feat: Add snippet VerifyMmrSuccessor

tasm-lib/src/mmr/verify_mmr_successor.rs

@@ -9,7 +9,7 @@ use triton_vm::{
 use crate::{
     arithmetic::u64::{
         add_u64::AddU64, log_2_floor_u64::Log2FloorU64, lt_u64::LtU64ConsumeArgs,
-        popcount_u64::PopCountU64, sub_u64::SubU64,
+        popcount_u64::PopCountU64, shift_right_u64::ShiftRightU64, sub_u64::SubU64,
     },
     data_type::DataType,
     field,
@@ -69,11 +69,12 @@ impl BasicSnippet for VerifyMmrSuccessor {
         );
         let ltu64 = library.import(Box::new(LtU64ConsumeArgs));
         let popcount_u64 = library.import(Box::new(PopCountU64));
-        let addu64 = library.import(Box::new(AddU64));
+        let add_u64 = library.import(Box::new(AddU64));
         let leaf_index_to_mti_and_pki = library.import(Box::new(MmrLeafIndexToMtIndexAndPeakIndex));
         let merkle_step_u64 = library.import(Box::new(MerkleStepU64Index));
         let ilog2_u64 = library.import(Box::new(Log2FloorU64));
         let sub_u64 = library.import(Box::new(SubU64));
+        let shr_u64 = library.import(Box::new(ShiftRightU64));
         let compare_digests = DataType::Digest.compare();
         let compare_u64 = DataType::U64.compare();
         let entrypoint = self.entrypoint();
@@ -86,16 +87,23 @@ impl BasicSnippet for VerifyMmrSuccessor {
             pop 1
             dup 1 dup 1 call {ilog2_u64}
             // [num_leafs_remaining] old_height
-            dup 2 push 2 pow split
+            hint old_height = stack[0]
+
+            dup 0 push 2 pow split
             // [num_leafs_remaining] old_height [1<<old_height]
+            hint two_pow_old_height = stack[0..2]
+
             dup 4 dup 4
             // [num_leafs_remaining] old_height [1<<old_height] [num_leafs_remaining]
             call {sub_u64}
             // [num_leafs_remaining] old_height [num_leafs_remaining*]
+            hint new_num_leafs_remaining = stack[0..2]
+
             swap 3 swap 1 swap 4
             // [num_leafs_remaining*] old_height [num_leafs_remaining]
             pop 2
             // [num_leafs_remaining*] old_height
+            hint num_leafs_remaining = stack[1..3]
         );
 
         triton_asm! {
@@ -138,18 +146,22 @@ impl BasicSnippet for VerifyMmrSuccessor {
 
             read_mem 1 push 2 add
             // _ *old_mmr *new_mmr num_old_peaks *old_peaks[0]
+            hint current_peak_ptr = stack[0]
 
             swap 1 push {Digest::LEN} mul
             // _ *old_mmr *new_mmr *old_peaks[0] (num_old_peaks*5)
 
             dup 1 add
             // _ *old_mmr *new_mmr *old_peaks[0] *end_of_memory
+            hint end_of_memory = stack[0]
 
             push 0 push 0
             // _ *old_mmr *new_mmr *old_peaks[0] *end_of_memory [0]
+            hint running_leaf_count = stack[0..2]
 
             dup 6 {&num_leafs}
             // _ *old_mmr *new_mmr *old_peaks[0] *end_of_memory [0] [old_num_leafs]
+            hint num_leafs_remaining = stack[0..2]
 
             push {0x455b00b5}
             // _ *old_mmr *new_mmr *old_peaks[0] *end_of_memory [0] [old_num_leafs] garbage
@@ -168,6 +180,18 @@ impl BasicSnippet for VerifyMmrSuccessor {
 
             // INVARIANT: _ *new_mmr *current_peak *end_of_memory [running_leaf_count] [num_leafs_remaining] garbage
             {main_loop}:
+                dup 2 dup 2
+                // _ *new_mmr *current_peak *end_of_memory [running_leaf_count] [num_leafs_remaining] garbage [num_leafs_remaining]
+
+                push 0 push 0 {&compare_u64}
+                // _ *new_mmr *current_peak *end_of_memory [running_leaf_count] [num_leafs_remaining] garbage (num_leafs_remaining == 0)
+
+                push 0 eq
+                // _ *new_mmr *current_peak *end_of_memory [running_leaf_count] [num_leafs_remaining] garbage (num_leafs_remaining != 0)
+
+                assert
+                // _ *new_mmr *current_peak *end_of_memory [running_leaf_count] [num_leafs_remaining] garbage
+
                 {&strip_top_bit}
                 // _ *new_mmr *current_peak *end_of_memory [running_leaf_count] [num_leafs_remaining*] old_height
 
@@ -179,22 +203,37 @@ impl BasicSnippet for VerifyMmrSuccessor {
 
                 call {leaf_index_to_mti_and_pki}
                 // _ *new_mmr *current_peak *end_of_memory [running_leaf_count] [num_leafs_remaining*] old_height [merkle_tree_index] peak_index
+                hint merkle_tree_index = stack[1..3]
+                hint peak_index = stack[0]
 
                 swap 2 swap 1
                 // _ *new_mmr *current_peak *end_of_memory [running_leaf_count] [num_leafs_remaining*] old_height peak_index [merkle_tree_index]
 
+                dup 3
+                // _ *new_mmr *current_peak *end_of_memory [running_leaf_count] [num_leafs_remaining*] old_height peak_index [merkle_tree_index] old_height
+
+                call {shr_u64}
+                // _ *new_mmr *current_peak *end_of_memory [running_leaf_count] [num_leafs_remaining*] old_height peak_index [merkle_tree_index >> old_height]
+
 
                 /* prepare & traverse */
-                dup 9 push {Digest::LEN-1} read_mem {Digest::LEN} pop 1
-                // _ *new_mmr *current_peak *end_of_memory [running_leaf_count] [num_leafs_remaining*] old_height peak_index [merkle_tree_index] [current_old_peak]
+                dup 9 push {Digest::LEN-1} add read_mem {Digest::LEN} pop 1
+                // _ *new_mmr *current_peak *end_of_memory [running_leaf_count] [num_leafs_remaining*] old_height peak_index [merkle_tree_index >> old_height] [current_old_peak]
+                hint merkle_node = stack[0..5]
 
                 call {traverse}
                 // _ *new_mmr *current_peak *end_of_memory [running_leaf_count] [num_leafs_remaining*] old_height peak_index [1] [some_new_peak]
+                hint landed_peak = stack[0..5]
+                hint merkle_index = stack[5..7]
+                hint peak_index = stack[7]
 
                 dup 15 {&field_peaks}
                 // _ *new_mmr *current_peak *end_of_memory [running_leaf_count] [num_leafs_remaining*] old_height peak_index [1] [some_new_peak] *new_peaks
 
-                push {1 + Digest::LEN - 1} dup 7 push {Digest::LEN} mul add
+                push {1 + Digest::LEN - 1} dup 9
+                hint peak_index = stack[0]
+
+                push {Digest::LEN} mul add
                 // _ *new_mmr *current_peak *end_of_memory [running_leaf_count] [num_leafs_remaining*] old_height peak_index [1] [some_new_peak] *new_peaks (5+peak_index*5)
 
                 add read_mem {Digest::LEN} pop 1
@@ -211,16 +250,18 @@ impl BasicSnippet for VerifyMmrSuccessor {
                 push 2 pow split
                 // _ *new_mmr *current_peak *end_of_memory [running_leaf_count] [num_leafs_remaining*] [1<<old_height]
 
-                dup 5 dup 5 call {addu64}
+                dup 5 dup 5 call {add_u64}
                 // _ *new_mmr *current_peak *end_of_memory [running_leaf_count] [num_leafs_remaining*] [running_leaf_count*]
+                hint running_leaf_count = stack[0..2]
 
-                swap 4 swap 1 swap 5 pop
+                swap 4 swap 1 swap 5 pop 1
                 // _ *new_mmr *current_peak *end_of_memory [running_leaf_count*] [num_leafs_remaining*] garbage
 
                 swap 6 push {Digest::LEN} add swap 6
                 // _ *new_mmr *next_peak *end_of_memory [running_leaf_count*] [num_leafs_remaining*] garbage
+                hint next_peak_ptr = stack[6]
 
-                return_or_recurse
+                recurse_or_return
 
             // INVARIANT: _ [merkle_tree_index] [current_node]
             {traverse}:
@@ -280,6 +321,7 @@ mod test {
     use crate::empty_stack;
     use crate::memory::encode_to_memory;
     use crate::memory::FIRST_NON_DETERMINISTICALLY_INITIALIZED_MEMORY_ADDRESS;
+    use crate::test_helpers::negative_test;
     use crate::traits::algorithm::ShadowedAlgorithm;
     use crate::traits::rust_shadow::RustShadow;
     use crate::{
@@ -313,6 +355,106 @@ mod test {
         number
     }
 
+    fn initial_state_from_mmr_tuple(
+        old_mmr: &MmrAccumulator,
+        new_mmr: &MmrAccumulator,
+        mmr_successor_proof: &MmrSuccessorProof,
+    ) -> AlgorithmInitialState {
+        let mut nondeterminism = NonDeterminism::new(vec![]);
+        VerifyMmrSuccessor::update_nondeterminism(&mut nondeterminism, &mmr_successor_proof);
+        let old_mmr_address = FIRST_NON_DETERMINISTICALLY_INITIALIZED_MEMORY_ADDRESS;
+        let new_mmr_address =
+            encode_to_memory(&mut nondeterminism.ram, old_mmr_address, old_mmr.clone());
+        let _garbage_address =
+            encode_to_memory(&mut nondeterminism.ram, new_mmr_address, new_mmr.clone());
+        let mut stack = empty_stack();
+        stack.push(old_mmr_address);
+        stack.push(new_mmr_address);
+        AlgorithmInitialState {
+            stack,
+            nondeterminism,
+        }
+    }
+
+    fn failing_initial_states() -> Vec<AlgorithmInitialState> {
+        let mut rng = thread_rng();
+        let mut initial_states = vec![];
+        for old_num_leafs in [0u64, 1, 8] {
+            for num_new_leafs in [0u64, 1, 8] {
+                let old_mmr = MmrAccumulator::init(
+                    (0..old_num_leafs.count_ones())
+                        .map(|_| rng.gen::<Digest>())
+                        .collect_vec(),
+                    old_num_leafs,
+                );
+                let new_leafs = (0..num_new_leafs)
+                    .map(|_| rng.gen::<Digest>())
+                    .collect_vec();
+                let mut new_mmr = old_mmr.clone();
+                for leaf in new_leafs.iter().copied() {
+                    new_mmr.append(leaf);
+                }
+                let mmr_successor_proof =
+                    MmrSuccessorProof::new_from_batch_append(&old_mmr, &new_leafs);
+
+                // rotate old num leafs
+                let old_mmr_fake_1 = MmrAccumulator::init(
+                    old_mmr.peaks(),
+                    (old_mmr.num_leafs() >> 1) ^ (old_mmr.num_leafs() << 63),
+                );
+                initial_states.push(initial_state_from_mmr_tuple(
+                    &old_mmr_fake_1,
+                    &new_mmr,
+                    &mmr_successor_proof,
+                ));
+
+                // rotate new num leafs
+                let new_mmr_fake_1 = MmrAccumulator::init(
+                    new_mmr.peaks(),
+                    (new_mmr.num_leafs() >> 1) ^ (new_mmr.num_leafs() << 63),
+                );
+                initial_states.push(initial_state_from_mmr_tuple(
+                    &old_mmr,
+                    &new_mmr_fake_1,
+                    &mmr_successor_proof,
+                ));
+
+                // modify path element
+                if !mmr_successor_proof.paths.is_empty() {
+                    let mut mmr_successor_proof_fake_1 = mmr_successor_proof.clone();
+                    mmr_successor_proof_fake_1.paths
+                        [rng.gen_range(0..mmr_successor_proof.paths.len())]
+                    .0[rng.gen_range(0..Digest::LEN)]
+                    .increment();
+                    initial_states.push(initial_state_from_mmr_tuple(
+                        &old_mmr,
+                        &new_mmr,
+                        &mmr_successor_proof_fake_1,
+                    ));
+                }
+
+                // flip old and new
+                initial_states.push(initial_state_from_mmr_tuple(
+                    &new_mmr,
+                    &old_mmr,
+                    &mmr_successor_proof,
+                ));
+
+                // inconsistent new mmr
+                let new_mmr_fake_2 = MmrAccumulator::init(
+                    [new_mmr.peaks(), vec![rng.gen::<Digest>()]].concat(),
+                    new_mmr.num_leafs(),
+                );
+                initial_states.push(initial_state_from_mmr_tuple(
+                    &old_mmr,
+                    &new_mmr_fake_2,
+                    &mmr_successor_proof,
+                ));
+            }
+        }
+        initial_states
+    }
+
     impl Algorithm for VerifyMmrSuccessor {
         fn rust_shadow(
             &self,
@@ -356,29 +498,46 @@ mod test {
                 .collect_vec();
             let mmr_successor_proof =
                 MmrSuccessorProof::new_from_batch_append(&old_mmr, &new_leafs);
-            println!(
-                "produced mmr successor proof of {} digests",
-                mmr_successor_proof.paths.len()
-            );
             let mut new_mmr = old_mmr.clone();
             for leaf in new_leafs {
                 new_mmr.append(leaf);
             }
 
-            let mut nondeterminism = NonDeterminism::new(vec![]);
-            VerifyMmrSuccessor::update_nondeterminism(&mut nondeterminism, &mmr_successor_proof);
-            let old_mmr_address = FIRST_NON_DETERMINISTICALLY_INITIALIZED_MEMORY_ADDRESS;
-            let new_mmr_address =
-                encode_to_memory(&mut nondeterminism.ram, old_mmr_address, old_mmr);
-            let _garbage_address =
-                encode_to_memory(&mut nondeterminism.ram, new_mmr_address, new_mmr);
-            let mut stack = empty_stack();
-            stack.push(old_mmr_address);
-            stack.push(new_mmr_address);
-            AlgorithmInitialState {
-                stack,
-                nondeterminism,
+            initial_state_from_mmr_tuple(&old_mmr, &new_mmr, &mmr_successor_proof)
+        }
+
+        fn corner_case_initial_states(&self) -> Vec<AlgorithmInitialState> {
+            let mut rng = thread_rng();
+            let mut initial_states = vec![];
+            for old_num_leafs in [0u64, 1, 8] {
+                for num_new_leafs in [0u64, 1, 8] {
+                    let old_mmr = MmrAccumulator::init(
+                        (0..old_num_leafs.count_ones())
+                            .map(|_| rng.gen::<Digest>())
+                            .collect_vec(),
+                        old_num_leafs,
+                    );
+                    let new_leafs = (0..num_new_leafs)
+                        .map(|_| rng.gen::<Digest>())
+                        .collect_vec();
+                    let mut new_mmr = old_mmr.clone();
+                    for leaf in new_leafs.iter().copied() {
+                        new_mmr.append(leaf);
+                    }
+                    let mmr_successor_proof =
+                        MmrSuccessorProof::new_from_batch_append(&old_mmr, &new_leafs);
+
+                    if old_num_leafs != 0 {
+                        // correct
+                        initial_states.push(initial_state_from_mmr_tuple(
+                            &old_mmr,
+                            &new_mmr,
+                            &mmr_successor_proof,
+                        ));
+                    }
+                }
             }
+            initial_states
         }
     }
 
@@ -387,6 +546,13 @@ mod test {
         ShadowedAlgorithm::new(VerifyMmrSuccessor).test();
     }
 
+    // #[test]
+    // fn verify_mmr_successor_negative_test() {
+    //     for init_state in failing_initial_states() {
+    //         negative_test(VerifyMmrSuccessor, init_state, &[]);
+    //     }
+    // }
+
     #[test]
     fn test_num_digests_in_proof() {
         let old_mmr_num_leafs = 116u64;