1use super::{
  2    BTreeMap, KvMap, MerkleError, MerkleStore, NodeIndex, RpoDigest, StoreNode, Vec, Word,
  3};
  4use crate::utils::collections::Diff;
  5
  6#[cfg(test)]
  7use super::{empty_roots::EMPTY_WORD, Felt, SimpleSmt};
  8
  9// MERKLE STORE DELTA
 10// ================================================================================================
 11
 12/// [MerkleStoreDelta] stores a vector of ([RpoDigest], [MerkleTreeDelta]) tuples where the
 13/// [RpoDigest] represents the root of the Merkle tree and [MerkleTreeDelta] represents the
 14/// differences between the initial and final Merkle tree states.
 15#[derive(Default, Debug, Clone, PartialEq, Eq)]
 16pub struct MerkleStoreDelta(pub Vec<(RpoDigest, MerkleTreeDelta)>);
 17
 18// MERKLE TREE DELTA
 19// ================================================================================================
 20
 21/// [MerkleDelta] stores the differences between the initial and final Merkle tree states.
 22///
 23/// The differences are represented as follows:
 24/// - depth: the depth of the merkle tree.
 25/// - cleared_slots: indexes of slots where values were set to [ZERO; 4].
 26/// - updated_slots: index-value pairs of slots where values were set to non [ZERO; 4] values.
 27#[cfg(not(test))]
 28#[derive(Debug, Clone, PartialEq, Eq)]
 29pub struct MerkleTreeDelta {
 30    depth: u8,
 31    cleared_slots: Vec<u64>,
 32    updated_slots: Vec<(u64, Word)>,
 33}
 34
 35impl MerkleTreeDelta {
 36    // CONSTRUCTOR
 37    // --------------------------------------------------------------------------------------------
 38    pub fn new(depth: u8) -> Self {
 39        Self {
 40            depth,
 41            cleared_slots: Vec::new(),
 42            updated_slots: Vec::new(),
 43        }
 44    }
 45
 46    // ACCESSORS
 47    // --------------------------------------------------------------------------------------------
 48    /// Returns the depth of the Merkle tree the [MerkleDelta] is associated with.
 49    pub fn depth(&self) -> u8 {
 50        self.depth
 51    }
 52
 53    /// Returns the indexes of slots where values were set to [ZERO; 4].
 54    pub fn cleared_slots(&self) -> &[u64] {
 55        &self.cleared_slots
 56    }
 57
 58    /// Returns the index-value pairs of slots where values were set to non [ZERO; 4] values.
 59    pub fn updated_slots(&self) -> &[(u64, Word)] {
 60        &self.updated_slots
 61    }
 62
 63    // MODIFIERS
 64    // --------------------------------------------------------------------------------------------
 65    /// Adds a slot index to the list of cleared slots.
 66    pub fn add_cleared_slot(&mut self, index: u64) {
 67        self.cleared_slots.push(index);
 68    }
 69
 70    /// Adds a slot index and a value to the list of updated slots.
 71    pub fn add_updated_slot(&mut self, index: u64, value: Word) {
 72        self.updated_slots.push((index, value));
 73    }
 74}
 75
 76/// Extracts a [MerkleDelta] object by comparing the leaves of two Merkle trees specifies by
 77/// their roots and depth.
 78pub fn merkle_tree_delta<T: KvMap<RpoDigest, StoreNode>>(
 79    tree_root_1: RpoDigest,
 80    tree_root_2: RpoDigest,
 81    depth: u8,
 82    merkle_store: &MerkleStore<T>,
 83) -> Result<MerkleTreeDelta, MerkleError> {
 84    if tree_root_1 == tree_root_2 {
 85        return Ok(MerkleTreeDelta::new(depth));
 86    }
 87
 88    let tree_1_leaves: BTreeMap<NodeIndex, RpoDigest> =
 89        merkle_store.non_empty_leaves(tree_root_1, depth).collect();
 90    let tree_2_leaves: BTreeMap<NodeIndex, RpoDigest> =
 91        merkle_store.non_empty_leaves(tree_root_2, depth).collect();
 92    let diff = tree_1_leaves.diff(&tree_2_leaves);
 93
 94    // TODO: Refactor this diff implementation to prevent allocation of both BTree and Vec.
 95    Ok(MerkleTreeDelta {
 96        depth,
 97        cleared_slots: diff.removed.into_iter().map(|index| index.value()).collect(),
 98        updated_slots: diff
 99            .updated
100            .into_iter()
101            .map(|(index, leaf)| (index.value(), *leaf))
102            .collect(),
103    })
104}
105
106// INTERNALS
107// --------------------------------------------------------------------------------------------
108#[cfg(test)]
109#[derive(Debug, Clone, PartialEq, Eq)]
110pub struct MerkleTreeDelta {
111    pub depth: u8,
112    pub cleared_slots: Vec<u64>,
113    pub updated_slots: Vec<(u64, Word)>,
114}
115
116// MERKLE DELTA
117// ================================================================================================
118#[test]
119fn test_compute_merkle_delta() {
120    let entries = vec![
121        (10, [Felt::new(0), Felt::new(1), Felt::new(2), Felt::new(3)]),
122        (15, [Felt::new(4), Felt::new(5), Felt::new(6), Felt::new(7)]),
123        (20, [Felt::new(8), Felt::new(9), Felt::new(10), Felt::new(11)]),
124        (31, [Felt::new(12), Felt::new(13), Felt::new(14), Felt::new(15)]),
125    ];
126    let simple_smt = SimpleSmt::with_leaves(30, entries.clone()).unwrap();
127    let mut store: MerkleStore = (&simple_smt).into();
128    let root = simple_smt.root();
129
130    // add a new node
131    let new_value = [Felt::new(16), Felt::new(17), Felt::new(18), Felt::new(19)];
132    let new_index = NodeIndex::new(simple_smt.depth(), 32).unwrap();
133    let root = store.set_node(root, new_index, new_value.into()).unwrap().root;
134
135    // update an existing node
136    let update_value = [Felt::new(20), Felt::new(21), Felt::new(22), Felt::new(23)];
137    let update_idx = NodeIndex::new(simple_smt.depth(), entries[0].0).unwrap();
138    let root = store.set_node(root, update_idx, update_value.into()).unwrap().root;
139
140    // remove a node
141    let remove_idx = NodeIndex::new(simple_smt.depth(), entries[1].0).unwrap();
142    let root = store.set_node(root, remove_idx, EMPTY_WORD.into()).unwrap().root;
143
144    let merkle_delta =
145        merkle_tree_delta(simple_smt.root(), root, simple_smt.depth(), &store).unwrap();
146    let expected_merkle_delta = MerkleTreeDelta {
147        depth: simple_smt.depth(),
148        cleared_slots: vec![remove_idx.value()],
149        updated_slots: vec![(update_idx.value(), update_value), (new_index.value(), new_value)],
150    };
151
152    assert_eq!(merkle_delta, expected_merkle_delta);
153}