Fix common prefix iterator to support string keys

sile · Dec 2, 2023 · 9d79d19 · 9d79d19
1 parent 0a249b6
commit 9d79d19
Show file tree

Hide file tree

Showing 4 changed files with 47 additions and 22 deletions.
diff --git a/src/lib.rs b/src/lib.rs
@@ -87,6 +87,9 @@ pub trait BorrowedBytes {
     fn is_empty(&self) -> bool {
         self.as_bytes().is_empty()
     }
+
+    /// Returns a suffix of this instance not containing the first `n` bytes.
+    fn strip_n_prefix(&self, n: usize) -> &Self;
 }
 
 impl BorrowedBytes for [u8] {
@@ -114,6 +117,10 @@ impl BorrowedBytes for [u8] {
     fn cmp_first_item(&self, bytes: &[u8]) -> Ordering {
         self.first().cmp(&bytes.first())
     }
+
+    fn strip_n_prefix(&self, n: usize) -> &Self {
+        &self[n..]
+    }
 }
 
 impl BorrowedBytes for str {
@@ -147,4 +154,8 @@ impl BorrowedBytes for str {
             .next()
             .cmp(&Self::from_bytes(bytes).chars().next())
     }
+
+    fn strip_n_prefix(&self, n: usize) -> &Self {
+        &self[n..]
+    }
 }
diff --git a/src/map.rs b/src/map.rs
@@ -275,14 +275,17 @@ impl<K: Bytes, V> GenericPatriciaMap<K, V> {
     ///     .flatten()
     ///     .eq(vec![&"a", &"b", &"c", &"d"].into_iter()));
     /// ```
-    pub fn common_prefixes<'a, 'b>(&'a self, key: &'b [u8]) -> CommonPrefixesIter<'a, 'b, K, V>
+    pub fn common_prefixes<'a, 'b, Q>(
+        &'a self,
+        key: &'b Q,
+    ) -> CommonPrefixesIter<'a, 'b, K::Borrowed, V>
     where
         'a: 'b,
+        Q: ?Sized + AsRef<K::Borrowed>,
     {
         CommonPrefixesIter {
-            key_bytes: key,
-            iterator: self.tree.common_prefixes(key),
-            _key: PhantomData,
+            key_bytes: key.as_ref().as_bytes(),
+            iterator: self.tree.common_prefixes(key.as_ref()),
         }
     }
 
@@ -304,9 +307,13 @@ impl<K: Bytes, V> GenericPatriciaMap<K, V> {
     ///     .flatten()
     ///     .eq(vec![&"a", &"b", &"c", &"d"].into_iter()));
     /// ```
-    pub fn common_prefix_values<'a>(&'a self, key: &[u8]) -> impl 'a + Iterator<Item = &'a V> {
+    pub fn common_prefix_values<'a, 'b, Q>(&'a self, key: &'b Q) -> impl 'a + Iterator<Item = &'a V>
+    where
+        'b: 'a,
+        Q: ?Sized + AsRef<K::Borrowed>,
+    {
         self.tree
-            .common_prefixes(key.to_owned())
+            .common_prefixes(key.as_ref())
             .filter_map(|(_, n)| n.value())
     }
 
@@ -676,18 +683,20 @@ impl<'a, V: 'a> Iterator for ValuesMut<'a, V> {
 /// An iterator over entries in a `PatriciaMap` that share a common prefix with
 /// a given key.
 #[derive(Debug)]
-pub struct CommonPrefixesIter<'a, 'b, K, V> {
+pub struct CommonPrefixesIter<'a, 'b, K: ?Sized, V> {
     key_bytes: &'b [u8],
-    iterator: node::CommonPrefixesIter<'a, &'b [u8], V>,
-    _key: PhantomData<K>,
+    iterator: node::CommonPrefixesIter<'a, 'b, K, V>,
 }
-impl<'a, 'b, K: 'b + Bytes, V> Iterator for CommonPrefixesIter<'a, 'b, K, V> {
-    type Item = (&'b K::Borrowed, &'a V);
+impl<'a, 'b, K, V> Iterator for CommonPrefixesIter<'a, 'b, K, V>
+where
+    K: 'b + ?Sized + BorrowedBytes,
+{
+    type Item = (&'b K, &'a V);
 
     fn next(&mut self) -> Option<Self::Item> {
         for (prefix_len, n) in self.iterator.by_ref() {
             if let Some(v) = n.value() {
-                return Some((K::Borrowed::from_bytes(&self.key_bytes[..prefix_len]), v));
+                return Some((K::from_bytes(&self.key_bytes[..prefix_len]), v));
             }
         }
 

diff --git a/src/node.rs b/src/node.rs
@@ -395,9 +395,12 @@ impl<V> Node<V> {
         }
     }
 
-    pub(crate) fn common_prefixes<K>(&self, key: K) -> CommonPrefixesIter<K, V>
+    pub(crate) fn common_prefixes<'a, 'b, K>(
+        &'a self,
+        key: &'b K,
+    ) -> CommonPrefixesIter<'a, 'b, K, V>
     where
-        K: AsRef<[u8]>,
+        K: ?Sized + BorrowedBytes,
     {
         CommonPrefixesIter {
             key,
@@ -849,20 +852,19 @@ impl<'a, V: 'a> Iterator for IterMut<'a, V> {
 /// An iterator over entries in that collects all values up to
 /// until the key stops matching.
 #[derive(Debug)]
-pub(crate) struct CommonPrefixesIter<'a, K, V> {
-    key: K,
+pub(crate) struct CommonPrefixesIter<'a, 'b, K: ?Sized, V> {
+    key: &'b K,
     stack: Vec<(usize, &'a Node<V>)>,
 }
 
-impl<'a, K, V> Iterator for CommonPrefixesIter<'a, K, V>
+impl<'a, 'b, K, V> Iterator for CommonPrefixesIter<'a, 'b, K, V>
 where
-    K: AsRef<[u8]>,
+    K: ?Sized + BorrowedBytes,
 {
     type Item = (usize, &'a Node<V>);
     fn next(&mut self) -> Option<Self::Item> {
         while let Some((offset, node)) = self.stack.pop() {
-            // TODO:
-            let key = &self.key.as_ref()[offset..];
+            let key = self.key.strip_n_prefix(offset);
             let (_next, common_prefix_len) = key.strip_common_prefix_and_len(node.label());
             if common_prefix_len == 0 && key.cmp_first_item(node.label()).is_ge() {
                 if let Some(sibling) = node.sibling() {

diff --git a/src/tree.rs b/src/tree.rs
@@ -82,9 +82,12 @@ impl<V> PatriciaTree<V> {
             None
         }
     }
-    pub(crate) fn common_prefixes<K>(&self, key: K) -> node::CommonPrefixesIter<K, V>
+    pub(crate) fn common_prefixes<'a, 'b, K>(
+        &'a self,
+        key: &'b K,
+    ) -> node::CommonPrefixesIter<'a, 'b, K, V>
     where
-        K: AsRef<[u8]>,
+        K: ?Sized + BorrowedBytes,
     {
         self.root.common_prefixes(key)
     }