diff --git a/src/node/mod.rs b/src/node/mod.rs index e5cfeb0..a7d0bd8 100644 --- a/src/node/mod.rs +++ b/src/node/mod.rs @@ -825,6 +825,19 @@ impl TreeNode { }) } + /// True if the key is a prefix of the set + fn is_prefix(&self, prefix: &[u8], store: &S) -> Result { + // if we find a tree at exactly the location, and it has a value, we have a hit + find_is_prefix(store, &TreeNodeRef(Ok(self)), prefix, |r| { + Ok(match r { + FindPrefixResult::Found(tree) => tree.value_opt().is_some(), + FindPrefixResult::IsPrefix { .. } => true, + FindPrefixResult::NotFound => false, + }) + }) + } + + fn get(&self, key: &[u8], store: &S) -> Result>, S::Error> { // if we find a tree at exactly the location, and it has a value, we have a hit find(store, &TreeNodeRef(Ok(self)), key, |r| { @@ -1484,6 +1497,62 @@ fn find( f(fr) } +enum FindPrefixResult { + // Found an exact match + Found(T), + // did not find anything, T is the closest match, with n remaining (unmatched) in the prefix of T + NotFound, + IsPrefix +} + +/// find a prefix in a tree. Will either return +/// - Found(tree) if we found the tree exactly, +/// - Prefix if we found a tree of which prefix is a prefix +/// - NotFound if there is no tree +fn find_is_prefix( + store: &S, + tree: &TreeNodeRef, + prefix: &[u8], + f: impl FnOnce(FindPrefixResult<&TreeNodeRef>) -> Result, +) -> Result { + let tree_prefix = tree.load_prefix(store)?; + let n = common_prefix(&tree_prefix, prefix); + // remaining in tree prefix + let rt = tree_prefix.len() - n; + // remaining in prefix + let rp = prefix.len() - n; + let fr = if rp == 0 && rt == 0 { + // direct hit, however need to check if this node has a value + FindPrefixResult::Found(tree) + } else if rp == 0 { + // tree does not contain a key that is a direct prefix of the string + FindPrefixResult::NotFound + } else if rt == 0 { + // Here we can check if the current node is has a key, + // If yes, FindResult is IsPrefix + if tree.value_opt().is_some() { + FindPrefixResult::IsPrefix + } else { + // else we keep looking + let c = prefix[n]; + let tree_children = tree.load_children(store)?; + if let Some(tree_children) = tree_children { + if let Some(child) = tree_children.find(c) { + return find_is_prefix(store, &child, &prefix[n..], f); + } else { + FindPrefixResult::NotFound + } + } else { + FindPrefixResult::NotFound + } + } + } else { + // disjoint, but we still need to store how far we matched + FindPrefixResult::NotFound + }; + f(fr) +} + /// Return the subtree with the given prefix. Will return an empty tree in case there is no match. fn filter_prefix( node: &TreeNodeRef, @@ -3509,6 +3578,11 @@ impl RadixTree { self.node.has_prefix(prefix.as_ref(), &self.store) } + #[cfg_attr(feature = "custom-store", visibility::make(pub))] + fn try_is_prefix(&self, key: impl AsRef<[u8]>) -> Result { + self.node.is_prefix(key.as_ref(), &self.store) + } + #[cfg_attr(feature = "custom-store", visibility::make(pub))] fn try_remove_prefix(&mut self, prefix: impl AsRef<[u8]>) -> Result<(), S::Error> { self.try_remove_prefix_with(&RadixTree::single(prefix, &[]), |_| Ok(true))