feat: add insert function to SMT (#5776)

Description --- Adds a new insert function to only add a leaf if it already exists. Motivation and Context --- If this is used for an unspent UTXO set we want it to disallow duplicate UTXOs How Has This Been Tested? --- unit tests --------- Co-authored-by: Aaron Feickert <[email protected]> Co-authored-by: Cayle Sharrock <[email protected]>
tari-project · Sep 20, 2023 · 5901b4a · 5901b4a
1 parent eb74081
commit 5901b4a
Show file tree

Hide file tree

Showing 2 changed files with 61 additions and 0 deletions.
diff --git a/base_layer/mmr/src/sparse_merkle_tree/error.rs b/base_layer/mmr/src/sparse_merkle_tree/error.rs
@@ -33,4 +33,6 @@ pub enum SMTError {
         "Cannot construct a proof. Either the key exists for an exclusion proof, or it does not for an inclusion proof"
     )]
     NonViableProof,
+    #[error("A duplicate key was found when trying to inserting")]
+    KeyExists,
 }
diff --git a/base_layer/mmr/src/sparse_merkle_tree/tree.rs b/base_layer/mmr/src/sparse_merkle_tree/tree.rs
@@ -253,6 +253,15 @@ impl<H: Digest<OutputSize = U32>> SparseMerkleTree<H> {
         Ok(result)
     }
 
+    /// This will only add new node when it does not exist
+    pub fn insert(&mut self, key: NodeKey, value: ValueHash) -> Result<UpdateResult, SMTError> {
+        if self.get(&key)?.is_some() {
+            return Err(SMTError::KeyExists);
+        }
+        // So we no know it does not exist, so lets add it.
+        self.upsert(key, value)
+    }
+
     /// Returns true if the tree contains the key `key`.
     pub fn contains(&self, key: &NodeKey) -> bool {
         match self.search_node(key) {
@@ -419,6 +428,7 @@ mod test {
     use crate::sparse_merkle_tree::{
         tree::{DeleteResult, SparseMerkleTree},
         NodeKey,
+        SMTError,
         UpdateResult,
         ValueHash,
         EMPTY_NODE_HASH,
@@ -859,4 +869,53 @@ mod test {
         // Sanity check that the tree is now empty
         assert!(tree.is_empty());
     }
+
+    #[test]
+    fn insert_only_if_exist() {
+        let mut tree = SparseMerkleTree::<Blake2b<U32>>::default();
+
+        // An empty tree contains no keys
+        assert!(!tree.contains(&short_key(0)));
+        assert!(!tree.contains(&short_key(1)));
+
+        // Add a key, which the tree must then contain
+        assert_eq!(
+            tree.insert(short_key(1), ValueHash::from([1u8; 32])).unwrap(),
+            UpdateResult::Inserted
+        );
+        assert!(!tree.contains(&short_key(0)));
+        assert!(tree.contains(&short_key(1)));
+        assert_eq!(
+            tree.insert(short_key(1), ValueHash::from([2u8; 32])).unwrap_err(),
+            SMTError::KeyExists
+        );
+        assert_eq!(tree.get(&short_key(1)).unwrap().unwrap(), &ValueHash::from([1u8; 32]));
+
+        // Delete the key, which the tree must not contain
+        tree.delete(&short_key(1)).unwrap();
+        assert!(!tree.contains(&short_key(0)));
+        assert!(!tree.contains(&short_key(1)));
+
+        // Build a more complex tree with two keys, which the tree must then contain
+        assert_eq!(
+            tree.insert(short_key(0), ValueHash::from([0u8; 32])).unwrap(),
+            UpdateResult::Inserted
+        );
+        assert_eq!(
+            tree.insert(short_key(1), ValueHash::from([1u8; 32])).unwrap(),
+            UpdateResult::Inserted
+        );
+        assert!(tree.contains(&short_key(0)));
+        assert!(tree.contains(&short_key(1)));
+        assert_eq!(
+            tree.insert(short_key(0), ValueHash::from([1u8; 32])).unwrap_err(),
+            SMTError::KeyExists
+        );
+        assert_eq!(tree.get(&short_key(0)).unwrap().unwrap(), &ValueHash::from([0u8; 32]));
+        assert_eq!(
+            tree.insert(short_key(1), ValueHash::from([2u8; 32])).unwrap_err(),
+            SMTError::KeyExists
+        );
+        assert_eq!(tree.get(&short_key(1)).unwrap().unwrap(), &ValueHash::from([1u8; 32]));
+    }
 }