make PrecomputedSubtrees more generic

src/merkle/smt/mod.rs

@@ -89,7 +89,7 @@ pub(crate) trait SparseMerkleTree<const DEPTH: u8> {
     };
 
     Self::path_and_leaf_to_opening(merkle_path, leaf)
-    }
+}
 
     /// Inserts a value at the specified key, returning the previous value associated with that key.
     /// Recall that by definition, any key that hasn't been updated is associated with
@@ -114,7 +114,7 @@ pub(crate) trait SparseMerkleTree<const DEPTH: u8> {
     self.recompute_nodes_from_index_to_root(node_index, Self::hash_leaf(&leaf));
 
     old_value
-    }
+}
 
     /// Recomputes the branch nodes (including the root) from `index` all the way to the root.
     /// `node_hash_at_index` is the hash of the node stored at index.
@@ -122,7 +122,7 @@ pub(crate) trait SparseMerkleTree<const DEPTH: u8> {
     &mut self,
     mut index: NodeIndex,
     node_hash_at_index: RpoDigest,
-    ) {
+) {
     let mut node_hash = node_hash_at_index;
     for node_depth in (0..index.depth()).rev() {
         let is_right = index.is_value_odd();
@@ -144,7 +144,7 @@ pub(crate) trait SparseMerkleTree<const DEPTH: u8> {
         }
     }
     self.set_root(node_hash);
-    }
+}
 
     /// Computes what changes are necessary to insert the specified key-value pairs into this Merkle
     /// tree, allowing for validation before applying those changes.
@@ -157,7 +157,7 @@ pub(crate) trait SparseMerkleTree<const DEPTH: u8> {
     fn compute_mutations(
     &self,
     kv_pairs: impl IntoIterator<Item = (Self::Key, Self::Value)>,
-    ) -> MutationSet<DEPTH, Self::Key, Self::Value> {
+) -> MutationSet<DEPTH, Self::Key, Self::Value> {
     use NodeMutation::*;
 
     let mut new_root = self.root();
@@ -242,7 +242,7 @@ pub(crate) trait SparseMerkleTree<const DEPTH: u8> {
         node_mutations,
         new_pairs,
     }
-    }
+}
 
     /// Apply the prospective mutations computed with [`SparseMerkleTree::compute_mutations()`] to
     /// this tree.
@@ -255,10 +255,10 @@ pub(crate) trait SparseMerkleTree<const DEPTH: u8> {
     fn apply_mutations(
     &mut self,
     mutations: MutationSet<DEPTH, Self::Key, Self::Value>,
-    ) -> Result<(), MerkleError>
+) -> Result<(), MerkleError>
-    where
+where
     Self: Sized,
-    {
+{
     use NodeMutation::*;
     let MutationSet {
         old_root,
@@ -287,7 +287,7 @@ pub(crate) trait SparseMerkleTree<const DEPTH: u8> {
     self.set_root(new_root);
 
     Ok(())
-    }
+}
 
     // REQUIRED METHODS
     // ---------------------------------------------------------------------------------------------
@@ -336,7 +336,7 @@ pub(crate) trait SparseMerkleTree<const DEPTH: u8> {
     existing_leaf: Self::Leaf,
     key: &Self::Key,
     value: &Self::Value,
-    ) -> Self::Leaf;
+) -> Self::Leaf;
 
     /// Maps a key to a leaf index
     fn key_to_leaf_index(key: &Self::Key) -> LeafIndex<DEPTH>;
@@ -352,8 +352,8 @@ pub(crate) trait SparseMerkleTree<const DEPTH: u8> {
 
     fn sorted_pairs_to_leaves(
         pairs: Vec<(Self::Key, Self::Value)>,
-    ) -> PairComputations<Self::Leaf> {
+    ) -> PairComputations<u64, Self::Leaf> {
-        let mut accumulator: PairComputations<Self::Leaf> = Default::default();
+        let mut accumulator: PairComputations<u64, Self::Leaf> = Default::default();
 
         // The kv-pairs we've seen so far that correspond to a single leaf.
         let mut current_leaf_buffer: Vec<(Self::Key, Self::Value)> = Default::default();
@@ -632,14 +632,14 @@ impl SubtreeLeaf {
 }
 
 #[derive(Debug, Clone, PartialEq, Eq)]
-pub struct PairComputations<L> {
+pub struct PairComputations<K, L> {
     /// Literal leaves to be added to the sparse Merkle tree's internal mapping.
-    pub nodes: BTreeMap<u64, L>,
+    pub nodes: BTreeMap<K, L>,
     /// "Conceptual" leaves that will be used for computations.
     pub leaves: Vec<Vec<SubtreeLeaf>>,
 }
 
-impl<L> PairComputations<L> {
+impl<K, L> PairComputations<K, L> {
     pub fn add_leaf(&mut self, leaf: SubtreeLeaf) {
         let last_subtree = match self.leaves.last_mut() {
             // Base case.
@@ -670,7 +670,7 @@ impl<L> PairComputations<L> {
 }
 
 // Derive requires `L` to impl Default, even though we don't actually need that.
-impl<L> Default for PairComputations<L> {
+impl<K, L> Default for PairComputations<K, L> {
     fn default() -> Self {
         Self {
             nodes: Default::default(),