WIP: Try another approach, no luck yet

Author: Avi-D-coder

src/stored.rs

@@ -21,7 +21,10 @@ pub trait Store<V> {
         hash_idx: Idx,
     ) -> Result<NodeHash, Self::Error>;
 
-    fn get_node(&self, hash_idx: Idx) -> Result<Node<&Branch<Idx>, &Leaf<V>>, Self::Error>;
+    fn get_node<'s>(
+        &'s self,
+        hash_idx: Idx,
+    ) -> Result<Node<&'s Branch<'s, Idx>, &'s Leaf<V>>, Self::Error>;
 }
 
 impl<V, S: Store<V>> Store<V> for &S {

src/stored/merkle.rs

@@ -5,7 +5,7 @@ use bumpalo::Bump;
 use ouroboros::self_referencing;
 
 use crate::{
-    transaction::nodes::{NodeRef, TrieRoot},
+    transaction::nodes::{NodeRef, PrefixesBuffer, TrieRoot},
     Branch, Leaf, PortableHash, PortableHasher, TrieError,
 };
 
@@ -20,12 +20,14 @@ type Result<T, E = TrieError> = core::result::Result<T, E>;
 #[derive(Clone, Debug, PartialEq, Eq)]
 pub struct Snapshot<V> {
     /// The last branch is the root of the trie if it exists.
-    branches: Box<[Branch<Idx>]>,
+    branches: Box<[Branch<'static, Idx>]>,
     /// A Snapshot containing only
     leaves: Box<[Leaf<V>]>,
 
     // we only store the hashes of the nodes that have not been visited.
     unvisited_nodes: Box<[NodeHash]>,
+
+    prefixies_buffer: PrefixesBuffer,
 }
 
 impl<V: PortableHash> Snapshot<V> {

src/transaction.rs

@@ -22,7 +22,6 @@ use self::nodes::{
 pub struct Transaction<S, V> {
     pub data_store: S,
     current_root: TrieRoot<NodeRef<V>>,
-    prefixes_buffer: PrefixesBuffer,
 }
 
 impl<Db: DatabaseSet<V>, V: Clone + PortableHash> Transaction<SnapshotBuilder<Db, V>, V> {
@@ -44,7 +43,6 @@ impl<Db: DatabaseSet<V>, V: Clone + PortableHash> Transaction<SnapshotBuilder<Db
                     left,
                     right,
                     mask: branch.mask,
-                    prior_word: branch.prior_word,
                     prefix: branch.prefix.clone(),
                 };
 
@@ -397,6 +395,7 @@ impl<S: Store<V>, V> Transaction<S, V> {
                                 mask: new_branch.mask,
                                 prior_word: new_branch.prior_word,
                                 prefix: new_branch.prefix.clone(),
+                                prefix,
                             }));
 
                             continue;
@@ -549,6 +548,7 @@ impl<Db, V: PortableHash + Clone> Transaction<SnapshotBuilder<Db, V>, V> {
         Transaction {
             current_root: builder.trie_root(),
             data_store: builder,
+            prefixes_buffer: builder.prefixes_buffer,
         }
     }
 }
@@ -582,6 +582,7 @@ impl<V: PortableHash + Clone> Transaction<Snapshot<V>, V> {
         Ok(Transaction {
             current_root: snapshot.trie_root()?,
             data_store: snapshot,
+            prefixes_buffer: snapshot.prefixes_buffer,
         })
     }
 }

src/transaction/nodes.rs

@@ -74,20 +74,20 @@ pub enum Node<B, L> {
 /// When executing against a `SnapshotBuilder`, it's a reference to a `NodeHash`,
 /// which can in turn be used to retrieve the `Node`.
 #[derive(Clone, PartialEq, Eq, PartialOrd, Ord)]
-pub enum NodeRef<V> {
-    ModBranch(Box<Branch<Self>>),
+pub enum NodeRef<'s, V> {
+    ModBranch(Box<Branch<'s, Self>>),
     ModLeaf(Box<Leaf<V>>),
     Stored(stored::Idx),
 }
 
-impl<V> NodeRef<V> {
+impl<V> NodeRef<'_, V> {
     #[inline(always)]
     pub fn temp_null_stored() -> Self {
         NodeRef::Stored(u32::MAX)
     }
 }
 
-impl<V> fmt::Debug for NodeRef<V> {
+impl<V> fmt::Debug for NodeRef<'_, V> {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         match self {
             Self::ModBranch(b) => f.debug_tuple("ModBranch").field(b).finish(),
@@ -97,7 +97,7 @@ impl<V> fmt::Debug for NodeRef<V> {
     }
 }
 
-impl<V> From<Box<Branch<NodeRef<V>>>> for NodeRef<V> {
+impl<'s, V> From<Box<Branch<'s, NodeRef<'s, V>>>> for NodeRef<'s, V> {
     #[inline]
     fn from(branch: Box<Branch<NodeRef<V>>>) -> Self {
         NodeRef::ModBranch(branch)
@@ -263,7 +263,7 @@ mod tests {
     }
 }
 
-pub struct Prefix {
+pub struct PrefixBufferRef {
     /// This value will be 0 if the branch occurs in the first word of the hash key.
     /// The value is the prior word if the branches parent's word index no more than 1 less.
     /// If the parent's word index is more than 1 word prior,
@@ -272,7 +272,7 @@ pub struct Prefix {
     prior_word_or_prefix_idx: u32,
 }
 
-impl Prefix {
+impl PrefixBufferRef {
     pub fn get_prefix<'s, 'txn: 's>(
         &'s self,
         prefixies: &'txn PrefixesBuffer,
@@ -296,7 +296,7 @@ impl Prefix {
     }
 }
 
-#[derive(Clone, PartialEq, Eq, PartialOrd, Ord)]
+#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
 pub struct PrefixesBuffer {
     buffer: Vec<u32>,
 }
@@ -318,21 +318,32 @@ impl PrefixesBuffer {
     }
 }
 
+#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
+#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
+pub enum PrefixCow<'a> {
+    StartOfKey,
+    PriorWord(u32),
+    Segment(&'a [u32]),
+    SegmentOwned(Box<[u32]>),
+}
+
+/// A branch node in the trie.
+/// `NR` is the type of the node references.
+/// `PR` is the type of reference to the prefix.
 #[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
 #[derive(Clone, PartialEq, Eq, PartialOrd, Ord)]
-pub struct Branch<NR> {
+pub struct Branch<'s, NR> {
     pub left: NR,
     pub right: NR,
     pub mask: BranchMask,
-    pub prefix: Prefix,
+    pub prefix: PrefixCow<'s>,
 }
 
-impl<NR> fmt::Debug for Branch<NR> {
+impl<NR> fmt::Debug for Branch<'_, NR> {
     #[inline]
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         f.debug_struct("Branch")
             .field("mask", &self.mask)
-            .field("prior_word", &self.prior_word)
             .field("prefix", &self.prefix)
             .finish()
     }
@@ -359,7 +370,7 @@ pub enum KeyPositionAdjacent {
     PrefixVec(usize),
 }
 
-impl<NR> Branch<NR> {
+impl<NR> Branch<NR, PrefixCow<'a>> {
     /// Returns the position of the key relative to the branch.
     #[inline(always)]
     pub fn key_position(&self, key_hash: &KeyHash) -> KeyPosition {
@@ -424,19 +435,7 @@ impl<NR> Branch<NR> {
     }
 }
 
-impl<V> Branch<NodeRef<V>> {
-    pub(crate) fn from_stored(branch: &Branch<stored::Idx>) -> Branch<NodeRef<V>> {
-        Branch {
-            left: NodeRef::Stored(branch.left),
-            right: NodeRef::Stored(branch.right),
-            mask: branch.mask,
-            prior_word: branch.prior_word,
-            // TODO remove the clone
-            // Maybe use a AsRef<[u32]> instead of Box<[u32]>
-            prefix: branch.prefix.clone(),
-        }
-    }
-
+impl<V, PR> Branch<NodeRef<V>, PR> {
     /// A wrapper around `new_at_branch_ret` which returns nothing.
     /// This exists to aid compiler inlining.
     ///
@@ -591,17 +590,19 @@ impl<V> Branch<NodeRef<V>> {
 
         debug_assert!(new_leaf.key_hash.0[..word_idx] == old_leaf.as_ref().key_hash.0[..word_idx]);
 
-        let prior_word_idx = word_idx.saturating_sub(1);
-        let prefix = new_leaf.key_hash.0[prefix_start_idx..prior_word_idx].into();
-        let prior_word = if word_idx == 0 {
-            0
-        } else {
+        let prefix = if word_idx == 0 {
+            PrefixCow::StartOfKey
+        } else if prefix_start_idx == word_idx {
+            let prior_word_idx = word_idx - 1;
             debug_assert_eq!(
                 new_leaf.key_hash.0[prior_word_idx],
                 old_leaf.as_ref().key_hash.0[prior_word_idx]
             );
-
-            new_leaf.key_hash.0[prior_word_idx]
+            PrefixCow::PriorWord(new_leaf.key_hash.0[prior_word_idx])
+        } else if prefix_start_idx == word_idx - 1 {
+            PrefixCow::PriorWord(new_leaf.key_hash.0[word_idx - 1])
+        } else {
+            PrefixCow::Segment(&new_leaf.key_hash.0[prefix_start_idx..word_idx])
         };
 
         let mask = BranchMask::new(word_idx as u32, a, b);
@@ -628,7 +629,6 @@ impl<V> Branch<NodeRef<V>> {
                 left,
                 right,
                 mask,
-                prior_word,
                 prefix,
             }),
             // TODO use an enum