From 3a179233a9411a830d98d536d8a6933da0ee8e84 Mon Sep 17 00:00:00 2001
From: wwweert123 <wwweert123>
Date: Thu, 29 Feb 2024 16:02:59 +0000
Subject: [PATCH 1/2] feat: add support for finding prefix in radixtree

---
 src/node/mod.rs | 76 ++++++++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 75 insertions(+), 1 deletion(-)

diff --git a/src/node/mod.rs b/src/node/mod.rs
index e5cfeb0..f6168fd 100644
--- a/src/node/mod.rs
+++ b/src/node/mod.rs
@@ -825,6 +825,19 @@ impl<S: BlobStore> TreeNode<S> {
         })
     }
 
+    /// True if the key is a prefix of the set
+    fn is_prefix(&self, prefix: &[u8], store: &S) -> Result<bool, S::Error> {
+        // 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<Option<Value<S>>, 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<S: BlobStore, T>(
     f(fr)
 }
 
+enum FindPrefixResult<T> {
+    // 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<S: BlobStore, T>(
+    store: &S,
+    tree: &TreeNodeRef<S>,
+    prefix: &[u8],
+    f: impl FnOnce(FindPrefixResult<&TreeNodeRef<S>>) -> Result<T, S::Error>,
+) -> Result<T, S::Error> {
+    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<S: BlobStore>(
     node: &TreeNodeRef<S>,
@@ -1491,7 +1560,7 @@ fn filter_prefix<S: BlobStore>(
     prefix: &[u8],
     substitution: &[u8],
 ) -> Result<TreeNode<S>, S::Error> {
-    find(store, node, prefix, |x| {
+    find(store, node, prefix, |x: FindResult<&TreeNodeRef<'_, S>>| {
         Ok(match x {
             FindResult::Found(res) => {
                 let mut res = res.to_owned();
@@ -3509,6 +3578,11 @@ impl<S: BlobStore + Clone> RadixTree<S> {
         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<bool, S::Error> {
+        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))

From d2d5bb7eb6fcd72d9cdde965b3c7cb3c2688bf6e Mon Sep 17 00:00:00 2001
From: wwweert123 <wwweert123>
Date: Fri, 1 Mar 2024 10:30:47 +0000
Subject: [PATCH 2/2] fix: fix code

---
 src/node/mod.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/node/mod.rs b/src/node/mod.rs
index f6168fd..a7d0bd8 100644
--- a/src/node/mod.rs
+++ b/src/node/mod.rs
@@ -1560,7 +1560,7 @@ fn filter_prefix<S: BlobStore>(
     prefix: &[u8],
     substitution: &[u8],
 ) -> Result<TreeNode<S>, S::Error> {
-    find(store, node, prefix, |x: FindResult<&TreeNodeRef<'_, S>>| {
+    find(store, node, prefix, |x| {
         Ok(match x {
             FindResult::Found(res) => {
                 let mut res = res.to_owned();