make PrecomputedSubtrees more generic

This commit is contained in:
Qyriad 2024-11-04 12:24:41 -07:00
parent 74ab46ca69
commit 49d88600c0

View file

@ -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>
where
) -> Result<(), MerkleError>
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> {
let mut accumulator: PairComputations<Self::Leaf> = Default::default();
) -> PairComputations<u64, Self::Leaf> {
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(),