feat(backend): add updates and deletions to Merkle tree state transition proofs (#383)

* Add updates and deletions to Merkle tree state transition proofs

* Code review

src/backends/plonky2/primitives/merkletree/circuit.rs

@@ -30,7 +30,9 @@ use crate::{
         basetypes::D,
         circuits::common::{CircuitBuilderPod, ValueTarget},
         error::Result,
-        primitives::merkletree::{MerkleClaimAndProof, MerkleProofStateTransition},
+        primitives::merkletree::{
+            MerkleClaimAndProof, MerkleTreeOp, MerkleTreeStateTransitionProof,
+        },
     },
     measure_gates_begin, measure_gates_end,
     middleware::{EMPTY_HASH, EMPTY_VALUE, F, HASH_SIZE},
@@ -396,19 +398,18 @@ fn kv_hash_target(
 /// Verifies that the merkletree state transition (from old_root to new_root)
 /// has been done correctly for the given new_key. This will allow verifying
 /// correct new leaf insertion, and leaf edition&deletion (if needed).
-pub struct MerkleProofStateTransitionTarget {
+/// See `MerkleTreeStateTransitionProof` struct for an explanation of the fields.
+pub struct MerkleTreeStateTransitionProofTarget {
     pub(crate) max_depth: usize,
     // `enabled` determines if the merkleproof state transition verification is enabled
     pub(crate) enabled: BoolTarget,
-    pub(crate) typ: Target,
+    pub(crate) op: Target,
-
     pub(crate) old_root: HashOutTarget,
-    pub(crate) proof_non_existence: MerkleClaimAndProofTarget,
+    pub(crate) op_proof: MerkleClaimAndProofTarget,
-
     pub(crate) new_root: HashOutTarget,
-    pub(crate) new_key: ValueTarget,
+    pub(crate) op_key: ValueTarget,
-    pub(crate) new_value: ValueTarget,
+    pub(crate) op_value: ValueTarget,
-    pub(crate) new_siblings: Vec<HashOutTarget>,
+    pub(crate) siblings: Vec<HashOutTarget>,
 
     // auxiliary witness
     pub(crate) divergence_level: Target,
@@ -416,85 +417,117 @@ pub struct MerkleProofStateTransitionTarget {
 /// creates the targets and defines the logic of the circuit
 pub fn verify_merkle_state_transition_circuit(
     builder: &mut CircuitBuilder<F, D>,
-    proof: &MerkleProofStateTransitionTarget,
+    proof: &MerkleTreeStateTransitionProofTarget,
 ) {
     let measure = measure_gates_begin!(
         builder,
         format!("MerkleProofStateTransition_{}", proof.max_depth)
     );
     let zero = builder.zero();
+    let one = builder.one();
+    let two = builder.constant(F::from_canonical_u64(2));
 
-    // for now, only type=0 (insertion proof) is supported
+    // Op type check
-    builder.connect(proof.typ, zero);
+    let is_insertion = builder.is_equal(proof.op, zero);
+    let is_update = builder.is_equal(proof.op, one);
+    let is_deletion = builder.is_equal(proof.op, two);
+    let op_type_check = {
+        let a = builder.or(is_insertion, is_update);
+        builder.or(a, is_deletion)
+    };
+    builder.assert_one(op_type_check.target);
 
-    // 1) check that for the old_root, the new_key does not exist in the tree
+    // 1) Verify the provided op proof.
-    verify_merkle_proof_circuit(builder, &proof.proof_non_existence);
+    verify_merkle_proof_circuit(builder, &proof.op_proof);
 
-    // 2) check that for the new_root, the new_key does exist in the tree
+    // 2) Check that the provided siblings yield an existence proof of
+    // (op_key, op_value) with root specified by op:
+    // Insert: `new_root`
+    // Update: `new_root`
+    // Delete: `old_root`
+    let root = HashOutTarget {
+        elements: std::array::from_fn(|i| {
+            builder.select(
+                is_deletion,
+                proof.old_root.elements[i],
+                proof.new_root.elements[i],
+            )
+        }),
+    };
     let new_key_proof = MerkleProofExistenceTarget {
         max_depth: proof.max_depth,
         enabled: proof.enabled,
-        root: proof.new_root,
+        root,
-        key: proof.new_key,
+        key: proof.op_key,
-        value: proof.new_value,
+        value: proof.op_value,
-        siblings: proof.new_siblings.clone(),
+        siblings: proof.siblings.clone(),
     };
     let new_leaf_path = verify_merkle_proof_existence_circuit(builder, &new_key_proof);
 
-    // 3.1) assert that proof_non_existence.existence==false
+    // 3.1) assert that op_proof.existence is of the appropriate type according to op:
+    // Insert/Delete: Non-existence
+    // Update: Existence
+    let proof_type = is_update;
     builder.conditional_assert_eq(
         proof.enabled.target,
-        proof.proof_non_existence.existence.target,
+        proof.op_proof.existence.target,
-        zero,
+        proof_type.target,
-    );
-    // 3.2) assert that proof.enabled matches with proof_non_existence.enabled
-    builder.connect(
-        proof.proof_non_existence.enabled.target,
-        proof.enabled.target,
     );
+    // 3.2) assert that proof.enabled matches with op_proof.enabled
+    builder.connect(proof.op_proof.enabled.target, proof.enabled.target);
 
-    // 4) assert proof_non_existence.root==old_root, and that it uses new_key
+    // 4) assert proof_non_existence.root corresponds to the root
+    // specified by the op (old_root for Insert/Update and new_root
+    // otherwise), and that it uses op_key
+    let claim_root = HashOutTarget {
+        elements: std::array::from_fn(|i| {
+            builder.select(
+                is_deletion,
+                proof.new_root.elements[i],
+                proof.old_root.elements[i],
+            )
+        }),
+    };
     for j in 0..HASH_SIZE {
         // 4.1) assert that proof.proof_non_existence.root == proof.old_root
         builder.conditional_assert_eq(
             proof.enabled.target,
-            proof.proof_non_existence.root.elements[j],
+            proof.op_proof.root.elements[j],
-            proof.old_root.elements[j],
+            claim_root.elements[j],
         );
-        // 4.2) assert that the non-existence proof uses the new_key (value not needed for
+        // 4.2) assert that the non-existence proof uses the op_key (value not needed).
-        //   non-existence)
         builder.conditional_assert_eq(
             proof.enabled.target,
-            proof.proof_non_existence.key.elements[j],
+            proof.op_proof.key.elements[j],
-            proof.new_key.elements[j],
+            proof.op_key.elements[j],
         );
     }
 
     // prepare value for check 5.2)
     let old_leaf_hash = kv_hash_target(
         builder,
-        &proof.proof_non_existence.other_key,
+        &proof.op_proof.other_key,
-        &proof.proof_non_existence.other_value,
+        &proof.op_proof.other_value,
     );
     // prepare values for check 5.3)
-    let old_leaf_path = keypath_target(
+    let old_leaf_path = keypath_target(proof.max_depth, builder, &proof.op_proof.other_key);
-        proof.max_depth,
-        builder,
-        &proof.proof_non_existence.other_key,
-    );
 
-    // 5) check that old_siblings & new_siblings match as expected. Let
+    // 5) check that old_siblings & new_siblings match as expected. They
-    //    d=divergence_level, assert that:
+    // should match at all levels for an update. Otherwise, let
+    //    d=divergence_level and assert that:
     // 5.1) old_siblings[i] == new_siblings[i] ∀ i \ {d}
     // 5.2) at i==d, if old_siblings[i] != new_siblings[i]: old_siblings[i] ==
     //   EMPTY_HASH new_siblings[i] == old_leaf_hash
     // 5.3) assert that if old_key!=empty, both old_leaf_path&new_leaf_path
     //   should diverge at the inputted divergence level
-    let old_siblings = proof.proof_non_existence.siblings.clone();
+    let old_siblings = proof.op_proof.siblings.clone();
-    let new_siblings = proof.new_siblings.clone();
+    let new_siblings = proof.siblings.clone();
     for i in 0..proof.max_depth {
         let i_targ = builder.constant(F::from_canonical_u64(i as u64));
         let is_divergence_level = builder.is_equal(i_targ, proof.divergence_level);
+        let is_not_update = builder.not(is_update);
+        // There is no divergence level for an update.
+        let is_divergence_level = builder.and(is_not_update, is_divergence_level);
 
         // 5.1) for all i except for i==divergence_level, assert that the
         //   siblings are the same
@@ -549,21 +582,20 @@ pub fn verify_merkle_state_transition_circuit(
     measure_gates_end!(builder, measure);
 }
 
-impl MerkleProofStateTransitionTarget {
+impl MerkleTreeStateTransitionProofTarget {
     pub fn new_virtual(max_depth: usize, builder: &mut CircuitBuilder<F, D>) -> Self {
         Self {
             max_depth,
             enabled: builder.add_virtual_bool_target_safe(),
-            typ: builder.add_virtual_target(),
+            op: builder.add_virtual_target(),
 
             old_root: builder.add_virtual_hash(),
-            proof_non_existence: MerkleClaimAndProofTarget::new_virtual(max_depth, builder),
+            op_proof: MerkleClaimAndProofTarget::new_virtual(max_depth, builder),
             new_root: builder.add_virtual_hash(),
-            new_key: builder.add_virtual_value(),
+            op_key: builder.add_virtual_value(),
-            new_value: builder.add_virtual_value(),
+            op_value: builder.add_virtual_value(),
-
             // siblings are padded till max_depth length
-            new_siblings: builder.add_virtual_hashes(max_depth),
+            siblings: builder.add_virtual_hashes(max_depth),
 
             divergence_level: builder.add_virtual_target(),
         }
@@ -575,26 +607,30 @@ impl MerkleProofStateTransitionTarget {
         &self,
         pw: &mut PartialWitness<F>,
         enabled: bool,
-        mp: &MerkleProofStateTransition,
+        mp: &MerkleTreeStateTransitionProof,
     ) -> Result<()> {
         pw.set_bool_target(self.enabled, enabled)?;
-        pw.set_target(self.typ, F::from_canonical_u8(mp.typ))?;
+        pw.set_target(self.op, F::from_canonical_u8(mp.op as u8))?;
 
         pw.set_hash_target(self.old_root, HashOut::from_vec(mp.old_root.0.to_vec()))?;
-        self.proof_non_existence.set_targets(
+        self.op_proof.set_targets(
             pw,
             enabled,
             &MerkleClaimAndProof {
-                root: mp.old_root,
+                root: if mp.op == MerkleTreeOp::Delete {
-                key: mp.new_key,
+                    mp.new_root
-                value: EMPTY_VALUE, // not needed for non-existence
+                } else {
-                proof: mp.proof_non_existence.clone(),
+                    mp.old_root
+                },
+                key: mp.op_key,
+                value: mp.value.unwrap_or(EMPTY_VALUE), // not needed for non-existence
+                proof: mp.op_proof.clone(),
             },
         )?;
 
         pw.set_hash_target(self.new_root, HashOut::from_vec(mp.new_root.0.to_vec()))?;
-        pw.set_target_arr(&self.new_key.elements, &mp.new_key.0)?;
+        pw.set_target_arr(&self.op_key.elements, &mp.op_key.0)?;
-        pw.set_target_arr(&self.new_value.elements, &mp.new_value.0)?;
+        pw.set_target_arr(&self.op_value.elements, &mp.op_value.0)?;
 
         let new_siblings = mp.siblings.clone();
 
@@ -605,7 +641,7 @@ impl MerkleProofStateTransitionTarget {
             .take(self.max_depth)
             .enumerate()
         {
-            pw.set_hash_target(self.new_siblings[i], HashOut::from_vec(sibling.0.to_vec()))?;
+            pw.set_hash_target(self.siblings[i], HashOut::from_vec(sibling.0.to_vec()))?;
         }
         pw.set_target(
             self.divergence_level,
@@ -961,7 +997,7 @@ pub mod tests {
     fn run_state_transition_circuit(
         expect_pass: bool,
         max_depth: usize,
-        state_transition_proof: &MerkleProofStateTransition,
+        state_transition_proof: &MerkleTreeStateTransitionProof,
     ) -> Result<()> {
         // sanity check, run the out-circuit proof verification
         if expect_pass {
@@ -975,7 +1011,7 @@ pub mod tests {
         let mut builder = CircuitBuilder::<F, D>::new(config);
         let mut pw = PartialWitness::<F>::new();
 
-        let targets = MerkleProofStateTransitionTarget::new_virtual(max_depth, &mut builder);
+        let targets = MerkleTreeStateTransitionProofTarget::new_virtual(max_depth, &mut builder);
         verify_merkle_state_transition_circuit(&mut builder, &targets);
         targets.set_targets(&mut pw, true, state_transition_proof)?;
 
@@ -1010,15 +1046,34 @@ pub mod tests {
         assert_eq!(state_transition_proof.old_root, old_root);
         assert_eq!(state_transition_proof.new_root, tree.root());
 
+        // Deleting this key should yield the old tree, and the proof
+        // should be the same (mutatis mutandis).
+        let mut tree_with_deleted_key = tree.clone();
+        let state_transition_proof1 = tree_with_deleted_key.delete(&key)?;
+        run_state_transition_circuit(true, max_depth, &state_transition_proof1)?;
+        assert_eq!(state_transition_proof1.old_root, tree.root());
+        assert_eq!(state_transition_proof1.new_root, old_root);
+
         // add a new leaf, which shares path with the previous one, but diverges
         // one level before, where there is no leaf yet to be pushed down.
-        let old_root = tree.root();
+        let mut tree_with_another_leaf = tree.clone();
+        let old_root = tree_with_another_leaf.root();
         let key = RawValue::from(21);
         let value = RawValue::from(1021);
-        let state_transition_proof = tree.insert(&key, &value)?;
+        let state_transition_proof = tree_with_another_leaf.insert(&key, &value)?;
         run_state_transition_circuit(true, max_depth, &state_transition_proof)?;
         assert_eq!(state_transition_proof.old_root, old_root);
-        assert_eq!(state_transition_proof.new_root, tree.root());
+        assert_eq!(
+            state_transition_proof.new_root,
+            tree_with_another_leaf.root()
+        );
+
+        // Alternatively add this key with another value then update.
+        let value1 = RawValue::from(99);
+        tree.insert(&key, &value1)?;
+        let state_transition_proof1 = tree.update(&key, &value)?;
+        run_state_transition_circuit(true, max_depth, &state_transition_proof1)?;
+        assert_eq!(tree.root(), tree_with_another_leaf.root());
 
         // another leaf which will push further down the leaf with key=37
         let old_root = tree.root();
@@ -1058,6 +1113,39 @@ pub mod tests {
         assert_eq!(state_transition_proof.old_root, old_root);
         assert_eq!(state_transition_proof.new_root, tree.root());
 
+        // ...and delete.
+        let state_transition_proof1 = tree.delete(&key)?;
+        run_state_transition_circuit(true, max_depth, &state_transition_proof1)?;
+        assert_eq!(
+            state_transition_proof1.old_root,
+            state_transition_proof.new_root
+        );
+        assert_eq!(
+            state_transition_proof1.new_root,
+            state_transition_proof.old_root
+        );
+
+        // ...and add a key-value pair with the same key but a
+        // different value, then update.
+        let value1 = RawValue::from(50);
+        let state_transition_proof1 = tree.insert(&key, &value1)?;
+        run_state_transition_circuit(true, max_depth, &state_transition_proof1)?;
+        assert_eq!(
+            state_transition_proof1.old_root,
+            state_transition_proof.old_root
+        );
+
+        let state_transition_proof2 = tree.update(&key, &value)?;
+        run_state_transition_circuit(true, max_depth, &state_transition_proof2)?;
+        assert_eq!(
+            state_transition_proof2.old_root,
+            state_transition_proof1.new_root
+        );
+        assert_eq!(
+            state_transition_proof2.new_root,
+            state_transition_proof.new_root
+        );
+
         Ok(())
     }
 
@@ -1099,14 +1187,11 @@ pub mod tests {
 
         // Tamper with state transition.
         const OFFSET: usize = 20;
-        let other_leaf = state_transition_proof
+        let other_leaf = state_transition_proof.op_proof.other_leaf.unwrap();
-            .proof_non_existence
-            .other_leaf
-            .unwrap();
         let altered_proof = MerkleProof {
             existence: true,
             siblings: [
-                state_transition_proof.proof_non_existence.siblings.clone(),
+                state_transition_proof.op_proof.siblings.clone(),
                 vec![EMPTY_HASH; OFFSET],
                 vec![kv_hash(&other_leaf.0, Some(other_leaf.1))],
             ]
@@ -1115,8 +1200,8 @@ pub mod tests {
         };
         let altered_root = altered_proof.compute_root_from_leaf(
             max_depth,
-            &state_transition_proof.new_key,
+            &state_transition_proof.op_key,
-            Some(state_transition_proof.new_value),
+            Some(state_transition_proof.op_value),
         )?;
         state_transition_proof.siblings = altered_proof.siblings;
         state_transition_proof.new_root = altered_root;
@@ -1165,13 +1250,13 @@ pub mod tests {
 
     fn run_circuit_disabled(
         max_depth: usize,
-        state_transition_proof: &MerkleProofStateTransition,
+        state_transition_proof: &MerkleTreeStateTransitionProof,
     ) -> Result<()> {
         let config = CircuitConfig::standard_recursion_config();
         let mut builder = CircuitBuilder::<F, D>::new(config);
         let mut pw = PartialWitness::<F>::new();
 
-        let targets = MerkleProofStateTransitionTarget::new_virtual(max_depth, &mut builder);
+        let targets = MerkleTreeStateTransitionProofTarget::new_virtual(max_depth, &mut builder);
         verify_merkle_state_transition_circuit(&mut builder, &targets);
         targets.set_targets(&mut pw, true, state_transition_proof)?;
 
@@ -1183,7 +1268,7 @@ pub mod tests {
         let mut builder = CircuitBuilder::<F, D>::new(config);
         let mut pw = PartialWitness::<F>::new();
 
-        let targets = MerkleProofStateTransitionTarget::new_virtual(max_depth, &mut builder);
+        let targets = MerkleTreeStateTransitionProofTarget::new_virtual(max_depth, &mut builder);
         verify_merkle_state_transition_circuit(&mut builder, &targets);
         targets.set_targets(&mut pw, false, state_transition_proof)?;
 

src/backends/plonky2/primitives/merkletree/error.rs

@@ -12,12 +12,12 @@ pub enum TreeInnerError {
     KeyExists,
     #[error("max depth reached")]
     MaxDepth,
-    #[error("reached empty node, should not have entered")]
-    EmptyNode,
     #[error("proof of {0} does not verify")]
     ProofFail(String), // inclusion / exclusion
     #[error("invalid {0} proof")]
     InvalidProof(String),
+    #[error("invalid state transition proof argument: {0}")]
+    InvalidStateTransitionProogArg(String),
     #[error("state transition proof does not verify, reason: {0}")]
     StateTransitionProofFail(String),
     #[error("key too short (key length: {0}) for the max_depth: {1}")]
@@ -66,15 +66,15 @@ impl TreeError {
     pub(crate) fn max_depth() -> Self {
         new!(MaxDepth)
     }
-    pub(crate) fn empty_node() -> Self {
-        new!(EmptyNode)
-    }
     pub(crate) fn proof_fail(obj: String) -> Self {
         new!(ProofFail(obj))
     }
     pub(crate) fn invalid_proof(obj: String) -> Self {
         new!(InvalidProof(obj))
     }
+    pub(crate) fn invalid_state_transition_proof_arg(reason: String) -> Self {
+        new!(InvalidStateTransitionProogArg(reason))
+    }
     pub(crate) fn state_transition_fail(reason: String) -> Self {
         new!(StateTransitionProofFail(reason))
     }

src/backends/plonky2/primitives/merkletree/mod.rs

@@ -24,18 +24,12 @@ pub struct MerkleTree {
 impl MerkleTree {
     /// builds a new `MerkleTree` where the leaves contain the given key-values
     pub fn new(max_depth: usize, kvs: &HashMap<RawValue, RawValue>) -> TreeResult<Self> {
-        // Construct leaves.
+        // Start with an empty node as root.
-        let mut leaves: Vec<_> = kvs
+        let mut root = Node::None;
-            .iter()
-            .map(|(k, v)| Leaf::new(max_depth, *k, *v))
-            .collect::<TreeResult<_>>()?;
 
-        // Start with a leaf or conclude with an empty node as root.
+        // Iterate over key-value pairs (if any) and add them.
-        let mut root = leaves.pop().map(Node::Leaf).unwrap_or(Node::None);
+        for (k, v) in kvs.iter() {
-
+            root.apply_op(max_depth, MerkleTreeOp::Insert, *k, Some(*v))?;
-        // Iterate over remaining leaves (if any) and add them.
-        for leaf in leaves.into_iter() {
-            root.add_leaf(0, max_depth, leaf)?;
         }
 
         // Fill in hashes.
@@ -82,12 +76,12 @@ impl MerkleTree {
         &mut self,
         key: &RawValue,
         value: &RawValue,
-    ) -> TreeResult<MerkleProofStateTransition> {
+    ) -> TreeResult<MerkleTreeStateTransitionProof> {
         let proof_non_existence = self.prove_nonexistence(key)?;
 
         let old_root: Hash = self.root.hash();
         self.root
-            .add_leaf(0, self.max_depth, Leaf::new(self.max_depth, *key, *value)?)?;
+            .apply_op(self.max_depth, MerkleTreeOp::Insert, *key, Some(*value))?;
         let new_root = self.root.compute_hash();
 
         let (v, proof) = self.prove(key)?;
@@ -95,17 +89,70 @@ impl MerkleTree {
         assert_eq!(v, *value);
         assert!(proof.other_leaf.is_none());
 
-        Ok(MerkleProofStateTransition {
+        Ok(MerkleTreeStateTransitionProof {
-            typ: 0, // insertion
+            op: MerkleTreeOp::Insert, // insertion
             old_root,
-            proof_non_existence,
+            op_proof: proof_non_existence,
             new_root,
-            new_key: *key,
+            op_key: *key,
-            new_value: *value,
+            op_value: *value,
+            value: None,
             siblings: proof.siblings,
         })
     }
 
+    pub fn update(
+        &mut self,
+        key: &RawValue,
+        value: &RawValue,
+    ) -> TreeResult<MerkleTreeStateTransitionProof> {
+        let (old_value, old_proof) = self.prove(key)?;
+
+        let old_root: Hash = self.root.hash();
+        self.root
+            .apply_op(self.max_depth, MerkleTreeOp::Update, *key, Some(*value))?;
+        let new_root = self.root.compute_hash();
+
+        let (v, proof) = self.prove(key)?;
+        assert!(proof.existence);
+        assert_eq!(v, *value);
+        assert!(proof.other_leaf.is_none());
+
+        Ok(MerkleTreeStateTransitionProof {
+            op: MerkleTreeOp::Update,
+            old_root,
+            op_proof: old_proof,
+            new_root,
+            op_key: *key,
+            op_value: *value,
+            value: Some(old_value),
+            siblings: proof.siblings,
+        })
+    }
+
+    pub fn delete(&mut self, key: &RawValue) -> TreeResult<MerkleTreeStateTransitionProof> {
+        let (value, proof_existence) = self.prove(key)?;
+
+        let old_root: Hash = self.root.hash();
+        self.root
+            .apply_op(self.max_depth, MerkleTreeOp::Delete, *key, None)?;
+        let new_root = self.root.compute_hash();
+
+        let proof = self.prove_nonexistence(key)?;
+        assert!(!proof.existence);
+
+        Ok(MerkleTreeStateTransitionProof {
+            op: MerkleTreeOp::Delete,
+            old_root,
+            op_proof: proof,
+            new_root,
+            op_key: *key,
+            op_value: value,
+            value: None,
+            siblings: proof_existence.siblings,
+        })
+    }
+
     /// returns a proof of existence, which proves that the given key exists in
     /// the tree. It returns the `value` of the leaf at the given `key`, and the
     /// `MerkleProof`.
@@ -206,86 +253,130 @@ impl MerkleTree {
 
     pub fn verify_state_transition(
         max_depth: usize,
-        proof: &MerkleProofStateTransition,
+        proof: &MerkleTreeStateTransitionProof,
     ) -> TreeResult<()> {
-        let mut old_siblings = proof.proof_non_existence.siblings.clone();
+        let mut old_siblings = proof.op_proof.siblings.clone();
         let new_siblings = proof.siblings.clone();
 
-        // check that for the old_root, the new_key does not exist in the tree
+        match proof.op {
-        Self::verify_nonexistence(
+            // A deletion is but an insertion subject to a time reversal.
-            max_depth,
+            MerkleTreeOp::Delete => {
-            proof.old_root,
+                let equivalent_insertion_proof = MerkleTreeStateTransitionProof {
-            &proof.proof_non_existence,
+                    op: MerkleTreeOp::Insert,
-            &proof.new_key,
+                    new_root: proof.old_root,
-        )?;
+                    old_root: proof.new_root,
-
+                    ..proof.clone()
-        // check that new_siblings verify with the new_root
+                };
-        Self::verify(
+                Self::verify_state_transition(max_depth, &equivalent_insertion_proof)
-            max_depth,
-            proof.new_root,
-            &MerkleProof {
-                existence: true,
-                siblings: new_siblings.clone(),
-                other_leaf: None,
-            },
-            &proof.new_key,
-            &proof.new_value,
-        )?;
-
-        // if other_leaf exists, check path divergence
-        if let Some((other_key, _)) = proof.proof_non_existence.other_leaf {
-            let old_path = keypath(max_depth, other_key)?;
-            let new_path = keypath(max_depth, proof.new_key)?;
-
-            let divergence_lvl: usize = match zip_eq(old_path, new_path).position(|(x, y)| x != y) {
-                Some(d) => d,
-                None => return Err(TreeError::max_depth()),
-            };
-
-            if divergence_lvl != new_siblings.len() - 1 {
-                return Err(TreeError::state_transition_fail(
-                    "paths divergence does not match".to_string(),
-                ));
             }
-        }
+            MerkleTreeOp::Update => {
+                // check that for the old_root, (op_key, value) *does* exist in the tree
+                Self::verify(
+                    max_depth,
+                    proof.old_root,
+                    &proof.op_proof,
+                    &proof.op_key,
+                    &proof.value.unwrap(),
+                )?;
+                // check that for the new_root, (op_key, op_value) *does* exist in the tree
+                Self::verify(
+                    max_depth,
+                    proof.new_root,
+                    &MerkleProof {
+                        existence: true,
+                        siblings: proof.siblings.clone(),
+                        other_leaf: None,
+                    },
+                    &proof.op_key,
+                    &proof.op_value,
+                )?;
 
-        // let d=divergence_level, assert that:
+                // All siblings should agree
-        // 1) old_siblings[i] == new_siblings[i] ∀ i \ {d}
+                (proof.siblings == proof.op_proof.siblings)
-        // 2) at i==d, if old_siblings[i] != new_siblings[i]:
+                    .then_some(())
-        //     old_siblings[i] == EMPTY_HASH
+                    .ok_or(TreeError::state_transition_fail(format!(
-        //     new_siblings[i] == old_leaf_hash
+                        "Invalid proof of update for key {}: Siblings don't match.",
-        let d = new_siblings.len() - 1;
+                        proof.op_key
-        old_siblings.resize(d + 1, EMPTY_HASH);
+                    )))
-        for i in 0..d {
-            if old_siblings[i] != new_siblings[i] {
-                return Err(TreeError::state_transition_fail(
-                    "siblings don't match: old[i]!=new[i] ∀ i (except at i==d)".to_string(),
-                ));
             }
-        }
+            MerkleTreeOp::Insert => {
-        if old_siblings[d] != new_siblings[d] {
+                // check that for the old_root, the new_key does not exist in the tree
-            if old_siblings[d] != EMPTY_HASH {
+                Self::verify_nonexistence(
-                return Err(TreeError::state_transition_fail(
+                    max_depth,
-                    "siblings don't match: old[d]!=empty".to_string(),
+                    proof.old_root,
-                ));
+                    &proof.op_proof,
-            }
+                    &proof.op_key,
-            let k = proof
+                )?;
-                .proof_non_existence
+
+                // check that new_siblings verify with the new_root
+                Self::verify(
+                    max_depth,
+                    proof.new_root,
+                    &MerkleProof {
+                        existence: true,
+                        siblings: new_siblings.clone(),
+                        other_leaf: None,
+                    },
+                    &proof.op_key,
+                    &proof.op_value,
+                )?;
+
+                // if other_leaf exists, check path divergence
+                if let Some((other_key, _)) = proof.op_proof.other_leaf {
+                    let old_path = keypath(max_depth, other_key)?;
+                    let new_path = keypath(max_depth, proof.op_key)?;
+
+                    let divergence_lvl: usize =
+                        match zip_eq(old_path, new_path).position(|(x, y)| x != y) {
+                            Some(d) => d,
+                            None => return Err(TreeError::max_depth()),
+                        };
+
+                    if divergence_lvl != new_siblings.len() - 1 {
+                        return Err(TreeError::state_transition_fail(
+                            "paths divergence does not match".to_string(),
+                        ));
+                    }
+                }
+
+                // let d=divergence_level, assert that:
+                // 1) old_siblings[i] == new_siblings[i] ∀ i \ {d}
+                // 2) at i==d, if old_siblings[i] != new_siblings[i]:
+                //     old_siblings[i] == EMPTY_HASH
+                //     new_siblings[i] == old_leaf_hash
+                let d = new_siblings.len() - 1;
+                old_siblings.resize(d + 1, EMPTY_HASH);
+                for i in 0..d {
+                    if old_siblings[i] != new_siblings[i] {
+                        return Err(TreeError::state_transition_fail(
+                            "siblings don't match: old[i]!=new[i] ∀ i (except at i==d)".to_string(),
+                        ));
+                    }
+                }
+                if old_siblings[d] != new_siblings[d] {
+                    if old_siblings[d] != EMPTY_HASH {
+                        return Err(TreeError::state_transition_fail(
+                            "siblings don't match: old[d]!=empty".to_string(),
+                        ));
+                    }
+                    let k = proof
+                .op_proof
                 .other_leaf
                 .map(|(k, _)| k)
                 .ok_or(TreeError::state_transition_fail(
                         "proof.proof_non_existence.other_leaf can not be empty for the case old_siblings[d]!=new_siblings[d]".to_string()
                         ))?;
-            let v: Option<RawValue> = proof.proof_non_existence.other_leaf.map(|(_, v)| v);
+                    let v: Option<RawValue> = proof.op_proof.other_leaf.map(|(_, v)| v);
-            let old_leaf_hash = kv_hash(&k, v);
+                    let old_leaf_hash = kv_hash(&k, v);
-            if new_siblings[d] != old_leaf_hash {
+                    if new_siblings[d] != old_leaf_hash {
-                return Err(TreeError::state_transition_fail(
+                        return Err(TreeError::state_transition_fail(
-                    "siblings don't match: new[d]!=old_leaf_hash".to_string(),
+                            "siblings don't match: new[d]!=old_leaf_hash".to_string(),
-                ));
+                        ));
+                    }
+                }
+                Ok(())
             }
         }
-
-        Ok(())
     }
 
     /// returns an iterator over the leaves of the tree
@@ -424,19 +515,38 @@ impl fmt::Display for MerkleClaimAndProof {
     }
 }
 
+#[derive(Copy, Clone, Debug, PartialEq, Serialize, Deserialize)]
+pub enum MerkleTreeOp {
+    Insert = 0,
+    Update,
+    Delete,
+}
+
 #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
-pub struct MerkleProofStateTransition {
+pub struct MerkleTreeStateTransitionProof {
-    // type: 0:insertion, 1:update, 2:deletion
+    pub(crate) op: MerkleTreeOp,
-    pub(crate) typ: u8,
 
     pub(crate) old_root: Hash,
-    // proof of non-existence of the new_key for the old_root
+
-    pub(crate) proof_non_existence: MerkleProof,
+    /// Insert: proof of non-existence of the op_key for the old_root
+    /// Update: proof of existence of (op_key, value) for the old_root
+    /// Delete: proof of non-existence of the op_key for the new_root
+    pub(crate) op_proof: MerkleProof,
 
     pub(crate) new_root: Hash,
-    pub(crate) new_key: RawValue,
-    pub(crate) new_value: RawValue,
 
+    /// Key & value relevant to transition proof. These are the
+    /// inserted/updated key-value pair for insertions and updates. For
+    /// deletions, these are the key-value pair that is deleted.
+    pub(crate) op_key: RawValue,
+    pub(crate) op_value: RawValue,
+
+    /// Update: value to be replaced.
+    pub(crate) value: Option<RawValue>,
+
+    /// Insert: siblings of inserted (op_key, op_value) leading to new_root
+    /// Update: siblings of updated (op_key, op_value) leading to new_root
+    /// Delete: siblings of deleted (op_key, op_value) leading to old_root
     pub(crate) siblings: Vec<Hash>,
 }
 
@@ -505,6 +615,13 @@ impl Node {
             Self::Intermediate(_n) => false,
         }
     }
+    fn is_intermediate(&self) -> bool {
+        match self {
+            Self::None => false,
+            Self::Leaf(_l) => false,
+            Self::Intermediate(_n) => true,
+        }
+    }
     fn compute_hash(&mut self) -> Hash {
         match self {
             Self::None => EMPTY_HASH,
@@ -564,53 +681,166 @@ impl Node {
         }
     }
 
-    // adds the leaf at the tree from the current node (self), without computing any hash
+    /// Applies given Merkle tree op without computing hashes.
-    pub(crate) fn add_leaf(&mut self, lvl: usize, max_depth: usize, leaf: Leaf) -> TreeResult<()> {
+    pub(crate) fn apply_op(
+        &mut self,
+        max_depth: usize,
+        op: MerkleTreeOp,
+        key: RawValue,
+        maybe_value: Option<RawValue>,
+    ) -> TreeResult<()> {
+        let key_path = keypath(max_depth, key)?;
+        // Rule out invalid arguments
+        match (op, maybe_value) {
+            (MerkleTreeOp::Insert, None) | (MerkleTreeOp::Update, None) => {
+                Err(TreeError::invalid_state_transition_proof_arg(format!(
+                    "{:?} op requires a value argument.",
+                    op
+                )))
+            }
+            (MerkleTreeOp::Delete, Some(_)) => {
+                Err(TreeError::invalid_state_transition_proof_arg(format!(
+                    "{:?} op requires no value argument, yet one was provided.",
+                    op
+                )))
+            }
+            _ => Ok(()),
+        }?;
+
+        // Loop through to leaf.
+        self.apply_op_loop(0, max_depth, op, key, &key_path, maybe_value)?;
+
+        // If we are dealing with a deletion, normalise along key
+        // path.
+        if let MerkleTreeOp::Delete = op {
+            self.normalise_path(&key_path);
+        }
+
+        Ok(())
+    }
+
+    /// Normalises a Merkle tree along a specified path. Useful
+    /// post-deletion.
+    fn normalise_path(&mut self, key_path: &[bool]) {
+        match self {
+            Self::Leaf(_) | Self::None => (),
+            Self::Intermediate(Intermediate {
+                hash: _h,
+                left,
+                right,
+            }) => {
+                if key_path[0] {
+                    right.normalise_path(&key_path[1..]);
+                } else {
+                    left.normalise_path(&key_path[1..]);
+                }
+
+                // If we have a branch with children (NIL, X) or (X,
+                // NIL) where X is not a branch, then replace with X.
+                if left.is_empty() && !right.is_intermediate() {
+                    *self = *right.clone();
+                } else if right.is_empty() && !left.is_intermediate() {
+                    *self = *left.clone();
+                }
+            }
+        }
+    }
+
+    fn apply_op_loop(
+        &mut self,
+        lvl: usize,
+        max_depth: usize,
+        op: MerkleTreeOp,
+        key: RawValue,
+        key_path: &[bool],
+        maybe_value: Option<RawValue>,
+    ) -> TreeResult<()> {
         if lvl >= max_depth {
             return Err(TreeError::max_depth());
         }
 
         match self {
             Self::Intermediate(n) => {
-                if leaf.path[lvl] {
+                if key_path[lvl] {
-                    if n.right.is_empty() {
+                    n.right
-                        // empty sub-node, add the leaf here
+                        .apply_op_loop(lvl + 1, max_depth, op, key, key_path, maybe_value)
-                        n.right = Box::new(Node::Leaf(leaf));
-                        return Ok(());
-                    }
-                    n.right.add_leaf(lvl + 1, max_depth, leaf)?;
                 } else {
-                    if n.left.is_empty() {
+                    n.left
-                        // empty sub-node, add the leaf here
+                        .apply_op_loop(lvl + 1, max_depth, op, key, key_path, maybe_value)
-                        n.left = Box::new(Node::Leaf(leaf));
-                        return Ok(());
-                    }
-                    n.left.add_leaf(lvl + 1, max_depth, leaf)?;
                 }
             }
-            Self::Leaf(l) => {
+            _ => {
+                *self = Self::op_node_check(max_depth, lvl, self, op, key, key_path, maybe_value)?;
+                Ok(())
+            }
+        }
+    }
+
+    /// Checks the terminal node against the desired op and returns a
+    /// suitable replacement.
+    ///
+    /// - Insertion => Node should be empty or contain a different
+    ///   key. A leaf is inserted in the right place.
+    /// - Update/Deletion => Node should contain the given key. The
+    ///   value is replaced in the case of an update and the leaf removed
+    ///   in the case of a deletion.
+    pub(crate) fn op_node_check(
+        max_depth: usize,
+        lvl: usize,
+        node: &Node,
+        op: MerkleTreeOp,
+        key: RawValue,
+        key_path: &[bool],
+        maybe_value: Option<RawValue>,
+    ) -> TreeResult<Node> {
+        use MerkleTreeOp::*;
+
+        // Invalid args are assumed to have been ruled out.
+        match (op, node, maybe_value) {
+            // Insertion case
+            (Insert, Node::None, Some(value)) => Ok(Node::Leaf(Leaf::new(max_depth, key, value)?)),
+            (Insert, Node::Leaf(l), Some(value)) => {
                 // in this case, it means that we found a leaf in the new-leaf
                 // path, thus we need to push both leaves (old-leaf and
                 // new-leaf) down the path till their paths diverge.
 
                 // first check that keys of both leaves are different
                 // (l=old-leaf, leaf=new-leaf)
-                if l.key == leaf.key {
+                if l.key == key {
                     // Note: current approach returns an error when trying to
                     // add to a leaf where the key already exists. We could also
                     // ignore it if needed.
-                    return Err(TreeError::key_exists());
+                    Err(TreeError::key_exists())
+                } else {
+                    let old_leaf = l.clone();
+                    // set new node as an intermediate node
+                    let mut new_node = Node::Intermediate(Intermediate::empty());
+                    new_node.down_till_divergence(
+                        lvl,
+                        max_depth,
+                        old_leaf,
+                        Leaf {
+                            hash: None,
+                            path: key_path.to_vec(),
+                            key,
+                            value,
+                        },
+                    )?;
+                    Ok(new_node)
                 }
-                let old_leaf = l.clone();
-                // set self as an intermediate node
-                *self = Node::Intermediate(Intermediate::empty());
-                return self.down_till_divergence(lvl, max_depth, old_leaf, leaf);
             }
-            Self::None => {
+            // Update case
-                return Err(TreeError::empty_node());
+            (Update, Node::Leaf(l), Some(value)) if l.key == key => {
+                Ok(Node::Leaf(Leaf::new(max_depth, key, value)?))
             }
+            // Deletion case
+            (Delete, Node::Leaf(l), None) if l.key == key => Ok(Node::None),
+            // Case of terminal node that does not match.
+            _ => Err(TreeError::state_transition_fail(format!(
+                "{:?} op requires key {} to be present in the tree, yet it is not.",
+                op, key
+            ))),
         }
-        Ok(())
     }
 
     /// goes down through a 'virtual' path till finding a divergence. This
@@ -691,11 +921,11 @@ impl Intermediate {
 }
 
 #[derive(Clone, Debug)]
-struct Leaf {
+pub(crate) struct Leaf {
-    hash: Option<Hash>,
+    pub(crate) hash: Option<Hash>,
-    path: Vec<bool>,
+    pub(crate) path: Vec<bool>,
-    key: RawValue,
+    pub(crate) key: RawValue,
-    value: RawValue,
+    pub(crate) value: RawValue,
 }
 impl Leaf {
     fn new(max_depth: usize, key: RawValue, value: RawValue) -> TreeResult<Self> {
@@ -874,16 +1104,58 @@ pub mod tests {
         MerkleTree::verify_state_transition(max_depth, &state_transition_proof)?;
         assert_eq!(state_transition_proof.old_root, old_root);
         assert_eq!(state_transition_proof.new_root, tree.root());
-        assert_eq!(state_transition_proof.new_key, key);
+        assert_eq!(state_transition_proof.op_key, key);
-        assert_eq!(state_transition_proof.new_value, value);
+        assert_eq!(state_transition_proof.op_value, value);
+        assert_eq!(state_transition_proof.value, None);
+
+        // Deleting this key should yield the old tree, and the proof
+        // should be the same (mutatis mutandis).
+        let mut tree_with_deleted_key = tree.clone();
+        let state_transition_proof1 = tree_with_deleted_key.delete(&key)?;
+        MerkleTree::verify_state_transition(max_depth, &state_transition_proof1)?;
+        assert_eq!(
+            state_transition_proof1.old_root,
+            state_transition_proof.new_root
+        );
+        assert_eq!(
+            state_transition_proof1.new_root,
+            state_transition_proof.old_root
+        );
+        assert_eq!(
+            state_transition_proof1.op_key,
+            state_transition_proof.op_key
+        );
+        assert_eq!(
+            state_transition_proof1.op_value,
+            state_transition_proof.op_value
+        );
+        assert_eq!(
+            state_transition_proof1.op_proof,
+            state_transition_proof.op_proof
+        );
+        assert_eq!(
+            state_transition_proof1.siblings,
+            state_transition_proof.siblings
+        );
 
         // 2nd part of the test. Add a new leaf
+        let mut tree_with_another_leaf = tree.clone();
         let key = RawValue::from(21);
         let value = RawValue::from(1021);
-        let state_transition_proof = tree.insert(&key, &value)?;
+        let state_transition_proof = tree_with_another_leaf.insert(&key, &value)?;
 
         MerkleTree::verify_state_transition(max_depth, &state_transition_proof)?;
 
+        // Alternatively add this key with another value then update.
+        let value1 = RawValue::from(99);
+        tree.insert(&key, &value1)?;
+        let state_transition_proof1 = tree.update(&key, &value)?;
+
+        MerkleTree::verify_state_transition(max_depth, &state_transition_proof1)?;
+
+        // `tree` and `tree_with_another_leaf` should coincide.
+        assert_eq!(tree.root(), tree_with_another_leaf.root());
+
         Ok(())
     }
 }