diff --git a/src/backends/plonky2/primitives/merkletree/circuit.rs b/src/backends/plonky2/primitives/merkletree/circuit.rs index f9a1c74..0cec246 100644 --- a/src/backends/plonky2/primitives/merkletree/circuit.rs +++ b/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) new_value: ValueTarget, - pub(crate) new_siblings: Vec, + pub(crate) op_key: ValueTarget, + pub(crate) op_value: ValueTarget, + pub(crate) siblings: Vec, // 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, - 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 - builder.connect(proof.typ, zero); + // Op type check + 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 - verify_merkle_proof_circuit(builder, &proof.proof_non_existence); + // 1) Verify the provided op proof. + 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, - key: proof.new_key, - value: proof.new_value, - siblings: proof.new_siblings.clone(), + root, + key: proof.op_key, + value: proof.op_value, + 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, - zero, - ); - // 3.2) assert that proof.enabled matches with proof_non_existence.enabled - builder.connect( - proof.proof_non_existence.enabled.target, - proof.enabled.target, + proof.op_proof.existence.target, + proof_type.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.old_root.elements[j], + proof.op_proof.root.elements[j], + claim_root.elements[j], ); - // 4.2) assert that the non-existence proof uses the new_key (value not needed for - // non-existence) + // 4.2) assert that the non-existence proof uses the op_key (value not needed). builder.conditional_assert_eq( proof.enabled.target, - proof.proof_non_existence.key.elements[j], - proof.new_key.elements[j], + proof.op_proof.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.proof_non_existence.other_value, + &proof.op_proof.other_key, + &proof.op_proof.other_value, ); // prepare values for check 5.3) - let old_leaf_path = keypath_target( - proof.max_depth, - builder, - &proof.proof_non_existence.other_key, - ); + let old_leaf_path = keypath_target(proof.max_depth, builder, &proof.op_proof.other_key); - // 5) check that old_siblings & new_siblings match as expected. Let - // d=divergence_level, assert that: + // 5) check that old_siblings & new_siblings match as expected. They + // 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 new_siblings = proof.new_siblings.clone(); + let old_siblings = proof.op_proof.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) -> 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(), - new_value: builder.add_virtual_value(), - + op_key: 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, 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, - key: mp.new_key, - value: EMPTY_VALUE, // not needed for non-existence - proof: mp.proof_non_existence.clone(), + root: if mp.op == MerkleTreeOp::Delete { + mp.new_root + } else { + 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.new_value.elements, &mp.new_value.0)?; + pw.set_target_arr(&self.op_key.elements, &mp.op_key.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::::new(config); let mut pw = PartialWitness::::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 - .proof_non_existence - .other_leaf - .unwrap(); + let other_leaf = state_transition_proof.op_proof.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, - Some(state_transition_proof.new_value), + &state_transition_proof.op_key, + 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::::new(config); let mut pw = PartialWitness::::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::::new(config); let mut pw = PartialWitness::::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)?; diff --git a/src/backends/plonky2/primitives/merkletree/error.rs b/src/backends/plonky2/primitives/merkletree/error.rs index 36f3ed0..e8fcfbc 100644 --- a/src/backends/plonky2/primitives/merkletree/error.rs +++ b/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)) } diff --git a/src/backends/plonky2/primitives/merkletree/mod.rs b/src/backends/plonky2/primitives/merkletree/mod.rs index afb118e..fa0e195 100644 --- a/src/backends/plonky2/primitives/merkletree/mod.rs +++ b/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) -> TreeResult { - // Construct leaves. - let mut leaves: Vec<_> = kvs - .iter() - .map(|(k, v)| Leaf::new(max_depth, *k, *v)) - .collect::>()?; + // Start with an empty node as root. + let mut root = Node::None; - // Start with a leaf or conclude with an empty node as root. - let mut root = leaves.pop().map(Node::Leaf).unwrap_or(Node::None); - - // Iterate over remaining leaves (if any) and add them. - for leaf in leaves.into_iter() { - root.add_leaf(0, max_depth, leaf)?; + // Iterate over key-value pairs (if any) and add them. + for (k, v) in kvs.iter() { + root.apply_op(max_depth, MerkleTreeOp::Insert, *k, Some(*v))?; } // Fill in hashes. @@ -82,12 +76,12 @@ impl MerkleTree { &mut self, key: &RawValue, value: &RawValue, - ) -> TreeResult { + ) -> TreeResult { 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 { - typ: 0, // insertion + Ok(MerkleTreeStateTransitionProof { + op: MerkleTreeOp::Insert, // insertion old_root, - proof_non_existence, + op_proof: proof_non_existence, new_root, - new_key: *key, - new_value: *value, + op_key: *key, + op_value: *value, + value: None, siblings: proof.siblings, }) } + pub fn update( + &mut self, + key: &RawValue, + value: &RawValue, + ) -> TreeResult { + 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 { + 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 - Self::verify_nonexistence( - max_depth, - proof.old_root, - &proof.proof_non_existence, - &proof.new_key, - )?; - - // 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.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(), - )); + match proof.op { + // A deletion is but an insertion subject to a time reversal. + MerkleTreeOp::Delete => { + let equivalent_insertion_proof = MerkleTreeStateTransitionProof { + op: MerkleTreeOp::Insert, + new_root: proof.old_root, + old_root: proof.new_root, + ..proof.clone() + }; + Self::verify_state_transition(max_depth, &equivalent_insertion_proof) } - } + 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: - // 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(), - )); + // All siblings should agree + (proof.siblings == proof.op_proof.siblings) + .then_some(()) + .ok_or(TreeError::state_transition_fail(format!( + "Invalid proof of update for key {}: Siblings don't match.", + proof.op_key + ))) } - } - 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 - .proof_non_existence + MerkleTreeOp::Insert => { + // check that for the old_root, the new_key does not exist in the tree + Self::verify_nonexistence( + max_depth, + proof.old_root, + &proof.op_proof, + &proof.op_key, + )?; + + // 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 = proof.proof_non_existence.other_leaf.map(|(_, v)| v); - let old_leaf_hash = kv_hash(&k, v); - if new_siblings[d] != old_leaf_hash { - return Err(TreeError::state_transition_fail( - "siblings don't match: new[d]!=old_leaf_hash".to_string(), - )); + let v: Option = proof.op_proof.other_leaf.map(|(_, v)| v); + let old_leaf_hash = kv_hash(&k, v); + if new_siblings[d] != old_leaf_hash { + return Err(TreeError::state_transition_fail( + "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 { - // type: 0:insertion, 1:update, 2:deletion - pub(crate) typ: u8, +pub struct MerkleTreeStateTransitionProof { + pub(crate) op: MerkleTreeOp, 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, + + /// 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, } @@ -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 - pub(crate) fn add_leaf(&mut self, lvl: usize, max_depth: usize, leaf: Leaf) -> TreeResult<()> { + /// Applies given Merkle tree op without computing hashes. + pub(crate) fn apply_op( + &mut self, + max_depth: usize, + op: MerkleTreeOp, + key: RawValue, + maybe_value: Option, + ) -> 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, + ) -> TreeResult<()> { if lvl >= max_depth { return Err(TreeError::max_depth()); } match self { Self::Intermediate(n) => { - if leaf.path[lvl] { - if n.right.is_empty() { - // empty sub-node, add the leaf here - n.right = Box::new(Node::Leaf(leaf)); - return Ok(()); - } - n.right.add_leaf(lvl + 1, max_depth, leaf)?; + if key_path[lvl] { + n.right + .apply_op_loop(lvl + 1, max_depth, op, key, key_path, maybe_value) } else { - if n.left.is_empty() { - // empty sub-node, add the leaf here - n.left = Box::new(Node::Leaf(leaf)); - return Ok(()); - } - n.left.add_leaf(lvl + 1, max_depth, leaf)?; + n.left + .apply_op_loop(lvl + 1, max_depth, op, key, key_path, maybe_value) } } - 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, + ) -> TreeResult { + 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 => { - return Err(TreeError::empty_node()); + // Update case + (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 { - hash: Option, - path: Vec, - key: RawValue, - value: RawValue, +pub(crate) struct Leaf { + pub(crate) hash: Option, + pub(crate) path: Vec, + pub(crate) key: RawValue, + pub(crate) value: RawValue, } impl Leaf { fn new(max_depth: usize, key: RawValue, value: RawValue) -> TreeResult { @@ -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.new_value, value); + assert_eq!(state_transition_proof.op_key, key); + 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(()) } }