diff --git a/src/lib.rs b/src/lib.rs index 7533d48..ed20ca2 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -82,6 +82,11 @@ pub trait BorrowedBytes { /// Compares the first item of this instance with the first item represented in the the given bytes. fn cmp_first_item(&self, bytes: &[u8]) -> Ordering; + + /// Returns `true` if this instance is empty, otherwise `false`. + fn is_empty(&self) -> bool { + self.as_bytes().is_empty() + } } impl BorrowedBytes for [u8] { diff --git a/src/map.rs b/src/map.rs index 0c093be..71d2720 100644 --- a/src/map.rs +++ b/src/map.rs @@ -129,7 +129,7 @@ impl GenericPatriciaMap { /// assert!(!map.contains_key("bar")); /// ``` pub fn contains_key>(&self, key: Q) -> bool { - self.tree.get(key.as_ref().as_bytes()).is_some() + self.tree.get(key.as_ref()).is_some() } /// Returns a reference to the value corresponding to the key. @@ -161,7 +161,7 @@ impl GenericPatriciaMap { /// assert_eq!(map.get("foo"), Some(&2)); /// ``` pub fn get_mut>(&mut self, key: Q) -> Option<&mut V> { - self.tree.get_mut(key.as_ref().as_bytes()) + self.tree.get_mut(key.as_ref()) } /// Finds the longest common prefix of `key` and the keys in this map, @@ -185,9 +185,7 @@ impl GenericPatriciaMap { where Q: ?Sized + AsRef, { - let (key, value) = self - .tree - .get_longest_common_prefix(key.as_ref().as_bytes())?; + let (key, value) = self.tree.get_longest_common_prefix(key.as_ref())?; Some((K::Borrowed::from_bytes(key), value)) } @@ -217,9 +215,7 @@ impl GenericPatriciaMap { where Q: ?Sized + AsRef, { - let (key, value) = self - .tree - .get_longest_common_prefix_mut(key.as_ref().as_bytes())?; + let (key, value) = self.tree.get_longest_common_prefix_mut(key.as_ref())?; Some((K::Borrowed::from_bytes(key), value)) } @@ -256,7 +252,7 @@ impl GenericPatriciaMap { /// assert_eq!(map.remove("foo"), None); /// ``` pub fn remove>(&mut self, key: Q) -> Option { - self.tree.remove(key.as_ref().as_bytes()) + self.tree.remove(key.as_ref()) } /// Returns an iterator that collects all entries in the map up to a certain key. @@ -338,7 +334,7 @@ impl GenericPatriciaMap { /// assert_eq!(b.keys().collect::>(), [b"elixir", b"erlang"]); /// ``` pub fn split_by_prefix>(&mut self, prefix: Q) -> Self { - let subtree = self.tree.split_by_prefix(prefix.as_ref().as_bytes()); + let subtree = self.tree.split_by_prefix(prefix.as_ref()); GenericPatriciaMap { tree: subtree, _key: PhantomData, @@ -455,7 +451,7 @@ impl GenericPatriciaMap { 'b: 'a, { self.tree - .iter_prefix(prefix.as_bytes()) + .iter_prefix(prefix) .into_iter() .flat_map(move |(prefix_len, nodes)| { Iter::::new(nodes, Vec::from(&prefix.as_bytes()[..prefix_len])) @@ -482,7 +478,7 @@ impl GenericPatriciaMap { 'b: 'a, { self.tree - .iter_prefix_mut(prefix.as_bytes()) + .iter_prefix_mut(prefix) .into_iter() .flat_map(move |(prefix_len, nodes)| { IterMut::::new(nodes, Vec::from(&prefix.as_bytes()[..prefix_len])) diff --git a/src/node.rs b/src/node.rs index 1b04447..16fa957 100644 --- a/src/node.rs +++ b/src/node.rs @@ -408,7 +408,7 @@ impl Node { pub(crate) fn get(&self, key: &K) -> Option<&V> { let (next, common_prefix_len) = key.strip_common_prefix_and_len(self.label()); if common_prefix_len == self.label().len() { - if next.as_bytes().is_empty() { + if next.is_empty() { self.value() } else { self.child().and_then(|child| child.get(next)) @@ -420,29 +420,27 @@ impl Node { } } - pub(crate) fn get_mut(&mut self, key: &[u8]) -> Option<&mut V> { - let common_prefix_len = self.skip_common_prefix(key); - let next = &key[common_prefix_len..]; + pub(crate) fn get_mut(&mut self, key: &K) -> Option<&mut V> { + let (next, common_prefix_len) = key.strip_common_prefix_and_len(self.label()); if common_prefix_len == self.label().len() { if next.is_empty() { self.value_mut() } else { self.child_mut().and_then(|child| child.get_mut(next)) } - } else if common_prefix_len == 0 && self.label().first() <= key.first() { + } else if common_prefix_len == 0 && key.cmp_first_item(self.label()).is_ge() { self.sibling_mut().and_then(|sibling| sibling.get_mut(next)) } else { None } } - pub(crate) fn get_longest_common_prefix( + pub(crate) fn get_longest_common_prefix( &self, - key: &[u8], + key: &K, offset: usize, ) -> Option<(usize, &V)> { - let common_prefix_len = self.skip_common_prefix(key); - let next = &key[common_prefix_len..]; + let (next, common_prefix_len) = key.strip_common_prefix_and_len(self.label()); if common_prefix_len == self.label().len() { let offset = offset + common_prefix_len; if next.is_empty() { @@ -452,20 +450,19 @@ impl Node { .and_then(|child| child.get_longest_common_prefix(next, offset)) .or_else(|| self.value().map(|v| (offset, v))) } - } else if common_prefix_len == 0 && self.label().first() <= key.first() { + } else if common_prefix_len == 0 && key.cmp_first_item(self.label()).is_ge() { self.sibling() .and_then(|sibling| sibling.get_longest_common_prefix(next, offset)) } else { None } } - pub(crate) fn get_longest_common_prefix_mut( + pub(crate) fn get_longest_common_prefix_mut( &mut self, - key: &[u8], + key: &K, offset: usize, ) -> Option<(usize, &mut V)> { - let common_prefix_len = self.skip_common_prefix(key); - let next = &key[common_prefix_len..]; + let (next, common_prefix_len) = key.strip_common_prefix_and_len(self.label()); if common_prefix_len == self.label().len() { let offset = offset + common_prefix_len; if next.is_empty() { @@ -476,7 +473,7 @@ impl Node { .and_then(|child| child.get_longest_common_prefix_mut(next, offset)) .or_else(|| this.value.map(|v| (offset, v))) } - } else if common_prefix_len == 0 && self.label().first() <= key.first() { + } else if common_prefix_len == 0 && key.cmp_first_item(self.label()).is_ge() { self.sibling_mut() .and_then(|sibling| sibling.get_longest_common_prefix_mut(next, offset)) } else { @@ -484,16 +481,19 @@ impl Node { } } - pub(crate) fn get_prefix_node(&self, key: &[u8], offset: usize) -> Option<(usize, &Self)> { - let common_prefix_len = self.skip_common_prefix(key); - let next = &key[common_prefix_len..]; + pub(crate) fn get_prefix_node( + &self, + key: &K, + offset: usize, + ) -> Option<(usize, &Self)> { + let (next, common_prefix_len) = key.strip_common_prefix_and_len(self.label()); if next.is_empty() { - Some((common_prefix_len, self)) + Some((common_prefix_len, self)) // TODO: return offset? } else if common_prefix_len == self.label().len() { let offset = offset + common_prefix_len; self.child() .and_then(|child| child.get_prefix_node(next, offset)) - } else if common_prefix_len == 0 && self.label().first() <= key.first() { + } else if common_prefix_len == 0 && key.cmp_first_item(self.label()).is_ge() { self.sibling() .and_then(|sibling| sibling.get_prefix_node(next, offset)) } else { @@ -501,20 +501,19 @@ impl Node { } } - pub(crate) fn get_prefix_node_mut( + pub(crate) fn get_prefix_node_mut( &mut self, - key: &[u8], + key: &K, offset: usize, ) -> Option<(usize, &mut Self)> { - let common_prefix_len = self.skip_common_prefix(key); - let next = &key[common_prefix_len..]; + let (next, common_prefix_len) = key.strip_common_prefix_and_len(self.label()); if next.is_empty() { - Some((common_prefix_len, self)) + Some((common_prefix_len, self)) // TODO: return offset? } else if common_prefix_len == self.label().len() { let offset = offset + common_prefix_len; self.child_mut() .and_then(|child| child.get_prefix_node_mut(next, offset)) - } else if common_prefix_len == 0 && self.label().first() <= key.first() { + } else if common_prefix_len == 0 && key.cmp_first_item(self.label()).is_ge() { self.sibling_mut() .and_then(|sibling| sibling.get_prefix_node_mut(next, offset)) } else { @@ -522,9 +521,13 @@ impl Node { } } - pub(crate) fn split_by_prefix(&mut self, prefix: &[u8], level: usize) -> Option { - let common_prefix_len = self.skip_common_prefix(prefix); - if common_prefix_len == prefix.len() { + pub(crate) fn split_by_prefix( + &mut self, + prefix: &K, + level: usize, + ) -> Option { + let (next, common_prefix_len) = prefix.strip_common_prefix_and_len(self.label()); + if common_prefix_len == prefix.as_bytes().len() { let value = self.take_value(); let child = self.take_child(); let node = Node::new(&self.label()[common_prefix_len..], value, child, None); @@ -533,7 +536,6 @@ impl Node { } Some(node) } else if common_prefix_len == self.label().len() { - let next = &prefix[common_prefix_len..]; self.child_mut() .and_then(|child| child.split_by_prefix(next, level + 1)) .map(|old| { @@ -541,8 +543,7 @@ impl Node { self.try_merge_with_child(level); old }) - } else if common_prefix_len == 0 && self.label().first() <= prefix.first() { - let next = &prefix[common_prefix_len..]; + } else if common_prefix_len == 0 && prefix.cmp_first_item(self.label()).is_ge() { self.sibling_mut() .and_then(|sibling| sibling.split_by_prefix(next, level)) .map(|old| { @@ -553,9 +554,8 @@ impl Node { None } } - pub(crate) fn remove(&mut self, key: &[u8], level: usize) -> Option { - let common_prefix_len = self.skip_common_prefix(key); - let next = &key[common_prefix_len..]; + pub(crate) fn remove(&mut self, key: &K, level: usize) -> Option { + let (next, common_prefix_len) = key.strip_common_prefix_and_len(self.label()); if common_prefix_len == self.label().len() { if next.is_empty() { self.take_value().map(|old| { @@ -571,7 +571,7 @@ impl Node { old }) } - } else if common_prefix_len == 0 && self.label().first() <= key.first() { + } else if common_prefix_len == 0 && key.cmp_first_item(self.label()).is_ge() { self.sibling_mut() .and_then(|sibling| sibling.remove(next, level)) .map(|old| { @@ -626,13 +626,6 @@ impl Node { None } } - fn skip_common_prefix(&self, key: &[u8]) -> usize { - self.label() - .iter() - .zip(key.iter()) - .take_while(|x| x.0 == x.1) - .count() - } pub(crate) fn flags(&self) -> Flags { Flags::from_bits_truncate(unsafe { *self.ptr }) } @@ -873,8 +866,10 @@ where type Item = (usize, &'a Node); fn next(&mut self) -> Option { while let Some((offset, node)) = self.stack.pop() { - let common_prefix_len = node.skip_common_prefix(&self.key.as_ref()[offset..]); - if common_prefix_len == 0 && node.label().first() <= self.key.as_ref().get(offset) { + // TODO: + let key = &self.key.as_ref()[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() { self.stack.push((offset, sibling)); } diff --git a/src/tree.rs b/src/tree.rs index bf06f94..2f53ac5 100644 --- a/src/tree.rs +++ b/src/tree.rs @@ -35,40 +35,49 @@ impl PatriciaTree { pub fn get(&self, key: &K) -> Option<&V> { self.root.get(key) } - pub fn get_mut(&mut self, key: &[u8]) -> Option<&mut V> { + pub fn get_mut(&mut self, key: &K) -> Option<&mut V> { self.root.get_mut(key) } - pub fn get_longest_common_prefix<'a>(&self, key: &'a [u8]) -> Option<(&'a [u8], &V)> { + pub fn get_longest_common_prefix<'a, K: ?Sized + BorrowedBytes>( + &self, + key: &'a K, + ) -> Option<(&'a [u8], &V)> { self.root .get_longest_common_prefix(key, 0) - .map(|(n, v)| (&key[..n], v)) + .map(|(n, v)| (&key.as_bytes()[..n], v)) } - pub fn get_longest_common_prefix_mut<'a>( + pub fn get_longest_common_prefix_mut<'a, K: ?Sized + BorrowedBytes>( &mut self, - key: &'a [u8], + key: &'a K, ) -> Option<(&'a [u8], &mut V)> { self.root .get_longest_common_prefix_mut(key, 0) - .map(|(n, v)| (&key[..n], v)) + .map(|(n, v)| (&key.as_bytes()[..n], v)) } - pub fn iter_prefix<'a, 'b>(&'a self, prefix: &'b [u8]) -> Option<(usize, Nodes)> { + pub fn iter_prefix<'a, 'b, K: ?Sized + BorrowedBytes>( + &'a self, + prefix: &'b K, + ) -> Option<(usize, Nodes)> { if let Some((common_prefix_len, node)) = self.root.get_prefix_node(prefix, 0) { let nodes = Nodes { nodes: node.iter_descendant(), label_lens: Vec::new(), }; - Some((prefix.len() - common_prefix_len, nodes)) + Some((prefix.as_bytes().len() - common_prefix_len, nodes)) } else { None } } - pub fn iter_prefix_mut<'a, 'b>(&'a mut self, prefix: &'b [u8]) -> Option<(usize, NodesMut)> { + pub fn iter_prefix_mut<'a, 'b, K: ?Sized + BorrowedBytes>( + &'a mut self, + prefix: &'b K, + ) -> Option<(usize, NodesMut)> { if let Some((common_prefix_len, node)) = self.root.get_prefix_node_mut(prefix, 0) { let nodes = NodesMut { nodes: node.iter_descendant_mut(), label_lens: Vec::new(), }; - Some((prefix.len() - common_prefix_len, nodes)) + Some((prefix.as_bytes().len() - common_prefix_len, nodes)) } else { None } @@ -79,7 +88,7 @@ impl PatriciaTree { { self.root.common_prefixes(key) } - pub fn remove(&mut self, key: &[u8]) -> Option { + pub fn remove(&mut self, key: &K) -> Option { if let Some(old) = self.root.remove(key, 0) { self.len -= 1; Some(old) @@ -87,9 +96,9 @@ impl PatriciaTree { None } } - pub fn split_by_prefix(&mut self, prefix: &[u8]) -> Self { + pub fn split_by_prefix(&mut self, prefix: &K) -> Self { if let Some(splitted_root) = self.root.split_by_prefix(prefix, 0) { - let mut splitted_root = Node::new(prefix, None, Some(splitted_root), None); + let mut splitted_root = Node::new(prefix.as_bytes(), 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();