From 5a3b81a722c36b820a195184067cea0ba172b8ac Mon Sep 17 00:00:00 2001 From: sword_smith Date: Fri, 9 Aug 2024 18:30:58 +0200 Subject: [PATCH 01/10] feat(mmr): root from auth struct snippet Add stub for calculating a Merkle root from an authentication struct. Also adds code to generate the witness. Cf. https://github.com/Neptune-Crypto/twenty-first/pull/228 --- tasm-lib/src/lib.rs | 2 +- tasm-lib/src/mmr.rs | 2 + tasm-lib/src/mmr/authentication_struct.rs | 1 + .../src/mmr/authentication_struct/shared.rs | 716 ++++++++++++ .../mmr/root_from_authentication_struct.rs | 1026 +++++++++++++++++ 5 files changed, 1746 insertions(+), 1 deletion(-) create mode 100644 tasm-lib/src/mmr/authentication_struct.rs create mode 100644 tasm-lib/src/mmr/authentication_struct/shared.rs create mode 100644 tasm-lib/src/mmr/root_from_authentication_struct.rs diff --git a/tasm-lib/src/lib.rs b/tasm-lib/src/lib.rs index 920727dd..106d6059 100644 --- a/tasm-lib/src/lib.rs +++ b/tasm-lib/src/lib.rs @@ -229,7 +229,7 @@ pub fn execute_test( terminal_state } -/// If the environment variable “TRITON_TUI” is set, write +/// If the environment variable TASMLIB_TRITON_TUI is set, write /// 1. the program to file `program.tasm`, and /// 2. the VM state to file `vm_state.json`. /// diff --git a/tasm-lib/src/mmr.rs b/tasm-lib/src/mmr.rs index cf240a9d..9ef3b582 100644 --- a/tasm-lib/src/mmr.rs +++ b/tasm-lib/src/mmr.rs @@ -1,7 +1,9 @@ +pub mod authentication_struct; pub mod bag_peaks; pub mod calculate_new_peaks_from_append; pub mod calculate_new_peaks_from_leaf_mutation; pub mod leaf_index_to_mt_index_and_peak_index; +pub mod root_from_authentication_struct; pub mod verify_from_memory; pub mod verify_from_secret_in_leaf_index_on_stack; pub mod verify_from_secret_in_secret_leaf_index; diff --git a/tasm-lib/src/mmr/authentication_struct.rs b/tasm-lib/src/mmr/authentication_struct.rs new file mode 100644 index 00000000..eec3c896 --- /dev/null +++ b/tasm-lib/src/mmr/authentication_struct.rs @@ -0,0 +1 @@ +pub mod shared; diff --git a/tasm-lib/src/mmr/authentication_struct/shared.rs b/tasm-lib/src/mmr/authentication_struct/shared.rs new file mode 100644 index 00000000..16758680 --- /dev/null +++ b/tasm-lib/src/mmr/authentication_struct/shared.rs @@ -0,0 +1,716 @@ +use std::collections::HashMap; +use std::collections::HashSet; + +use itertools::Itertools; +use num_traits::One; + +use crate::twenty_first::bfe; +use crate::twenty_first::prelude::*; +use crate::twenty_first::util_types::mmr::mmr_accumulator::MmrAccumulator; +use crate::twenty_first::util_types::mmr::shared_advanced::get_peak_heights; +use crate::twenty_first::util_types::mmr::shared_basic::leaf_index_to_mt_index_and_peak_index; + +const ROOT_MT_INDEX: u64 = 1; + +/// A witness to facilitate the proving of the authenticity of a Merkle +/// authentication struct. +#[derive(Debug, Clone)] +pub struct AuthStructIntegrityProof { + // All indices are Merkle tree node indices + pub nd_auth_struct_indices: Vec, + pub nd_sibling_indices: Vec<(u64, u64)>, + pub nd_siblings: Vec<(Digest, Digest)>, +} + +/// An authentication structure that can be used to prove membership of a list +/// of leaves in a Merkle tree, along with the indexed leaves in question, and +/// the witness necessary to prove membership in a ZK program. +#[derive(Debug, Clone)] +pub struct AuthenticatedMerkleAuthStruct { + pub auth_struct: Vec, + pub indexed_leafs: Vec<(u64, Digest)>, + pub witness: AuthStructIntegrityProof, +} + +impl AuthStructIntegrityProof { + /// Return the Merkle tree node indices of the digests required to prove + /// membership for the specified leaf indices, as well as the node indices + /// that can be derived from the leaf indices and their authentication + /// path. + fn auth_struct_and_nd_indices( + num_leafs: u64, + leaf_indices: &[u64], + ) -> (Vec, Vec<(u64, u64)>) { + // The set of indices of nodes that need to be included in the authentications + // structure. In principle, every node of every authentication path is needed. + // The root is never needed. Hence, it is not considered below. + let mut node_is_needed = HashSet::new(); + + // The set of indices of nodes that can be computed from other nodes in the + // authentication structure or the leafs that are explicitly supplied during + // verification. Every node on the direct path from the leaf to the root can + // be computed by the very nature of “authentication path”. + let mut node_can_be_computed = HashSet::new(); + + for &leaf_index in leaf_indices { + assert!(num_leafs > leaf_index, "Leaf index must be less than number of leafs. Got leaf_index = {leaf_index}; num_leafs = {num_leafs}"); + + let mut node_index = leaf_index + num_leafs; + while node_index > ROOT_MT_INDEX { + let sibling_index = node_index ^ 1; + node_can_be_computed.insert(node_index); + node_is_needed.insert(sibling_index); + node_index /= 2; + } + } + + let set_difference = node_is_needed.difference(&node_can_be_computed).copied(); + let set_union = node_is_needed + .union(&node_can_be_computed) + .sorted_unstable() + .rev(); + + let mut set_union = set_union.peekable(); + + let mut set_union_as_ordered_pairs = Vec::new(); + while set_union.peek().is_some() { + let right_index = *set_union.next().unwrap(); + + // Crashes on odd-length of input list, which is what we want, as + // this acts as a sanity check. + let left_index = *set_union.next().unwrap(); + set_union_as_ordered_pairs.push((left_index, right_index)); + } + + ( + set_difference.sorted_unstable().rev().collect(), + set_union_as_ordered_pairs, + ) + } + + pub fn root_from_authentication_struct( + &self, + tree_height: u32, + auth_struct: Vec, + indexed_leafs: Vec<(u64, Digest)>, + ) -> Digest { + fn digest_to_xfe(digest: Digest, challenge: XFieldElement) -> XFieldElement { + let leaf_xfe_lo = XFieldElement::new([digest.0[0], digest.0[1], digest.0[2]]); + let leaf_xfe_hi = + challenge * XFieldElement::new([digest.0[3], digest.0[4], BFieldElement::one()]); + + leaf_xfe_lo + leaf_xfe_hi + } + + fn node_index_to_bfe(node_index: u64) -> BFieldElement { + BFieldElement::new(node_index) + } + + // Sanity check + assert_eq!( + self.nd_auth_struct_indices.len(), + auth_struct.len(), + "Provided auth struct length must match that specified in receiver" + ); + + // Get challenges + let (alpha, beta, gamma) = { + let mut sponge = Tip5::init(); + sponge.pad_and_absorb_all(&indexed_leafs.encode()); + sponge.pad_and_absorb_all(&auth_struct.encode()); + let challenges = sponge.sample_scalars(3); + (challenges[0], challenges[1], challenges[2]) + }; + + // Accumulate `p` from public data + let mut p = XFieldElement::one(); + for i in (0..indexed_leafs.len()).rev() { + let node_index_as_bfe = node_index_to_bfe((1 << tree_height) ^ indexed_leafs[i].0); + let leaf_as_xfe = digest_to_xfe(indexed_leafs[i].1, alpha); + let fact = leaf_as_xfe - beta + gamma * node_index_as_bfe; + p *= fact; + } + + let mut prev = 0; + for i in (0..auth_struct.len()).rev() { + let auth_struct_index = self.nd_auth_struct_indices[i]; + + // `auth_struct` must be sorted high-to-low by node-index. But since + // we're traversing in reverse order, the inequality is flipped. + assert!(auth_struct_index > prev); + prev = auth_struct_index; + + let auth_struct_index_as_bfe = node_index_to_bfe(auth_struct_index); + + let auth_str_elem_as_xfe = digest_to_xfe(auth_struct[i], alpha); + let fact = auth_str_elem_as_xfe - beta + gamma * auth_struct_index_as_bfe; + p *= fact; + } + + // Use secret data to invert `p` back and to calculate the root + let mut t = auth_struct + .first() + .copied() + .unwrap_or_else(|| indexed_leafs.first().unwrap().1); + let mut t_xfe = digest_to_xfe(t, alpha); + let mut parent_index_bfe = BFieldElement::one(); + for ((l, r), (left_index, right_index)) in self + .nd_siblings + .iter() + .zip_eq(self.nd_sibling_indices.clone()) + { + assert_eq!(left_index + 1, right_index); + + t = Tip5::hash_pair(*l, *r); + + let l_xfe = digest_to_xfe(*l, alpha); + let r_xfe = digest_to_xfe(*r, alpha); + t_xfe = digest_to_xfe(t, alpha); + + let left_index_bfe = node_index_to_bfe(left_index); + let right_index_bfe = node_index_to_bfe(right_index); + parent_index_bfe = left_index_bfe / bfe!(2); + + let fact1 = l_xfe - beta + gamma * left_index_bfe; + let fact2 = r_xfe - beta + gamma * right_index_bfe; + let fact_parent = t_xfe - beta + gamma * parent_index_bfe; + + p *= fact1.inverse() * fact2.inverse() * fact_parent; + } + + assert_eq!(t_xfe - beta + gamma, p); + + assert!(parent_index_bfe.is_one()); + + t + } + + /// Return the authentication structure authenticity witness, + /// authentication structure, and the (leaf-index, leaf-digest) pairs + /// from a list of MMR membership proofs. All MMR membership proofs must + /// belong under the same peak, i.e., be part of the same Merkle tree in + /// the list of Merkle trees that the MMR contains. + /// + /// Panics if the input list of MMR-membership proofs is empty, or if they + /// do not all belong under the same peak. + pub fn new_from_mmr_membership_proofs( + mmra: &MmrAccumulator, + indexed_mmr_mps: Vec<(u64, Digest, MmrMembershipProof)>, + ) -> HashMap { + #[derive(Clone, Debug)] + struct IndexedAuthenticatedMmrLeaf { + merkle_tree_node_index: u64, + merkle_tree_leaf_index: u64, + leaf_digest: Digest, + membership_proof: MmrMembershipProof, + } + + // Split indexed MMR-mps into a hashmap with one entry for each + // referenced peak in the MMR. + let num_mmr_leafs = mmra.num_leafs(); + let mut peak_index_to_indexed_mmr_mp: HashMap> = + HashMap::default(); + let peak_heights = get_peak_heights(num_mmr_leafs); + for (mmr_leaf_index, leaf, mmr_mp) in indexed_mmr_mps { + let (mt_index, peak_index) = + leaf_index_to_mt_index_and_peak_index(mmr_leaf_index, num_mmr_leafs); + let peak_index_as_usize: usize = peak_index.try_into().unwrap(); + let num_leafs_local_mt = 1 << peak_heights[peak_index_as_usize]; + let mt_leaf_index = mt_index - num_leafs_local_mt; + peak_index_to_indexed_mmr_mp + .entry(peak_index) + .or_default() + .push(IndexedAuthenticatedMmrLeaf { + merkle_tree_node_index: mt_index, + merkle_tree_leaf_index: mt_leaf_index, + leaf_digest: leaf, + membership_proof: mmr_mp, + }); + } + + // Loop over all peaks and collect an authentication witness struct + // for each peak. + let mut peak_index_to_authenticated_auth_struct = HashMap::default(); + for (peak_index, indexed_mmr_mp_structs) in peak_index_to_indexed_mmr_mp { + let peak_index_as_usize: usize = peak_index.try_into().unwrap(); + let num_leafs_in_local_mt = 1 << peak_heights[peak_index_as_usize]; + let local_mt_leaf_indices = indexed_mmr_mp_structs + .iter() + .map(|x| x.merkle_tree_leaf_index) + .collect_vec(); + + let (nd_auth_struct_indices, nd_sibling_indices) = + Self::auth_struct_and_nd_indices(num_leafs_in_local_mt, &local_mt_leaf_indices); + let peak = mmra.peaks()[peak_index_as_usize]; + + let mut node_digests: HashMap = HashMap::default(); + node_digests.insert(ROOT_MT_INDEX, peak); + + // Loop over all indexed leafs for this peak + for indexed_mmr_mp in indexed_mmr_mp_structs.iter() { + let mut mt_node_index = indexed_mmr_mp.merkle_tree_node_index; + let mut node = indexed_mmr_mp.leaf_digest; + + // Loop over all authentication path elements for this indexed leaf + for ap_elem in indexed_mmr_mp.membership_proof.authentication_path.iter() { + node_digests.insert(mt_node_index, node); + node_digests.insert(mt_node_index ^ 1, *ap_elem); + node = if mt_node_index & 1 == 0 { + Tip5::hash_pair(node, *ap_elem) + } else { + Tip5::hash_pair(*ap_elem, node) + }; + + mt_node_index /= 2; + } + + // Sanity check that MMR-MPs are valid + assert_eq!(peak, node, "Derived peak must match provided peak"); + } + let nd_siblings = nd_sibling_indices + .iter() + .map(|(left_idx, right_idx)| (node_digests[left_idx], node_digests[right_idx])) + .collect_vec(); + let auth_struct = nd_auth_struct_indices + .iter() + .map(|idx| node_digests[idx]) + .collect_vec(); + let indexed_leafs = indexed_mmr_mp_structs + .into_iter() + .map(|indexed_mmr_mp| { + ( + indexed_mmr_mp.merkle_tree_leaf_index, + indexed_mmr_mp.leaf_digest, + ) + }) + .collect_vec(); + + let witness = Self { + nd_auth_struct_indices, + nd_sibling_indices, + nd_siblings, + }; + + peak_index_to_authenticated_auth_struct.insert( + peak_index, + AuthenticatedMerkleAuthStruct { + auth_struct, + indexed_leafs, + witness, + }, + ); + } + + peak_index_to_authenticated_auth_struct + } + + /// Return the authentication structure witness, authentication structure, + /// and the (leaf-index, leaf-digest) pairs. + pub fn new_from_merkle_tree( + tree: &MerkleTree, + mut revealed_leaf_indices: Vec, + ) -> AuthenticatedMerkleAuthStruct { + revealed_leaf_indices.sort_unstable(); + revealed_leaf_indices.dedup(); + revealed_leaf_indices.reverse(); + let num_leafs: u64 = tree.num_leafs() as u64; + + let (mut nd_auth_struct_indices, nd_sibling_indices) = + Self::auth_struct_and_nd_indices(num_leafs, &revealed_leaf_indices); + if revealed_leaf_indices.is_empty() { + nd_auth_struct_indices = vec![ROOT_MT_INDEX]; + } + + let nd_siblings = nd_sibling_indices + .iter() + .map(|&(l, r)| { + let l: usize = l.try_into().unwrap(); + let r: usize = r.try_into().unwrap(); + (tree.node(l).unwrap(), tree.node(r).unwrap()) + }) + .collect_vec(); + + let revealed_leafs = revealed_leaf_indices + .iter() + .map(|j| tree.node((*j + num_leafs) as usize).unwrap()) + .collect_vec(); + let indexed_leafs = revealed_leaf_indices + .clone() + .into_iter() + .zip_eq(revealed_leafs) + .collect_vec(); + + let auth_struct = nd_auth_struct_indices + .iter() + .map(|node_index| tree.node(*node_index as usize).unwrap()) + .collect_vec(); + + let witness = Self { + nd_auth_struct_indices, + nd_sibling_indices, + nd_siblings, + }; + + AuthenticatedMerkleAuthStruct { + auth_struct, + indexed_leafs, + witness, + } + } +} + +#[cfg(test)] +mod tests { + use proptest::collection::vec; + use proptest::prop_assert_eq; + use rand::random; + use test_strategy::proptest; + + use crate::twenty_first::math::other::random_elements; + use crate::twenty_first::util_types::mmr::mmr_accumulator::util::mmra_with_mps; + + use super::*; + + #[proptest(cases = 20)] + fn root_from_authentication_struct_mmr_prop_test( + #[strategy(0..u64::MAX / 2)] mmr_leaf_count: u64, + #[strategy(0usize..20)] _num_revealed_leafs: usize, + #[strategy(vec(0u64..#mmr_leaf_count, #_num_revealed_leafs))] + mmr_revealed_leaf_indices: Vec, + ) { + let indexed_leafs_input: Vec<(u64, Digest)> = mmr_revealed_leaf_indices + .iter() + .map(|idx| (*idx, random())) + .collect_vec(); + let (mmra, mmr_mps) = mmra_with_mps(mmr_leaf_count, indexed_leafs_input.clone()); + let indexed_mmr_mps = mmr_mps + .into_iter() + .zip_eq(indexed_leafs_input) + .map(|(mmr_mp, (idx, leaf))| (idx, leaf, mmr_mp)) + .collect_vec(); + + let authenticated_auth_structs = + AuthStructIntegrityProof::new_from_mmr_membership_proofs(&mmra, indexed_mmr_mps); + + let peak_heights = get_peak_heights(mmr_leaf_count); + for (peak_index, authentication_auth_struct) in authenticated_auth_structs { + let AuthenticatedMerkleAuthStruct { + auth_struct, + indexed_leafs, + witness, + } = &authentication_auth_struct; + let tree_height: u32 = peak_heights[peak_index as usize]; + let computed_root = witness.root_from_authentication_struct( + tree_height, + auth_struct.to_owned(), + indexed_leafs.to_owned(), + ); + + prop_assert_eq!(mmra.peaks()[peak_index as usize], computed_root); + } + } + + #[test] + fn auth_struct_on_empty_mmr() { + let empty_mmra = MmrAccumulator::init(vec![], 0); + let authenticated_auth_structs = + AuthStructIntegrityProof::new_from_mmr_membership_proofs(&empty_mmra, vec![]); + assert!(authenticated_auth_structs.is_empty()); + } + + #[test] + fn auth_struct_non_empty_mmr_empty_leaf_list() { + let mmra_10_leafs = MmrAccumulator::new_from_leafs(vec![Digest::default(); 10]); + let authenticated_auth_structs = + AuthStructIntegrityProof::new_from_mmr_membership_proofs(&mmra_10_leafs, vec![]); + assert!(authenticated_auth_structs.is_empty()); + } + + #[should_panic] + #[test] + fn panics_on_missing_nd_digests() { + let tree_height = 3u32; + let num_leafs = 1u64 << tree_height; + let leafs: Vec = random_elements(num_leafs.try_into().unwrap()); + let tree = MerkleTree::::new::(&leafs).unwrap(); + + let authenticated_auth_struct = + AuthStructIntegrityProof::new_from_merkle_tree(&tree, vec![0]); + let AuthenticatedMerkleAuthStruct { + auth_struct, + indexed_leafs, + mut witness, + } = authenticated_auth_struct; + witness.nd_sibling_indices.clear(); + witness.nd_siblings.clear(); + + witness.root_from_authentication_struct(tree_height, auth_struct, indexed_leafs); + } + + #[test] + fn auth_struct_from_mmr_mps_test_height_5_9_indices() { + let local_tree_height = 5; + let mmr_leaf_indices = [0, 1, 2, 16, 17, 18, 27, 29, 31]; + let indexed_leafs_input: Vec<(u64, Digest)> = mmr_leaf_indices + .iter() + .map(|idx| (*idx, random())) + .collect_vec(); + let (mmra, mmr_mps) = mmra_with_mps(1 << local_tree_height, indexed_leafs_input.clone()); + let indexed_mmr_mps = mmr_mps + .into_iter() + .zip_eq(indexed_leafs_input) + .map(|(mmr_mp, (idx, leaf))| (idx, leaf, mmr_mp)) + .collect_vec(); + + let authenticity_witnesses = + AuthStructIntegrityProof::new_from_mmr_membership_proofs(&mmra, indexed_mmr_mps); + assert!( + authenticity_witnesses.len().is_one(), + "All indices belong to first peak" + ); + let AuthenticatedMerkleAuthStruct { + auth_struct, + indexed_leafs, + witness, + } = &authenticity_witnesses[&0]; + + let tree_height: u32 = local_tree_height.try_into().unwrap(); + let computed_root = witness.root_from_authentication_struct( + tree_height, + auth_struct.to_owned(), + indexed_leafs.to_owned(), + ); + + let peak_index = 0; + let expected_root = mmra.peaks()[peak_index]; + assert_eq!(expected_root, computed_root); + } + + #[test] + fn auth_struct_from_mmr_mps_test_height_4_2_indices() { + let local_tree_height = 4; + let mmr_leaf_indices = [0, 1]; + let indexed_leafs_input: Vec<(u64, Digest)> = mmr_leaf_indices + .iter() + .map(|idx| (*idx, random())) + .collect_vec(); + let (mmra, mmr_mps) = mmra_with_mps(1 << local_tree_height, indexed_leafs_input.clone()); + let indexed_mmr_mps = mmr_mps + .into_iter() + .zip_eq(indexed_leafs_input) + .map(|(mmr_mp, (idx, leaf))| (idx, leaf, mmr_mp)) + .collect_vec(); + + let authenticity_witnesses = + AuthStructIntegrityProof::new_from_mmr_membership_proofs(&mmra, indexed_mmr_mps); + assert!( + authenticity_witnesses.len().is_one(), + "All indices belong to first peak" + ); + let AuthenticatedMerkleAuthStruct { + auth_struct, + indexed_leafs, + witness, + } = &authenticity_witnesses[&0]; + + let tree_height: u32 = local_tree_height.try_into().unwrap(); + let computed_root = witness.root_from_authentication_struct( + tree_height, + auth_struct.to_owned(), + indexed_leafs.to_owned(), + ); + + let peak_index = 0; + let expected_root = mmra.peaks()[peak_index]; + assert_eq!(expected_root, computed_root); + } + + #[proptest(cases = 20)] + fn root_from_authentication_struct_prop_test( + #[strategy(0..12u64)] tree_height: u64, + #[strategy(0usize..100)] _num_revealed_leafs: usize, + #[strategy(vec(0u64..1<<#tree_height, #_num_revealed_leafs))] revealed_leaf_indices: Vec< + u64, + >, + ) { + let num_leafs = 1u64 << tree_height; + let leafs: Vec = random_elements(num_leafs.try_into().unwrap()); + let tree = MerkleTree::::new::(&leafs).unwrap(); + + let authenticated_auth_struct = + AuthStructIntegrityProof::new_from_merkle_tree(&tree, revealed_leaf_indices); + let AuthenticatedMerkleAuthStruct { + auth_struct, + indexed_leafs, + witness, + } = authenticated_auth_struct; + + let tree_height: u32 = tree_height.try_into().unwrap(); + let computed_root = + witness.root_from_authentication_struct(tree_height, auth_struct, indexed_leafs); + let expected_root = tree.root(); + prop_assert_eq!(expected_root, computed_root); + } + + fn prop_from_merkle_tree( + tree_height: usize, + leaf_indices: Vec, + nd_auth_struct_indices: Vec, + nd_sibling_indices: Vec<(u64, u64)>, + ) { + let leafs: Vec = random_elements(1 << tree_height); + let tree = MerkleTree::::new::(&leafs).unwrap(); + + let auth_struct = nd_auth_struct_indices + .iter() + .map(|i| tree.node(*i as usize).unwrap()) + .collect_vec(); + let revealed_leafs = leaf_indices + .iter() + .map(|i| tree.leaf(*i as usize).unwrap()) + .collect_vec(); + let revealed_leafs = leaf_indices + .into_iter() + .zip_eq(revealed_leafs) + .collect_vec(); + let nd_siblings = nd_sibling_indices + .iter() + .map(|(left_idx, right_idx)| { + ( + tree.node(*left_idx as usize).unwrap(), + tree.node(*right_idx as usize).unwrap(), + ) + }) + .collect_vec(); + + let mmr_auth_struct = AuthStructIntegrityProof { + nd_auth_struct_indices, + nd_sibling_indices, + nd_siblings, + }; + let tree_height: u32 = tree_height.try_into().unwrap(); + let calculated_root = mmr_auth_struct.root_from_authentication_struct( + tree_height, + auth_struct, + revealed_leafs, + ); + assert_eq!(tree.root(), calculated_root); + } + + #[test] + fn root_from_authentication_struct_tree_height_0_no_revealed_leafs() { + let tree_height = 0; + let leaf_indices = vec![]; + let nd_auth_struct_indices = vec![1]; + let nd_sibling_indices = vec![]; + prop_from_merkle_tree( + tree_height, + leaf_indices, + nd_auth_struct_indices, + nd_sibling_indices, + ) + } + + #[test] + fn root_from_authentication_struct_tree_height_0_1_revealed() { + let tree_height = 0; + let leaf_indices = vec![0]; + let nd_auth_struct_indices = vec![]; + let nd_sibling_indices = vec![]; + prop_from_merkle_tree( + tree_height, + leaf_indices, + nd_auth_struct_indices, + nd_sibling_indices, + ) + } + + #[test] + fn root_from_authentication_struct_tree_height_1_1_revealed() { + let tree_height = 1; + let leaf_indices = vec![0u64]; + let nd_auth_struct_indices = vec![3]; + let nd_sibling_indices = vec![(2u64, 3u64)]; + prop_from_merkle_tree( + tree_height, + leaf_indices, + nd_auth_struct_indices, + nd_sibling_indices, + ) + } + + #[test] + fn root_from_authentication_struct_tree_height_1_2_revealed() { + let tree_height = 1; + let leaf_indices = vec![0u64, 1]; + let nd_auth_struct_indices = vec![]; + let nd_sibling_indices = vec![(2u64, 3u64)]; + prop_from_merkle_tree( + tree_height, + leaf_indices, + nd_auth_struct_indices, + nd_sibling_indices, + ) + } + + #[test] + fn root_from_authentication_struct_tree_height_2_0_revealed() { + let tree_height = 2; + let leaf_indices = vec![]; + let auth_struct_indices = vec![1]; + let nd_sibling_indices = vec![]; + prop_from_merkle_tree( + tree_height, + leaf_indices, + auth_struct_indices, + nd_sibling_indices, + ) + } + + #[test] + fn root_from_authentication_struct_tree_height_2_2_revealed() { + let tree_height = 2; + let leaf_indices = vec![0u64, 1]; + let auth_struct_indices = vec![3]; + let nd_sibling_indices = vec![(4u64, 5u64), (2, 3)]; + prop_from_merkle_tree( + tree_height, + leaf_indices, + auth_struct_indices, + nd_sibling_indices, + ) + } + + #[test] + fn root_from_authentication_struct_tree_height_4_4_revealed() { + let tree_height = 4; + let leaf_indices = vec![14u64, 12, 10, 8]; + let num_leafs = 1 << tree_height; + let auth_struct_indices = vec![ + num_leafs + 15, + num_leafs + 13, + num_leafs + 11, + num_leafs + 9, + 2, + ]; + let nd_sibling_indices_layer_height_0 = [(14u64, 15u64), (12, 13), (10, 11), (8, 9)] + .map(|(l, r)| (l + num_leafs, r + num_leafs)); + let nd_sibling_indices_layer_height_1 = [(14u64, 15u64), (12u64, 13u64)]; + let nd_sibling_indices_layer_height_2 = [(6u64, 7u64)]; + let nd_sibling_indices_layer_height_3 = [(2u64, 3u64)]; + let nd_sibling_indices = [ + nd_sibling_indices_layer_height_0.to_vec(), + nd_sibling_indices_layer_height_1.to_vec(), + nd_sibling_indices_layer_height_2.to_vec(), + nd_sibling_indices_layer_height_3.to_vec(), + ] + .concat(); + + prop_from_merkle_tree( + tree_height, + leaf_indices, + auth_struct_indices, + nd_sibling_indices, + ) + } +} diff --git a/tasm-lib/src/mmr/root_from_authentication_struct.rs b/tasm-lib/src/mmr/root_from_authentication_struct.rs new file mode 100644 index 00000000..4a8a16a8 --- /dev/null +++ b/tasm-lib/src/mmr/root_from_authentication_struct.rs @@ -0,0 +1,1026 @@ +use triton_vm::prelude::*; +use twenty_first::math::x_field_element::EXTENSION_DEGREE; + +use crate::data_type::DataType; +use crate::hashing::absorb_multiple::AbsorbMultiple; +use crate::library::Library; +use crate::prelude::BasicSnippet; + +pub struct RootFromAuthenticationStruct; + +impl RootFromAuthenticationStruct { + fn indexed_leaf_element_type() -> DataType { + DataType::Tuple(vec![DataType::U64, DataType::Digest]) + } +} + +impl BasicSnippet for RootFromAuthenticationStruct { + fn inputs(&self) -> Vec<(DataType, String)> { + vec![ + (DataType::U32, "tree_height".to_owned()), + ( + DataType::List(Box::new(DataType::Digest)), + "auth_struct".to_owned(), + ), + ( + DataType::List(Box::new(Self::indexed_leaf_element_type())), + "indexed_leafs".to_owned(), + ), + ] + } + + fn outputs(&self) -> Vec<(DataType, String)> { + vec![(DataType::Digest, "root".to_owned())] + } + + fn entrypoint(&self) -> String { + "tasmlib_mmr_root_from_authentication_struct".to_owned() + } + + fn code(&self, library: &mut Library) -> Vec { + let absorb_multiple = library.import(Box::new(AbsorbMultiple)); + let alpha_challenge_pointer_write = library.kmalloc(EXTENSION_DEGREE as u32); + let alpha_challenge_pointer_read = + alpha_challenge_pointer_write + bfe!(EXTENSION_DEGREE as u64 - 1); + let beta_challenge_pointer_write = library.kmalloc(EXTENSION_DEGREE as u32); + let beta_challenge_pointer_read = + beta_challenge_pointer_write + bfe!(EXTENSION_DEGREE as u64 - 1); + let gamma_challenge_pointer_write = library.kmalloc(EXTENSION_DEGREE as u32); + let gamma_challenge_pointer_read = + gamma_challenge_pointer_write + bfe!(EXTENSION_DEGREE as u64 - 1); + + let indexed_leaf_element_size = Self::indexed_leaf_element_type().stack_size(); + let calculate_and_store_challenges = triton_asm!( + // _ *auth_struct *indexed_leafs + + sponge_init + // _ *auth_struct *indexed_leafs + + read_mem 1 + push 1 + add + // _ *auth_struct indexed_leafs_len *indexed_leafs + + swap 1 + // _ *auth_struct *indexed_leafs indexed_leafs_len + + push {indexed_leaf_element_size} + mul + push 1 + add + // _ *auth_struct *indexed_leafs indexed_leafs_size + + call {absorb_multiple} + // _ *auth_struct + + read_mem 1 + push 1 + add + // _ auth_struct_len *auth_struct + + swap 1 + push {Digest::LEN} + mul + push 1 + add + // _ *auth_struct auth_struct_size + + call {absorb_multiple} + // _ + + sponge_squeeze + // _ w9 w8 w7 w6 w5 w4 w3 w2 w1 w0 + + pop 1 + // _ w9 w8 w7 w6 w5 w4 w3 w2 w1 + hint alpha: XFieldElement = stack[0..3] + hint minus_beta: XFieldElement = stack[3..6] + hint gamma: XFieldElement = stack[6..9] + // _ [gamma] [-beta] [alpha] <- rename + + push {alpha_challenge_pointer_write} + write_mem {EXTENSION_DEGREE} + pop 1 + // _ [gamma] [-beta] + + push {beta_challenge_pointer_write} + write_mem {EXTENSION_DEGREE} + pop 1 + // _ [gamma] + + push {gamma_challenge_pointer_write} + write_mem {EXTENSION_DEGREE} + pop 1 + // _ + ); + + let u64_size = DataType::U64.stack_size(); + const TWO_POW_32: u64 = 1 << 32; + let u64_to_bfe = triton_asm!( + // _ leaf_idx_hi leaf_idx_lo + + swap 1 + push {TWO_POW_32} + mul + // _ leaf_idx_lo (leaf_idx_hi << 32) + + add + + // _ leaf_idx_as_bfe + ); + + let digest_to_xfe = triton_asm!( + // _ [digest] + + push 1 + push {alpha_challenge_pointer_read} + read_mem {EXTENSION_DEGREE} + pop 1 + // _ l4 l3 l2 l1 l0 1 [α; 3] + + xx_mul + // _ l4 l3 l2 [(l1 l0 1) * α] + + xx_add + // _ xfe + ); + + let entrypoint = self.entrypoint(); + let accumulate_indexed_leafs_loop_label = format!("{entrypoint}_acc_indexed_leafs"); + let accumulated_indexed_leafs_loop = triton_asm!( + // INVARIANT: _ num_leafs *auth_struct *idx_leafs *idx_leafs[n]_lw [0; 2] [p; 3] + {accumulate_indexed_leafs_loop_label}: + // _ num_leafs *auth_struct *idx_leafs *idx_leafs[n]_lw [0; 2] [p; 3] + + /* Read leaf-index, convert it to BFE, and multiply it with `gamma` challenge */ + push {gamma_challenge_pointer_read} + read_mem {EXTENSION_DEGREE} + pop 1 + // _ num_leafs *auth_struct *idx_leafs *idx_leafs[n]_lw [0; 2] [p; 3] [gamma] + // _ num_leafs *auth_struct *idx_leafs *idx_leafs[n]_lw [0; 2] [p; 3] [γ] <-- rename + + dup 8 + read_mem {u64_size} + swap 11 + pop 1 + // _ num_leafs *auth_struct *idx_leafs (*idx_leafs[n]_lw - 2) [0; 2] [p; 3] [γ] [leaf_idx; 2] + + {&u64_to_bfe} + // _ num_leafs *auth_struct *idx_leafs (*idx_leafs[n]_lw - 2) [0; 2] [p; 3] [γ] leaf_idx_bfe + + dup 12 + add + // _ num_leafs *auth_struct *idx_leafs (*idx_leafs[n]_lw - 2) [0; 2] [p; 3] [γ] node_idx_bfe + + xb_mul + // _ num_leafs *auth_struct *idx_leafs (*idx_leafs[n]_lw - 2) [0; 2] [p; 3] [γ * node_idx; 3] + + dup 8 + read_mem {Digest::LEN} + swap 14 + pop 1 + // _ num_leafs *auth_struct *idx_leafs *idx_leafs[n-1]_lw [0; 2] [p; 3] [γ * node_idx; 3] [leaf; 5] + // _ num_leafs *auth_struct *idx_leafs *idx_leafs[n-1]_lw [0; 2] [p; 3] [γ * node_idx; 3] l4 l3 l2 l1 l0 + + /* Convert `leaf` to XFE, using challenge */ + {&digest_to_xfe} + hint leaf_as_xfe: XFieldElement = stack[0..2] + // _ num_leafs *auth_struct *idx_leafs *idx_leafs[n-1]_lw [0; 2] [p; 3] [γ * node_idx; 3] [(l1 l0 1) * α + (l4 l3 l2)] + // _ num_leafs *auth_struct *idx_leafs *idx_leafs[n-1]_lw [0; 2] [p; 3] [γ * node_idx; 3] [leaf_as_xfe] <-- rename + + push {beta_challenge_pointer_read} + read_mem {EXTENSION_DEGREE} + pop 1 + // _ num_leafs *auth_struct *idx_leafs *idx_leafs[n-1]_lw [0; 2] [p; 3] [γ * node_idx; 3] [leaf_as_xfe] [-beta] + + xx_add + // _ num_leafs *auth_struct *idx_leafs *idx_leafs[n-1]_lw [0; 2] [p; 3] [γ * node_idx; 3] [leaf_as_xfe - beta] + + xx_add + // _ num_leafs *auth_struct *idx_leafs *idx_leafs[n-1]_lw [0; 2] [p; 3] [γ * node_idx + leaf_as_xfe - beta] + + xx_mul + // _ num_leafs *auth_struct *idx_leafs *idx_leafs[n-1]_lw [0; 2] [p'; 3] + + recurse_or_return + ); + let accumulate_indexed_leafs_from_public_data = triton_asm!( + // _ num_leafs *auth_struct *indexed_leafs + + dup 0 + read_mem 1 + push 1 + add + swap 1 + // _ num_leafs *auth_struct *indexed_leafs *indexed_leafs indexed_leafs_len + + push {indexed_leaf_element_size} + mul + // _ num_leafs *auth_struct *indexed_leafs *indexed_leafs (indexed_leafs_size - 1) + + add + // _ num_leafs *auth_struct *indexed_leafs *indexed_leafs_last_word + + push 0 + push 0 + push 0 + push 0 + push 1 + hint prev = stack[3..5] + hint p: XFieldElement = stack[0..3] + // _ num_leafs *auth_struct *indexed_leafs *indexed_leafs_last_word [0u64; 2] [p; 3] + + dup 6 + dup 6 + eq + push 0 + eq + skiz + call {accumulate_indexed_leafs_loop_label} + // _ num_leafs *auth_struct *indexed_leafs *indexed_leafs [0u64; 2] [p; 3] + ); + + let accumulate_auth_struct_leafs_from_public_data_label = + format!("{entrypoint}_auth_struct_loop"); + + // let u64_lt = library.import(Box::new(LtU64PreserveArgs)); + let u64_lt = triton_asm!( + // _ lhs_hi lhs_lo rhs_hi rhs_lo + + /* calculate rhs_hi < lhs_hi || rhs_hi == lhs_hi && rhs_lo < lhs_lo */ + dup 2 + swap 1 + // _ lhs_hi lhs_lo rhs_hi lhs_lo rhs_lo + + lt + // _ lhs_hi lhs_lo rhs_hi (lhs_lo > rhs_lo) + // _ lhs_hi lhs_lo rhs_hi (rhs_lo < lhs_lo) + + dup 1 + dup 4 + eq + // _ lhs_hi lhs_lo rhs_hi (rhs_lo < lhs_lo) (rhs_hi == lhs_hi) + + mul + // _ lhs_hi lhs_lo rhs_hi (rhs_lo < lhs_lo && rhs_hi == lhs_hi) + + dup 3 + swap 1 + swap 2 + // _ lhs_hi lhs_lo (rhs_lo < lhs_lo && rhs_hi == lhs_hi) lhs_hi rhs_hi + + lt + // _ lhs_hi lhs_lo (rhs_lo < lhs_lo && rhs_hi == lhs_hi) (lhs_hi > rhs_hi) + // _ lhs_hi lhs_lo (rhs_lo < lhs_lo && rhs_hi == lhs_hi) (rhs_hi < lhs_hi) + + add + // _ lhs_hi lhs_lo (rhs_lo < lhs_lo && rhs_hi == lhs_hi || rhs_hi < lhs_hi) + ); + + let accumulate_auth_struct_leafs_from_public_data = triton_asm!( + // INVARIANT: _ *auth_struct *auth_struct[n]_lw [prev; 2] [p] + {accumulate_auth_struct_leafs_from_public_data_label}: + /* Divine in auth-struct node-index and verify ordering */ + + divine 2 + hint auth_struct_elem_node_index: u64 = stack[0..2] + // _ *auth_struct *auth_struct[n]_lw [prev; 2] [p] [node_index] + + + /* Notice that the u64-lt snippet crashes if divined node-index + words are not valid u32s. So no need to check explicitly. */ + dup 6 + dup 6 + {&u64_lt} + // _ *auth_struct *auth_struct[n]_lw [prev; 2] [p] [node_index] (prev < nodex_index) + + assert + // _ *auth_struct *auth_struct[n]_lw [prev; 2] [p] [node_index] + // _ *auth_struct *auth_struct[n]_lw prev_hi prev_lo p2 p1 p0 node_index_hi node_index_lo + + swap 5 + pop 1 + swap 5 + pop 1 + // _ *auth_struct *auth_struct[n]_lw [node_index] [p] + + /* Calculate `node_index * challenge` */ + push {gamma_challenge_pointer_read} + read_mem {EXTENSION_DEGREE} + pop 1 + // _ *auth_struct *auth_struct[n]_lw [node_index] [p] [γ] + + dup 7 + push {TWO_POW_32} + mul + dup 7 + add + // _ *auth_struct *auth_struct[n]_lw [node_index] [p] [γ] node_index_bfe + + xb_mul + // _ *auth_struct *auth_struct[n]_lw [node_index] [p] [node_index * γ] + + /* Read auth-struct element and convert to XFE */ + dup 8 + read_mem {Digest::LEN} + swap 14 + pop 1 + // _ *auth_struct *auth_struct[n]_lw [node_index] [p] [node_index * γ] [auth_struct_digest] + + {&digest_to_xfe} + // _ *auth_struct *auth_struct[n]_lw [node_index] [p] [node_index * γ] [auth_struct_xfe] + + push {beta_challenge_pointer_read} + read_mem {EXTENSION_DEGREE} + pop 1 + xx_add + // _ *auth_struct *auth_struct[n]_lw [node_index] [p] [node_index * γ] [auth_struct_xfe - beta] + + xx_add + // _ *auth_struct *auth_struct[n]_lw [node_index] [p] [node_index * γ + auth_struct_xfe - beta] + + xx_mul + // _ *auth_struct *auth_struct[n]_lw [node_index] [p * (node_index * γ + auth_struct_xfe - beta)] + // _ *auth_struct *auth_struct[n]_lw [prev] [p'] <-- rename + + recurse_or_return + ); + + let set_initial_t_value = triton_asm!( + // _ num_leafs *indexed_leafs *auth_struct [p] + + /* Set `t` digest to according to the rules: + let mut t = auth_struct + .first() + .copied() + .unwrap_or_else(|| indexed_leafs.first().unwrap().1); + */ + dup 3 + read_mem 1 + pop 1 + // _ num_leafs *indexed_leafs *auth_struct [p] auth_struct_len + ); + + triton_asm!( + {entrypoint}: + // _ tree_height *auth_struct *indexed_leafs + + dup 1 + dup 1 + // _ tree_height *auth_struct *indexed_leafs *auth_struct *indexed_leafs + + {&calculate_and_store_challenges} + // _ tree_height *auth_struct *indexed_leafs + + /* Calculate number of leafs in Merkle tree */ + swap 2 + push 2 + pow + swap 2 + hint num_leafs = stack[2] + // _ num_leafs *auth_struct *indexed_leafs + + {&accumulate_indexed_leafs_from_public_data} + // _ num_leafs *auth_struct *indexed_leafs *indexed_leafs [garbage; 2] [p; 3] + // _ num_leafs *auth_struct *indexed_leafs *indexed_leafs [garbage; 2] p2 p1 p0 <-- rename + + /* Prepare for next loop, absorption of auth-struct digests into accumulator */ + swap 7 + swap 6 + pop 1 + // _ num_leafs p0 *auth_struct *indexed_leafs [0; 2] p2 p1 + + dup 5 + read_mem 1 + push 1 + add + swap 1 + // _ num_leafs p0 *auth_struct *indexed_leafs [0; 2] p2 p1 *auth_struct auth_struct_len + + push {Digest::LEN} + mul + add + // _ num_leafs p0 *auth_struct *indexed_leafs [0; 2] p2 p1 *auth_struct_last_word + + swap 5 + swap 7 + // _ num_leafs *indexed_leafs *auth_struct *auth_struct_last_word [0; 2] p2 p1 p0 + // _ num_leafs *indexed_leafs *auth_struct *auth_struct_last_word [prev; 2] [p] <-- rename + + dup 6 + dup 6 + eq + push 0 + eq + // _ num_leafs *indexed_leafs *auth_struct *auth_struct_last_word [prev; 2] [p] (*auth_struct_last_word != *auth_struct) + + skiz + call {accumulate_auth_struct_leafs_from_public_data_label} + // _ num_leafs *indexed_leafs *auth_struct *auth_struct_last_word [prev; 2] [p] + + /* Cleanup stack before next loop */ + swap 3 + pop 1 + swap 3 + pop 1 + swap 3 + pop 1 + // _ num_leafs *indexed_leafs *auth_struct [p] + + return + + {&accumulated_indexed_leafs_loop} + {&accumulate_auth_struct_leafs_from_public_data} + ) + } +} + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + use std::collections::VecDeque; + + use itertools::Itertools; + use num::One; + use rand::rngs::StdRng; + use rand::Rng; + use rand::SeedableRng; + use twenty_first::prelude::CpuParallel; + use twenty_first::prelude::MerkleTree; + use twenty_first::prelude::Sponge; + + use crate::rust_shadowing_helper_functions::list::list_insert; + use crate::rust_shadowing_helper_functions::list::load_list_with_copy_elements; + use crate::snippet_bencher::BenchmarkCase; + use crate::traits::procedure::Procedure; + use crate::traits::procedure::ProcedureInitialState; + use crate::traits::procedure::ShadowedProcedure; + use crate::traits::rust_shadow::RustShadow; + use crate::VmHasher; + + use super::*; + + const SIZE_OF_INDEXED_LEAFS_ELEMENT: usize = Digest::LEN + 2; + + #[test] + fn test() { + ShadowedProcedure::new(RootFromAuthenticationStruct).test(); + } + + impl Procedure for RootFromAuthenticationStruct { + fn rust_shadow( + &self, + stack: &mut Vec, + memory: &mut std::collections::HashMap, + nondeterminism: &NonDeterminism, + public_input: &[BFieldElement], + sponge: &mut Option, + ) -> Vec { + fn digest_to_xfe(digest: Digest, challenge: XFieldElement) -> XFieldElement { + let [l0, l1, l2, l3, l4] = digest.0; + let leaf_xfe_lo = XFieldElement::new([BFieldElement::new(1), l0, l1]); + let leaf_xfe_hi = XFieldElement::new([l2, l3, l4]); + + challenge * leaf_xfe_lo + leaf_xfe_hi + } + + assert_eq!( + SIZE_OF_INDEXED_LEAFS_ELEMENT, + Self::indexed_leaf_element_type().stack_size() + ); + let indexed_leafs_pointer = stack.pop().unwrap(); + let auth_struct_pointer = stack.pop().unwrap(); + let tree_height: u32 = stack.pop().unwrap().try_into().unwrap(); + + let bfes_to_indexed_leaf = + |bfes: [BFieldElement; SIZE_OF_INDEXED_LEAFS_ELEMENT]| -> (u64, Digest) { + *<(u64, Digest)>::decode(&bfes).unwrap() + }; + let bfes_to_digest = |bfes: [BFieldElement; Digest::LEN]| -> Digest { + println!("bfes_to_digest:\nbfes: {bfes:?}"); + *Digest::decode(&bfes[0..Digest::LEN]).unwrap() + }; + + let indexed_leafs: Vec<[BFieldElement; SIZE_OF_INDEXED_LEAFS_ELEMENT]> = + load_list_with_copy_elements(indexed_leafs_pointer, memory); + println!("indexed_leafs: {indexed_leafs:?}"); + let indexed_leafs = indexed_leafs + .into_iter() + .map(bfes_to_indexed_leaf) + .collect_vec(); + let auth_struct: Vec<[BFieldElement; Digest::LEN]> = + load_list_with_copy_elements(auth_struct_pointer, memory); + println!("auth_struct: {auth_struct:?}"); + let auth_struct = auth_struct.into_iter().map(bfes_to_digest).collect_vec(); + + let sponge = sponge.as_mut().expect("sponge must be initialized"); + + sponge.pad_and_absorb_all(&indexed_leafs.encode()); + sponge.pad_and_absorb_all(&auth_struct.encode()); + + let sponge_output = sponge.squeeze(); + let alpha = XFieldElement::new([sponge_output[1], sponge_output[2], sponge_output[3]]); + let beta = -XFieldElement::new([sponge_output[4], sponge_output[5], sponge_output[6]]); + let gamma = XFieldElement::new([sponge_output[7], sponge_output[8], sponge_output[9]]); + println!("alpha {alpha}"); + println!("-beta {}", -beta); + println!("gamma {}", -gamma); + + let num_leafs = 1 << tree_height; + let mut p = XFieldElement::one(); + for (leaf_idx, leaf) in indexed_leafs.iter().copied().rev() { + let leaf_idx_as_bfe = bfe!(leaf_idx); + let node_idx_as_bfe = leaf_idx_as_bfe + bfe!(num_leafs); + + let leaf_as_xfe = digest_to_xfe(leaf, alpha); + + println!("gamma * node_idx_as_bfe = {}", gamma * node_idx_as_bfe); + println!("leaf_as_xfe = {}", leaf_as_xfe); + + let fact = leaf_as_xfe - beta + gamma * node_idx_as_bfe; + println!("indexed-leaf fact: {fact}"); + p *= fact; + } + + println!("Rust-shadow, p, after indexed-leafs absorption: {p}"); + + let mut prev = 0u64; + let mut individual_tokens: VecDeque = + nondeterminism.individual_tokens.to_owned().into(); + for auth_struct_elem in auth_struct.iter().copied().rev() { + let auth_struct_elem_node_index_hi: u32 = + individual_tokens.pop_front().unwrap().try_into().unwrap(); + let auth_struct_elem_node_index_lo: u32 = + individual_tokens.pop_front().unwrap().try_into().unwrap(); + let auth_struct_elem_node_index = ((auth_struct_elem_node_index_hi as u64) << 32) + + auth_struct_elem_node_index_lo as u64; + println!("auth_struct_elem_node_index: {auth_struct_elem_node_index}"); + assert!(auth_struct_elem_node_index > prev); + prev = auth_struct_elem_node_index; + + let auth_struct_index_as_bfe = bfe!(auth_struct_elem_node_index); + + let auth_struct_elem_xfe = digest_to_xfe(auth_struct_elem, alpha); + let fact = auth_struct_elem_xfe - beta + gamma * auth_struct_index_as_bfe; + + println!("auth struct fact: {fact}"); + p *= fact; + } + + println!("Rust-shadow, p, after auth-struct absorption: {p}"); + + vec![] + } + + fn pseudorandom_initial_state( + &self, + seed: [u8; 32], + bench_case: Option, + ) -> ProcedureInitialState { + let mut rng: StdRng = SeedableRng::from_seed(seed); + + // TODO: use real `mmr_authentication_struct` code here + let tree_height = rng.gen_range(1..5); + let leaf_count = 1 << tree_height; + let num_revealed_leafs = rng.gen_range(1..leaf_count); + let revealed_leaf_indices = (0..num_revealed_leafs) + .map(|_| rng.gen_range(0..leaf_count)) + .unique() + .collect_vec(); + let leafs = (0..leaf_count).map(|_| rng.gen()).collect_vec(); + let tree = MerkleTree::::new::(&leafs).unwrap(); + let mmr_authentication_struct = + mmr_authentication_struct::AuthStructIntegrityProof::new_from_merkle_tree( + &tree, + revealed_leaf_indices, + ); + + let mut memory = HashMap::new(); + let authentication_structure_ptr = rng.gen(); + let indexed_leafs_ptr = rng.gen(); + + println!( + "indexed_leafs.len(): {}", + mmr_authentication_struct.indexed_leafs.len() + ); + println!( + "authentication_structure.len(): {}", + mmr_authentication_struct.auth_struct.len() + ); + + list_insert( + authentication_structure_ptr, + mmr_authentication_struct.auth_struct.clone(), + &mut memory, + ); + list_insert( + indexed_leafs_ptr, + mmr_authentication_struct.indexed_leafs, + &mut memory, + ); + + let stack = [ + self.init_stack_for_isolated_run(), + vec![ + bfe!(tree_height), + authentication_structure_ptr, + indexed_leafs_ptr, + ], + ] + .concat(); + + println!( + "node indices: {}", + mmr_authentication_struct + .witness + .nd_auth_struct_indices + .iter() + .join(", ") + ); + let individual_tokens = mmr_authentication_struct + .witness + .nd_auth_struct_indices + .into_iter() + .rev() + .flat_map(|node_index| node_index.encode().into_iter().rev().collect_vec()) + .collect_vec(); + println!("individual_tokens: {}", individual_tokens.iter().join(", ")); + let nondeterminism = NonDeterminism::new(individual_tokens).with_ram(memory); + ProcedureInitialState { + stack, + nondeterminism, + public_input: vec![], + sponge: Some(Tip5::init()), + } + } + + fn corner_case_initial_states(&self) -> Vec { + vec![] + } + } +} + +// TODO: Use this logic from `twenty-first` instead! +mod mmr_authentication_struct { + use std::collections::{HashMap, HashSet}; + + use itertools::Itertools; + use num::One; + use twenty_first::{ + prelude::{AlgebraicHasher, Inverse, MerkleTree, Mmr, MmrMembershipProof, Sponge}, + util_types::mmr::{ + mmr_accumulator::MmrAccumulator, shared_advanced::get_peak_heights, + shared_basic::leaf_index_to_mt_index_and_peak_index, + }, + }; + + use super::*; + + const ROOT_MT_INDEX: u64 = 1; + + /// A witness to facilitate the proving of the authenticity of a Merkle + /// authentication struct. + #[derive(Debug, Clone)] + pub struct AuthStructIntegrityProof { + // All indices are Merkle tree node indices + pub nd_auth_struct_indices: Vec, + pub nd_sibling_indices: Vec<(u64, u64)>, + pub nd_siblings: Vec<(Digest, Digest)>, + } + + /// An authentication structure that can be used to prove membership of a list + /// of leaves in a Merkle tree, along with the indexed leaves in question, and + /// the witness necessary to prove membership in a ZK program. + #[derive(Debug, Clone)] + pub struct AuthenticatedMerkleAuthStruct { + pub auth_struct: Vec, + pub indexed_leafs: Vec<(u64, Digest)>, + pub witness: AuthStructIntegrityProof, + } + + impl AuthStructIntegrityProof { + /// Return the Merkle tree node indices of the digests required to prove + /// membership for the specified leaf indices, as well as the node indices + /// that can be derived from the leaf indices and their authentication + /// path. + fn auth_struct_and_nd_indices( + num_leafs: u64, + leaf_indices: &[u64], + ) -> (Vec, Vec<(u64, u64)>) { + // The set of indices of nodes that need to be included in the authentications + // structure. In principle, every node of every authentication path is needed. + // The root is never needed. Hence, it is not considered below. + let mut node_is_needed = HashSet::new(); + + // The set of indices of nodes that can be computed from other nodes in the + // authentication structure or the leafs that are explicitly supplied during + // verification. Every node on the direct path from the leaf to the root can + // be computed by the very nature of “authentication path”. + let mut node_can_be_computed = HashSet::new(); + + for &leaf_index in leaf_indices { + assert!(num_leafs > leaf_index, "Leaf index must be less than number of leafs. Got leaf_index = {leaf_index}; num_leafs = {num_leafs}"); + + let mut node_index = leaf_index + num_leafs; + while node_index > ROOT_MT_INDEX { + let sibling_index = node_index ^ 1; + node_can_be_computed.insert(node_index); + node_is_needed.insert(sibling_index); + node_index /= 2; + } + } + + let set_difference = node_is_needed.difference(&node_can_be_computed).copied(); + let set_union = node_is_needed + .union(&node_can_be_computed) + .sorted_unstable() + .rev(); + + let mut set_union = set_union.peekable(); + + let mut set_union_as_ordered_pairs = Vec::new(); + while set_union.peek().is_some() { + let right_index = *set_union.next().unwrap(); + + // Crashes on odd-length of input list, which is what we want, as + // this acts as a sanity check. + let left_index = *set_union.next().unwrap(); + set_union_as_ordered_pairs.push((left_index, right_index)); + } + + ( + set_difference.sorted_unstable().rev().collect(), + set_union_as_ordered_pairs, + ) + } + + pub fn root_from_authentication_struct( + &self, + tree_height: u32, + auth_struct: Vec, + indexed_leafs: Vec<(u64, Digest)>, + ) -> Digest { + fn digest_to_xfe(digest: Digest, challenge: XFieldElement) -> XFieldElement { + let leaf_xfe_lo = XFieldElement::new([digest.0[0], digest.0[1], digest.0[2]]); + let leaf_xfe_hi = challenge + * XFieldElement::new([digest.0[3], digest.0[4], BFieldElement::one()]); + + leaf_xfe_lo + leaf_xfe_hi + } + + fn node_index_to_bfe(node_index: u64) -> BFieldElement { + BFieldElement::new(node_index) + } + + // Sanity check + assert_eq!( + self.nd_auth_struct_indices.len(), + auth_struct.len(), + "Provided auth struct length must match that specified in receiver" + ); + + // Get challenges + let (alpha, beta, gamma) = { + let mut sponge = Tip5::init(); + sponge.pad_and_absorb_all(&indexed_leafs.encode()); + sponge.pad_and_absorb_all(&auth_struct.encode()); + let challenges = sponge.sample_scalars(3); + (challenges[0], challenges[1], challenges[2]) + }; + + // Accumulate `p` from public data + let mut p = XFieldElement::one(); + for i in (0..indexed_leafs.len()).rev() { + let node_index_as_bfe = node_index_to_bfe((1 << tree_height) ^ indexed_leafs[i].0); + let leaf_as_xfe = digest_to_xfe(indexed_leafs[i].1, alpha); + let fact = leaf_as_xfe - beta + gamma * node_index_as_bfe; + p *= fact; + } + + let mut prev = 0; + for i in (0..auth_struct.len()).rev() { + let auth_struct_index = self.nd_auth_struct_indices[i]; + + // `auth_struct` must be sorted high-to-low by node-index. But since + // we're traversing in reverse order, the inequality is flipped. + assert!(auth_struct_index > prev); + prev = auth_struct_index; + + let auth_struct_index_as_bfe = node_index_to_bfe(auth_struct_index); + + let auth_str_elem_as_xfe = digest_to_xfe(auth_struct[i], alpha); + let fact = auth_str_elem_as_xfe - beta + gamma * auth_struct_index_as_bfe; + p *= fact; + } + + // Use secret data to invert `p` back and to calculate the root + let mut t = auth_struct + .first() + .copied() + .unwrap_or_else(|| indexed_leafs.first().unwrap().1); + let mut t_xfe = digest_to_xfe(t, alpha); + let mut parent_index_bfe = BFieldElement::one(); + for ((l, r), (left_index, right_index)) in self + .nd_siblings + .iter() + .zip_eq(self.nd_sibling_indices.clone()) + { + assert_eq!(left_index + 1, right_index); + + t = Tip5::hash_pair(*l, *r); + + let l_xfe = digest_to_xfe(*l, alpha); + let r_xfe = digest_to_xfe(*r, alpha); + t_xfe = digest_to_xfe(t, alpha); + + let left_index_bfe = node_index_to_bfe(left_index); + let right_index_bfe = node_index_to_bfe(right_index); + parent_index_bfe = left_index_bfe / bfe!(2); + + let fact1 = l_xfe - beta + gamma * left_index_bfe; + let fact2 = r_xfe - beta + gamma * right_index_bfe; + let fact_parent = t_xfe - beta + gamma * parent_index_bfe; + + p *= fact1.inverse() * fact2.inverse() * fact_parent; + } + + assert_eq!(t_xfe - beta + gamma, p); + assert!(parent_index_bfe.is_one()); + + t + } + + /// Return the authentication structure authenticity witness, + /// authentication structure, and the (leaf-index, leaf-digest) pairs + /// from a list of MMR membership proofs. All MMR membership proofs must + /// belong under the same peak, i.e., be part of the same Merkle tree in + /// the list of Merkle trees that the MMR contains. + /// + /// Panics if the input list of MMR-membership proofs is empty, or if they + /// do not all belong under the same peak. + pub fn new_from_mmr_membership_proofs( + mmra: &MmrAccumulator, + indexed_mmr_mps: Vec<(u64, Digest, MmrMembershipProof)>, + ) -> HashMap { + #[derive(Clone, Debug)] + struct IndexedAuthenticatedMmrLeaf { + merkle_tree_node_index: u64, + merkle_tree_leaf_index: u64, + leaf_digest: Digest, + membership_proof: MmrMembershipProof, + } + + // Split indexed MMR-mps into a hashmap with one entry for each + // referenced peak in the MMR. + let num_mmr_leafs = mmra.num_leafs(); + let mut peak_index_to_indexed_mmr_mp: HashMap> = + HashMap::default(); + let peak_heights = get_peak_heights(num_mmr_leafs); + for (mmr_leaf_index, leaf, mmr_mp) in indexed_mmr_mps { + let (mt_index, peak_index) = + leaf_index_to_mt_index_and_peak_index(mmr_leaf_index, num_mmr_leafs); + let peak_index_as_usize: usize = peak_index.try_into().unwrap(); + let num_leafs_local_mt = 1 << peak_heights[peak_index_as_usize]; + let mt_leaf_index = mt_index - num_leafs_local_mt; + peak_index_to_indexed_mmr_mp + .entry(peak_index) + .or_default() + .push(IndexedAuthenticatedMmrLeaf { + merkle_tree_node_index: mt_index, + merkle_tree_leaf_index: mt_leaf_index, + leaf_digest: leaf, + membership_proof: mmr_mp, + }); + } + + // Loop over all peaks and collect an authentication witness struct + // for each peak. + let mut peak_index_to_authenticated_auth_struct = HashMap::default(); + for (peak_index, indexed_mmr_mp_structs) in peak_index_to_indexed_mmr_mp { + let peak_index_as_usize: usize = peak_index.try_into().unwrap(); + let num_leafs_in_local_mt = 1 << peak_heights[peak_index_as_usize]; + let local_mt_leaf_indices = indexed_mmr_mp_structs + .iter() + .map(|x| x.merkle_tree_leaf_index) + .collect_vec(); + + let (nd_auth_struct_indices, nd_sibling_indices) = + Self::auth_struct_and_nd_indices(num_leafs_in_local_mt, &local_mt_leaf_indices); + let peak = mmra.peaks()[peak_index_as_usize]; + + let mut node_digests: HashMap = HashMap::default(); + node_digests.insert(ROOT_MT_INDEX, peak); + + // Loop over all indexed leafs for this peak + for indexed_mmr_mp in indexed_mmr_mp_structs.iter() { + let mut mt_node_index = indexed_mmr_mp.merkle_tree_node_index; + let mut node = indexed_mmr_mp.leaf_digest; + + // Loop over all authentication path elements for this indexed leaf + for ap_elem in indexed_mmr_mp.membership_proof.authentication_path.iter() { + node_digests.insert(mt_node_index, node); + node_digests.insert(mt_node_index ^ 1, *ap_elem); + node = if mt_node_index & 1 == 0 { + Tip5::hash_pair(node, *ap_elem) + } else { + Tip5::hash_pair(*ap_elem, node) + }; + + mt_node_index /= 2; + } + + // Sanity check that MMR-MPs are valid + assert_eq!(peak, node, "Derived peak must match provided peak"); + } + let nd_siblings = nd_sibling_indices + .iter() + .map(|(left_idx, right_idx)| (node_digests[left_idx], node_digests[right_idx])) + .collect_vec(); + let auth_struct = nd_auth_struct_indices + .iter() + .map(|idx| node_digests[idx]) + .collect_vec(); + let indexed_leafs = indexed_mmr_mp_structs + .into_iter() + .map(|indexed_mmr_mp| { + ( + indexed_mmr_mp.merkle_tree_leaf_index, + indexed_mmr_mp.leaf_digest, + ) + }) + .collect_vec(); + + let witness = Self { + nd_auth_struct_indices, + nd_sibling_indices, + nd_siblings, + }; + + peak_index_to_authenticated_auth_struct.insert( + peak_index, + AuthenticatedMerkleAuthStruct { + auth_struct, + indexed_leafs, + witness, + }, + ); + } + + peak_index_to_authenticated_auth_struct + } + + /// Return the authentication structure witness, authentication structure, + /// and the (leaf-index, leaf-digest) pairs. + pub fn new_from_merkle_tree( + tree: &MerkleTree, + mut revealed_leaf_indices: Vec, + ) -> AuthenticatedMerkleAuthStruct { + revealed_leaf_indices.sort_unstable(); + revealed_leaf_indices.dedup(); + revealed_leaf_indices.reverse(); + let num_leafs: u64 = tree.num_leafs() as u64; + + let (mut nd_auth_struct_indices, nd_sibling_indices) = + Self::auth_struct_and_nd_indices(num_leafs, &revealed_leaf_indices); + if revealed_leaf_indices.is_empty() { + nd_auth_struct_indices = vec![ROOT_MT_INDEX]; + } + + let nd_siblings = nd_sibling_indices + .iter() + .map(|&(l, r)| { + let l: usize = l.try_into().unwrap(); + let r: usize = r.try_into().unwrap(); + (tree.node(l).unwrap(), tree.node(r).unwrap()) + }) + .collect_vec(); + + let revealed_leafs = revealed_leaf_indices + .iter() + .map(|j| tree.node((*j + num_leafs) as usize).unwrap()) + .collect_vec(); + let indexed_leafs = revealed_leaf_indices + .clone() + .into_iter() + .zip_eq(revealed_leafs) + .collect_vec(); + + let auth_struct = nd_auth_struct_indices + .iter() + .map(|node_index| tree.node(*node_index as usize).unwrap()) + .collect_vec(); + + let witness = Self { + nd_auth_struct_indices, + nd_sibling_indices, + nd_siblings, + }; + + AuthenticatedMerkleAuthStruct { + auth_struct, + indexed_leafs, + witness, + } + } + } +} From 7f55fa70c58eb4baa52a71a6e22a7ab90e2a4976 Mon Sep 17 00:00:00 2001 From: sword-smith Date: Mon, 12 Aug 2024 23:24:52 +0200 Subject: [PATCH 02/10] mmr::authentication_struct: implement snipet and disallow empty list of leafs When calculating the root, don't allow the list of leafs to be empty. We think we never need the snippet with those inputs, so that should be fine. It makes the initialization of the `t` value (that should end up containing the root) simpler, as we can do: ```rust let mut t = indexed_leafs.first().unwrap().1; ``` instead of ```rust let mut t = auth_struct .first() .copied() .unwrap_or_else(|| indexed_leafs.first().unwrap().1); ``` Code works with sane test. --- ...thentication_struct_derive_challenges.json | 24 + tasm-lib/src/mmr.rs | 1 - tasm-lib/src/mmr/authentication_struct.rs | 2 + .../derive_challenges.rs | 248 +++++ .../root_from_authentication_struct.rs | 863 ++++++++---------- .../src/mmr/authentication_struct/shared.rs | 53 +- .../src/rust_shadowing_helper_functions.rs | 1 + .../rust_shadowing_helper_functions/input.rs | 10 +- .../rust_shadowing_helper_functions/memory.rs | 15 + 9 files changed, 709 insertions(+), 508 deletions(-) create mode 100644 tasm-lib/benchmarks/tasmlib_mmr_authentication_struct_derive_challenges.json create mode 100644 tasm-lib/src/mmr/authentication_struct/derive_challenges.rs rename tasm-lib/src/mmr/{ => authentication_struct}/root_from_authentication_struct.rs (52%) create mode 100644 tasm-lib/src/rust_shadowing_helper_functions/memory.rs diff --git a/tasm-lib/benchmarks/tasmlib_mmr_authentication_struct_derive_challenges.json b/tasm-lib/benchmarks/tasmlib_mmr_authentication_struct_derive_challenges.json new file mode 100644 index 00000000..e317de7a --- /dev/null +++ b/tasm-lib/benchmarks/tasmlib_mmr_authentication_struct_derive_challenges.json @@ -0,0 +1,24 @@ +[ + { + "name": "tasmlib_mmr_authentication_struct_derive_challenges", + "benchmark_result": { + "clock_cycle_count": 1196, + "hash_table_height": 985, + "u32_table_height": 25, + "op_stack_table_height": 833, + "ram_table_height": 1454 + }, + "case": "CommonCase" + }, + { + "name": "tasmlib_mmr_authentication_struct_derive_challenges", + "benchmark_result": { + "clock_cycle_count": 2102, + "hash_table_height": 1891, + "u32_table_height": 26, + "op_stack_table_height": 1437, + "ram_table_height": 2964 + }, + "case": "WorstCase" + } +] \ No newline at end of file diff --git a/tasm-lib/src/mmr.rs b/tasm-lib/src/mmr.rs index 9ef3b582..803819ec 100644 --- a/tasm-lib/src/mmr.rs +++ b/tasm-lib/src/mmr.rs @@ -3,7 +3,6 @@ pub mod bag_peaks; pub mod calculate_new_peaks_from_append; pub mod calculate_new_peaks_from_leaf_mutation; pub mod leaf_index_to_mt_index_and_peak_index; -pub mod root_from_authentication_struct; pub mod verify_from_memory; pub mod verify_from_secret_in_leaf_index_on_stack; pub mod verify_from_secret_in_secret_leaf_index; diff --git a/tasm-lib/src/mmr/authentication_struct.rs b/tasm-lib/src/mmr/authentication_struct.rs index eec3c896..2406a8f5 100644 --- a/tasm-lib/src/mmr/authentication_struct.rs +++ b/tasm-lib/src/mmr/authentication_struct.rs @@ -1 +1,3 @@ +pub mod derive_challenges; +pub mod root_from_authentication_struct; pub mod shared; diff --git a/tasm-lib/src/mmr/authentication_struct/derive_challenges.rs b/tasm-lib/src/mmr/authentication_struct/derive_challenges.rs new file mode 100644 index 00000000..a4920f92 --- /dev/null +++ b/tasm-lib/src/mmr/authentication_struct/derive_challenges.rs @@ -0,0 +1,248 @@ +use triton_vm::prelude::*; + +use crate::data_type::DataType; +use crate::hashing::absorb_multiple::AbsorbMultiple; +use crate::mmr::authentication_struct::shared; +use crate::prelude::BasicSnippet; +use crate::Library; + +/// Derive and return the challenges that the authentication structure verification +/// program uses. +pub struct DeriveChallenges; + +impl BasicSnippet for DeriveChallenges { + fn inputs(&self) -> Vec<(DataType, String)> { + vec![ + ( + DataType::List(Box::new(DataType::Digest)), + "auth_struct".to_owned(), + ), + ( + DataType::List(Box::new(shared::indexed_leaf_element_type())), + "indexed_leafs".to_owned(), + ), + ] + } + + fn outputs(&self) -> Vec<(DataType, String)> { + vec![ + (DataType::Xfe, "alpha".to_owned()), + (DataType::Xfe, "-beta".to_owned()), + (DataType::Xfe, "gamma".to_owned()), + ] + } + + fn entrypoint(&self) -> String { + "tasmlib_mmr_authentication_struct_derive_challenges".to_owned() + } + + fn code(&self, library: &mut Library) -> Vec { + let absorb_multiple = library.import(Box::new(AbsorbMultiple)); + + let entrypoint = self.entrypoint(); + + let indexed_leaf_element_size = shared::indexed_leaf_element_type().stack_size(); + triton_asm!( + {entrypoint}: + // _ *auth_struct *indexed_leafs + + sponge_init + // _ *auth_struct *indexed_leafs + + read_mem 1 + push 1 + add + // _ *auth_struct indexed_leafs_len *indexed_leafs + + swap 1 + // _ *auth_struct *indexed_leafs indexed_leafs_len + + push {indexed_leaf_element_size} + mul + push 1 + add + // _ *auth_struct *indexed_leafs indexed_leafs_size + + call {absorb_multiple} + // _ *auth_struct + + read_mem 1 + push 1 + add + // _ auth_struct_len *auth_struct + + swap 1 + push {Digest::LEN} + mul + push 1 + add + // _ *auth_struct auth_struct_size + + call {absorb_multiple} + // _ + + sponge_squeeze + // _ w9 w8 w7 w6 w5 w4 w3 w2 w1 w0 + + pop 1 + // _ w9 w8 w7 w6 w5 w4 w3 w2 w1 + // _ [gamma] [-beta] [alpha] <- rename + + return + ) + } +} + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + + use itertools::Itertools; + use num::One; + use rand::rngs::StdRng; + use rand::Rng; + use rand::SeedableRng; + use shared::AuthenticatedMerkleAuthStruct; + use twenty_first::prelude::Sponge; + use twenty_first::util_types::mmr::mmr_accumulator::util::mmra_with_mps; + + use crate::mmr::authentication_struct::shared::AuthStructIntegrityProof; + use crate::rust_shadowing_helper_functions::list::list_insert; + use crate::rust_shadowing_helper_functions::list::load_list_with_copy_elements; + use crate::snippet_bencher::BenchmarkCase; + use crate::traits::procedure::Procedure; + use crate::traits::procedure::ProcedureInitialState; + use crate::traits::procedure::ShadowedProcedure; + use crate::traits::rust_shadow::RustShadow; + use crate::VmHasher; + + use super::*; + + const SIZE_OF_INDEXED_LEAFS_ELEMENT: usize = Digest::LEN + 2; + + #[test] + fn test() { + ShadowedProcedure::new(DeriveChallenges).test(); + } + + impl Procedure for DeriveChallenges { + fn rust_shadow( + &self, + stack: &mut Vec, + memory: &mut HashMap, + _nondeterminism: &NonDeterminism, + _public_input: &[BFieldElement], + sponge: &mut Option, + ) -> Vec { + let indexed_leafs_pointer = stack.pop().unwrap(); + let auth_struct_pointer = stack.pop().unwrap(); + let bfes_to_indexed_leaf = + |bfes: [BFieldElement; SIZE_OF_INDEXED_LEAFS_ELEMENT]| -> (u64, Digest) { + *<(u64, Digest)>::decode(&bfes).unwrap() + }; + let bfes_to_digest = |bfes: [BFieldElement; Digest::LEN]| -> Digest { + *Digest::decode(&bfes[0..Digest::LEN]).unwrap() + }; + let indexed_leafs: Vec<[BFieldElement; SIZE_OF_INDEXED_LEAFS_ELEMENT]> = + load_list_with_copy_elements(indexed_leafs_pointer, memory); + let indexed_leafs = indexed_leafs + .into_iter() + .map(bfes_to_indexed_leaf) + .collect_vec(); + let auth_struct: Vec<[BFieldElement; Digest::LEN]> = + load_list_with_copy_elements(auth_struct_pointer, memory); + let auth_struct = auth_struct.into_iter().map(bfes_to_digest).collect_vec(); + + let sponge = sponge.as_mut().expect("sponge must be initialized"); + + sponge.pad_and_absorb_all(&indexed_leafs.encode()); + sponge.pad_and_absorb_all(&auth_struct.encode()); + + let sponge_output = sponge.squeeze(); + for elem in sponge_output.into_iter().skip(1).rev() { + stack.push(elem); + } + + vec![] + } + + fn pseudorandom_initial_state( + &self, + seed: [u8; 32], + bench_case: Option, + ) -> ProcedureInitialState { + let mut rng: StdRng = SeedableRng::from_seed(seed); + + let (tree_height, num_revealed_leafs) = match bench_case { + Some(BenchmarkCase::CommonCase) => (32, 10), + Some(BenchmarkCase::WorstCase) => (62, 10), + None => (rng.gen_range(0..62), 10), + }; + + let leaf_count = 1 << tree_height; + let revealed_leaf_indices = (0..num_revealed_leafs) + .map(|_| rng.gen_range(0..leaf_count)) + .unique() + .collect_vec(); + let indexed_leafs = revealed_leaf_indices + .into_iter() + .map(|leaf_idx: u64| (leaf_idx, rng.gen())) + .collect_vec(); + let (mmra, mps) = mmra_with_mps(leaf_count, indexed_leafs.clone()); + let indexed_mmr_mps = indexed_leafs + .into_iter() + .zip_eq(mps) + .map(|((leaf_idx, leaf), mp)| (leaf_idx, leaf, mp)) + .collect_vec(); + let authenticity_witnesses = + AuthStructIntegrityProof::new_from_mmr_membership_proofs(&mmra, indexed_mmr_mps); + assert!( + authenticity_witnesses.len().is_one(), + "All indices belong to first peak" + ); + let AuthenticatedMerkleAuthStruct { + auth_struct, + indexed_leafs, + .. + } = &authenticity_witnesses[&0]; + + let mut memory = HashMap::new(); + let authentication_structure_ptr = rng.gen(); + let indexed_leafs_ptr = rng.gen(); + + list_insert( + authentication_structure_ptr, + auth_struct.to_owned(), + &mut memory, + ); + list_insert(indexed_leafs_ptr, indexed_leafs.to_owned(), &mut memory); + + let stack = [ + self.init_stack_for_isolated_run(), + vec![authentication_structure_ptr, indexed_leafs_ptr], + ] + .concat(); + + let nondeterminism = NonDeterminism::default().with_ram(memory); + ProcedureInitialState { + stack, + nondeterminism, + public_input: vec![], + sponge: Some(Tip5::init()), + } + } + } +} + +#[cfg(test)] +mod benches { + use crate::traits::procedure::ShadowedProcedure; + use crate::traits::rust_shadow::RustShadow; + + use super::*; + + #[test] + fn bag_peaks_benchmark() { + ShadowedProcedure::new(DeriveChallenges).bench(); + } +} diff --git a/tasm-lib/src/mmr/root_from_authentication_struct.rs b/tasm-lib/src/mmr/authentication_struct/root_from_authentication_struct.rs similarity index 52% rename from tasm-lib/src/mmr/root_from_authentication_struct.rs rename to tasm-lib/src/mmr/authentication_struct/root_from_authentication_struct.rs index 4a8a16a8..404ba496 100644 --- a/tasm-lib/src/mmr/root_from_authentication_struct.rs +++ b/tasm-lib/src/mmr/authentication_struct/root_from_authentication_struct.rs @@ -1,9 +1,10 @@ use triton_vm::prelude::*; use twenty_first::math::x_field_element::EXTENSION_DEGREE; +use twenty_first::prelude::Inverse; use crate::data_type::DataType; -use crate::hashing::absorb_multiple::AbsorbMultiple; use crate::library::Library; +use crate::mmr::authentication_struct::derive_challenges::DeriveChallenges; use crate::prelude::BasicSnippet; pub struct RootFromAuthenticationStruct; @@ -38,7 +39,6 @@ impl BasicSnippet for RootFromAuthenticationStruct { } fn code(&self, library: &mut Library) -> Vec { - let absorb_multiple = library.import(Box::new(AbsorbMultiple)); let alpha_challenge_pointer_write = library.kmalloc(EXTENSION_DEGREE as u32); let alpha_challenge_pointer_read = alpha_challenge_pointer_write + bfe!(EXTENSION_DEGREE as u64 - 1); @@ -48,51 +48,19 @@ impl BasicSnippet for RootFromAuthenticationStruct { let gamma_challenge_pointer_write = library.kmalloc(EXTENSION_DEGREE as u32); let gamma_challenge_pointer_read = gamma_challenge_pointer_write + bfe!(EXTENSION_DEGREE as u64 - 1); + let t_digest_pointer_write = library.kmalloc(Digest::LEN as u32); + let t_digest_pointer_read = t_digest_pointer_write + bfe!(Digest::LEN as u64 - 1); + let t_xfe_pointer_write = library.kmalloc(EXTENSION_DEGREE as u32); + let t_xfe_pointer_read = t_xfe_pointer_write + bfe!(EXTENSION_DEGREE as u64 - 1); + let p_pointer_write = library.kmalloc(EXTENSION_DEGREE as u32); + let p_pointer_read = p_pointer_write + bfe!(EXTENSION_DEGREE as u64 - 1); let indexed_leaf_element_size = Self::indexed_leaf_element_type().stack_size(); + let derive_challenges = library.import(Box::new(DeriveChallenges)); let calculate_and_store_challenges = triton_asm!( // _ *auth_struct *indexed_leafs - sponge_init - // _ *auth_struct *indexed_leafs - - read_mem 1 - push 1 - add - // _ *auth_struct indexed_leafs_len *indexed_leafs - - swap 1 - // _ *auth_struct *indexed_leafs indexed_leafs_len - - push {indexed_leaf_element_size} - mul - push 1 - add - // _ *auth_struct *indexed_leafs indexed_leafs_size - - call {absorb_multiple} - // _ *auth_struct - - read_mem 1 - push 1 - add - // _ auth_struct_len *auth_struct - - swap 1 - push {Digest::LEN} - mul - push 1 - add - // _ *auth_struct auth_struct_size - - call {absorb_multiple} - // _ - - sponge_squeeze - // _ w9 w8 w7 w6 w5 w4 w3 w2 w1 w0 - - pop 1 - // _ w9 w8 w7 w6 w5 w4 w3 w2 w1 + call {derive_challenges} hint alpha: XFieldElement = stack[0..3] hint minus_beta: XFieldElement = stack[3..6] hint gamma: XFieldElement = stack[6..9] @@ -165,6 +133,8 @@ impl BasicSnippet for RootFromAuthenticationStruct { pop 1 // _ num_leafs *auth_struct *idx_leafs (*idx_leafs[n]_lw - 2) [0; 2] [p; 3] [γ] [leaf_idx; 2] + // TODO: Assert that `leaf_idx < num_leafs`? + {&u64_to_bfe} // _ num_leafs *auth_struct *idx_leafs (*idx_leafs[n]_lw - 2) [0; 2] [p; 3] [γ] leaf_idx_bfe @@ -214,6 +184,15 @@ impl BasicSnippet for RootFromAuthenticationStruct { swap 1 // _ num_leafs *auth_struct *indexed_leafs *indexed_leafs indexed_leafs_len + /* Disallow empty list of indexed leafs */ + dup 0 + push 0 + eq + push 0 + eq + assert + // _ num_leafs *auth_struct *indexed_leafs *indexed_leafs indexed_leafs_len + push {indexed_leaf_element_size} mul // _ num_leafs *auth_struct *indexed_leafs *indexed_leafs (indexed_leafs_size - 1) @@ -346,21 +325,187 @@ impl BasicSnippet for RootFromAuthenticationStruct { recurse_or_return ); - let set_initial_t_value = triton_asm!( - // _ num_leafs *indexed_leafs *auth_struct [p] + let nd_loop_label = format!("{entrypoint}_nd_loop"); + let dup_top_two_digests = triton_asm![dup 9; Digest::LEN * 2]; + let dup_top_digest = triton_asm![dup 4; Digest::LEN]; + let dup_top_xfe = triton_asm![dup 2; EXTENSION_DEGREE]; + let one_half = BFieldElement::new(2).inverse(); + let nd_loop = triton_asm!( + // _ INVARIANT: _ + {nd_loop_label}: + divine 1 + // _ left_index - /* Set `t` digest to according to the rules: - let mut t = auth_struct - .first() - .copied() - .unwrap_or_else(|| indexed_leafs.first().unwrap().1); - */ - dup 3 - read_mem 1 - pop 1 - // _ num_leafs *indexed_leafs *auth_struct [p] auth_struct_len + divine 1 + // _ left_index right_index + + dup 1 + push 1 + add + eq + // _ l_index_bfe (left_index + 1 == right_index) + + assert + // _ l_index_bfe + + /* Update parent index */ + push {one_half} + mul + hint parent_index: BFieldElement = stack[0..1] + // _ (l_index_bfe / 2) + // _ parent_index <-- rename + + /* Calculate parent digest, preserving child digests */ + divine {Digest::LEN} + hint right: Digest = stack[0..5] + + divine {Digest::LEN} + hint left: Digest = stack[0..5] + + {&dup_top_two_digests} + hash + hint t: Digest = stack[0..5] + // _ parent_index [right] [left] [t] + + {&dup_top_digest} + push {t_digest_pointer_write} + write_mem {Digest::LEN} + pop 1 + // _ parent_index [right] [left] [t] + + {&digest_to_xfe} + hint t_xfe: XFieldElement = stack[0..3] + // _ parent_index [right] [left] [t_xfe] + + {&dup_top_xfe} + push {t_xfe_pointer_write} + write_mem {EXTENSION_DEGREE} + pop 1 + // _ parent_index [right] [left] [t_xfe] + + push {beta_challenge_pointer_read} + read_mem {EXTENSION_DEGREE} + pop 1 + xx_add + // _ parent_index [right] [left] [t_xfe - β] + + dup 13 + push {gamma_challenge_pointer_read} + read_mem {EXTENSION_DEGREE} + pop 1 + // _ parent_index [right] [left] [t_xfe - β] parent_index [γ] + + swap 1 + swap 2 + swap 3 + xb_mul + // _ parent_index [right] [left] [t_xfe - β] [γ * parent_index] + + xx_add + // _ parent_index [right] [left] [t_xfe - β + γ * parent_index] + // _ parent_index [right] [left] [fact_parent] <-- rename + + /* Accumulate `fact_parent` into `p` */ + push {p_pointer_read} + read_mem {EXTENSION_DEGREE} + pop 1 + xx_mul + push {p_pointer_write} + write_mem {EXTENSION_DEGREE} + pop 1 + // _ parent_index [right] [left] + + /* Claculate `fact_1` */ + {&digest_to_xfe} + // _ parent_index [right] [left_xfe] + + push {beta_challenge_pointer_read} + read_mem {EXTENSION_DEGREE} + pop 1 + xx_add + // _ parent_index [right] [left_xfe - β] + + push {gamma_challenge_pointer_read} + read_mem {EXTENSION_DEGREE} + pop 1 + // _ parent_index [right] [left_xfe - β] [γ] + + dup 11 + push 2 + mul + // _ parent_index [right] [left_xfe - β] [γ] left_index + + xb_mul + // _ parent_index [right] [left_xfe - β] [γ * left_index] + + xx_add + // _ parent_index [right] [left_xfe - β + γ * left_index] + // _ parent_index [right] [fact_1] <-- rename + + x_invert + // _ parent_index [right] [fact_1^{-1}] + + /* Divide `fact_1` out of `p` */ + push {p_pointer_read} + read_mem {EXTENSION_DEGREE} + pop 1 + xx_mul + push {p_pointer_write} + write_mem {EXTENSION_DEGREE} + pop 1 + // _ parent_index [right] + + /* Calculate `fact_2` */ + {&digest_to_xfe} + // _ parent_index [right_xfe] + + push {beta_challenge_pointer_read} + read_mem {EXTENSION_DEGREE} + pop 1 + xx_add + // _ parent_index [right_xfe - β] + + push {gamma_challenge_pointer_read} + read_mem {EXTENSION_DEGREE} + pop 1 + // _ parent_index [right_xfe - β] [γ] + + dup 6 + push 2 + mul + push 1 + add + // _ parent_index [right_xfe - β] [γ] right_index + + xb_mul + // _ parent_index [right_xfe - β] [right_index * γ] + + xx_add + // _ parent_index [right_xfe - β + right_index * γ] + // _ parent_index [fact_2] <-- rename + + x_invert + // _ parent_index [fact_2^{-1}] + + /* Divide `fact_2` out of `p` */ + push {p_pointer_read} + read_mem {EXTENSION_DEGREE} + pop 1 + xx_mul + push {p_pointer_write} + write_mem {EXTENSION_DEGREE} + pop 1 + // _ parent_index + + push 1 + eq + skiz + return + + recurse ); + let compare_xfes = DataType::Xfe.compare(); triton_asm!( {entrypoint}: // _ tree_height *auth_struct *indexed_leafs @@ -372,65 +517,146 @@ impl BasicSnippet for RootFromAuthenticationStruct { {&calculate_and_store_challenges} // _ tree_height *auth_struct *indexed_leafs - /* Calculate number of leafs in Merkle tree */ + /* Calculate number of leafs in Merkle tree + Notice that `tree_num_leafs` is not necessarily a u32, but is a + BFE and a power of two whose log_2 value is in the range + [0,63] */ swap 2 push 2 pow swap 2 - hint num_leafs = stack[2] - // _ num_leafs *auth_struct *indexed_leafs + hint tree_num_leafs: BFieldElement = stack[2] + // _ tree_num_leafs *auth_struct *indexed_leafs {&accumulate_indexed_leafs_from_public_data} - // _ num_leafs *auth_struct *indexed_leafs *indexed_leafs [garbage; 2] [p; 3] - // _ num_leafs *auth_struct *indexed_leafs *indexed_leafs [garbage; 2] p2 p1 p0 <-- rename + // _ tree_num_leafs *auth_struct *indexed_leafs *indexed_leafs [garbage; 2] [p; 3] + // _ tree_num_leafs *auth_struct *indexed_leafs *indexed_leafs [garbage; 2] p2 p1 p0 <-- rename /* Prepare for next loop, absorption of auth-struct digests into accumulator */ swap 7 swap 6 pop 1 - // _ num_leafs p0 *auth_struct *indexed_leafs [0; 2] p2 p1 + // _ tree_num_leafs p0 *auth_struct *indexed_leafs [0; 2] p2 p1 dup 5 read_mem 1 push 1 add swap 1 - // _ num_leafs p0 *auth_struct *indexed_leafs [0; 2] p2 p1 *auth_struct auth_struct_len + // _ tree_num_leafs p0 *auth_struct *indexed_leafs [0; 2] p2 p1 *auth_struct auth_struct_len push {Digest::LEN} mul add - // _ num_leafs p0 *auth_struct *indexed_leafs [0; 2] p2 p1 *auth_struct_last_word + // _ tree_num_leafs p0 *auth_struct *indexed_leafs [0; 2] p2 p1 *auth_struct_last_word swap 5 swap 7 - // _ num_leafs *indexed_leafs *auth_struct *auth_struct_last_word [0; 2] p2 p1 p0 - // _ num_leafs *indexed_leafs *auth_struct *auth_struct_last_word [prev; 2] [p] <-- rename + // _ tree_num_leafs *indexed_leafs *auth_struct *auth_struct_last_word [0; 2] p2 p1 p0 + // _ tree_num_leafs *indexed_leafs *auth_struct *auth_struct_last_word [prev; 2] [p] <-- rename dup 6 dup 6 eq push 0 eq - // _ num_leafs *indexed_leafs *auth_struct *auth_struct_last_word [prev; 2] [p] (*auth_struct_last_word != *auth_struct) + // _ tree_num_leafs *indexed_leafs *auth_struct *auth_struct_last_word [prev; 2] [p] (*auth_struct_last_word != *auth_struct) skiz call {accumulate_auth_struct_leafs_from_public_data_label} - // _ num_leafs *indexed_leafs *auth_struct *auth_struct_last_word [prev; 2] [p] + // _ tree_num_leafs *indexed_leafs *auth_struct *auth_struct_last_word [prev; 2] [p] /* Cleanup stack before next loop */ - swap 3 + swap 4 pop 1 - swap 3 + swap 4 pop 1 - swap 3 + swap 4 + pop 2 + // _ tree_num_leafs *indexed_leafs [p] + + /* Set initial t values, from indexed_leafs[0] */ + dup 3 + push {Digest::LEN} + add + read_mem {Digest::LEN} + pop 1 + // _ tree_num_leafs *indexed_leafs [p] [t; 5] + + dup 4 + dup 4 + dup 4 + dup 4 + dup 4 + {&digest_to_xfe} + // _ tree_num_leafs *indexed_leafs [p] [t; 5] [t_xfe] + + /* Write t values, and `p` to static memory */ + push {t_xfe_pointer_write} + write_mem {EXTENSION_DEGREE} pop 1 - // _ num_leafs *indexed_leafs *auth_struct [p] + push {t_digest_pointer_write} + write_mem {Digest::LEN} + pop 1 + // _ tree_num_leafs *indexed_leafs [p] + + push {p_pointer_write} + write_mem {EXTENSION_DEGREE} + pop 2 + // _ tree_num_leafs + + /* Call the ND-loop if tree_num_leafs != 1 */ + push 1 + eq + push 0 + eq + // (tree_num_leafs != 1) + + skiz + call {nd_loop_label} + // _ + + /* Assert that p == t_xfe - beta + gamma */ + break + push {p_pointer_read} + read_mem {EXTENSION_DEGREE} + pop 1 + // _ [p] + + push {t_xfe_pointer_read} + read_mem {EXTENSION_DEGREE} + pop 1 + // _ [p] [t_xfe] + + push {beta_challenge_pointer_read} + read_mem {EXTENSION_DEGREE} + pop 1 + xx_add + // _ [p] [t_xfe - β] + + push {gamma_challenge_pointer_read} + read_mem {EXTENSION_DEGREE} + pop 1 + xx_add + // _ [p] [t_xfe - β + γ] + + {&compare_xfes} + // _ (p == t_xfe - β + γ) + + assert + // _ + + /* Return `t` (digest) */ + push {t_digest_pointer_read} + read_mem {Digest::LEN} + pop 1 + // _ [t] return {&accumulated_indexed_leafs_loop} {&accumulate_auth_struct_leafs_from_public_data} + {&nd_loop} ) } } @@ -445,12 +671,16 @@ mod tests { use rand::rngs::StdRng; use rand::Rng; use rand::SeedableRng; + use twenty_first::prelude::AlgebraicHasher; use twenty_first::prelude::CpuParallel; use twenty_first::prelude::MerkleTree; use twenty_first::prelude::Sponge; + use crate::mmr::authentication_struct::shared::AuthStructIntegrityProof; + use crate::rust_shadowing_helper_functions::input::consume_digest_from_secret_in; use crate::rust_shadowing_helper_functions::list::list_insert; use crate::rust_shadowing_helper_functions::list::load_list_with_copy_elements; + use crate::rust_shadowing_helper_functions::memory::write_to_memory; use crate::snippet_bencher::BenchmarkCase; use crate::traits::procedure::Procedure; use crate::traits::procedure::ProcedureInitialState; @@ -484,6 +714,33 @@ mod tests { challenge * leaf_xfe_lo + leaf_xfe_hi } + fn mimic_use_of_static_memory( + memory: &mut HashMap, + alpha: XFieldElement, + beta: XFieldElement, + gamma: XFieldElement, + t: Digest, + t_xfe: XFieldElement, + p: XFieldElement, + ) { + const ALPHA_POINTER_WRITE: BFieldElement = BFieldElement::new(BFieldElement::P - 4); + const BETA_POINTER_WRITE: BFieldElement = BFieldElement::new(BFieldElement::P - 7); + const GAMMA_POINTER_WRITE: BFieldElement = + BFieldElement::new(BFieldElement::P - 10); + const T_DIGEST_POINTER_WRITE: BFieldElement = + BFieldElement::new(BFieldElement::P - 15); + const T_XFE_POINTER_WRITE: BFieldElement = + BFieldElement::new(BFieldElement::P - 18); + const P_POINTER_WRITE: BFieldElement = BFieldElement::new(BFieldElement::P - 21); + + write_to_memory(ALPHA_POINTER_WRITE, alpha, memory); + write_to_memory(BETA_POINTER_WRITE, beta, memory); + write_to_memory(GAMMA_POINTER_WRITE, gamma, memory); + write_to_memory(T_DIGEST_POINTER_WRITE, t, memory); + write_to_memory(T_XFE_POINTER_WRITE, t_xfe, memory); + write_to_memory(P_POINTER_WRITE, p, memory); + } + assert_eq!( SIZE_OF_INDEXED_LEAFS_ELEMENT, Self::indexed_leaf_element_type().stack_size() @@ -497,20 +754,17 @@ mod tests { *<(u64, Digest)>::decode(&bfes).unwrap() }; let bfes_to_digest = |bfes: [BFieldElement; Digest::LEN]| -> Digest { - println!("bfes_to_digest:\nbfes: {bfes:?}"); *Digest::decode(&bfes[0..Digest::LEN]).unwrap() }; let indexed_leafs: Vec<[BFieldElement; SIZE_OF_INDEXED_LEAFS_ELEMENT]> = load_list_with_copy_elements(indexed_leafs_pointer, memory); - println!("indexed_leafs: {indexed_leafs:?}"); let indexed_leafs = indexed_leafs .into_iter() .map(bfes_to_indexed_leaf) .collect_vec(); let auth_struct: Vec<[BFieldElement; Digest::LEN]> = load_list_with_copy_elements(auth_struct_pointer, memory); - println!("auth_struct: {auth_struct:?}"); let auth_struct = auth_struct.into_iter().map(bfes_to_digest).collect_vec(); let sponge = sponge.as_mut().expect("sponge must be initialized"); @@ -522,28 +776,17 @@ mod tests { let alpha = XFieldElement::new([sponge_output[1], sponge_output[2], sponge_output[3]]); let beta = -XFieldElement::new([sponge_output[4], sponge_output[5], sponge_output[6]]); let gamma = XFieldElement::new([sponge_output[7], sponge_output[8], sponge_output[9]]); - println!("alpha {alpha}"); - println!("-beta {}", -beta); - println!("gamma {}", -gamma); - let num_leafs = 1 << tree_height; + let tree_num_leafs = 1 << tree_height; let mut p = XFieldElement::one(); for (leaf_idx, leaf) in indexed_leafs.iter().copied().rev() { let leaf_idx_as_bfe = bfe!(leaf_idx); - let node_idx_as_bfe = leaf_idx_as_bfe + bfe!(num_leafs); - + let node_idx_as_bfe = leaf_idx_as_bfe + bfe!(tree_num_leafs); let leaf_as_xfe = digest_to_xfe(leaf, alpha); - - println!("gamma * node_idx_as_bfe = {}", gamma * node_idx_as_bfe); - println!("leaf_as_xfe = {}", leaf_as_xfe); - let fact = leaf_as_xfe - beta + gamma * node_idx_as_bfe; - println!("indexed-leaf fact: {fact}"); p *= fact; } - println!("Rust-shadow, p, after indexed-leafs absorption: {p}"); - let mut prev = 0u64; let mut individual_tokens: VecDeque = nondeterminism.individual_tokens.to_owned().into(); @@ -554,7 +797,6 @@ mod tests { individual_tokens.pop_front().unwrap().try_into().unwrap(); let auth_struct_elem_node_index = ((auth_struct_elem_node_index_hi as u64) << 32) + auth_struct_elem_node_index_lo as u64; - println!("auth_struct_elem_node_index: {auth_struct_elem_node_index}"); assert!(auth_struct_elem_node_index > prev); prev = auth_struct_elem_node_index; @@ -563,11 +805,45 @@ mod tests { let auth_struct_elem_xfe = digest_to_xfe(auth_struct_elem, alpha); let fact = auth_struct_elem_xfe - beta + gamma * auth_struct_index_as_bfe; - println!("auth struct fact: {fact}"); p *= fact; } - println!("Rust-shadow, p, after auth-struct absorption: {p}"); + let mut t = indexed_leafs[0].1; + let mut t_xfe = digest_to_xfe(t, alpha); + if tree_num_leafs != 1 { + loop { + let left_index = individual_tokens.pop_front().unwrap(); + let right_index = individual_tokens.pop_front().unwrap(); + assert_eq!(left_index + bfe!(1), right_index); + + let parent_index = left_index / bfe!(2); + + let right = consume_digest_from_secret_in(&mut individual_tokens); + let left = consume_digest_from_secret_in(&mut individual_tokens); + + t = Tip5::hash_pair(left, right); + t_xfe = digest_to_xfe(t, alpha); + let l_xfe = digest_to_xfe(left, alpha); + let r_xfe = digest_to_xfe(right, alpha); + let fact1 = l_xfe - beta + gamma * left_index; + let fact2 = r_xfe - beta + gamma * right_index; + let fact_parent = t_xfe - beta + gamma * parent_index; + + p *= fact1.inverse() * fact2.inverse() * fact_parent; + + if parent_index.is_one() { + break; + } + } + } + + assert_eq!(t_xfe - beta + gamma, p); + + for elem in t.encode().into_iter().rev() { + stack.push(elem); + } + + mimic_use_of_static_memory(memory, alpha, -beta, gamma, t, t_xfe, p); vec![] } @@ -580,25 +856,28 @@ mod tests { let mut rng: StdRng = SeedableRng::from_seed(seed); // TODO: use real `mmr_authentication_struct` code here - let tree_height = rng.gen_range(1..5); + let tree_height = rng.gen_range(0..5); let leaf_count = 1 << tree_height; - let num_revealed_leafs = rng.gen_range(1..leaf_count); + let num_revealed_leafs = rng.gen_range(1..=leaf_count); let revealed_leaf_indices = (0..num_revealed_leafs) .map(|_| rng.gen_range(0..leaf_count)) .unique() .collect_vec(); + println!( + "revealed_leaf_indices: [{}]", + revealed_leaf_indices.iter().join(", ") + ); let leafs = (0..leaf_count).map(|_| rng.gen()).collect_vec(); let tree = MerkleTree::::new::(&leafs).unwrap(); let mmr_authentication_struct = - mmr_authentication_struct::AuthStructIntegrityProof::new_from_merkle_tree( - &tree, - revealed_leaf_indices, - ); + AuthStructIntegrityProof::new_from_merkle_tree(&tree, revealed_leaf_indices); + println!("tree.root() = {}", tree.root()); let mut memory = HashMap::new(); let authentication_structure_ptr = rng.gen(); let indexed_leafs_ptr = rng.gen(); + println!("leaf_count: {leaf_count}"); println!( "indexed_leafs.len(): {}", mmr_authentication_struct.indexed_leafs.len() @@ -637,13 +916,40 @@ mod tests { .iter() .join(", ") ); - let individual_tokens = mmr_authentication_struct + let nd_auth_struct_indices = mmr_authentication_struct .witness .nd_auth_struct_indices .into_iter() .rev() .flat_map(|node_index| node_index.encode().into_iter().rev().collect_vec()) .collect_vec(); + println!( + "nd_auth_struct_indices (encoded as BFEs): {}", + nd_auth_struct_indices.iter().join(", ") + ); + let nd_loop_nd = mmr_authentication_struct + .witness + .nd_sibling_indices + .iter() + .copied() + .zip_eq( + mmr_authentication_struct + .witness + .nd_siblings + .iter() + .copied(), + ) + .flat_map(|((left_index, right_index), (left_node, right_node))| { + [ + vec![bfe!(left_index), bfe!(right_index)], + right_node.encode().into_iter().rev().collect_vec(), + left_node.encode().into_iter().rev().collect_vec(), + ] + .concat() + }) + .collect_vec(); + + let individual_tokens = [nd_auth_struct_indices, nd_loop_nd].concat(); println!("individual_tokens: {}", individual_tokens.iter().join(", ")); let nondeterminism = NonDeterminism::new(individual_tokens).with_ram(memory); ProcedureInitialState { @@ -659,368 +965,3 @@ mod tests { } } } - -// TODO: Use this logic from `twenty-first` instead! -mod mmr_authentication_struct { - use std::collections::{HashMap, HashSet}; - - use itertools::Itertools; - use num::One; - use twenty_first::{ - prelude::{AlgebraicHasher, Inverse, MerkleTree, Mmr, MmrMembershipProof, Sponge}, - util_types::mmr::{ - mmr_accumulator::MmrAccumulator, shared_advanced::get_peak_heights, - shared_basic::leaf_index_to_mt_index_and_peak_index, - }, - }; - - use super::*; - - const ROOT_MT_INDEX: u64 = 1; - - /// A witness to facilitate the proving of the authenticity of a Merkle - /// authentication struct. - #[derive(Debug, Clone)] - pub struct AuthStructIntegrityProof { - // All indices are Merkle tree node indices - pub nd_auth_struct_indices: Vec, - pub nd_sibling_indices: Vec<(u64, u64)>, - pub nd_siblings: Vec<(Digest, Digest)>, - } - - /// An authentication structure that can be used to prove membership of a list - /// of leaves in a Merkle tree, along with the indexed leaves in question, and - /// the witness necessary to prove membership in a ZK program. - #[derive(Debug, Clone)] - pub struct AuthenticatedMerkleAuthStruct { - pub auth_struct: Vec, - pub indexed_leafs: Vec<(u64, Digest)>, - pub witness: AuthStructIntegrityProof, - } - - impl AuthStructIntegrityProof { - /// Return the Merkle tree node indices of the digests required to prove - /// membership for the specified leaf indices, as well as the node indices - /// that can be derived from the leaf indices and their authentication - /// path. - fn auth_struct_and_nd_indices( - num_leafs: u64, - leaf_indices: &[u64], - ) -> (Vec, Vec<(u64, u64)>) { - // The set of indices of nodes that need to be included in the authentications - // structure. In principle, every node of every authentication path is needed. - // The root is never needed. Hence, it is not considered below. - let mut node_is_needed = HashSet::new(); - - // The set of indices of nodes that can be computed from other nodes in the - // authentication structure or the leafs that are explicitly supplied during - // verification. Every node on the direct path from the leaf to the root can - // be computed by the very nature of “authentication path”. - let mut node_can_be_computed = HashSet::new(); - - for &leaf_index in leaf_indices { - assert!(num_leafs > leaf_index, "Leaf index must be less than number of leafs. Got leaf_index = {leaf_index}; num_leafs = {num_leafs}"); - - let mut node_index = leaf_index + num_leafs; - while node_index > ROOT_MT_INDEX { - let sibling_index = node_index ^ 1; - node_can_be_computed.insert(node_index); - node_is_needed.insert(sibling_index); - node_index /= 2; - } - } - - let set_difference = node_is_needed.difference(&node_can_be_computed).copied(); - let set_union = node_is_needed - .union(&node_can_be_computed) - .sorted_unstable() - .rev(); - - let mut set_union = set_union.peekable(); - - let mut set_union_as_ordered_pairs = Vec::new(); - while set_union.peek().is_some() { - let right_index = *set_union.next().unwrap(); - - // Crashes on odd-length of input list, which is what we want, as - // this acts as a sanity check. - let left_index = *set_union.next().unwrap(); - set_union_as_ordered_pairs.push((left_index, right_index)); - } - - ( - set_difference.sorted_unstable().rev().collect(), - set_union_as_ordered_pairs, - ) - } - - pub fn root_from_authentication_struct( - &self, - tree_height: u32, - auth_struct: Vec, - indexed_leafs: Vec<(u64, Digest)>, - ) -> Digest { - fn digest_to_xfe(digest: Digest, challenge: XFieldElement) -> XFieldElement { - let leaf_xfe_lo = XFieldElement::new([digest.0[0], digest.0[1], digest.0[2]]); - let leaf_xfe_hi = challenge - * XFieldElement::new([digest.0[3], digest.0[4], BFieldElement::one()]); - - leaf_xfe_lo + leaf_xfe_hi - } - - fn node_index_to_bfe(node_index: u64) -> BFieldElement { - BFieldElement::new(node_index) - } - - // Sanity check - assert_eq!( - self.nd_auth_struct_indices.len(), - auth_struct.len(), - "Provided auth struct length must match that specified in receiver" - ); - - // Get challenges - let (alpha, beta, gamma) = { - let mut sponge = Tip5::init(); - sponge.pad_and_absorb_all(&indexed_leafs.encode()); - sponge.pad_and_absorb_all(&auth_struct.encode()); - let challenges = sponge.sample_scalars(3); - (challenges[0], challenges[1], challenges[2]) - }; - - // Accumulate `p` from public data - let mut p = XFieldElement::one(); - for i in (0..indexed_leafs.len()).rev() { - let node_index_as_bfe = node_index_to_bfe((1 << tree_height) ^ indexed_leafs[i].0); - let leaf_as_xfe = digest_to_xfe(indexed_leafs[i].1, alpha); - let fact = leaf_as_xfe - beta + gamma * node_index_as_bfe; - p *= fact; - } - - let mut prev = 0; - for i in (0..auth_struct.len()).rev() { - let auth_struct_index = self.nd_auth_struct_indices[i]; - - // `auth_struct` must be sorted high-to-low by node-index. But since - // we're traversing in reverse order, the inequality is flipped. - assert!(auth_struct_index > prev); - prev = auth_struct_index; - - let auth_struct_index_as_bfe = node_index_to_bfe(auth_struct_index); - - let auth_str_elem_as_xfe = digest_to_xfe(auth_struct[i], alpha); - let fact = auth_str_elem_as_xfe - beta + gamma * auth_struct_index_as_bfe; - p *= fact; - } - - // Use secret data to invert `p` back and to calculate the root - let mut t = auth_struct - .first() - .copied() - .unwrap_or_else(|| indexed_leafs.first().unwrap().1); - let mut t_xfe = digest_to_xfe(t, alpha); - let mut parent_index_bfe = BFieldElement::one(); - for ((l, r), (left_index, right_index)) in self - .nd_siblings - .iter() - .zip_eq(self.nd_sibling_indices.clone()) - { - assert_eq!(left_index + 1, right_index); - - t = Tip5::hash_pair(*l, *r); - - let l_xfe = digest_to_xfe(*l, alpha); - let r_xfe = digest_to_xfe(*r, alpha); - t_xfe = digest_to_xfe(t, alpha); - - let left_index_bfe = node_index_to_bfe(left_index); - let right_index_bfe = node_index_to_bfe(right_index); - parent_index_bfe = left_index_bfe / bfe!(2); - - let fact1 = l_xfe - beta + gamma * left_index_bfe; - let fact2 = r_xfe - beta + gamma * right_index_bfe; - let fact_parent = t_xfe - beta + gamma * parent_index_bfe; - - p *= fact1.inverse() * fact2.inverse() * fact_parent; - } - - assert_eq!(t_xfe - beta + gamma, p); - assert!(parent_index_bfe.is_one()); - - t - } - - /// Return the authentication structure authenticity witness, - /// authentication structure, and the (leaf-index, leaf-digest) pairs - /// from a list of MMR membership proofs. All MMR membership proofs must - /// belong under the same peak, i.e., be part of the same Merkle tree in - /// the list of Merkle trees that the MMR contains. - /// - /// Panics if the input list of MMR-membership proofs is empty, or if they - /// do not all belong under the same peak. - pub fn new_from_mmr_membership_proofs( - mmra: &MmrAccumulator, - indexed_mmr_mps: Vec<(u64, Digest, MmrMembershipProof)>, - ) -> HashMap { - #[derive(Clone, Debug)] - struct IndexedAuthenticatedMmrLeaf { - merkle_tree_node_index: u64, - merkle_tree_leaf_index: u64, - leaf_digest: Digest, - membership_proof: MmrMembershipProof, - } - - // Split indexed MMR-mps into a hashmap with one entry for each - // referenced peak in the MMR. - let num_mmr_leafs = mmra.num_leafs(); - let mut peak_index_to_indexed_mmr_mp: HashMap> = - HashMap::default(); - let peak_heights = get_peak_heights(num_mmr_leafs); - for (mmr_leaf_index, leaf, mmr_mp) in indexed_mmr_mps { - let (mt_index, peak_index) = - leaf_index_to_mt_index_and_peak_index(mmr_leaf_index, num_mmr_leafs); - let peak_index_as_usize: usize = peak_index.try_into().unwrap(); - let num_leafs_local_mt = 1 << peak_heights[peak_index_as_usize]; - let mt_leaf_index = mt_index - num_leafs_local_mt; - peak_index_to_indexed_mmr_mp - .entry(peak_index) - .or_default() - .push(IndexedAuthenticatedMmrLeaf { - merkle_tree_node_index: mt_index, - merkle_tree_leaf_index: mt_leaf_index, - leaf_digest: leaf, - membership_proof: mmr_mp, - }); - } - - // Loop over all peaks and collect an authentication witness struct - // for each peak. - let mut peak_index_to_authenticated_auth_struct = HashMap::default(); - for (peak_index, indexed_mmr_mp_structs) in peak_index_to_indexed_mmr_mp { - let peak_index_as_usize: usize = peak_index.try_into().unwrap(); - let num_leafs_in_local_mt = 1 << peak_heights[peak_index_as_usize]; - let local_mt_leaf_indices = indexed_mmr_mp_structs - .iter() - .map(|x| x.merkle_tree_leaf_index) - .collect_vec(); - - let (nd_auth_struct_indices, nd_sibling_indices) = - Self::auth_struct_and_nd_indices(num_leafs_in_local_mt, &local_mt_leaf_indices); - let peak = mmra.peaks()[peak_index_as_usize]; - - let mut node_digests: HashMap = HashMap::default(); - node_digests.insert(ROOT_MT_INDEX, peak); - - // Loop over all indexed leafs for this peak - for indexed_mmr_mp in indexed_mmr_mp_structs.iter() { - let mut mt_node_index = indexed_mmr_mp.merkle_tree_node_index; - let mut node = indexed_mmr_mp.leaf_digest; - - // Loop over all authentication path elements for this indexed leaf - for ap_elem in indexed_mmr_mp.membership_proof.authentication_path.iter() { - node_digests.insert(mt_node_index, node); - node_digests.insert(mt_node_index ^ 1, *ap_elem); - node = if mt_node_index & 1 == 0 { - Tip5::hash_pair(node, *ap_elem) - } else { - Tip5::hash_pair(*ap_elem, node) - }; - - mt_node_index /= 2; - } - - // Sanity check that MMR-MPs are valid - assert_eq!(peak, node, "Derived peak must match provided peak"); - } - let nd_siblings = nd_sibling_indices - .iter() - .map(|(left_idx, right_idx)| (node_digests[left_idx], node_digests[right_idx])) - .collect_vec(); - let auth_struct = nd_auth_struct_indices - .iter() - .map(|idx| node_digests[idx]) - .collect_vec(); - let indexed_leafs = indexed_mmr_mp_structs - .into_iter() - .map(|indexed_mmr_mp| { - ( - indexed_mmr_mp.merkle_tree_leaf_index, - indexed_mmr_mp.leaf_digest, - ) - }) - .collect_vec(); - - let witness = Self { - nd_auth_struct_indices, - nd_sibling_indices, - nd_siblings, - }; - - peak_index_to_authenticated_auth_struct.insert( - peak_index, - AuthenticatedMerkleAuthStruct { - auth_struct, - indexed_leafs, - witness, - }, - ); - } - - peak_index_to_authenticated_auth_struct - } - - /// Return the authentication structure witness, authentication structure, - /// and the (leaf-index, leaf-digest) pairs. - pub fn new_from_merkle_tree( - tree: &MerkleTree, - mut revealed_leaf_indices: Vec, - ) -> AuthenticatedMerkleAuthStruct { - revealed_leaf_indices.sort_unstable(); - revealed_leaf_indices.dedup(); - revealed_leaf_indices.reverse(); - let num_leafs: u64 = tree.num_leafs() as u64; - - let (mut nd_auth_struct_indices, nd_sibling_indices) = - Self::auth_struct_and_nd_indices(num_leafs, &revealed_leaf_indices); - if revealed_leaf_indices.is_empty() { - nd_auth_struct_indices = vec![ROOT_MT_INDEX]; - } - - let nd_siblings = nd_sibling_indices - .iter() - .map(|&(l, r)| { - let l: usize = l.try_into().unwrap(); - let r: usize = r.try_into().unwrap(); - (tree.node(l).unwrap(), tree.node(r).unwrap()) - }) - .collect_vec(); - - let revealed_leafs = revealed_leaf_indices - .iter() - .map(|j| tree.node((*j + num_leafs) as usize).unwrap()) - .collect_vec(); - let indexed_leafs = revealed_leaf_indices - .clone() - .into_iter() - .zip_eq(revealed_leafs) - .collect_vec(); - - let auth_struct = nd_auth_struct_indices - .iter() - .map(|node_index| tree.node(*node_index as usize).unwrap()) - .collect_vec(); - - let witness = Self { - nd_auth_struct_indices, - nd_sibling_indices, - nd_siblings, - }; - - AuthenticatedMerkleAuthStruct { - auth_struct, - indexed_leafs, - witness, - } - } - } -} diff --git a/tasm-lib/src/mmr/authentication_struct/shared.rs b/tasm-lib/src/mmr/authentication_struct/shared.rs index 16758680..d17f4920 100644 --- a/tasm-lib/src/mmr/authentication_struct/shared.rs +++ b/tasm-lib/src/mmr/authentication_struct/shared.rs @@ -4,6 +4,7 @@ use std::collections::HashSet; use itertools::Itertools; use num_traits::One; +use crate::data_type::DataType; use crate::twenty_first::bfe; use crate::twenty_first::prelude::*; use crate::twenty_first::util_types::mmr::mmr_accumulator::MmrAccumulator; @@ -12,6 +13,10 @@ use crate::twenty_first::util_types::mmr::shared_basic::leaf_index_to_mt_index_a const ROOT_MT_INDEX: u64 = 1; +pub(super) fn indexed_leaf_element_type() -> DataType { + DataType::Tuple(vec![DataType::U64, DataType::Digest]) +} + /// A witness to facilitate the proving of the authenticity of a Merkle /// authentication struct. #[derive(Debug, Clone)] @@ -88,6 +93,9 @@ impl AuthStructIntegrityProof { ) } + /// Calculate a root from an authentication structure, indexed leafs, and + /// additional witness data on `self`. Crashes if the witness data is + /// incorrect and if the list of indexed leafs is empty. pub fn root_from_authentication_struct( &self, tree_height: u32, @@ -148,10 +156,7 @@ impl AuthStructIntegrityProof { } // Use secret data to invert `p` back and to calculate the root - let mut t = auth_struct - .first() - .copied() - .unwrap_or_else(|| indexed_leafs.first().unwrap().1); + let mut t = indexed_leafs.first().unwrap().1; let mut t_xfe = digest_to_xfe(t, alpha); let mut parent_index_bfe = BFieldElement::one(); for ((l, r), (left_index, right_index)) in self @@ -159,7 +164,10 @@ impl AuthStructIntegrityProof { .iter() .zip_eq(self.nd_sibling_indices.clone()) { + let left_index_bfe = node_index_to_bfe(left_index); + let right_index_bfe = node_index_to_bfe(right_index); assert_eq!(left_index + 1, right_index); + parent_index_bfe = left_index_bfe / bfe!(2); t = Tip5::hash_pair(*l, *r); @@ -167,10 +175,6 @@ impl AuthStructIntegrityProof { let r_xfe = digest_to_xfe(*r, alpha); t_xfe = digest_to_xfe(t, alpha); - let left_index_bfe = node_index_to_bfe(left_index); - let right_index_bfe = node_index_to_bfe(right_index); - parent_index_bfe = left_index_bfe / bfe!(2); - let fact1 = l_xfe - beta + gamma * left_index_bfe; let fact2 = r_xfe - beta + gamma * right_index_bfe; let fact_parent = t_xfe - beta + gamma * parent_index_bfe; @@ -315,11 +319,8 @@ impl AuthStructIntegrityProof { revealed_leaf_indices.reverse(); let num_leafs: u64 = tree.num_leafs() as u64; - let (mut nd_auth_struct_indices, nd_sibling_indices) = + let (nd_auth_struct_indices, nd_sibling_indices) = Self::auth_struct_and_nd_indices(num_leafs, &revealed_leaf_indices); - if revealed_leaf_indices.is_empty() { - nd_auth_struct_indices = vec![ROOT_MT_INDEX]; - } let nd_siblings = nd_sibling_indices .iter() @@ -597,20 +598,6 @@ mod tests { assert_eq!(tree.root(), calculated_root); } - #[test] - fn root_from_authentication_struct_tree_height_0_no_revealed_leafs() { - let tree_height = 0; - let leaf_indices = vec![]; - let nd_auth_struct_indices = vec![1]; - let nd_sibling_indices = vec![]; - prop_from_merkle_tree( - tree_height, - leaf_indices, - nd_auth_struct_indices, - nd_sibling_indices, - ) - } - #[test] fn root_from_authentication_struct_tree_height_0_1_revealed() { let tree_height = 0; @@ -653,20 +640,6 @@ mod tests { ) } - #[test] - fn root_from_authentication_struct_tree_height_2_0_revealed() { - let tree_height = 2; - let leaf_indices = vec![]; - let auth_struct_indices = vec![1]; - let nd_sibling_indices = vec![]; - prop_from_merkle_tree( - tree_height, - leaf_indices, - auth_struct_indices, - nd_sibling_indices, - ) - } - #[test] fn root_from_authentication_struct_tree_height_2_2_revealed() { let tree_height = 2; diff --git a/tasm-lib/src/rust_shadowing_helper_functions.rs b/tasm-lib/src/rust_shadowing_helper_functions.rs index 306b551a..c3a1a225 100644 --- a/tasm-lib/src/rust_shadowing_helper_functions.rs +++ b/tasm-lib/src/rust_shadowing_helper_functions.rs @@ -3,6 +3,7 @@ pub mod claim; pub mod dyn_malloc; pub mod input; pub mod list; +pub mod memory; /// Count the number of non-leaf nodes that were inserted *prior* to /// the insertion of this leaf. diff --git a/tasm-lib/src/rust_shadowing_helper_functions/input.rs b/tasm-lib/src/rust_shadowing_helper_functions/input.rs index f7f997ad..58552535 100644 --- a/tasm-lib/src/rust_shadowing_helper_functions/input.rs +++ b/tasm-lib/src/rust_shadowing_helper_functions/input.rs @@ -1,3 +1,5 @@ +use std::collections::VecDeque; + use num::Zero; use triton_vm::prelude::*; use twenty_first::math::other::random_elements; @@ -42,15 +44,11 @@ pub fn read_digest_from_std_in(std_in: &[BFieldElement], std_in_cursor: &mut usi Digest::new(values) } -pub fn read_digest_from_secret_in( - secret_in: &[BFieldElement], - secret_in_cursor: &mut usize, -) -> Digest { +pub fn consume_digest_from_secret_in(secret_in: &mut VecDeque) -> Digest { let mut values = [BFieldElement::zero(); Digest::LEN]; let mut i = 0; while i < Digest::LEN { - values[Digest::LEN - 1 - i] = secret_in[*secret_in_cursor]; - *secret_in_cursor += 1; + values[Digest::LEN - 1 - i] = secret_in.pop_front().unwrap(); i += 1; } diff --git a/tasm-lib/src/rust_shadowing_helper_functions/memory.rs b/tasm-lib/src/rust_shadowing_helper_functions/memory.rs new file mode 100644 index 00000000..aaeda8e7 --- /dev/null +++ b/tasm-lib/src/rust_shadowing_helper_functions/memory.rs @@ -0,0 +1,15 @@ +use std::collections::HashMap; + +use triton_vm::prelude::BFieldCodec; +use triton_vm::prelude::BFieldElement; + +pub fn write_to_memory( + mut pointer: BFieldElement, + value: T, + memory: &mut HashMap, +) { + for word in value.encode() { + memory.insert(pointer, word); + pointer.increment(); + } +} From b55b1d09e43adf577c170af015157a5bf41ece06 Mon Sep 17 00:00:00 2001 From: sword_smith Date: Thu, 15 Aug 2024 12:44:05 +0200 Subject: [PATCH 03/10] mmr::auhtentication_struct: Refactor tests Make Rust-shadowing and state initialization nicer. --- .../root_from_authentication_struct.rs | 164 ++++++++++-------- .../src/mmr/authentication_struct/shared.rs | 2 +- .../rust_shadowing_helper_functions/input.rs | 4 +- 3 files changed, 97 insertions(+), 73 deletions(-) diff --git a/tasm-lib/src/mmr/authentication_struct/root_from_authentication_struct.rs b/tasm-lib/src/mmr/authentication_struct/root_from_authentication_struct.rs index 404ba496..2d156f4e 100644 --- a/tasm-lib/src/mmr/authentication_struct/root_from_authentication_struct.rs +++ b/tasm-lib/src/mmr/authentication_struct/root_from_authentication_struct.rs @@ -668,16 +668,16 @@ mod tests { use itertools::Itertools; use num::One; + use num::Zero; use rand::rngs::StdRng; use rand::Rng; use rand::SeedableRng; use twenty_first::prelude::AlgebraicHasher; - use twenty_first::prelude::CpuParallel; - use twenty_first::prelude::MerkleTree; use twenty_first::prelude::Sponge; + use twenty_first::util_types::mmr::mmr_accumulator::util::mmra_with_mps; use crate::mmr::authentication_struct::shared::AuthStructIntegrityProof; - use crate::rust_shadowing_helper_functions::input::consume_digest_from_secret_in; + use crate::rust_shadowing_helper_functions::input::read_digest_from_input; use crate::rust_shadowing_helper_functions::list::list_insert; use crate::rust_shadowing_helper_functions::list::load_list_with_copy_elements; use crate::rust_shadowing_helper_functions::memory::write_to_memory; @@ -703,7 +703,7 @@ mod tests { stack: &mut Vec, memory: &mut std::collections::HashMap, nondeterminism: &NonDeterminism, - public_input: &[BFieldElement], + _public_input: &[BFieldElement], sponge: &mut Option, ) -> Vec { fn digest_to_xfe(digest: Digest, challenge: XFieldElement) -> XFieldElement { @@ -741,10 +741,63 @@ mod tests { write_to_memory(P_POINTER_WRITE, p, memory); } + fn accumulate_indexed_leafs( + indexed_leafs: &[(u64, Digest)], + alpha: XFieldElement, + beta: XFieldElement, + gamma: XFieldElement, + tree_num_leafs: u64, + ) -> XFieldElement { + let mut p = XFieldElement::one(); + for (leaf_idx, leaf) in indexed_leafs.iter().copied().rev() { + let leaf_idx_as_bfe = bfe!(leaf_idx); + let node_idx_as_bfe = leaf_idx_as_bfe + bfe!(tree_num_leafs); + let leaf_as_xfe = digest_to_xfe(leaf, alpha); + let fact = leaf_as_xfe - beta + gamma * node_idx_as_bfe; + p *= fact; + } + + p + } + + fn accumulate_auth_struct( + mut p: XFieldElement, + auth_struct: Vec, + individual_tokens: &mut VecDeque, + alpha: XFieldElement, + beta: XFieldElement, + gamma: XFieldElement, + ) -> XFieldElement { + let mut prev = 0u64; + + for auth_struct_elem in auth_struct.iter().copied().rev() { + let auth_struct_elem_node_index_hi: u32 = + individual_tokens.pop_front().unwrap().try_into().unwrap(); + let auth_struct_elem_node_index_lo: u32 = + individual_tokens.pop_front().unwrap().try_into().unwrap(); + let auth_struct_elem_node_index = ((auth_struct_elem_node_index_hi as u64) + << 32) + + auth_struct_elem_node_index_lo as u64; + assert!(auth_struct_elem_node_index > prev); + prev = auth_struct_elem_node_index; + + let auth_struct_index_as_bfe = bfe!(auth_struct_elem_node_index); + + let auth_struct_elem_xfe = digest_to_xfe(auth_struct_elem, alpha); + let fact = auth_struct_elem_xfe - beta + gamma * auth_struct_index_as_bfe; + + p *= fact; + } + + p + } + assert_eq!( SIZE_OF_INDEXED_LEAFS_ELEMENT, Self::indexed_leaf_element_type().stack_size() ); + + // declare input-arguments let indexed_leafs_pointer = stack.pop().unwrap(); let auth_struct_pointer = stack.pop().unwrap(); let tree_height: u32 = stack.pop().unwrap().try_into().unwrap(); @@ -767,8 +820,8 @@ mod tests { load_list_with_copy_elements(auth_struct_pointer, memory); let auth_struct = auth_struct.into_iter().map(bfes_to_digest).collect_vec(); + // Calculate challenges let sponge = sponge.as_mut().expect("sponge must be initialized"); - sponge.pad_and_absorb_all(&indexed_leafs.encode()); sponge.pad_and_absorb_all(&auth_struct.encode()); @@ -777,37 +830,17 @@ mod tests { let beta = -XFieldElement::new([sponge_output[4], sponge_output[5], sponge_output[6]]); let gamma = XFieldElement::new([sponge_output[7], sponge_output[8], sponge_output[9]]); - let tree_num_leafs = 1 << tree_height; - let mut p = XFieldElement::one(); - for (leaf_idx, leaf) in indexed_leafs.iter().copied().rev() { - let leaf_idx_as_bfe = bfe!(leaf_idx); - let node_idx_as_bfe = leaf_idx_as_bfe + bfe!(tree_num_leafs); - let leaf_as_xfe = digest_to_xfe(leaf, alpha); - let fact = leaf_as_xfe - beta + gamma * node_idx_as_bfe; - p *= fact; - } + let tree_num_leafs = 1u64 << tree_height; + + // Accumulate into `p` from public data + let mut p = + accumulate_indexed_leafs(&indexed_leafs, alpha, beta, gamma, tree_num_leafs); - let mut prev = 0u64; let mut individual_tokens: VecDeque = nondeterminism.individual_tokens.to_owned().into(); - for auth_struct_elem in auth_struct.iter().copied().rev() { - let auth_struct_elem_node_index_hi: u32 = - individual_tokens.pop_front().unwrap().try_into().unwrap(); - let auth_struct_elem_node_index_lo: u32 = - individual_tokens.pop_front().unwrap().try_into().unwrap(); - let auth_struct_elem_node_index = ((auth_struct_elem_node_index_hi as u64) << 32) - + auth_struct_elem_node_index_lo as u64; - assert!(auth_struct_elem_node_index > prev); - prev = auth_struct_elem_node_index; - - let auth_struct_index_as_bfe = bfe!(auth_struct_elem_node_index); - - let auth_struct_elem_xfe = digest_to_xfe(auth_struct_elem, alpha); - let fact = auth_struct_elem_xfe - beta + gamma * auth_struct_index_as_bfe; - - p *= fact; - } + p = accumulate_auth_struct(p, auth_struct, &mut individual_tokens, alpha, beta, gamma); + // "Unaccumulate" into `p` from secret data, and calculate Merkle root let mut t = indexed_leafs[0].1; let mut t_xfe = digest_to_xfe(t, alpha); if tree_num_leafs != 1 { @@ -818,8 +851,8 @@ mod tests { let parent_index = left_index / bfe!(2); - let right = consume_digest_from_secret_in(&mut individual_tokens); - let left = consume_digest_from_secret_in(&mut individual_tokens); + let right = read_digest_from_input(&mut individual_tokens); + let left = read_digest_from_input(&mut individual_tokens); t = Tip5::hash_pair(left, right); t_xfe = digest_to_xfe(t, alpha); @@ -839,6 +872,7 @@ mod tests { assert_eq!(t_xfe - beta + gamma, p); + // Return the Merkle root on the stack for elem in t.encode().into_iter().rev() { stack.push(elem); } @@ -855,38 +889,39 @@ mod tests { ) -> ProcedureInitialState { let mut rng: StdRng = SeedableRng::from_seed(seed); - // TODO: use real `mmr_authentication_struct` code here let tree_height = rng.gen_range(0..5); - let leaf_count = 1 << tree_height; - let num_revealed_leafs = rng.gen_range(1..=leaf_count); + let num_leafs_in_merkle_tree = 1 << tree_height; + let num_revealed_leafs = rng.gen_range(1..=num_leafs_in_merkle_tree); let revealed_leaf_indices = (0..num_revealed_leafs) - .map(|_| rng.gen_range(0..leaf_count)) + .map(|_| rng.gen_range(0..num_leafs_in_merkle_tree)) .unique() .collect_vec(); - println!( - "revealed_leaf_indices: [{}]", - revealed_leaf_indices.iter().join(", ") - ); - let leafs = (0..leaf_count).map(|_| rng.gen()).collect_vec(); - let tree = MerkleTree::::new::(&leafs).unwrap(); + let num_revealed_leafs = revealed_leaf_indices.len(); + assert!(!num_revealed_leafs.is_zero()); + + let revealed_leafs: Vec = + (0..num_revealed_leafs).map(|_| rng.gen()).collect_vec(); + let indexed_leafs = revealed_leaf_indices + .into_iter() + .zip_eq(revealed_leafs) + .collect_vec(); + + let (mmra, mps) = mmra_with_mps(num_leafs_in_merkle_tree, indexed_leafs.clone()); + let indexed_mmr_mps = indexed_leafs + .into_iter() + .zip_eq(mps) + .map(|((idx, leaf), mp)| (idx, leaf, mp)) + .collect_vec(); + let mmr_authentication_struct = - AuthStructIntegrityProof::new_from_merkle_tree(&tree, revealed_leaf_indices); - println!("tree.root() = {}", tree.root()); + AuthStructIntegrityProof::new_from_mmr_membership_proofs(&mmra, indexed_mmr_mps); + assert!(mmr_authentication_struct.len().is_one()); + let mmr_authentication_struct = &mmr_authentication_struct[&0]; let mut memory = HashMap::new(); let authentication_structure_ptr = rng.gen(); let indexed_leafs_ptr = rng.gen(); - println!("leaf_count: {leaf_count}"); - println!( - "indexed_leafs.len(): {}", - mmr_authentication_struct.indexed_leafs.len() - ); - println!( - "authentication_structure.len(): {}", - mmr_authentication_struct.auth_struct.len() - ); - list_insert( authentication_structure_ptr, mmr_authentication_struct.auth_struct.clone(), @@ -894,7 +929,7 @@ mod tests { ); list_insert( indexed_leafs_ptr, - mmr_authentication_struct.indexed_leafs, + mmr_authentication_struct.indexed_leafs.clone(), &mut memory, ); @@ -908,25 +943,13 @@ mod tests { ] .concat(); - println!( - "node indices: {}", - mmr_authentication_struct - .witness - .nd_auth_struct_indices - .iter() - .join(", ") - ); let nd_auth_struct_indices = mmr_authentication_struct .witness .nd_auth_struct_indices - .into_iter() + .iter() .rev() .flat_map(|node_index| node_index.encode().into_iter().rev().collect_vec()) .collect_vec(); - println!( - "nd_auth_struct_indices (encoded as BFEs): {}", - nd_auth_struct_indices.iter().join(", ") - ); let nd_loop_nd = mmr_authentication_struct .witness .nd_sibling_indices @@ -950,7 +973,6 @@ mod tests { .collect_vec(); let individual_tokens = [nd_auth_struct_indices, nd_loop_nd].concat(); - println!("individual_tokens: {}", individual_tokens.iter().join(", ")); let nondeterminism = NonDeterminism::new(individual_tokens).with_ram(memory); ProcedureInitialState { stack, diff --git a/tasm-lib/src/mmr/authentication_struct/shared.rs b/tasm-lib/src/mmr/authentication_struct/shared.rs index d17f4920..58f371df 100644 --- a/tasm-lib/src/mmr/authentication_struct/shared.rs +++ b/tasm-lib/src/mmr/authentication_struct/shared.rs @@ -529,7 +529,7 @@ mod tests { #[proptest(cases = 20)] fn root_from_authentication_struct_prop_test( #[strategy(0..12u64)] tree_height: u64, - #[strategy(0usize..100)] _num_revealed_leafs: usize, + #[strategy(1usize..100)] _num_revealed_leafs: usize, #[strategy(vec(0u64..1<<#tree_height, #_num_revealed_leafs))] revealed_leaf_indices: Vec< u64, >, diff --git a/tasm-lib/src/rust_shadowing_helper_functions/input.rs b/tasm-lib/src/rust_shadowing_helper_functions/input.rs index 58552535..a766bddf 100644 --- a/tasm-lib/src/rust_shadowing_helper_functions/input.rs +++ b/tasm-lib/src/rust_shadowing_helper_functions/input.rs @@ -44,7 +44,9 @@ pub fn read_digest_from_std_in(std_in: &[BFieldElement], std_in_cursor: &mut usi Digest::new(values) } -pub fn consume_digest_from_secret_in(secret_in: &mut VecDeque) -> Digest { +/// Read and consume a digest from an input source, either public input or +/// secret input. Returns the digest. +pub fn read_digest_from_input(secret_in: &mut VecDeque) -> Digest { let mut values = [BFieldElement::zero(); Digest::LEN]; let mut i = 0; while i < Digest::LEN { From d6fc65d37f0958d6c5872ff5fbfd396ea8c1f134 Mon Sep 17 00:00:00 2001 From: sword_smith Date: Thu, 15 Aug 2024 13:00:47 +0200 Subject: [PATCH 04/10] mmr::authentication_struct: Add benchmarks for TASM program --- ...b_mmr_root_from_authentication_struct.json | 24 ++++++++++++++++ .../root_from_authentication_struct.rs | 28 +++++++++++++++++-- 2 files changed, 49 insertions(+), 3 deletions(-) create mode 100644 tasm-lib/benchmarks/tasmlib_mmr_root_from_authentication_struct.json diff --git a/tasm-lib/benchmarks/tasmlib_mmr_root_from_authentication_struct.json b/tasm-lib/benchmarks/tasmlib_mmr_root_from_authentication_struct.json new file mode 100644 index 00000000..a7a1a6c9 --- /dev/null +++ b/tasm-lib/benchmarks/tasmlib_mmr_root_from_authentication_struct.json @@ -0,0 +1,24 @@ +[ + { + "name": "tasmlib_mmr_root_from_authentication_struct", + "benchmark_result": { + "clock_cycle_count": 93344, + "hash_table_height": 5437, + "u32_table_height": 11235, + "op_stack_table_height": 140194, + "ram_table_height": 39822 + }, + "case": "CommonCase" + }, + { + "name": "tasmlib_mmr_root_from_authentication_struct", + "benchmark_result": { + "clock_cycle_count": 193548, + "hash_table_height": 10771, + "u32_table_height": 41115, + "op_stack_table_height": 290804, + "ram_table_height": 82518 + }, + "case": "WorstCase" + } +] \ No newline at end of file diff --git a/tasm-lib/src/mmr/authentication_struct/root_from_authentication_struct.rs b/tasm-lib/src/mmr/authentication_struct/root_from_authentication_struct.rs index 2d156f4e..198f0935 100644 --- a/tasm-lib/src/mmr/authentication_struct/root_from_authentication_struct.rs +++ b/tasm-lib/src/mmr/authentication_struct/root_from_authentication_struct.rs @@ -617,7 +617,6 @@ impl BasicSnippet for RootFromAuthenticationStruct { // _ /* Assert that p == t_xfe - beta + gamma */ - break push {p_pointer_read} read_mem {EXTENSION_DEGREE} pop 1 @@ -663,6 +662,7 @@ impl BasicSnippet for RootFromAuthenticationStruct { #[cfg(test)] mod tests { + use std::cmp::min; use std::collections::HashMap; use std::collections::VecDeque; @@ -889,9 +889,18 @@ mod tests { ) -> ProcedureInitialState { let mut rng: StdRng = SeedableRng::from_seed(seed); - let tree_height = rng.gen_range(0..5); + let (tree_height, num_revealed_leafs) = match bench_case { + None => { + let tree_height = rng.gen_range(0..32); + let num_leafs_in_merkle_tree = 1 << tree_height; + let num_revealed_leafs = rng.gen_range(1..=min(num_leafs_in_merkle_tree, 20)); + (tree_height, num_revealed_leafs) + } + Some(BenchmarkCase::CommonCase) => (32, 20), + Some(BenchmarkCase::WorstCase) => (62, 20), + }; let num_leafs_in_merkle_tree = 1 << tree_height; - let num_revealed_leafs = rng.gen_range(1..=num_leafs_in_merkle_tree); + let revealed_leaf_indices = (0..num_revealed_leafs) .map(|_| rng.gen_range(0..num_leafs_in_merkle_tree)) .unique() @@ -987,3 +996,16 @@ mod tests { } } } + +#[cfg(test)] +mod benches { + use crate::traits::procedure::ShadowedProcedure; + use crate::traits::rust_shadow::RustShadow; + + use super::*; + + #[test] + fn bench_root_from_auth_struct() { + ShadowedProcedure::new(RootFromAuthenticationStruct).bench(); + } +} From ea1c289ff23cf3a6a5263bcd0ee4394cdc277c97 Mon Sep 17 00:00:00 2001 From: sword_smith Date: Thu, 15 Aug 2024 13:41:14 +0200 Subject: [PATCH 05/10] perf(mmr::authentication_struct): Read ND indices in one go --- .../tasmlib_mmr_root_from_authentication_struct.json | 8 ++++---- .../root_from_authentication_struct.rs | 5 +---- 2 files changed, 5 insertions(+), 8 deletions(-) diff --git a/tasm-lib/benchmarks/tasmlib_mmr_root_from_authentication_struct.json b/tasm-lib/benchmarks/tasmlib_mmr_root_from_authentication_struct.json index a7a1a6c9..6a8f9829 100644 --- a/tasm-lib/benchmarks/tasmlib_mmr_root_from_authentication_struct.json +++ b/tasm-lib/benchmarks/tasmlib_mmr_root_from_authentication_struct.json @@ -2,8 +2,8 @@ { "name": "tasmlib_mmr_root_from_authentication_struct", "benchmark_result": { - "clock_cycle_count": 93344, - "hash_table_height": 5437, + "clock_cycle_count": 92793, + "hash_table_height": 5431, "u32_table_height": 11235, "op_stack_table_height": 140194, "ram_table_height": 39822 @@ -13,8 +13,8 @@ { "name": "tasmlib_mmr_root_from_authentication_struct", "benchmark_result": { - "clock_cycle_count": 193548, - "hash_table_height": 10771, + "clock_cycle_count": 192404, + "hash_table_height": 10765, "u32_table_height": 41115, "op_stack_table_height": 290804, "ram_table_height": 82518 diff --git a/tasm-lib/src/mmr/authentication_struct/root_from_authentication_struct.rs b/tasm-lib/src/mmr/authentication_struct/root_from_authentication_struct.rs index 198f0935..626719d6 100644 --- a/tasm-lib/src/mmr/authentication_struct/root_from_authentication_struct.rs +++ b/tasm-lib/src/mmr/authentication_struct/root_from_authentication_struct.rs @@ -333,10 +333,7 @@ impl BasicSnippet for RootFromAuthenticationStruct { let nd_loop = triton_asm!( // _ INVARIANT: _ {nd_loop_label}: - divine 1 - // _ left_index - - divine 1 + divine 2 // _ left_index right_index dup 1 From 030782c292a5166a1a65aef674a569691a0a079d Mon Sep 17 00:00:00 2001 From: sword_smith Date: Thu, 15 Aug 2024 14:12:32 +0200 Subject: [PATCH 06/10] perf(mmr::authentication_struct): Don't store t_xfe to memory Instead of storing this value to memory in each loop iteration, we just calculate it from `t: Digest` when we need it. --- ...b_mmr_root_from_authentication_struct.json | 16 ++++----- .../root_from_authentication_struct.rs | 35 ++++--------------- 2 files changed, 14 insertions(+), 37 deletions(-) diff --git a/tasm-lib/benchmarks/tasmlib_mmr_root_from_authentication_struct.json b/tasm-lib/benchmarks/tasmlib_mmr_root_from_authentication_struct.json index 6a8f9829..6e012149 100644 --- a/tasm-lib/benchmarks/tasmlib_mmr_root_from_authentication_struct.json +++ b/tasm-lib/benchmarks/tasmlib_mmr_root_from_authentication_struct.json @@ -2,22 +2,22 @@ { "name": "tasmlib_mmr_root_from_authentication_struct", "benchmark_result": { - "clock_cycle_count": 92793, - "hash_table_height": 5431, + "clock_cycle_count": 89479, + "hash_table_height": 5419, "u32_table_height": 11235, - "op_stack_table_height": 140194, - "ram_table_height": 39822 + "op_stack_table_height": 135778, + "ram_table_height": 38168 }, "case": "CommonCase" }, { "name": "tasmlib_mmr_root_from_authentication_struct", "benchmark_result": { - "clock_cycle_count": 192404, - "hash_table_height": 10765, + "clock_cycle_count": 185532, + "hash_table_height": 10753, "u32_table_height": 41115, - "op_stack_table_height": 290804, - "ram_table_height": 82518 + "op_stack_table_height": 281644, + "ram_table_height": 79085 }, "case": "WorstCase" } diff --git a/tasm-lib/src/mmr/authentication_struct/root_from_authentication_struct.rs b/tasm-lib/src/mmr/authentication_struct/root_from_authentication_struct.rs index 626719d6..b4adf98d 100644 --- a/tasm-lib/src/mmr/authentication_struct/root_from_authentication_struct.rs +++ b/tasm-lib/src/mmr/authentication_struct/root_from_authentication_struct.rs @@ -50,8 +50,6 @@ impl BasicSnippet for RootFromAuthenticationStruct { gamma_challenge_pointer_write + bfe!(EXTENSION_DEGREE as u64 - 1); let t_digest_pointer_write = library.kmalloc(Digest::LEN as u32); let t_digest_pointer_read = t_digest_pointer_write + bfe!(Digest::LEN as u64 - 1); - let t_xfe_pointer_write = library.kmalloc(EXTENSION_DEGREE as u32); - let t_xfe_pointer_read = t_xfe_pointer_write + bfe!(EXTENSION_DEGREE as u64 - 1); let p_pointer_write = library.kmalloc(EXTENSION_DEGREE as u32); let p_pointer_read = p_pointer_write + bfe!(EXTENSION_DEGREE as u64 - 1); @@ -328,7 +326,6 @@ impl BasicSnippet for RootFromAuthenticationStruct { let nd_loop_label = format!("{entrypoint}_nd_loop"); let dup_top_two_digests = triton_asm![dup 9; Digest::LEN * 2]; let dup_top_digest = triton_asm![dup 4; Digest::LEN]; - let dup_top_xfe = triton_asm![dup 2; EXTENSION_DEGREE]; let one_half = BFieldElement::new(2).inverse(); let nd_loop = triton_asm!( // _ INVARIANT: _ @@ -374,12 +371,6 @@ impl BasicSnippet for RootFromAuthenticationStruct { hint t_xfe: XFieldElement = stack[0..3] // _ parent_index [right] [left] [t_xfe] - {&dup_top_xfe} - push {t_xfe_pointer_write} - write_mem {EXTENSION_DEGREE} - pop 1 - // _ parent_index [right] [left] [t_xfe] - push {beta_challenge_pointer_read} read_mem {EXTENSION_DEGREE} pop 1 @@ -580,18 +571,7 @@ impl BasicSnippet for RootFromAuthenticationStruct { pop 1 // _ tree_num_leafs *indexed_leafs [p] [t; 5] - dup 4 - dup 4 - dup 4 - dup 4 - dup 4 - {&digest_to_xfe} - // _ tree_num_leafs *indexed_leafs [p] [t; 5] [t_xfe] - - /* Write t values, and `p` to static memory */ - push {t_xfe_pointer_write} - write_mem {EXTENSION_DEGREE} - pop 1 + /* Write t value, and `p` to static memory */ push {t_digest_pointer_write} write_mem {Digest::LEN} pop 1 @@ -619,9 +599,10 @@ impl BasicSnippet for RootFromAuthenticationStruct { pop 1 // _ [p] - push {t_xfe_pointer_read} - read_mem {EXTENSION_DEGREE} + push {t_digest_pointer_read} + read_mem {Digest::LEN} pop 1 + {&digest_to_xfe} // _ [p] [t_xfe] push {beta_challenge_pointer_read} @@ -717,7 +698,6 @@ mod tests { beta: XFieldElement, gamma: XFieldElement, t: Digest, - t_xfe: XFieldElement, p: XFieldElement, ) { const ALPHA_POINTER_WRITE: BFieldElement = BFieldElement::new(BFieldElement::P - 4); @@ -726,15 +706,12 @@ mod tests { BFieldElement::new(BFieldElement::P - 10); const T_DIGEST_POINTER_WRITE: BFieldElement = BFieldElement::new(BFieldElement::P - 15); - const T_XFE_POINTER_WRITE: BFieldElement = - BFieldElement::new(BFieldElement::P - 18); - const P_POINTER_WRITE: BFieldElement = BFieldElement::new(BFieldElement::P - 21); + const P_POINTER_WRITE: BFieldElement = BFieldElement::new(BFieldElement::P - 18); write_to_memory(ALPHA_POINTER_WRITE, alpha, memory); write_to_memory(BETA_POINTER_WRITE, beta, memory); write_to_memory(GAMMA_POINTER_WRITE, gamma, memory); write_to_memory(T_DIGEST_POINTER_WRITE, t, memory); - write_to_memory(T_XFE_POINTER_WRITE, t_xfe, memory); write_to_memory(P_POINTER_WRITE, p, memory); } @@ -874,7 +851,7 @@ mod tests { stack.push(elem); } - mimic_use_of_static_memory(memory, alpha, -beta, gamma, t, t_xfe, p); + mimic_use_of_static_memory(memory, alpha, -beta, gamma, t, p); vec![] } From 6e9f0e0ec3d3611a110a12b9a5d960b9d90b6222 Mon Sep 17 00:00:00 2001 From: sword_smith Date: Thu, 15 Aug 2024 14:32:04 +0200 Subject: [PATCH 07/10] perf(mmr::authentication_path): Save two `mul` by buffering left_index, not parent_index It's a bit more performant to store `left_index` on the stack and then calculate from it `parent_index` and `right_index` then to store `parent_index` as we did before and calculate the left and right index from that. --- ...b_mmr_root_from_authentication_struct.json | 12 +-- .../root_from_authentication_struct.rs | 79 +++++++++---------- .../src/mmr/authentication_struct/shared.rs | 6 +- 3 files changed, 45 insertions(+), 52 deletions(-) diff --git a/tasm-lib/benchmarks/tasmlib_mmr_root_from_authentication_struct.json b/tasm-lib/benchmarks/tasmlib_mmr_root_from_authentication_struct.json index 6e012149..295148f3 100644 --- a/tasm-lib/benchmarks/tasmlib_mmr_root_from_authentication_struct.json +++ b/tasm-lib/benchmarks/tasmlib_mmr_root_from_authentication_struct.json @@ -2,10 +2,10 @@ { "name": "tasmlib_mmr_root_from_authentication_struct", "benchmark_result": { - "clock_cycle_count": 89479, - "hash_table_height": 5419, + "clock_cycle_count": 87275, + "hash_table_height": 5413, "u32_table_height": 11235, - "op_stack_table_height": 135778, + "op_stack_table_height": 133574, "ram_table_height": 38168 }, "case": "CommonCase" @@ -13,10 +13,10 @@ { "name": "tasmlib_mmr_root_from_authentication_struct", "benchmark_result": { - "clock_cycle_count": 185532, - "hash_table_height": 10753, + "clock_cycle_count": 180956, + "hash_table_height": 10747, "u32_table_height": 41115, - "op_stack_table_height": 281644, + "op_stack_table_height": 277068, "ram_table_height": 79085 }, "case": "WorstCase" diff --git a/tasm-lib/src/mmr/authentication_struct/root_from_authentication_struct.rs b/tasm-lib/src/mmr/authentication_struct/root_from_authentication_struct.rs index b4adf98d..32d672a6 100644 --- a/tasm-lib/src/mmr/authentication_struct/root_from_authentication_struct.rs +++ b/tasm-lib/src/mmr/authentication_struct/root_from_authentication_struct.rs @@ -340,15 +340,9 @@ impl BasicSnippet for RootFromAuthenticationStruct { // _ l_index_bfe (left_index + 1 == right_index) assert + hint left_index: BFieldElement = stack[0..1] // _ l_index_bfe - /* Update parent index */ - push {one_half} - mul - hint parent_index: BFieldElement = stack[0..1] - // _ (l_index_bfe / 2) - // _ parent_index <-- rename - /* Calculate parent digest, preserving child digests */ divine {Digest::LEN} hint right: Digest = stack[0..5] @@ -359,39 +353,41 @@ impl BasicSnippet for RootFromAuthenticationStruct { {&dup_top_two_digests} hash hint t: Digest = stack[0..5] - // _ parent_index [right] [left] [t] + // _ left_index [right] [left] [t] {&dup_top_digest} push {t_digest_pointer_write} write_mem {Digest::LEN} pop 1 - // _ parent_index [right] [left] [t] + // _ left_index [right] [left] [t] {&digest_to_xfe} hint t_xfe: XFieldElement = stack[0..3] - // _ parent_index [right] [left] [t_xfe] + // _ left_index [right] [left] [t_xfe] push {beta_challenge_pointer_read} read_mem {EXTENSION_DEGREE} pop 1 xx_add - // _ parent_index [right] [left] [t_xfe - β] + // _ left_index [right] [left] [t_xfe - β] dup 13 + push {one_half} + mul push {gamma_challenge_pointer_read} read_mem {EXTENSION_DEGREE} pop 1 - // _ parent_index [right] [left] [t_xfe - β] parent_index [γ] + // _ left_index [right] [left] [t_xfe - β] parent_index [γ] swap 1 swap 2 swap 3 xb_mul - // _ parent_index [right] [left] [t_xfe - β] [γ * parent_index] + // _ left_index [right] [left] [t_xfe - β] [γ * parent_index] xx_add - // _ parent_index [right] [left] [t_xfe - β + γ * parent_index] - // _ parent_index [right] [left] [fact_parent] <-- rename + // _ left_index [right] [left] [t_xfe - β + γ * parent_index] + // _ left_index [right] [left] [fact_parent] <-- rename /* Accumulate `fact_parent` into `p` */ push {p_pointer_read} @@ -401,39 +397,37 @@ impl BasicSnippet for RootFromAuthenticationStruct { push {p_pointer_write} write_mem {EXTENSION_DEGREE} pop 1 - // _ parent_index [right] [left] + // _ left_index [right] [left] - /* Claculate `fact_1` */ + /* Calculate `fact_left` */ {&digest_to_xfe} - // _ parent_index [right] [left_xfe] + // _ left_index [right] [left_xfe] push {beta_challenge_pointer_read} read_mem {EXTENSION_DEGREE} pop 1 xx_add - // _ parent_index [right] [left_xfe - β] + // _ left_index [right] [left_xfe - β] push {gamma_challenge_pointer_read} read_mem {EXTENSION_DEGREE} pop 1 - // _ parent_index [right] [left_xfe - β] [γ] + // _ left_index [right] [left_xfe - β] [γ] dup 11 - push 2 - mul - // _ parent_index [right] [left_xfe - β] [γ] left_index + // _ left_index [right] [left_xfe - β] [γ] left_index xb_mul - // _ parent_index [right] [left_xfe - β] [γ * left_index] + // _ left_index [right] [left_xfe - β] [γ * left_index] xx_add - // _ parent_index [right] [left_xfe - β + γ * left_index] - // _ parent_index [right] [fact_1] <-- rename + // _ left_index [right] [left_xfe - β + γ * left_index] + // _ left_index [right] [fact_left] <-- rename x_invert - // _ parent_index [right] [fact_1^{-1}] + // _ left_index [right] [fact_left^{-1}] - /* Divide `fact_1` out of `p` */ + /* Divide `fact_left` out of `p` */ push {p_pointer_read} read_mem {EXTENSION_DEGREE} pop 1 @@ -441,41 +435,39 @@ impl BasicSnippet for RootFromAuthenticationStruct { push {p_pointer_write} write_mem {EXTENSION_DEGREE} pop 1 - // _ parent_index [right] + // _ left_index [right] - /* Calculate `fact_2` */ + /* Calculate `fact_right` */ {&digest_to_xfe} - // _ parent_index [right_xfe] + // _ left_index [right_xfe] push {beta_challenge_pointer_read} read_mem {EXTENSION_DEGREE} pop 1 xx_add - // _ parent_index [right_xfe - β] + // _ left_index [right_xfe - β] push {gamma_challenge_pointer_read} read_mem {EXTENSION_DEGREE} pop 1 - // _ parent_index [right_xfe - β] [γ] + // _ left_index [right_xfe - β] [γ] dup 6 - push 2 - mul push 1 add - // _ parent_index [right_xfe - β] [γ] right_index + // _ left_index [right_xfe - β] [γ] right_index xb_mul - // _ parent_index [right_xfe - β] [right_index * γ] + // _ left_index [right_xfe - β] [right_index * γ] xx_add - // _ parent_index [right_xfe - β + right_index * γ] - // _ parent_index [fact_2] <-- rename + // _ left_index [right_xfe - β + right_index * γ] + // _ left_index [fact_right] <-- rename x_invert - // _ parent_index [fact_2^{-1}] + // _ left_index [fact_right^{-1}] - /* Divide `fact_2` out of `p` */ + /* Divide `fact_right` out of `p` */ push {p_pointer_read} read_mem {EXTENSION_DEGREE} pop 1 @@ -483,9 +475,10 @@ impl BasicSnippet for RootFromAuthenticationStruct { push {p_pointer_write} write_mem {EXTENSION_DEGREE} pop 1 - // _ parent_index + // _ left_index - push 1 + /* Terminate loop when left_index == 2 <=> parent_index == 1 */ + push 2 eq skiz return diff --git a/tasm-lib/src/mmr/authentication_struct/shared.rs b/tasm-lib/src/mmr/authentication_struct/shared.rs index 58f371df..5418c063 100644 --- a/tasm-lib/src/mmr/authentication_struct/shared.rs +++ b/tasm-lib/src/mmr/authentication_struct/shared.rs @@ -175,11 +175,11 @@ impl AuthStructIntegrityProof { let r_xfe = digest_to_xfe(*r, alpha); t_xfe = digest_to_xfe(t, alpha); - let fact1 = l_xfe - beta + gamma * left_index_bfe; - let fact2 = r_xfe - beta + gamma * right_index_bfe; + let fact_left = l_xfe - beta + gamma * left_index_bfe; + let fact_right = r_xfe - beta + gamma * right_index_bfe; let fact_parent = t_xfe - beta + gamma * parent_index_bfe; - p *= fact1.inverse() * fact2.inverse() * fact_parent; + p *= fact_left.inverse() * fact_right.inverse() * fact_parent; } assert_eq!(t_xfe - beta + gamma, p); From 5c6d43600f003e3027b8b48a776d097434f14777 Mon Sep 17 00:00:00 2001 From: sword_smith Date: Thu, 15 Aug 2024 15:16:32 +0200 Subject: [PATCH 08/10] perf(mmr::authentication_struct): Keep `p` on stack in nd-loop --- ...b_mmr_root_from_authentication_struct.json | 16 +- .../root_from_authentication_struct.rs | 198 +++++++++--------- 2 files changed, 110 insertions(+), 104 deletions(-) diff --git a/tasm-lib/benchmarks/tasmlib_mmr_root_from_authentication_struct.json b/tasm-lib/benchmarks/tasmlib_mmr_root_from_authentication_struct.json index 295148f3..e6ea6a91 100644 --- a/tasm-lib/benchmarks/tasmlib_mmr_root_from_authentication_struct.json +++ b/tasm-lib/benchmarks/tasmlib_mmr_root_from_authentication_struct.json @@ -2,22 +2,22 @@ { "name": "tasmlib_mmr_root_from_authentication_struct", "benchmark_result": { - "clock_cycle_count": 87275, - "hash_table_height": 5413, + "clock_cycle_count": 85070, + "hash_table_height": 5407, "u32_table_height": 11235, - "op_stack_table_height": 133574, - "ram_table_height": 38168 + "op_stack_table_height": 132462, + "ram_table_height": 39264 }, "case": "CommonCase" }, { "name": "tasmlib_mmr_root_from_authentication_struct", "benchmark_result": { - "clock_cycle_count": 180956, - "hash_table_height": 10747, + "clock_cycle_count": 176379, + "hash_table_height": 10741, "u32_table_height": 41115, - "op_stack_table_height": 277068, - "ram_table_height": 79085 + "op_stack_table_height": 274770, + "ram_table_height": 81367 }, "case": "WorstCase" } diff --git a/tasm-lib/src/mmr/authentication_struct/root_from_authentication_struct.rs b/tasm-lib/src/mmr/authentication_struct/root_from_authentication_struct.rs index 32d672a6..b94ebdc4 100644 --- a/tasm-lib/src/mmr/authentication_struct/root_from_authentication_struct.rs +++ b/tasm-lib/src/mmr/authentication_struct/root_from_authentication_struct.rs @@ -50,8 +50,10 @@ impl BasicSnippet for RootFromAuthenticationStruct { gamma_challenge_pointer_write + bfe!(EXTENSION_DEGREE as u64 - 1); let t_digest_pointer_write = library.kmalloc(Digest::LEN as u32); let t_digest_pointer_read = t_digest_pointer_write + bfe!(Digest::LEN as u64 - 1); - let p_pointer_write = library.kmalloc(EXTENSION_DEGREE as u32); - let p_pointer_read = p_pointer_write + bfe!(EXTENSION_DEGREE as u64 - 1); + let right_digest_pointer_write = library.kmalloc(Digest::LEN as u32); + let right_digest_pointer_read = right_digest_pointer_write + bfe!(Digest::LEN as u64 - 1); + let left_digest_pointer_write = library.kmalloc(Digest::LEN as u32); + let left_digest_pointer_read = left_digest_pointer_write + bfe!(Digest::LEN as u64 - 1); let indexed_leaf_element_size = Self::indexed_leaf_element_type().stack_size(); let derive_challenges = library.import(Box::new(DeriveChallenges)); @@ -324,164 +326,163 @@ impl BasicSnippet for RootFromAuthenticationStruct { ); let nd_loop_label = format!("{entrypoint}_nd_loop"); - let dup_top_two_digests = triton_asm![dup 9; Digest::LEN * 2]; let dup_top_digest = triton_asm![dup 4; Digest::LEN]; let one_half = BFieldElement::new(2).inverse(); let nd_loop = triton_asm!( - // _ INVARIANT: _ + // _ INVARIANT: _ [p] {nd_loop_label}: divine 2 - // _ left_index right_index + // _ [p] left_index right_index dup 1 push 1 add eq - // _ l_index_bfe (left_index + 1 == right_index) + // _ [p] l_index_bfe (left_index + 1 == right_index) assert hint left_index: BFieldElement = stack[0..1] - // _ l_index_bfe + // _ [p] l_index_bfe + + swap 3 + swap 2 + swap 1 + // _ l_index_bfe [p] /* Calculate parent digest, preserving child digests */ divine {Digest::LEN} hint right: Digest = stack[0..5] - - divine {Digest::LEN} - hint left: Digest = stack[0..5] - - {&dup_top_two_digests} - hash - hint t: Digest = stack[0..5] - // _ left_index [right] [left] [t] + // _ l_index_bfe [p] [right] {&dup_top_digest} - push {t_digest_pointer_write} + push {right_digest_pointer_write} write_mem {Digest::LEN} pop 1 - // _ left_index [right] [left] [t] + // _ l_index_bfe [p] [right] {&digest_to_xfe} - hint t_xfe: XFieldElement = stack[0..3] - // _ left_index [right] [left] [t_xfe] + // _ l_index_bfe [p] [right_xfe] push {beta_challenge_pointer_read} read_mem {EXTENSION_DEGREE} pop 1 xx_add - // _ left_index [right] [left] [t_xfe - β] + // _ l_index_bfe [p] [right_xfe - β] - dup 13 - push {one_half} - mul push {gamma_challenge_pointer_read} read_mem {EXTENSION_DEGREE} pop 1 - // _ left_index [right] [left] [t_xfe - β] parent_index [γ] + // _ l_index_bfe [p] [right_xfe - β] [γ] + + dup 9 + push 1 + add + // _ l_index_bfe [p] [right_xfe - β] [γ] r_index_bfe - swap 1 - swap 2 - swap 3 xb_mul - // _ left_index [right] [left] [t_xfe - β] [γ * parent_index] + // _ l_index_bfe [p] [right_xfe - β] [γ * r_index_bfe] xx_add - // _ left_index [right] [left] [t_xfe - β + γ * parent_index] - // _ left_index [right] [left] [fact_parent] <-- rename + // _ l_index_bfe [p] [t_xfe - β + γ * parent_index] + // _ l_index_bfe [p] [fact_right] - /* Accumulate `fact_parent` into `p` */ - push {p_pointer_read} - read_mem {EXTENSION_DEGREE} - pop 1 - xx_mul - push {p_pointer_write} - write_mem {EXTENSION_DEGREE} + divine {Digest::LEN} + hint left: Digest = stack[0..5] + // _ l_index_bfe [p] [fact_right] [left] + + {&dup_top_digest} + push {left_digest_pointer_write} + write_mem {Digest::LEN} pop 1 - // _ left_index [right] [left] + // _ l_index_bfe [p] [fact_right] [left] - /* Calculate `fact_left` */ {&digest_to_xfe} - // _ left_index [right] [left_xfe] + // _ l_index_bfe [p] [fact_right] [left_xfe] push {beta_challenge_pointer_read} read_mem {EXTENSION_DEGREE} pop 1 xx_add - // _ left_index [right] [left_xfe - β] + // _ l_index_bfe [p] [fact_right] [left_xfe - β] push {gamma_challenge_pointer_read} read_mem {EXTENSION_DEGREE} pop 1 - // _ left_index [right] [left_xfe - β] [γ] - - dup 11 - // _ left_index [right] [left_xfe - β] [γ] left_index + // _ l_index_bfe [p] [fact_right] [left_xfe - β] [γ] + dup 12 xb_mul - // _ left_index [right] [left_xfe - β] [γ * left_index] + // _ l_index_bfe [p] [fact_right] [left_xfe - β] [l_index_bfe * γ] xx_add - // _ left_index [right] [left_xfe - β + γ * left_index] - // _ left_index [right] [fact_left] <-- rename + // _ l_index_bfe [p] [fact_right] [fact_left] + + xx_mul + // _ l_index_bfe [p] [fact_right * fact_left] x_invert - // _ left_index [right] [fact_left^{-1}] + // _ l_index_bfe [p] [(fact_right*fact_left)^{-1}] - /* Divide `fact_left` out of `p` */ - push {p_pointer_read} - read_mem {EXTENSION_DEGREE} - pop 1 xx_mul - push {p_pointer_write} - write_mem {EXTENSION_DEGREE} + // _ l_index_bfe [p * (fact_right*fact_left)^{-1}] + + /* Calculate t = hash(left, right) */ + push {right_digest_pointer_read} + read_mem {Digest::LEN} pop 1 - // _ left_index [right] + push {left_digest_pointer_read} + read_mem {Digest::LEN} + pop 1 + hash + // _ l_index_bfe [p * (fact_right*fact_left)^{-1}] [t] + + // TODO: We only need to store `t` here if l_index_bfe == 2 + {&dup_top_digest} + push {t_digest_pointer_write} + write_mem {Digest::LEN} + pop 1 + // _ l_index_bfe [p * (fact_right*fact_left)^{-1}] [t] - /* Calculate `fact_right` */ {&digest_to_xfe} - // _ left_index [right_xfe] + // _ l_index_bfe [p * (fact_right*fact_left)^{-1}] [t_xfe] push {beta_challenge_pointer_read} read_mem {EXTENSION_DEGREE} pop 1 xx_add - // _ left_index [right_xfe - β] + // _ l_index_bfe [p * (fact_right*fact_left)^{-1}] [t_xfe - β] push {gamma_challenge_pointer_read} read_mem {EXTENSION_DEGREE} pop 1 - // _ left_index [right_xfe - β] [γ] + // _ l_index_bfe [p * (fact_right*fact_left)^{-1}] [t_xfe - β] [γ] - dup 6 - push 1 - add - // _ left_index [right_xfe - β] [γ] right_index + dup 9 + push {one_half} + mul + // _ l_index_bfe [p * (fact_right*fact_left)^{-1}] [t_xfe - β] [γ] parent_index xb_mul - // _ left_index [right_xfe - β] [right_index * γ] + // _ l_index_bfe [p * (fact_right*fact_left)^{-1}] [t_xfe - β] [parent_index * γ] xx_add - // _ left_index [right_xfe - β + right_index * γ] - // _ left_index [fact_right] <-- rename + // _ l_index_bfe [p * (fact_right*fact_left)^{-1}] [fact_parent] - x_invert - // _ left_index [fact_right^{-1}] - - /* Divide `fact_right` out of `p` */ - push {p_pointer_read} - read_mem {EXTENSION_DEGREE} - pop 1 xx_mul - push {p_pointer_write} - write_mem {EXTENSION_DEGREE} - pop 1 - // _ left_index + // _ l_index_bfe [p * (fact_right*fact_left)^{-1} * fact_parent] + // _ l_index_bfe [p'] + + swap 1 + swap 2 + swap 3 + // _ [p'] l_index_bfe /* Terminate loop when left_index == 2 <=> parent_index == 1 */ push 2 eq skiz return + // _ [p'] recurse ); @@ -564,34 +565,32 @@ impl BasicSnippet for RootFromAuthenticationStruct { pop 1 // _ tree_num_leafs *indexed_leafs [p] [t; 5] - /* Write t value, and `p` to static memory */ + /* Write t value */ push {t_digest_pointer_write} write_mem {Digest::LEN} pop 1 // _ tree_num_leafs *indexed_leafs [p] - push {p_pointer_write} - write_mem {EXTENSION_DEGREE} - pop 2 - // _ tree_num_leafs + // _ tree_num_leafs *indexed_leafs p2 p1 p0 + swap 2 + swap 4 + swap 1 + swap 3 + pop 1 + // _ [p] tree_num_leafs /* Call the ND-loop if tree_num_leafs != 1 */ push 1 eq push 0 eq - // (tree_num_leafs != 1) + // [p] (tree_num_leafs != 1) skiz call {nd_loop_label} - // _ - - /* Assert that p == t_xfe - beta + gamma */ - push {p_pointer_read} - read_mem {EXTENSION_DEGREE} - pop 1 // _ [p] + /* Assert that p == t_xfe - beta + gamma */ push {t_digest_pointer_read} read_mem {Digest::LEN} pop 1 @@ -691,7 +690,8 @@ mod tests { beta: XFieldElement, gamma: XFieldElement, t: Digest, - p: XFieldElement, + right: Digest, + left: Digest, ) { const ALPHA_POINTER_WRITE: BFieldElement = BFieldElement::new(BFieldElement::P - 4); const BETA_POINTER_WRITE: BFieldElement = BFieldElement::new(BFieldElement::P - 7); @@ -699,13 +699,17 @@ mod tests { BFieldElement::new(BFieldElement::P - 10); const T_DIGEST_POINTER_WRITE: BFieldElement = BFieldElement::new(BFieldElement::P - 15); - const P_POINTER_WRITE: BFieldElement = BFieldElement::new(BFieldElement::P - 18); + const RIGHT_DIGEST_POINTER_WRITE: BFieldElement = + BFieldElement::new(BFieldElement::P - 20); + const LEFT_DIGEST_POINTER_WRITE: BFieldElement = + BFieldElement::new(BFieldElement::P - 25); write_to_memory(ALPHA_POINTER_WRITE, alpha, memory); write_to_memory(BETA_POINTER_WRITE, beta, memory); write_to_memory(GAMMA_POINTER_WRITE, gamma, memory); write_to_memory(T_DIGEST_POINTER_WRITE, t, memory); - write_to_memory(P_POINTER_WRITE, p, memory); + write_to_memory(RIGHT_DIGEST_POINTER_WRITE, right, memory); + write_to_memory(LEFT_DIGEST_POINTER_WRITE, left, memory); } fn accumulate_indexed_leafs( @@ -810,6 +814,8 @@ mod tests { // "Unaccumulate" into `p` from secret data, and calculate Merkle root let mut t = indexed_leafs[0].1; let mut t_xfe = digest_to_xfe(t, alpha); + let mut right = Digest::default(); + let mut left = Digest::default(); if tree_num_leafs != 1 { loop { let left_index = individual_tokens.pop_front().unwrap(); @@ -818,8 +824,8 @@ mod tests { let parent_index = left_index / bfe!(2); - let right = read_digest_from_input(&mut individual_tokens); - let left = read_digest_from_input(&mut individual_tokens); + right = read_digest_from_input(&mut individual_tokens); + left = read_digest_from_input(&mut individual_tokens); t = Tip5::hash_pair(left, right); t_xfe = digest_to_xfe(t, alpha); @@ -844,7 +850,7 @@ mod tests { stack.push(elem); } - mimic_use_of_static_memory(memory, alpha, -beta, gamma, t, p); + mimic_use_of_static_memory(memory, alpha, -beta, gamma, t, right, left); vec![] } From f3e2c0c085fcdbfa8473726fc0ebdfdba5c34842 Mon Sep 17 00:00:00 2001 From: sword_smith Date: Thu, 15 Aug 2024 15:24:08 +0200 Subject: [PATCH 09/10] perf(mmr::authentication_struct): Only store `t` digest in last loop iteration --- ...b_mmr_root_from_authentication_struct.json | 16 +++++------ .../root_from_authentication_struct.rs | 28 +++++++++++++++---- 2 files changed, 30 insertions(+), 14 deletions(-) diff --git a/tasm-lib/benchmarks/tasmlib_mmr_root_from_authentication_struct.json b/tasm-lib/benchmarks/tasmlib_mmr_root_from_authentication_struct.json index e6ea6a91..1708cd5f 100644 --- a/tasm-lib/benchmarks/tasmlib_mmr_root_from_authentication_struct.json +++ b/tasm-lib/benchmarks/tasmlib_mmr_root_from_authentication_struct.json @@ -2,22 +2,22 @@ { "name": "tasmlib_mmr_root_from_authentication_struct", "benchmark_result": { - "clock_cycle_count": 85070, - "hash_table_height": 5407, + "clock_cycle_count": 82876, + "hash_table_height": 5413, "u32_table_height": 11235, - "op_stack_table_height": 132462, - "ram_table_height": 39264 + "op_stack_table_height": 128066, + "ram_table_height": 36514 }, "case": "CommonCase" }, { "name": "tasmlib_mmr_root_from_authentication_struct", "benchmark_result": { - "clock_cycle_count": 176379, - "hash_table_height": 10741, + "clock_cycle_count": 171813, + "hash_table_height": 10747, "u32_table_height": 41115, - "op_stack_table_height": 274770, - "ram_table_height": 81367 + "op_stack_table_height": 265630, + "ram_table_height": 75652 }, "case": "WorstCase" } diff --git a/tasm-lib/src/mmr/authentication_struct/root_from_authentication_struct.rs b/tasm-lib/src/mmr/authentication_struct/root_from_authentication_struct.rs index b94ebdc4..259b7e74 100644 --- a/tasm-lib/src/mmr/authentication_struct/root_from_authentication_struct.rs +++ b/tasm-lib/src/mmr/authentication_struct/root_from_authentication_struct.rs @@ -325,8 +325,22 @@ impl BasicSnippet for RootFromAuthenticationStruct { recurse_or_return ); - let nd_loop_label = format!("{entrypoint}_nd_loop"); let dup_top_digest = triton_asm![dup 4; Digest::LEN]; + let store_t_digest_in_memory_label = format!("{entrypoint}_store_t_digest"); + let store_t_digest_in_memory = triton_asm!( + {store_t_digest_in_memory_label}: + // _ [t] + + {&dup_top_digest} + push {t_digest_pointer_write} + write_mem {Digest::LEN} + pop 1 + // _ [t] + + return + ); + + let nd_loop_label = format!("{entrypoint}_nd_loop"); let one_half = BFieldElement::new(2).inverse(); let nd_loop = triton_asm!( // _ INVARIANT: _ [p] @@ -436,11 +450,12 @@ impl BasicSnippet for RootFromAuthenticationStruct { hash // _ l_index_bfe [p * (fact_right*fact_left)^{-1}] [t] - // TODO: We only need to store `t` here if l_index_bfe == 2 - {&dup_top_digest} - push {t_digest_pointer_write} - write_mem {Digest::LEN} - pop 1 + /* Store [t] digest in memory if this is last loop iteration */ + dup 8 + push 2 + eq + skiz + call {store_t_digest_in_memory_label} // _ l_index_bfe [p * (fact_right*fact_left)^{-1}] [t] {&digest_to_xfe} @@ -626,6 +641,7 @@ impl BasicSnippet for RootFromAuthenticationStruct { {&accumulated_indexed_leafs_loop} {&accumulate_auth_struct_leafs_from_public_data} {&nd_loop} + {&store_t_digest_in_memory} ) } } From 3b645722a22b5e001b0a1892663895a2de1af51e Mon Sep 17 00:00:00 2001 From: sword_smith Date: Thu, 15 Aug 2024 16:33:31 +0200 Subject: [PATCH 10/10] bench(mmr:authentication_struct): More realistic bench params In a mutator-set context for which this is developed, the leaf-indices will be grouped together. The benchmark now reflects that. --- ...b_mmr_root_from_authentication_struct.json | 20 +++---- .../root_from_authentication_struct.rs | 57 +++++++++++++++---- 2 files changed, 56 insertions(+), 21 deletions(-) diff --git a/tasm-lib/benchmarks/tasmlib_mmr_root_from_authentication_struct.json b/tasm-lib/benchmarks/tasmlib_mmr_root_from_authentication_struct.json index 1708cd5f..b723f98f 100644 --- a/tasm-lib/benchmarks/tasmlib_mmr_root_from_authentication_struct.json +++ b/tasm-lib/benchmarks/tasmlib_mmr_root_from_authentication_struct.json @@ -2,22 +2,22 @@ { "name": "tasmlib_mmr_root_from_authentication_struct", "benchmark_result": { - "clock_cycle_count": 82876, - "hash_table_height": 5413, - "u32_table_height": 11235, - "op_stack_table_height": 128066, - "ram_table_height": 36514 + "clock_cycle_count": 21218, + "hash_table_height": 1753, + "u32_table_height": 2932, + "op_stack_table_height": 32926, + "ram_table_height": 9604 }, "case": "CommonCase" }, { "name": "tasmlib_mmr_root_from_authentication_struct", "benchmark_result": { - "clock_cycle_count": 171813, - "hash_table_height": 10747, - "u32_table_height": 41115, - "op_stack_table_height": 265630, - "ram_table_height": 75652 + "clock_cycle_count": 26162, + "hash_table_height": 2053, + "u32_table_height": 1736, + "op_stack_table_height": 40578, + "ram_table_height": 11786 }, "case": "WorstCase" } diff --git a/tasm-lib/src/mmr/authentication_struct/root_from_authentication_struct.rs b/tasm-lib/src/mmr/authentication_struct/root_from_authentication_struct.rs index 259b7e74..b6cfc303 100644 --- a/tasm-lib/src/mmr/authentication_struct/root_from_authentication_struct.rs +++ b/tasm-lib/src/mmr/authentication_struct/root_from_authentication_struct.rs @@ -580,7 +580,7 @@ impl BasicSnippet for RootFromAuthenticationStruct { pop 1 // _ tree_num_leafs *indexed_leafs [p] [t; 5] - /* Write t value */ + /* Write t value (in case we're not entering the loop) */ push {t_digest_pointer_write} write_mem {Digest::LEN} pop 1 @@ -878,22 +878,57 @@ mod tests { ) -> ProcedureInitialState { let mut rng: StdRng = SeedableRng::from_seed(seed); - let (tree_height, num_revealed_leafs) = match bench_case { + let num_chunks = 45; + let num_accessible_chunk_indices = 1 << 8; + let (tree_height, revealed_leaf_indices) = match bench_case { None => { - let tree_height = rng.gen_range(0..32); + let tree_height = rng.gen_range(0..62); let num_leafs_in_merkle_tree = 1 << tree_height; - let num_revealed_leafs = rng.gen_range(1..=min(num_leafs_in_merkle_tree, 20)); - (tree_height, num_revealed_leafs) + let num_revealed_leafs = + rng.gen_range(1..=min(num_leafs_in_merkle_tree, num_chunks)); + + let revealed_leaf_indices = (0..num_revealed_leafs) + .map(|_| { + rng.gen_range( + 0..min(num_accessible_chunk_indices, num_leafs_in_merkle_tree), + ) + }) + .unique() + .collect_vec(); + + (tree_height, revealed_leaf_indices) + } + + // In both benchmarks, we leafs from the middle of the Merkle + // tree. Were we pick the indices is not so relevant for + // performance, as long as they're grouped together in a + // realistic way for the mutator set. + Some(BenchmarkCase::CommonCase) => { + let tree_height = 32; + let midpoint = 1 << (tree_height - 1); + let revealed_leaf_indices = (0..num_chunks) + .map(|_| rng.gen_range(midpoint..num_accessible_chunk_indices + midpoint)) + .unique() + .collect_vec(); + + (tree_height, revealed_leaf_indices) + } + Some(BenchmarkCase::WorstCase) => { + let tree_height = 62; + let midpoint = 1 << (tree_height - 1); + let revealed_leaf_indices = (0..num_chunks) + .map(|_| rng.gen_range(midpoint..num_accessible_chunk_indices + midpoint)) + .unique() + .collect_vec(); + + (tree_height, revealed_leaf_indices) } - Some(BenchmarkCase::CommonCase) => (32, 20), - Some(BenchmarkCase::WorstCase) => (62, 20), }; let num_leafs_in_merkle_tree = 1 << tree_height; - let revealed_leaf_indices = (0..num_revealed_leafs) - .map(|_| rng.gen_range(0..num_leafs_in_merkle_tree)) - .unique() - .collect_vec(); + // This picks leaf-indices with low values but I don't think that + // matters for performance. + let num_revealed_leafs = revealed_leaf_indices.len(); assert!(!num_revealed_leafs.is_zero());