refactor(BagPeaks): Commit to leaf count

Synchronize with upstream behavior in twenty-first.

Co-authored-by: Thorkil Schmidiger <[email protected]>

Author: aszepieniec

tasm-lib/src/mmr/bag_peaks.rs

@@ -1,41 +1,51 @@
 use std::collections::HashMap;
 
 use triton_vm::prelude::*;
-use twenty_first::util_types::shared::bag_peaks;
+use triton_vm::twenty_first::util_types::mmr::mmr_accumulator::MmrAccumulator;
 
-use crate::mmr::MAX_MMR_HEIGHT;
+use crate::arithmetic;
 use crate::prelude::*;
 use crate::traits::basic_snippet::Reviewer;
 use crate::traits::basic_snippet::SignOffFingerprint;
 
-/// [Bag peaks](bag_peaks) into a single [`Digest`].
+/// [Bag peaks](bag_peaks) of an MMR into a single [`Digest`].
 ///
 /// # Behavior
 ///
 /// ```text
-/// BEFORE: _ *peaks
+/// BEFORE: _ *mmr_accumulator
 /// AFTER:  _ [bagged_peaks: Digest]
 /// ```
 ///
 /// # Preconditions
 ///
 /// - the input argument points to a properly [`BFieldCodec`]-encoded list of
 ///   [`Digest`]s in memory
-/// - the pointed-to list contains fewer than [`MAX_MMR_HEIGHT`] elements
+/// - the pointed-to MMR accumulator is consistent, *i.e.*, the number of peaks
+///   matches with the number of set bits in the leaf count.
 ///
 /// # Postconditions
 ///
 /// - the output is a single [`Digest`] computed like in [`bag_peaks`]
 /// - the output is properly [`BFieldCodec`] encoded
 ///
+/// # Crashes
+///
+///  - if the MMR accumulator is inconsistent, *i.e.*, if the number of peaks
+///    does not match the number of set bits in the leaf count
+///
 #[derive(Debug, Copy, Clone, Eq, PartialEq, Hash)]
 pub struct BagPeaks;
 
+impl BagPeaks {
+    const INCONSISTENT_PEAKS_NUM: usize = 11;
+}
+
 impl BasicSnippet for BagPeaks {
     fn inputs(&self) -> Vec<(DataType, String)> {
-        let list_of_digests = DataType::List(Box::new(DataType::Digest));
+        let mmr_accumulator = DataType::List(Box::new(DataType::Digest));
 
-        vec![(list_of_digests, "*peaks".to_string())]
+        vec![(mmr_accumulator, "*mmra".to_string())]
     }
 
     fn outputs(&self) -> Vec<(DataType, String)> {
@@ -46,107 +56,77 @@ impl BasicSnippet for BagPeaks {
         "tasmlib_mmr_bag_peaks".to_string()
     }
 
-    fn code(&self, _: &mut Library) -> Vec<LabelledInstruction> {
+    fn code(&self, library: &mut Library) -> Vec<LabelledInstruction> {
         let entrypoint = self.entrypoint();
-
-        let length_is_zero = format!("{entrypoint}_length_is_zero");
-        let length_is_not_zero = format!("{entrypoint}_length_is_not_zero");
-        let length_is_one = format!("{entrypoint}_length_is_one");
-        let length_is_gt_one = format!("{entrypoint}_length_is_gt_one");
         let bagging_loop = format!("{entrypoint}_loop");
 
-        let Digest(zero_digest) = bag_peaks(&[]);
+        let destructure_mmra = MmrAccumulator::destructure();
+        let pop_count = library.import(Box::new(arithmetic::u64::popcount::PopCount));
 
         triton_asm!(
-        // BEFORE: _ *peaks
+        // BEFORE: _ *mmra
         // AFTER:  _ [bagged_peaks: Digest]
         {entrypoint}:
-            dup 0 read_mem 1 pop 1
-            // _ *peaks length
+            {&destructure_mmra}
+            // _ *peaks *leaf_count
 
-            push {MAX_MMR_HEIGHT}
-            dup 1
-            lt
-            assert
+            addi 1 read_mem 2 pop 1
+            hint leaf_count : u64 = stack[0..2]
+            // _ *peaks [leaf_count]
 
-            /* special case: length is 0 */
-            push 1
-            dup 1 push 0 eq
-            // _ *peaks length 1 (length==0)
+            dup 1 dup 1
+            call {pop_count}
+            // _ *peaks [leaf_count] popcount
 
-            skiz call {length_is_zero}
-            skiz call {length_is_not_zero}
+            dup 3 read_mem 1 pop 1
+            // _ *peaks [leaf_count] popcount num_peaks
 
-            // _ [bagged_peaks: Digest]
-            return
+            dup 1 eq
+            // _ *peaks [leaf_count] pop_count (num_peaks==pop_count)
 
-        // BEFORE: _ *peaks length 1
-        // AFTER:  _ [bagged_peaks: Digest] 0
-        {length_is_zero}:
-            pop 3
+            assert error_id {Self::INCONSISTENT_PEAKS_NUM}
+            hint num_peaks: u32 = stack[0]
+            // _ *peaks [leaf_count] num_peaks
 
-            push {zero_digest[4]}
-            push {zero_digest[3]}
-            push {zero_digest[2]}
-            push {zero_digest[1]}
-            push {zero_digest[0]}
-            hint bag_of_no_peaks: Digest = stack[0..5]
+            place 2
+            // _ *peaks num_peaks [leaf_count]
+            // _ *peaks len [leaf_count] <-- rename
 
             push 0
-            return
-
-        // BEFORE: _ *peaks length
-        // AFTER:  _ [bagged_peaks: Digest]
-        {length_is_not_zero}:
-            /* special case: length is 1 */
-            push 1
-            dup 1 push 1 eq
-            // _ *peaks length 1 (length==1)
-
-            skiz call {length_is_one}
-            skiz call {length_is_gt_one}
-            return
-
-        // BEFORE: _ *peaks length 1
-        // AFTER:  _ [bagged_peaks: Digest] 0
-        {length_is_one}:
-            pop 2
-
-            push {Digest::LEN} add
-            // _ *peaks[0]_lw
-
-            read_mem {Digest::LEN}
-            pop 1
-
             push 0
-            return
+            push 0
+            push 0
+            push 0
+            push 0
+            push 0
+            push 0
+            pick 9
+            pick 9
+            // _ *peaks len 0 0 0 0 0 0 0 0 [leaf_count; 2]
 
-        // BEFORE: _ *peaks length
-        // AFTER:  _ [bagged_peaks: Digest]
-        {length_is_gt_one}:
-            /* Get pointer to last word of peaks list */
+            hash
+            // _ *peaks len [hash_of_leaf_count]
+
+            pick 5
             push {Digest::LEN}
             mul
-            // _ *peaks offset
+            // _ *peaks [hash_of_leaf_count] size_of_peaks_list
 
-            dup 1
+            dup 6
             add
-            // _ *peaks *peaks[last]_lw
-
-            read_mem {Digest::LEN}
-            // _ *peaks [peaks[last]: Digest] *peaks[last-1]_lw
+            // _ *peaks [hash_of_leaf_count] *peaks[last]_lw
 
             place 5
-            // _ *peaks *peaks[last-1]_lw [peaks[last]: Digest]
-            // _ *peaks *peaks[last-1]_lw [acc: Digest]
+            // _ *peaks *peaks[last]_lw [hash_of_leaf_count]
+
+            dup 6 dup 6 eq push 0 eq
+            // _ *peaks *peaks[last]_lw [hash_of_leaf_count] (num_peaks == 0)
 
-            call {bagging_loop}
-            // *peaks *peaks [acc: Digest]
-            // *peaks *peaks [bagged_peaks: Digest]
+            skiz call {bagging_loop}
+            // _ *peaks *peaks [bag_hash]
 
-            pick 6
-            pick 6
-            pop 2
+            pick 6 pick 6 pop 2
+            // _ [bagged_peaks: Digest]
 
             return
 
@@ -165,29 +145,31 @@ impl BasicSnippet for BagPeaks {
     }
 
     fn sign_offs(&self) -> HashMap<Reviewer, SignOffFingerprint> {
-        let mut sign_offs = HashMap::new();
-        sign_offs.insert(Reviewer("ferdinand"), 0xaf79abb21cb46bf.into());
-        sign_offs
+        [].into()
     }
 }
 
 #[cfg(test)]
 mod tests {
     use std::collections::HashMap;
 
+    use triton_vm::twenty_first::prelude::Mmr;
     use twenty_first::math::other::random_elements;
 
     use super::*;
     use crate::test_prelude::*;
 
     impl BagPeaks {
-        fn set_up_initial_state(&self, num_peaks: usize) -> FunctionInitialState {
-            let address = random();
+        fn set_up_initial_state(&self, leaf_count: u64) -> FunctionInitialState {
+            let address = rand::random();
             let mut stack = self.init_stack_for_isolated_run();
             push_encodable(&mut stack, &address);
 
             let mut memory = HashMap::new();
-            encode_to_memory(&mut memory, address, &random_elements::<Digest>(num_peaks));
+            let num_peaks = leaf_count.count_ones();
+            let mmra =
+                MmrAccumulator::init(random_elements::<Digest>(num_peaks as usize), leaf_count);
+            encode_to_memory(&mut memory, address, &mmra);
 
             FunctionInitialState { stack, memory }
         }
@@ -200,28 +182,45 @@ mod tests {
             memory: &mut HashMap<BFieldElement, BFieldElement>,
         ) {
             let address = pop_encodable(stack);
-            let peaks = *Vec::<Digest>::decode_from_memory(memory, address).unwrap();
-            push_encodable(stack, &bag_peaks(&peaks));
+            let mmra = *MmrAccumulator::decode_from_memory(memory, address).unwrap();
+
+            fn bag_peaks(peaks: &[Digest], leaf_count: u64) -> Digest {
+                // use `hash_10` over `hash` or `hash_varlen` to simplify hashing in Triton VM
+                let [lo_limb, hi_limb] = leaf_count.encode()[..] else {
+                    panic!("internal error: unknown encoding of type `u64`")
+                };
+                let padded_leaf_count = bfe_array![lo_limb, hi_limb, 0, 0, 0, 0, 0, 0, 0, 0];
+                let hashed_leaf_count = Digest::new(Tip5::hash_10(&padded_leaf_count));
+
+                peaks
+                    .iter()
+                    .rev()
+                    .fold(hashed_leaf_count, |acc, &peak| Tip5::hash_pair(peak, acc))
+            }
+
+            let bag = bag_peaks(&mmra.peaks(), mmra.num_leafs());
+            println!("bag: {bag}");
+            push_encodable(stack, &bag);
         }
 
         fn pseudorandom_initial_state(
             &self,
             seed: [u8; 32],
             bench_case: Option<BenchmarkCase>,
         ) -> FunctionInitialState {
-            let num_peaks = match bench_case {
-                Some(BenchmarkCase::CommonCase) => 30,
-                Some(BenchmarkCase::WorstCase) => 60,
-                None => StdRng::from_seed(seed).gen_range(0..=63),
+            let num_leafs = match bench_case {
+                Some(BenchmarkCase::CommonCase) => 348753,
+                Some(BenchmarkCase::WorstCase) => 843759843768,
+                None => StdRng::from_seed(seed).random_range(0u64..(u64::MAX >> 1)),
             };
 
-            self.set_up_initial_state(num_peaks)
+            self.set_up_initial_state(num_leafs)
         }
 
         fn corner_case_initial_states(&self) -> Vec<FunctionInitialState> {
             (0..=5)
                 .chain([63])
-                .map(|num_peaks| self.set_up_initial_state(num_peaks))
+                .map(|num_peaks| self.set_up_initial_state((1 << num_peaks) - 1))
                 .collect()
         }
     }