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