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