Skip to content

Commit

Permalink
Use K: BorrowedBytes instead of [u8] for node operations
Browse files Browse the repository at this point in the history
  • Loading branch information
@sile
sile committed Dec 1, 2023
1 parent c41401d commit d130ccf
Show file tree
Hide file tree
Showing 4 changed files with 75 additions and 70 deletions.
5 changes: 5 additions & 0 deletions src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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] {
Expand Down
20 changes: 8 additions & 12 deletions src/map.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -129,7 +129,7 @@ impl<K: Bytes, V> GenericPatriciaMap<K, V> {
/// assert!(!map.contains_key("bar"));
/// ```
pub fn contains_key<Q: AsRef<K::Borrowed>>(&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.
Expand Down Expand Up @@ -161,7 +161,7 @@ impl<K: Bytes, V> GenericPatriciaMap<K, V> {
/// assert_eq!(map.get("foo"), Some(&2));
/// ```
pub fn get_mut<Q: AsRef<K::Borrowed>>(&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,
Expand All @@ -185,9 +185,7 @@ impl<K: Bytes, V> GenericPatriciaMap<K, V> {
where
Q: ?Sized + AsRef<K::Borrowed>,
{
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))
}

Expand Down Expand Up @@ -217,9 +215,7 @@ impl<K: Bytes, V> GenericPatriciaMap<K, V> {
where
Q: ?Sized + AsRef<K::Borrowed>,
{
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))
}

Expand Down Expand Up @@ -256,7 +252,7 @@ impl<K: Bytes, V> GenericPatriciaMap<K, V> {
/// assert_eq!(map.remove("foo"), None);
/// ```
pub fn remove<Q: AsRef<K::Borrowed>>(&mut self, key: Q) -> Option<V> {
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.
Expand Down Expand Up @@ -338,7 +334,7 @@ impl<K: Bytes, V> GenericPatriciaMap<K, V> {
/// assert_eq!(b.keys().collect::<Vec<_>>(), [b"elixir", b"erlang"]);
/// ```
pub fn split_by_prefix<Q: AsRef<K::Borrowed>>(&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,
Expand Down Expand Up @@ -455,7 +451,7 @@ impl<K: Bytes, V> GenericPatriciaMap<K, V> {
'b: 'a,
{
self.tree
.iter_prefix(prefix.as_bytes())
.iter_prefix(prefix)
.into_iter()
.flat_map(move |(prefix_len, nodes)| {
Iter::<K, V>::new(nodes, Vec::from(&prefix.as_bytes()[..prefix_len]))
Expand All @@ -482,7 +478,7 @@ impl<K: Bytes, V> GenericPatriciaMap<K, V> {
'b: 'a,
{
self.tree
.iter_prefix_mut(prefix.as_bytes())
.iter_prefix_mut(prefix)
.into_iter()
.flat_map(move |(prefix_len, nodes)| {
IterMut::<K, V>::new(nodes, Vec::from(&prefix.as_bytes()[..prefix_len]))
Expand Down
85 changes: 40 additions & 45 deletions src/node.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -408,7 +408,7 @@ impl<V> Node<V> {
pub(crate) fn get<K: ?Sized + BorrowedBytes>(&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))
Expand All @@ -420,29 +420,27 @@ impl<V> Node<V> {
}
}

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<K: ?Sized + BorrowedBytes>(&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<K: ?Sized + BorrowedBytes>(
&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() {
Expand All @@ -452,20 +450,19 @@ impl<V> Node<V> {
.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<K: ?Sized + BorrowedBytes>(
&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() {
Expand All @@ -476,55 +473,61 @@ impl<V> Node<V> {
.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 {
None
}
}

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<K: ?Sized + BorrowedBytes>(
&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 {
None
}
}

pub(crate) fn get_prefix_node_mut(
pub(crate) fn get_prefix_node_mut<K: ?Sized + BorrowedBytes>(
&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 {
None
}
}

pub(crate) fn split_by_prefix(&mut self, prefix: &[u8], level: usize) -> Option<Self> {
let common_prefix_len = self.skip_common_prefix(prefix);
if common_prefix_len == prefix.len() {
pub(crate) fn split_by_prefix<K: ?Sized + BorrowedBytes>(
&mut self,
prefix: &K,
level: usize,
) -> Option<Self> {
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);
Expand All @@ -533,16 +536,14 @@ impl<V> Node<V> {
}
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| {
self.try_reclaim_child();
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| {
Expand All @@ -553,9 +554,8 @@ impl<V> Node<V> {
None
}
}
pub(crate) fn remove(&mut self, key: &[u8], level: usize) -> Option<V> {
let common_prefix_len = self.skip_common_prefix(key);
let next = &key[common_prefix_len..];
pub(crate) fn remove<K: ?Sized + BorrowedBytes>(&mut self, key: &K, level: usize) -> Option<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.take_value().map(|old| {
Expand All @@ -571,7 +571,7 @@ impl<V> Node<V> {
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| {
Expand Down Expand Up @@ -626,13 +626,6 @@ impl<V> Node<V> {
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 })
}
Expand Down Expand Up @@ -873,8 +866,10 @@ where
type Item = (usize, &'a Node<V>);
fn next(&mut self) -> Option<Self::Item> {
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));
}
Expand Down
35 changes: 22 additions & 13 deletions src/tree.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -35,40 +35,49 @@ impl<V> PatriciaTree<V> {
pub fn get<K: ?Sized + BorrowedBytes>(&self, key: &K) -> Option<&V> {
self.root.get(key)
}
pub fn get_mut(&mut self, key: &[u8]) -> Option<&mut V> {
pub fn get_mut<K: ?Sized + BorrowedBytes>(&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<V>)> {
pub fn iter_prefix<'a, 'b, K: ?Sized + BorrowedBytes>(
&'a self,
prefix: &'b K,
) -> Option<(usize, Nodes<V>)> {
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<V>)> {
pub fn iter_prefix_mut<'a, 'b, K: ?Sized + BorrowedBytes>(
&'a mut self,
prefix: &'b K,
) -> Option<(usize, NodesMut<V>)> {
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
}
Expand All @@ -79,17 +88,17 @@ impl<V> PatriciaTree<V> {
{
self.root.common_prefixes(key)
}
pub fn remove(&mut self, key: &[u8]) -> Option<V> {
pub fn remove<K: ?Sized + BorrowedBytes>(&mut self, key: &K) -> Option<V> {
if let Some(old) = self.root.remove(key, 0) {
self.len -= 1;
Some(old)
} else {
None
}
}
pub fn split_by_prefix(&mut self, prefix: &[u8]) -> Self {
pub fn split_by_prefix<K: ?Sized + BorrowedBytes>(&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();
Expand Down

0 comments on commit d130ccf

Please sign in to comment.