use std::ops::Bound;
#[cfg(feature = "loom")]
use loom::sync::atomic::Ordering;
#[cfg(not(feature = "loom"))]
use std::sync::atomic::Ordering;
use crate::key::Key;
use super::SkipList;
use super::back_finger::BackFinger;
use super::node::SkipNode;
use super::strip_mark;
macro_rules! skip_key_iter {
($vis:vis) => {
$vis struct SkipKeyIter<'g, K: Key, N: SkipNode> {
list: &'g SkipList<N>,
front: *mut N,
back: Option<*mut N>,
back_finger: BackFinger<N>,
end: Bound<K>,
start: Bound<K>,
against_storage: bool,
reversed: bool,
done: bool,
_guard: seize::LocalGuard<'g>,
}
};
}
#[cfg(any(test, feature = "loom", miri))]
skip_key_iter!(pub);
#[cfg(not(any(test, feature = "loom", miri)))]
skip_key_iter!(pub(crate));
impl<'g, K: Key, N: SkipNode> SkipKeyIter<'g, K, N> {
pub(crate) fn bounded_dir(
list: &'g SkipList<N>,
start: Bound<K>,
end: Bound<K>,
front: *mut N,
against_storage: bool,
reversed: bool,
guard: seize::LocalGuard<'g>,
) -> Self {
Self {
list,
front,
back: None,
back_finger: BackFinger::empty(),
end,
start,
against_storage,
reversed,
done: false,
_guard: guard,
}
}
fn resolve_back(&mut self) -> *mut N {
match &self.end {
Bound::Unbounded => self
.list
.back_seek_last(&mut self.back_finger, &self._guard),
Bound::Excluded(key) => {
self.list
.back_step(&mut self.back_finger, key.as_bytes(), &self._guard)
}
Bound::Included(key) => {
let ge = self.list.find_first_ge(key.as_bytes(), &self._guard);
if !ge.is_null()
&& !unsafe { &*ge }.is_marked()
&& unsafe { &*ge }.key_bytes() == key.as_bytes()
{
ge
} else {
self.list
.back_step(&mut self.back_finger, key.as_bytes(), &self._guard)
}
}
}
}
#[inline(always)]
fn check_end(&self, key: &[u8]) -> bool {
match &self.end {
Bound::Unbounded => true,
Bound::Excluded(end) => {
if self.reversed {
key > end.as_bytes()
} else {
key < end.as_bytes()
}
}
Bound::Included(end) => {
if self.reversed {
key >= end.as_bytes()
} else {
key <= end.as_bytes()
}
}
}
}
#[inline(always)]
fn check_start(&self, key: &[u8]) -> bool {
match &self.start {
Bound::Unbounded => true,
Bound::Excluded(start) => {
if self.reversed {
key < start.as_bytes()
} else {
key > start.as_bytes()
}
}
Bound::Included(start) => {
if self.reversed {
key <= start.as_bytes()
} else {
key >= start.as_bytes()
}
}
}
}
fn next_with_storage(&mut self) -> Option<&'g N> {
loop {
if self.done || self.front.is_null() {
return None;
}
let node = unsafe { &*self.front };
self.front = strip_mark(node.tower(0).load(Ordering::Acquire));
if node.is_marked() {
continue;
}
if !self.check_end(node.key_bytes()) {
self.done = true;
return None;
}
return Some(node);
}
}
fn next_against_storage(&mut self) -> Option<&'g N> {
debug_assert!(
self.front.is_null(),
"against-storage walk assumes a null front; add a convergence check first"
);
if self.back.is_none() {
self.back = Some(self.resolve_back());
if self.front.is_null() && self.back == Some(std::ptr::null_mut()) {
self.done = true;
}
}
loop {
let back = self.back.unwrap_or(std::ptr::null_mut());
if self.done || back.is_null() {
return None;
}
let node = unsafe { &*back };
let key = node.key_bytes();
self.back = Some(
self.list
.back_step(&mut self.back_finger, key, &self._guard),
);
if node.is_marked() {
continue;
}
if !self.check_start(key) {
self.done = true;
return None;
}
return Some(node);
}
}
}
impl<'g, K: Key, N: SkipNode> Iterator for SkipKeyIter<'g, K, N> {
type Item = &'g N;
fn next(&mut self) -> Option<Self::Item> {
if self.against_storage {
self.next_against_storage()
} else {
self.next_with_storage()
}
}
}
pub(crate) struct SkipKeyKeys<'a, K: Key, N: SkipNode> {
inner: SkipKeyIter<'a, K, N>,
}
impl<'a, K: Key, N: SkipNode> SkipKeyKeys<'a, K, N> {
pub(crate) fn new(inner: SkipKeyIter<'a, K, N>) -> Self {
Self { inner }
}
}
impl<K: Key, N: SkipNode> Iterator for SkipKeyKeys<'_, K, N> {
type Item = K;
fn next(&mut self) -> Option<Self::Item> {
self.inner
.next()
.map(|node| K::from_bytes(node.key_bytes()))
}
}
#[cfg(any(feature = "loom", miri))]
#[allow(private_interfaces)]
pub fn test_all<'g, N: SkipNode>(
list: &'g SkipList<N>,
guard: seize::LocalGuard<'g>,
) -> SkipKeyIter<'g, [u8; 4], N> {
SkipKeyIter::bounded_dir(
list,
Bound::Unbounded,
Bound::Unbounded,
strip_mark(unsafe { (*list.head_ptr()).tower(0).load(Ordering::Acquire) }),
false,
false,
guard,
)
}
#[cfg(any(feature = "loom", miri))]
#[allow(private_interfaces)]
pub fn test_all_back<'g, N: SkipNode>(
list: &'g SkipList<N>,
guard: seize::LocalGuard<'g>,
) -> SkipKeyIter<'g, [u8; 4], N> {
SkipKeyIter::bounded_dir(
list,
Bound::Unbounded,
Bound::Unbounded,
std::ptr::null_mut(),
true,
false,
guard,
)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::disk_loc::DiskLoc;
use crate::skiplist::SkipList;
use crate::skiplist::node::ConstNode;
type TestNode = ConstNode<[u8; 8], 1, DiskLoc>;
type TestList = SkipList<TestNode>;
fn disk() -> DiskLoc {
DiskLoc::new(0, 0, 0)
}
fn key_byte(i: u8) -> [u8; 8] {
let mut k = [0u8; 8];
k[7] = i;
k
}
fn alloc_node(key: [u8; 8]) -> *mut TestNode {
TestNode::alloc(key, [0u8; 1], disk(), 1)
}
fn insert_keys<'a>(list: &'a TestList, keys: &[u8]) -> seize::LocalGuard<'a> {
let guard = list.collector().enter();
for &i in keys {
let node = alloc_node(key_byte(i));
list.insert(node, &guard);
}
guard
}
fn collect_keys(iter: SkipKeyIter<'_, [u8; 8], TestNode>) -> Vec<u8> {
iter.map(|n| n.key_bytes()[7]).collect()
}
#[test]
fn skip_key_iter_yields_keys_in_order() {
let list: TestList = SkipList::new(false);
let guard = insert_keys(&list, &[1, 3, 5]);
let iter = SkipKeyIter::bounded_dir(
&list,
Bound::Unbounded,
Bound::Unbounded,
strip_mark(unsafe { (*list.head_ptr()).tower(0).load(Ordering::Acquire) }),
false,
false,
guard,
);
assert_eq!(collect_keys(iter), vec![1, 3, 5]);
}
#[test]
fn skip_key_iter_against_storage_yields_desc_on_asc_list() {
let list: TestList = SkipList::new(false);
let guard = insert_keys(&list, &[1, 3, 5]);
let iter = SkipKeyIter::bounded_dir(
&list,
Bound::Unbounded,
Bound::Unbounded,
std::ptr::null_mut(),
true,
false,
guard,
);
assert_eq!(collect_keys(iter), vec![5, 3, 1]);
}
#[test]
fn skip_key_iter_range_bounds_respected() {
let list: TestList = SkipList::new(false);
let guard = insert_keys(&list, &[1, 2, 3, 4, 5]);
let start_key = key_byte(2);
let end_key = key_byte(5);
let front = list.find_first_ge(&start_key, &guard);
let iter = SkipKeyIter::bounded_dir(
&list,
Bound::Included(start_key),
Bound::Excluded(end_key),
front,
false,
false,
guard,
);
assert_eq!(collect_keys(iter), vec![2, 3, 4]);
}
}