diff --git a/src/set.rs b/src/set.rs index b676682..720576b 100644 --- a/src/set.rs +++ b/src/set.rs @@ -402,10 +402,13 @@ mod tests { assert_eq!(splitted_set.iter().collect::>(), [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::>(), [b"bar", b"foo"]); assert_eq!(splitted_set.iter().collect::>(), [b"baz"]); + splitted_set.insert("aaa"); + assert_eq!(splitted_set.iter().collect::>(), [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::>(), [b"bar", b"baz", b"foo"]); diff --git a/src/tree.rs b/src/tree.rs index ac6a109..9131296 100644 --- a/src/tree.rs +++ b/src/tree.rs @@ -84,8 +84,9 @@ impl PatriciaTree { } pub fn split_by_prefix>(&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 {