diff --git a/src/lib.rs b/src/lib.rs index ed20ca2..cb7441a 100644 --- 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 index 71d2720..c9d7983 100644 --- a/src/map.rs +++ b/src/map.rs @@ -275,14 +275,17 @@ impl GenericPatriciaMap { /// .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, { 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 GenericPatriciaMap { /// .flatten() /// .eq(vec![&"a", &"b", &"c", &"d"].into_iter())); /// ``` - pub fn common_prefix_values<'a>(&'a self, key: &[u8]) -> impl 'a + Iterator { + pub fn common_prefix_values<'a, 'b, Q>(&'a self, key: &'b Q) -> impl 'a + Iterator + where + 'b: 'a, + Q: ?Sized + AsRef, + { 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, + 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 { 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 index 97ee21e..ef2254e 100644 --- a/src/node.rs +++ b/src/node.rs @@ -395,9 +395,12 @@ impl Node { } } - pub(crate) fn common_prefixes(&self, key: K) -> CommonPrefixesIter + 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)>, } -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); fn next(&mut self) -> Option { 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 index d7bec84..236952f 100644 --- a/src/tree.rs +++ b/src/tree.rs @@ -82,9 +82,12 @@ impl PatriciaTree { None } } - pub(crate) fn common_prefixes(&self, key: K) -> node::CommonPrefixesIter + 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) }