miden_objects/block/
partial_nullifier_tree.rs

1use crate::{
2    Digest, EMPTY_WORD, Felt, FieldElement, Word,
3    block::{BlockNumber, NullifierWitness},
4    crypto::merkle::PartialSmt,
5    errors::NullifierTreeError,
6    note::Nullifier,
7};
8
9/// The partial sparse merkle tree containing the nullifiers of consumed notes.
10#[derive(Debug, Clone, PartialEq, Eq)]
11pub struct PartialNullifierTree(PartialSmt);
12
13impl PartialNullifierTree {
14    /// The leaf value of an unspent nullifier.
15    pub const UNSPENT_NULLIFIER: Word = EMPTY_WORD;
16
17    /// Creates a new, empty partial nullifier tree.
18    pub fn new() -> Self {
19        PartialNullifierTree(PartialSmt::new())
20    }
21
22    /// Adds the given nullifier witness to the partial tree and tracks it. Once a nullifier has
23    /// been added to the tree, it can be marked as spent using [`Self::mark_spent`].
24    ///
25    /// # Errors
26    ///
27    /// Returns an error if:
28    /// - after the witness' merkle path was added the partial nullifier tree has a different root
29    ///   than before it was added.
30    pub fn add_nullifier_witness(
31        &mut self,
32        witness: NullifierWitness,
33    ) -> Result<(), NullifierTreeError> {
34        let (path, leaf) = witness.into_proof().into_parts();
35        self.0.add_path(leaf, path).map_err(NullifierTreeError::TreeRootConflict)
36    }
37
38    /// Marks the given nullifiers as spent at the given block number.
39    ///
40    /// # Errors
41    ///
42    /// Returns an error if:
43    /// - a nullifier was already spent.
44    /// - a nullifier is not tracked by this partial nullifier tree, that is, its
45    ///   [`NullifierWitness`] was not added to the tree previously.
46    pub fn mark_spent(
47        &mut self,
48        nullifiers: impl Iterator<Item = Nullifier>,
49        block_num: BlockNumber,
50    ) -> Result<(), NullifierTreeError> {
51        for nullifier in nullifiers {
52            self.mark_spent_single(nullifier, block_num)?;
53        }
54
55        Ok(())
56    }
57
58    /// Returns the root of the tree.
59    pub fn root(&self) -> Digest {
60        self.0.root()
61    }
62
63    /// Marks the given nullifier as spent at the given block number.
64    ///
65    /// # Errors
66    ///
67    /// See [`Self::mark_spent`] for the possible error conditions.
68    fn mark_spent_single(
69        &mut self,
70        nullifier: Nullifier,
71        block_num: BlockNumber,
72    ) -> Result<(), NullifierTreeError> {
73        let prev_nullifier_value = self
74            .0
75            .insert(nullifier.inner(), block_num_to_leaf_value(block_num))
76            .map_err(|source| NullifierTreeError::UntrackedNullifier { nullifier, source })?;
77
78        if prev_nullifier_value != Self::UNSPENT_NULLIFIER {
79            Err(NullifierTreeError::NullifierAlreadySpent(nullifier))
80        } else {
81            Ok(())
82        }
83    }
84}
85
86impl Default for PartialNullifierTree {
87    fn default() -> Self {
88        Self::new()
89    }
90}
91
92// HELPER FUNCTIONS
93// ================================================================================================
94
95/// Returns the nullifier's leaf value in the SMT by its block number.
96fn block_num_to_leaf_value(block: BlockNumber) -> Word {
97    [Felt::from(block), Felt::ZERO, Felt::ZERO, Felt::ZERO]
98}
99
100#[cfg(test)]
101mod tests {
102    use assert_matches::assert_matches;
103    use miden_crypto::merkle::Smt;
104    use winter_rand_utils::rand_array;
105
106    use super::*;
107
108    /// Test that using a stale nullifier witness together with a current one results in a different
109    /// tree root and thus an error.
110    #[test]
111    fn partial_nullifier_tree_root_mismatch() {
112        let key0 = Digest::from(Word::from(rand_array()));
113        let key1 = Digest::from(Word::from(rand_array()));
114        let key2 = Digest::from(Word::from(rand_array()));
115
116        let value0 = EMPTY_WORD;
117        let value1 = Word::from(rand_array());
118        let value2 = EMPTY_WORD;
119
120        let kv_pairs = vec![(key0, value0)];
121
122        let mut full = Smt::with_entries(kv_pairs).unwrap();
123        let stale_proof0 = full.open(&key0);
124        // Insert a non-empty value so the nullifier tree's root changes.
125        full.insert(key1, value1);
126        full.insert(key2, value2);
127        let proof2 = full.open(&key2);
128
129        assert_ne!(stale_proof0.compute_root(), proof2.compute_root());
130
131        let mut partial = PartialNullifierTree::new();
132
133        partial.add_nullifier_witness(NullifierWitness::new(stale_proof0)).unwrap();
134        let error = partial.add_nullifier_witness(NullifierWitness::new(proof2)).unwrap_err();
135
136        assert_matches!(error, NullifierTreeError::TreeRootConflict(_));
137    }
138}