Skip to content

Commit

Permalink
test: Add example program for merkle_step_mem
Browse files Browse the repository at this point in the history
  • Loading branch information
jan-ferdinand committed Jul 3, 2024
1 parent 4fc1cfd commit 5335b85
Show file tree
Hide file tree
Showing 2 changed files with 89 additions and 0 deletions.
42 changes: 42 additions & 0 deletions triton-vm/src/example_programs.rs
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@ lazy_static! {
calculate_new_mmr_peaks_from_append_with_safe_lists();
pub static ref MERKLE_TREE_AUTHENTICATION_PATH_VERIFY: Program =
merkle_tree_authentication_path_verify();
pub static ref MERKLE_TREE_UPDATE: Program = merkle_tree_update();
}

fn fibonacci_sequence() -> Program {
Expand Down Expand Up @@ -174,6 +175,47 @@ fn merkle_tree_authentication_path_verify() -> Program {
)
}

/// Triton program to verifiably change a Merkle tree's leaf. That is:
/// 1. verify that the supplied `old_leaf` is indeed a leaf in the Merkle tree
/// defined by the `merkle_root` and the `tree_height`,
/// 2. update the leaf at the specified `leaf_index` with the `new_leaf`, and
/// 3. return the new Merkle root.
///
/// The authentication path for the leaf to update has to be supplied via RAM.
///
/// - input:
/// - RAM address of leaf's authentication path
/// - leaf index to update
/// - Merkle tree's height
/// - old leaf
/// - (current) merkle root
/// - new leaf
/// - output:
/// - new root
fn merkle_tree_update() -> Program {
triton_program! {
read_io 3 // _ *ap leaf_index tree_height
push 2 pow add // _ *ap node_index
dup 1 push 1 dup 2 // _ *ap node_index *ap 1 node_index
read_io 5 // _ *ap node_index *ap 1 node_index [old_leaf; 5]
call compute_root // _ *ap node_index *ap' 1 1 [root; 5]
read_io 5 // _ *ap node_index *ap' 1 1 [root; 5] [presumed_root; 5]
assert_vector // _ *ap node_index *ap' 1 1 [root; 5]
pop 5 pop 3 // _ *ap node_index
push 1 swap 1 // _ *ap 1 node_index
read_io 5 // _ *ap 1 node_index [new_leaf; 5]
call compute_root // _ *ap' 1 1 [new_root; 5]
write_io 5 // _ *ap' 1 1
pop 3 halt // _

// BEFORE: _ *ap 1 node_index [leaf; 5]
// AFTER: _ (*ap + 5 * tree_height) 1 1 [root; 5]
compute_root:
merkle_step_mem
recurse_or_return
}
}

fn verify_sudoku() -> Program {
// RAM layout:
// 0..=8: primes for mapping digits 1..=9
Expand Down
47 changes: 47 additions & 0 deletions triton-vm/src/vm.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2365,6 +2365,53 @@ pub(crate) mod tests {
assert!(let Ok(_) = program.run(public_input.into(), non_determinism));
}

#[proptest]
fn triton_assembly_merkle_tree_update(
leaved_merkle_tree: LeavedMerkleTreeTestData,
#[strategy(arb())] new_leaf: Digest,
#[strategy(arb())] auth_path_address: BFieldElement,
) {
let merkle_tree = leaved_merkle_tree.merkle_tree;
let leaf_index = leaved_merkle_tree.revealed_indices[0];
let auth_path = merkle_tree.authentication_structure(&[leaf_index]).unwrap();

let mut ram = HashMap::new();
{
let mut auth_path_address = auth_path_address;
for &Digest(digest) in &auth_path {
for element in digest {
ram.insert(auth_path_address, element);
auth_path_address.increment();
}
}
}
let non_determinism = NonDeterminism::default().with_ram(ram);

let old_leaf = Digest::from(leaved_merkle_tree.leaves[leaf_index]);
let old_root = merkle_tree.root();
let mut public_input = vec![
auth_path_address,
u64::try_from(leaf_index).unwrap().into(),
u64::try_from(merkle_tree.height()).unwrap().into(),
];
public_input.extend(old_leaf.reversed().values());
public_input.extend(old_root.reversed().values());
public_input.extend(new_leaf.reversed().values());
let public_input = PublicInput::new(public_input);

let program = MERKLE_TREE_UPDATE.clone();
let_assert!(Ok(output) = program.run(public_input, non_determinism));
let new_root = Digest(output.try_into().unwrap());

let inclusion_proof = MerkleTreeInclusionProof::<Tip5> {
tree_height: merkle_tree.height(),
indexed_leaves: vec![(leaf_index, new_leaf)],
authentication_structure: auth_path,
_hasher: std::marker::PhantomData,
};
assert!(inclusion_proof.verify(new_root));
}

#[proptest]
fn run_tvm_get_collinear_y(
#[strategy(arb())] p0: (BFieldElement, BFieldElement),
Expand Down

0 comments on commit 5335b85

Please sign in to comment.