Merge pull request #23 from sile/split

Make split trees have empty root nodes
sile · Jan 23, 2023 · 3656cd2 · 3656cd2
2 parents 56b5e3c + 043e01f
commit 3656cd2
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 3 deletions.
diff --git a/src/set.rs b/src/set.rs
@@ -402,10 +402,13 @@ mod tests {
         assert_eq!(splitted_set.iter().collect::<Vec<_>>(), [b"bar"]);
 
         let mut set: PatriciaSet = vec!["foo", "bar", "baz"].into_iter().collect();
-        let splitted_set = set.split_by_prefix("baz");
+        let mut splitted_set = set.split_by_prefix("baz");
         assert_eq!(set.iter().collect::<Vec<_>>(), [b"bar", b"foo"]);
         assert_eq!(splitted_set.iter().collect::<Vec<_>>(), [b"baz"]);
 
+        splitted_set.insert("aaa");
+        assert_eq!(splitted_set.iter().collect::<Vec<_>>(), [b"aaa", b"baz"]);
+
         let mut set: PatriciaSet = vec!["foo", "bar", "baz"].into_iter().collect();
         let splitted_set = set.split_by_prefix("bazz");
         assert_eq!(set.iter().collect::<Vec<_>>(), [b"bar", b"baz", b"foo"]);

diff --git a/src/tree.rs b/src/tree.rs
@@ -84,8 +84,9 @@ impl<V> PatriciaTree<V> {
     }
     pub fn split_by_prefix<K: AsRef<[u8]>>(&mut self, prefix: K) -> Self {
         if let Some(splitted_root) = self.root.split_by_prefix(prefix.as_ref(), 0) {
-            let splitted_root = Node::new(prefix.as_ref(), None, Some(splitted_root), None);
-            let splitted = Self::from(splitted_root);
+            let mut splitted_root = Node::new(prefix.as_ref(), None, Some(splitted_root), None);
+            splitted_root.try_merge_with_child(1);
+            let splitted = Self::from(Node::new(b"", None, Some(splitted_root), None));
             self.len -= splitted.len();
             splitted
         } else {