Fix #40 (ugly code)

Author: stefnotch

src/node.rs

@@ -76,6 +76,16 @@ impl<K: Borrow<[u8]>, V> Branch<K, V> {
         self.entries.contains(index)
     }
 
+    /// Corresponds to the key-value pair at this position.
+    #[inline]
+    pub fn head_entry(&self) -> Option<&Leaf<K, V>> {
+        match self.entries.get(0) {
+            Some(Node::Leaf(leaf)) => Some(leaf),
+            None => None,
+            _ => unsafe { debug_unreachable!() },
+        }
+    }
+
     #[inline]
     pub fn entry_mut(&mut self, index: u8) -> &mut Node<K, V> {
         let entry = self.entries.get_mut(index);
@@ -89,6 +99,13 @@ impl<K: Borrow<[u8]>, V> Branch<K, V> {
         self.entries.get(nybble_index(self.choice, key))
     }
 
+    // Get the child node corresponding to the given key.
+    #[inline]
+    pub fn child_with_offsetted_key(&self, key: &[u8], key_offset: usize) -> Option<&Node<K, V>> {
+        self.entries
+            .get(nybble_index(self.choice.checked_sub(key_offset * 2)?, key))
+    }
+
     // Mutable version of `Branch::child`.
     #[inline]
     pub fn child_mut(&mut self, key: &[u8]) -> Option<&mut Node<K, V>> {
@@ -119,6 +136,18 @@ impl<K: Borrow<[u8]>, V> Branch<K, V> {
         self.entries.get_or_any(nybble_index(self.choice, key))
     }
 
+    // Retrieve the node which contains the exemplar. This does not recurse and return the actual
+    // exemplar - just the node which might be or contain it.
+    #[inline]
+    pub fn exemplar_with_offset(&self, key: &[u8], key_offset: usize) -> &Node<K, V> {
+        self.entries.get_or_any(
+            self.choice
+                .checked_sub(key_offset * 2)
+                .map(|choice| nybble_index(choice, key))
+                .unwrap_or_default(),
+        )
+    }
+
     // As `Branch::exemplar` but for mutable borrows.
     #[inline]
     pub fn exemplar_mut(&mut self, key: &[u8]) -> &mut Node<K, V> {
@@ -132,6 +161,12 @@ impl<K: Borrow<[u8]>, V> Branch<K, V> {
         self.exemplar(key).get_exemplar(key)
     }
 
+    #[inline]
+    pub fn get_exemplar_with_offset(&self, key: &[u8], key_offset: usize) -> &Leaf<K, V> {
+        self.exemplar_with_offset(key, key_offset)
+            .get_exemplar_with_offset(key, key_offset)
+    }
+
     // Mutably borrow the exemplar for the given key, mutually recursing through
     // `Node::get_exemplar_mut`.
     #[inline]
@@ -307,6 +342,13 @@ impl<K: Borrow<[u8]>, V> Node<K, V> {
         }
     }
 
+    pub fn get_exemplar_with_offset(&self, key: &[u8], key_offset: usize) -> &Leaf<K, V> {
+        match *self {
+            Node::Leaf(ref leaf) => leaf,
+            Node::Branch(ref branch) => branch.get_exemplar_with_offset(key, key_offset),
+        }
+    }
+
     // Mutably borrow the exemplar for a given key.
     pub fn get_exemplar_mut(&mut self, key: &[u8]) -> &mut Leaf<K, V> {
         match *self {
@@ -320,19 +362,23 @@ impl<K: Borrow<[u8]>, V> Node<K, V> {
     //
     // PRECONDITION:
     // - There exists at least one node in the trie with the given prefix.
-    pub fn get_prefix_validated<'a>(&'a self, prefix: &[u8]) -> &'a Node<K, V> {
+    pub fn get_prefix_validated<'a>(
+        &'a self,
+        prefix: &[u8],
+        prefix_offset: usize,
+    ) -> &'a Node<K, V> {
         match *self {
             Node::Leaf(..) => self,
             Node::Branch(ref branch) => {
-                if branch.choice >= prefix.len() * 2 {
+                if branch.choice >= (prefix.len() + prefix_offset) * 2 {
                     self
                 } else {
-                    let child_opt = branch.child(prefix);
+                    let child_opt = branch.child_with_offsetted_key(prefix, prefix_offset);
 
                     // unsafe: child must exist in the trie - prefix'd nodes must exist.
                     let child = unsafe { child_opt.unwrap_unchecked() };
 
-                    child.get_prefix_validated(prefix)
+                    child.get_prefix_validated(prefix, prefix_offset)
                 }
             }
         }
@@ -346,7 +392,31 @@ impl<K: Borrow<[u8]>, V> Node<K, V> {
             Node::Branch(ref branch)
                 if branch.get_exemplar(prefix).key_slice().starts_with(prefix) =>
             {
-                Some(self.get_prefix_validated(prefix))
+                Some(self.get_prefix_validated(prefix, 0))
+            }
+
+            _ => None,
+        }
+    }
+
+    // Borrow the node which contains all and only entries with keys continuing with
+    // `prefix`.
+    pub fn get_prefix_with_offset<'a>(
+        &'a self,
+        prefix: &[u8],
+        prefix_offset: usize,
+    ) -> Option<&'a Node<K, V>> {
+        match *self {
+            Node::Leaf(ref leaf) if leaf.key_slice()[prefix_offset..].starts_with(prefix) => {
+                Some(self)
+            }
+            Node::Branch(ref branch)
+                if branch
+                    .get_exemplar_with_offset(prefix, prefix_offset)
+                    .key_slice()[prefix_offset..]
+                    .starts_with(prefix) =>
+            {
+                Some(self.get_prefix_validated(prefix, prefix_offset))
             }
 
             _ => None,

src/subtrie.rs

@@ -6,7 +6,9 @@ use crate::iter::Iter;
 use crate::node::Node;
 
 pub struct SubTrie<'a, K: 'a, V: 'a> {
-    pub(crate) root: Option<&'a Node<K, V>>,
+    /// The index of the next byte to compare.
+    key_byte_index: usize,
+    root: Option<&'a Node<K, V>>,
 }
 
 impl<'a, K: fmt::Debug, V: fmt::Debug> fmt::Debug for SubTrie<'a, K, V> {
@@ -28,6 +30,17 @@ impl<'a, K: 'a, V: 'a> IntoIterator for SubTrie<'a, K, V> {
 }
 
 impl<'a, K: 'a, V: 'a> SubTrie<'a, K, V> {
+    pub fn new(root: Option<&'a Node<K, V>>, key_byte_index: usize) -> SubTrie<'a, K, V> {
+        SubTrie {
+            key_byte_index,
+            root,
+        }
+    }
+
+    pub fn empty() -> SubTrie<'a, K, V> {
+        SubTrie::new(None, 0)
+    }
+
     /// Returns true if the subtrie has no entries.
     pub fn is_empty(&self) -> bool {
         self.root.is_none()
@@ -49,10 +62,41 @@ impl<'a, K: Borrow<[u8]>, V> SubTrie<'a, K, V> {
         }
     }
 
-    pub fn subtrie<L: Borrow<[u8]>>(&self, prefix: L) -> SubTrie<'a, K, V> {
-        SubTrie {
-            root: self.root.and_then(|node| node.get_prefix(prefix.borrow())),
-        }
+    /// Takes the next step in the trie, returning a new subtrie.
+    pub fn subtrie<L: Borrow<[u8]>>(&self, next_key_part: L) -> SubTrie<'a, K, V> {
+        let root = match self.root {
+            Some(node) => node,
+            None => return SubTrie::empty(),
+        };
+        let node = root.get_prefix_with_offset(next_key_part.borrow(), self.key_byte_index);
+        SubTrie::new(node, self.key_byte_index + next_key_part.borrow().len())
+    }
+
+    /// Gets the value at the root of the subtrie.
+    pub fn get_value(&self) -> Option<&'a V> {
+        self.root
+            .and_then(|node| match node {
+                Node::Leaf(leaf) => Some(leaf),
+                Node::Branch(v) => v.head_entry(),
+            })
+            .and_then(|leaf| {
+                if self.key_byte_index == leaf.key.borrow().len() {
+                    Some(&leaf.val)
+                } else {
+                    None
+                }
+            })
+    }
+
+    /// Gets a subtrie rooted at the given prefix.
+    /// Is slightly less efficient than `subtrie`, since it re-compares the prefix.
+    pub fn subtrie_with_prefix<L: Borrow<[u8]>>(&self, prefix: L) -> SubTrie<'a, K, V> {
+        let root = match self.root {
+            Some(node) => node,
+            None => return SubTrie::empty(),
+        };
+        let node = root.get_prefix(prefix.borrow());
+        SubTrie::new(node, prefix.borrow().len())
     }
 
     pub fn get<L: Borrow<[u8]>>(&self, key: L) -> Option<&'a V> {

src/trie.rs

@@ -226,12 +226,12 @@ impl<K: Borrow<[u8]>, V> Trie<K, V> {
         K: Borrow<Q>,
         Q: Borrow<[u8]>,
     {
-        SubTrie {
-            root: self
-                .root
-                .as_ref()
-                .and_then(|node| node.get_prefix(prefix.borrow())),
-        }
+        let root = match self.root {
+            Some(ref node) => node,
+            None => return SubTrie::empty(),
+        };
+        let node = root.get_prefix(prefix.borrow());
+        SubTrie::new(node, prefix.borrow().len())
     }
 
     /// Get the longest common prefix of all the nodes in the trie and the given key.

src/wrapper.rs

@@ -90,6 +90,13 @@ impl Break for BString {
 #[derive(PartialEq, Eq, PartialOrd, Ord)]
 pub struct BStr(str);
 
+impl BStr {
+    #[inline]
+    pub fn ref_str(s: &str) -> &BStr {
+        <&BStr>::from(s)
+    }
+}
+
 impl fmt::Debug for BStr {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         self.0.fmt(f)
@@ -118,6 +125,13 @@ impl Borrow<[u8]> for BStr {
     }
 }
 
+impl Borrow<[u8]> for &'_ BStr {
+    #[inline]
+    fn borrow(&self) -> &[u8] {
+        self.0.as_bytes()
+    }
+}
+
 #[allow(clippy::derived_hash_with_manual_eq)]
 impl Hash for BStr {
     #[inline]

tests/lib.rs

@@ -498,3 +498,71 @@ fn issue_36_node_count_after_clear() {
     trie.clear();
     assert_eq!(0, trie.count());
 }
+
+#[test]
+fn issue_40_subtree_api_extension_base() {
+    use wrapper::{BStr, BString};
+
+    let mut trie = Trie::<BString, u32>::new();
+    trie.insert("abc".into(), 1);
+    trie.insert("a".into(), 99);
+    trie.insert("abcd".into(), 2);
+    trie.insert("abcde".into(), 3);
+    assert_eq!(
+        trie.subtrie(BStr::ref_str("abc")).get(BStr::ref_str("abc")),
+        Some(&1)
+    );
+    assert_eq!(
+        trie.subtrie(BStr::ref_str("abc"))
+            .get(BStr::ref_str("abcd")),
+        Some(&2)
+    );
+    assert_eq!(
+        trie.subtrie(BStr::ref_str("abcd"))
+            .get(BStr::ref_str("abc")),
+        None
+    );
+    assert_eq!(
+        trie.subtrie(BStr::ref_str("abcd"))
+            .get(BStr::ref_str("abcd")),
+        Some(&2)
+    );
+    assert_eq!(
+        trie.subtrie(BStr::ref_str("abcde"))
+            .get(BStr::ref_str("abcde")),
+        Some(&3)
+    );
+    assert_eq!(trie.subtrie(BStr::ref_str("abcde")).is_empty(), false);
+    assert_eq!(trie.subtrie(BStr::ref_str("abcdef")).is_empty(), true);
+}
+
+#[test]
+fn issue_40_subtree_api_extension_step() {
+    use wrapper::{BStr, BString};
+
+    let mut trie = Trie::<BString, u32>::new();
+    trie.insert("abc".into(), 1);
+    trie.insert("a".into(), 99);
+    trie.insert("abcd".into(), 2);
+    trie.insert("abcde".into(), 3);
+    assert_eq!(trie.subtrie(BStr::ref_str("a")).get_value(), Some(&99));
+    assert_eq!(
+        trie.subtrie(BStr::ref_str("a"))
+            .subtrie(BStr::ref_str("b"))
+            .get_value(),
+        None
+    );
+    assert_eq!(
+        trie.subtrie(BStr::ref_str("a"))
+            .subtrie(BStr::ref_str("bcd"))
+            .get_value(),
+        Some(&2)
+    );
+    assert_eq!(
+        trie.subtrie(BStr::ref_str("a"))
+            .subtrie(BStr::ref_str("b"))
+            .subtrie(BStr::ref_str("c"))
+            .get(BStr::ref_str("abcd")),
+        Some(&2)
+    );
+}