Skip to content

Commit

Permalink
feat(tree_index, #157): implement TreeIndex::peek_entry (#160)
Browse files Browse the repository at this point in the history
  • Loading branch information
@qthree
qthree authored Sep 24, 2024
1 parent b801f1d commit a298240
Show file tree
Hide file tree
Showing 5 changed files with 79 additions and 34 deletions.
35 changes: 34 additions & 1 deletion src/tree_index.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -618,7 +618,7 @@ where
Q: Ord + ?Sized,
{
if let Some(root_ref) = self.root.load(Acquire, guard).as_ref() {
return root_ref.search(key, guard);
return root_ref.search(key, guard).map(|(_k, v)| v);
}
None
}
Expand All @@ -645,6 +645,39 @@ where
self.peek(key, &guard).map(|v| reader(key, v))
}

/// Returns a guarded reference to the key-value pair for the specified key without acquiring locks.
///
/// Returns `None` if the key does not exist. The returned reference can survive as long as the
/// associated [`Guard`] is alive.
///
/// # Examples
///
/// ```
/// use std::sync::Arc;
/// use scc::TreeIndex;
/// use scc::ebr::Guard;
///
/// let treeindex: TreeIndex<Arc<str>, u32> = TreeIndex::new();
/// treeindex.insert("foo".into(), 1).expect("insert in empty TreeIndex");
///
/// let guard = Guard::new();
/// let key: Arc<str> = treeindex
/// .peek_entry("foo", &guard)
/// .map(|(k, _v)| Arc::clone(k))
/// .expect("peek by borrowed value");
/// ```
#[inline]
pub fn peek_entry<'g, Q>(&self, key: &Q, guard: &'g Guard) -> Option<(&'g K, &'g V)>
where
K: Borrow<Q>,
Q: Ord + ?Sized,
{
if let Some(root_ref) = self.root.load(Acquire, guard).as_ref() {
return root_ref.search(key, guard);
}
None
}

/// Returns `true` if the [`TreeIndex`] contains the key.
///
/// # Examples
Expand Down
18 changes: 9 additions & 9 deletions src/tree_index/internal_node.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -146,7 +146,7 @@ where
{
/// Searches for an entry associated with the given key.
#[inline]
pub(super) fn search<'g, Q>(&self, key: &Q, guard: &'g Guard) -> Option<&'g V>
pub(super) fn search<'g, Q>(&self, key: &Q, guard: &'g Guard) -> Option<(&'g K, &'g V)>
where
K: 'g + Borrow<Q>,
Q: Ord + ?Sized,
Expand Down Expand Up @@ -1145,15 +1145,15 @@ mod test {
match internal_node.insert(k, k, &mut (), &guard) {
Ok(result) => match result {
InsertResult::Success => {
assert_eq!(internal_node.search(&k, &guard), Some(&k));
assert_eq!(internal_node.search(&k, &guard), Some((&k, &k)));
}
InsertResult::Duplicate(..)
| InsertResult::Frozen(..)
| InsertResult::Retired(..) => unreachable!(),
InsertResult::Full(_, _) => {
internal_node.rollback(&guard);
for j in 0..k {
assert_eq!(internal_node.search(&j, &guard), Some(&j));
assert_eq!(internal_node.search(&j, &guard), Some((&j, &j)));
if j == k - 1 {
assert!(matches!(
internal_node.remove_if::<_, _, _>(
Expand All @@ -1176,13 +1176,13 @@ mod test {
InsertResult::Retry(k, v) => {
let result = internal_node.insert(k, v, &mut (), &guard);
assert!(result.is_ok());
assert_eq!(internal_node.search(&k, &guard), Some(&k));
assert_eq!(internal_node.search(&k, &guard), Some((&k, &k)));
}
},
Err((k, v)) => {
let result = internal_node.insert(k, v, &mut (), &guard);
assert!(result.is_ok());
assert_eq!(internal_node.search(&k, &guard), Some(&k));
assert_eq!(internal_node.search(&k, &guard), Some((&k, &k)));
}
}
}
Expand Down Expand Up @@ -1239,7 +1239,7 @@ mod test {
if max_key.map_or(false, |m| m == id) {
break;
}
assert_eq!(internal_node_clone.search(&id, &guard), Some(&id));
assert_eq!(internal_node_clone.search(&id, &guard), Some((&id, &id)));
}
for id in range {
if max_key.map_or(false, |m| m == id) {
Expand Down Expand Up @@ -1319,7 +1319,7 @@ mod test {
};
assert_eq!(
internal_node_clone.search(&fixed_point, &guard).unwrap(),
&fixed_point
(&fixed_point, &fixed_point)
);
}
{
Expand All @@ -1336,7 +1336,7 @@ mod test {
}
assert_eq!(
internal_node_clone.search(&fixed_point, &guard).unwrap(),
&fixed_point
(&fixed_point, &fixed_point)
);
}
for i in 0..workload_size {
Expand All @@ -1362,7 +1362,7 @@ mod test {
);
assert_eq!(
internal_node_clone.search(&fixed_point, &guard).unwrap(),
&fixed_point
(&fixed_point, &fixed_point)
);
}
}
Expand Down
34 changes: 20 additions & 14 deletions src/tree_index/leaf.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -482,13 +482,14 @@ where

/// Returns a value associated with the key.
#[inline]
pub(super) fn search<Q>(&self, key: &Q) -> Option<&V>
pub(super) fn search<Q>(&self, key: &Q) -> Option<(&K, &V)>
where
K: Borrow<Q>,
Q: Ord + ?Sized,
{
let metadata = self.metadata.load(Acquire);
self.search_slot(key, metadata).map(|i| self.value_at(i))
self.search_slot(key, metadata)
.map(|i| (self.key_at(i), self.value_at(i)))
}

/// Returns the index of the key-value pair that is smaller than the given key.
Expand Down Expand Up @@ -551,7 +552,8 @@ where
let rank = mutable_metadata % (1_usize << DIMENSION.num_bits_per_entry);
if rank < min_max_rank && rank > max_min_rank {
let k = self.key_at(i);
match k.borrow().cmp(key) {
let kq: &Q = k.borrow();
match kq.cmp(key) {
Ordering::Less => {
if max_min_rank < rank {
max_min_rank = rank;
Expand Down Expand Up @@ -711,7 +713,8 @@ where
K: Borrow<Q>,
Q: Ord + ?Sized,
{
self.key_at(index).borrow().cmp(key)
let kq: &Q = self.key_at(index).borrow();
kq.cmp(key)
}
}

Expand Down Expand Up @@ -998,8 +1001,8 @@ mod test {
leaf.insert("GOOD DAY".to_owned(), "OH MY GOD!!".to_owned()),
InsertResult::Success
));
assert_eq!(leaf.search("MY GOODNESS!").unwrap(), "OH MY GOD!!");
assert_eq!(leaf.search("GOOD DAY").unwrap(), "OH MY GOD!!");
assert_eq!(leaf.search("MY GOODNESS!").unwrap().1, "OH MY GOD!!");
assert_eq!(leaf.search("GOOD DAY").unwrap().1, "OH MY GOD!!");

for i in 0..DIMENSION.num_entries {
if let InsertResult::Full(k, v) = leaf.insert(i.to_string(), i.to_string()) {
Expand All @@ -1008,7 +1011,10 @@ mod test {
assert_eq!(v, i.to_string());
break;
}
assert_eq!(leaf.search(&i.to_string()).unwrap(), &i.to_string());
assert_eq!(
leaf.search(&i.to_string()).unwrap(),
(&i.to_string(), &i.to_string())
);
}

for i in 0..DIMENSION.num_entries {
Expand Down Expand Up @@ -1117,8 +1123,8 @@ mod test {
let mut leaf1 = None;
let mut leaf2 = None;
leaf.freeze_and_distribute(&mut leaf1, &mut leaf2);
assert_eq!(leaf1.as_ref().and_then(|l| l.search(&11)), Some(&17));
assert_eq!(leaf1.as_ref().and_then(|l| l.search(&17)), Some(&11));
assert_eq!(leaf1.as_ref().and_then(|l| l.search(&11)), Some((&11, &17)));
assert_eq!(leaf1.as_ref().and_then(|l| l.search(&17)), Some((&17, &11)));
assert!(leaf2.is_none());
assert!(matches!(leaf.insert(1, 7), InsertResult::Frozen(..)));
assert_eq!(leaf.remove_if(&17, &mut |_| true), RemoveResult::Frozen);
Expand Down Expand Up @@ -1162,7 +1168,7 @@ mod test {
assert_eq!(result.0, Some((&i, &i)));
}
for i in 0..insert {
assert_eq!(*leaf.search(&i).unwrap(), i);
assert_eq!(leaf.search(&i).unwrap(), (&i, &i));
}
if insert == DIMENSION.num_entries {
assert!(matches!(leaf.insert(usize::MAX, usize::MAX), InsertResult::Full(..)));
Expand Down Expand Up @@ -1220,7 +1226,7 @@ mod test {
barrier_clone.wait().await;
let inserted = match leaf_clone.insert(t, t) {
InsertResult::Success => {
assert_eq!(*leaf_clone.search(&t).unwrap(), t);
assert_eq!(leaf_clone.search(&t).unwrap(), (&t, &t));
true
}
InsertResult::Duplicate(_, _)
Expand Down Expand Up @@ -1249,7 +1255,7 @@ mod test {
barrier_clone.wait().await;
assert_eq!((*full_clone).load(Relaxed), num_excess);
if inserted {
assert_eq!(*leaf_clone.search(&t).unwrap(), t);
assert_eq!(leaf_clone.search(&t).unwrap(), (&t, &t));
}
{
let scanner = Scanner::new(&leaf_clone);
Expand Down Expand Up @@ -1303,11 +1309,11 @@ mod test {
let _result = leaf_clone.insert(i, i);
}
assert!(!leaf_clone.is_retired());
assert_eq!(leaf_clone.search(&k).unwrap(), &k);
assert_eq!(leaf_clone.search(&k).unwrap(), (&k, &k));
}
for i in 0..workload_size {
let _result = leaf_clone.remove_if(&i, &mut |v| *v != k);
assert_eq!(leaf_clone.search(&k).unwrap(), &k);
assert_eq!(leaf_clone.search(&k).unwrap(), (&k, &k));
}
}
}));
Expand Down
24 changes: 15 additions & 9 deletions src/tree_index/leaf_node.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -160,7 +160,7 @@ where
{
/// Searches for an entry associated with the given key.
#[inline]
pub(super) fn search<'g, Q>(&self, key: &Q, guard: &'g Guard) -> Option<&'g V>
pub(super) fn search<'g, Q>(&self, key: &Q, guard: &'g Guard) -> Option<(&'g K, &'g V)>
where
K: 'g + Borrow<Q>,
Q: Ord + ?Sized,
Expand Down Expand Up @@ -1169,10 +1169,13 @@ mod test {
Ok(InsertResult::Success)
));
assert_eq!(
leaf_node.search("MY GOODNESS!", &guard).unwrap(),
leaf_node.search("MY GOODNESS!", &guard).unwrap().1,
"OH MY GOD!!"
);
assert_eq!(
leaf_node.search("GOOD DAY", &guard).unwrap().1,
"OH MY GOD!!"
);
assert_eq!(leaf_node.search("GOOD DAY", &guard).unwrap(), "OH MY GOD!!");
assert!(matches!(
leaf_node.remove_if::<_, _, _>("GOOD DAY", &mut |v| v == "OH MY", &mut (), &guard),
Ok(RemoveResult::Fail)
Expand Down Expand Up @@ -1211,7 +1214,7 @@ mod test {
}
match result.unwrap() {
InsertResult::Success => {
assert_eq!(leaf_node.search(&k, &guard), Some(&k));
assert_eq!(leaf_node.search(&k, &guard), Some((&k, &k)));
continue;
}
InsertResult::Duplicate(..)
Expand All @@ -1220,13 +1223,16 @@ mod test {
InsertResult::Full(_, _) => {
leaf_node.rollback(&guard);
for r in 0..(k - 1) {
assert_eq!(leaf_node.search(&r, &guard), Some(&r));
assert_eq!(leaf_node.search(&r, &guard), Some((&r, &r)));
assert!(leaf_node
.remove_if::<_, _, _>(&r, &mut |_| true, &mut (), &guard)
.is_ok());
assert_eq!(leaf_node.search(&r, &guard), None);
}
assert_eq!(leaf_node.search(&(k - 1), &guard), Some(&(k - 1)));
assert_eq!(
leaf_node.search(&(k - 1), &guard),
Some((&(k - 1), &(k - 1)))
);
assert_eq!(
leaf_node.remove_if::<_, _, _>(&(k - 1), &mut |_| true, &mut (), &guard),
Ok(RemoveResult::Retired)
Expand Down Expand Up @@ -1292,7 +1298,7 @@ mod test {
if max_key.map_or(false, |m| m == id) {
break;
}
assert_eq!(leaf_node_clone.search(&id, &guard), Some(&id));
assert_eq!(leaf_node_clone.search(&id, &guard), Some((&id, &id)));
}
for id in range {
if max_key.map_or(false, |m| m == id) {
Expand Down Expand Up @@ -1376,7 +1382,7 @@ mod test {
leaf_node_clone.rollback(&guard);
}
}
assert_eq!(leaf_node_clone.search(&k, &guard).unwrap(), &k);
assert_eq!(leaf_node_clone.search(&k, &guard).unwrap(), (&k, &k));
}
for i in 0..workload_size {
let max_scanner = leaf_node_clone.max_le_appr(&k, &guard).unwrap();
Expand All @@ -1397,7 +1403,7 @@ mod test {
&mut (),
&guard,
);
assert_eq!(leaf_node_clone.search(&k, &guard).unwrap(), &k);
assert_eq!(leaf_node_clone.search(&k, &guard).unwrap(), (&k, &k));
}
}
}));
Expand Down
2 changes: 1 addition & 1 deletion src/tree_index/node.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -65,7 +65,7 @@ where
{
/// Searches for an entry associated with the given key.
#[inline]
pub(super) fn search<'g, Q>(&self, key: &Q, guard: &'g Guard) -> Option<&'g V>
pub(super) fn search<'g, Q>(&self, key: &Q, guard: &'g Guard) -> Option<(&'g K, &'g V)>
where
K: 'g + Borrow<Q>,
Q: Ord + ?Sized,
Expand Down

0 comments on commit a298240

Please sign in to comment.