From 0aebfa0596c875cb55da640b68ee284e2275dbe4 Mon Sep 17 00:00:00 2001 From: Victor Lopez Date: Tue, 28 Mar 2023 22:38:17 +0200 Subject: [PATCH] feat: add stdlib smt collection This commit introduces a `collections` module for the `stdlib`. It contains, initially, functions to support Sparse Merkle Tree functionality. Initially, the `smt::get` is available; it will fetch the value of a key from a Sparse Merkle tree. It adds a `adv.smtget` that will push to the advice stack information about a Sparse Merkle tree keyed leaf. --- CHANGELOG.md | 6 + assembly/src/assembler/instruction/mod.rs | 1 + assembly/src/parsers/adv_ops.rs | 9 +- assembly/src/parsers/nodes.rs | 2 + assembly/src/parsers/serde/deserialization.rs | 1 + assembly/src/parsers/serde/mod.rs | 25 +- assembly/src/parsers/serde/serialization.rs | 1 + assembly/src/parsers/tests.rs | 3 +- core/src/lib.rs | 3 +- core/src/operations/decorators/advice.rs | 17 ++ processor/Cargo.toml | 1 + processor/src/advice/mem_provider.rs | 15 +- processor/src/advice/mod.rs | 26 +- processor/src/decorators/mod.rs | 141 +++++----- processor/src/decorators/tests.rs | 187 +++++++++++++ processor/src/errors.rs | 4 + stdlib/asm/collections/smt.masm | 245 ++++++++++++++++++ stdlib/docs/smt_collections.md | 5 + stdlib/tests/collections/mod.rs | 1 + stdlib/tests/collections/smt.rs | 162 ++++++++++++ test-utils/src/crypto.rs | 12 +- test-utils/src/lib.rs | 10 +- test-utils/src/rand.rs | 35 +++ 23 files changed, 813 insertions(+), 99 deletions(-) create mode 100644 processor/src/decorators/tests.rs create mode 100644 stdlib/asm/collections/smt.masm create mode 100644 stdlib/docs/smt_collections.md create mode 100644 stdlib/tests/collections/smt.rs create mode 100644 test-utils/src/rand.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index e8dc30c290..6748659b8c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,12 @@ ## 0.6.0 (TBD) +#### Assembly +- Added new instructions: `adv.smtget`. + +#### Stdlib +- Added new module: `collections::smt` with `smt::get`. + ## 0.5.0 (2023-03-29) #### CLI diff --git a/assembly/src/assembler/instruction/mod.rs b/assembly/src/assembler/instruction/mod.rs index b908897bb9..063dd5a959 100644 --- a/assembly/src/assembler/instruction/mod.rs +++ b/assembly/src/assembler/instruction/mod.rs @@ -290,6 +290,7 @@ impl Assembler { Instruction::AdvMem(a, n) => adv_ops::adv_mem(span, *a, *n), Instruction::AdvExt2Inv => span.add_decorator(Decorator::Advice(Ext2Inv)), Instruction::AdvExt2INTT => span.add_decorator(Decorator::Advice(Ext2INTT)), + Instruction::AdvSmtGet => span.add_decorator(Decorator::Advice(SmtGet)), // ----- cryptographic instructions --------------------------------------------------- Instruction::Hash => crypto_ops::hash(span), diff --git a/assembly/src/parsers/adv_ops.rs b/assembly/src/parsers/adv_ops.rs index 03f71d3b20..447ee22e7c 100644 --- a/assembly/src/parsers/adv_ops.rs +++ b/assembly/src/parsers/adv_ops.rs @@ -8,7 +8,8 @@ use super::{ // INSTRUCTION PARSERS // ================================================================================================ -/// Returns `AdvU64Div`, `AdvKeyval`, or `AdvMem` instruction node. +/// Returns `AdvU64Div`, `AdvKeyval`, `AdvMem`, `AdvExt2Inv`, `AdvExt2INTT`, or `AdvSmtGet` +/// instruction node. /// /// # Errors /// Returns an error if: @@ -55,6 +56,12 @@ pub fn parse_adv_inject(op: &Token) -> Result { } Ok(Instruction(AdvExt2INTT)) } + "smtget" => { + if op.num_parts() > 2 { + return Err(ParsingError::extra_param(op)); + } + Ok(Instruction(AdvSmtGet)) + } _ => Err(ParsingError::invalid_op(op)), } } diff --git a/assembly/src/parsers/nodes.rs b/assembly/src/parsers/nodes.rs index 5e8deef161..a18c5afbac 100644 --- a/assembly/src/parsers/nodes.rs +++ b/assembly/src/parsers/nodes.rs @@ -261,6 +261,7 @@ pub enum Instruction { AdvMem(u32, u32), AdvExt2Inv, AdvExt2INTT, + AdvSmtGet, // ----- cryptographic operations ------------------------------------------------------------- Hash, @@ -536,6 +537,7 @@ impl fmt::Display for Instruction { Self::AdvMem(start_addr, num_words) => write!(f, "adv.mem.{start_addr}.{num_words}"), Self::AdvExt2Inv => write!(f, "adv.ext2inv"), Self::AdvExt2INTT => write!(f, "adv.ext2intt"), + Self::AdvSmtGet => write!(f, "adv.smtget"), // ----- cryptographic operations ----------------------------------------------------- Self::Hash => write!(f, "hash"), diff --git a/assembly/src/parsers/serde/deserialization.rs b/assembly/src/parsers/serde/deserialization.rs index ca62c3c290..9bccbbc3d7 100644 --- a/assembly/src/parsers/serde/deserialization.rs +++ b/assembly/src/parsers/serde/deserialization.rs @@ -335,6 +335,7 @@ impl Deserializable for Instruction { OpCode::AdvLoadW => Ok(Instruction::AdvLoadW), OpCode::AdvExt2Inv => Ok(Instruction::AdvExt2Inv), OpCode::AdvExt2INTT => Ok(Instruction::AdvExt2INTT), + OpCode::AdvSmtGet => Ok(Instruction::AdvSmtGet), // ----- cryptographic operations ----------------------------------------------------- OpCode::Hash => Ok(Instruction::Hash), diff --git a/assembly/src/parsers/serde/mod.rs b/assembly/src/parsers/serde/mod.rs index c5135c04f6..c552c9bb3d 100644 --- a/assembly/src/parsers/serde/mod.rs +++ b/assembly/src/parsers/serde/mod.rs @@ -257,22 +257,23 @@ pub enum OpCode { AdvMem = 227, AdvExt2Inv = 228, AdvExt2INTT = 229, + AdvSmtGet = 230, // ----- cryptographic operations ------------------------------------------------------------- - Hash = 230, - HMerge = 231, - HPerm = 232, - MTreeGet = 233, - MTreeSet = 234, - MTreeMerge = 235, - FriExt2Fold4 = 236, + Hash = 231, + HMerge = 232, + HPerm = 233, + MTreeGet = 234, + MTreeSet = 235, + MTreeMerge = 236, + FriExt2Fold4 = 237, // ----- exec / call -------------------------------------------------------------------------- - ExecLocal = 237, - ExecImported = 238, - CallLocal = 239, - CallImported = 240, - SysCall = 241, + ExecLocal = 238, + ExecImported = 239, + CallLocal = 240, + CallImported = 241, + SysCall = 242, // ----- control flow ------------------------------------------------------------------------- IfElse = 253, diff --git a/assembly/src/parsers/serde/serialization.rs b/assembly/src/parsers/serde/serialization.rs index 1d3600122c..010c5ec1cf 100644 --- a/assembly/src/parsers/serde/serialization.rs +++ b/assembly/src/parsers/serde/serialization.rs @@ -446,6 +446,7 @@ impl Serializable for Instruction { Self::AdvLoadW => OpCode::AdvLoadW.write_into(target), Self::AdvExt2Inv => OpCode::AdvExt2Inv.write_into(target), Self::AdvExt2INTT => OpCode::AdvExt2INTT.write_into(target), + Self::AdvSmtGet => OpCode::AdvSmtGet.write_into(target), // ----- cryptographic operations ----------------------------------------------------- Self::Hash => OpCode::Hash.write_into(target), diff --git a/assembly/src/parsers/tests.rs b/assembly/src/parsers/tests.rs index 7c95439d3c..34fcf0ae8b 100644 --- a/assembly/src/parsers/tests.rs +++ b/assembly/src/parsers/tests.rs @@ -184,11 +184,12 @@ fn test_ast_parsing_adv_ops() { #[test] fn test_ast_parsing_adv_injection() { - let source = "begin adv.u64div adv.keyval adv.mem.1.1 end"; + let source = "begin adv.u64div adv.keyval adv.mem.1.1 adv.smtget end"; let nodes: Vec = vec![ Node::Instruction(Instruction::AdvU64Div), Node::Instruction(Instruction::AdvKeyval), Node::Instruction(Instruction::AdvMem(1, 1)), + Node::Instruction(Instruction::AdvSmtGet), ]; assert_program_output(source, BTreeMap::new(), nodes); diff --git a/core/src/lib.rs b/core/src/lib.rs index 36472f42b3..dea405b28d 100644 --- a/core/src/lib.rs +++ b/core/src/lib.rs @@ -14,7 +14,8 @@ pub use ::crypto::{Word, ONE, WORD_SIZE, ZERO}; pub mod crypto { pub mod merkle { pub use ::crypto::merkle::{ - MerkleError, MerklePath, MerklePathSet, MerkleStore, MerkleTree, NodeIndex, SimpleSmt, + EmptySubtreeRoots, MerkleError, MerklePath, MerklePathSet, MerkleStore, MerkleTree, + NodeIndex, SimpleSmt, }; } diff --git a/core/src/operations/decorators/advice.rs b/core/src/operations/decorators/advice.rs index 6a00ee5e66..2535b9eda8 100644 --- a/core/src/operations/decorators/advice.rs +++ b/core/src/operations/decorators/advice.rs @@ -51,6 +51,22 @@ pub enum AdviceInjector { /// routine interpolates ( using inverse NTT ) the evaluations into a polynomial in /// coefficient form and pushes the result into the advice stack. Ext2INTT, + + /// Pushes the value and depth flags of a leaf indexed by `key` on a Sparse Merkle tree with + /// the provided `root`. + /// + /// The Sparse Merkle tree is tiered, meaning it will have leaf depths in `{16, 32, 48, 64}`. + /// The depth flags define the tier on which the leaf is located. + /// + /// The operand stack is expected to be arranged as follows: + /// - key, 4 elements. + /// - root of the Sparse Merkle tree, 4 elements. + /// + /// After a successful operation, the advice stack will look as follows: + /// - boolean flag set to `1` if the depth is `16` or `32`. + /// - boolean flag set to `1` if the depth is `16` or `48`. + /// - value word; will be zeroed if the tree don't contain a mapped value for the key. + SmtGet, } impl fmt::Display for AdviceInjector { @@ -63,6 +79,7 @@ impl fmt::Display for AdviceInjector { Self::Memory(start_addr, num_words) => write!(f, "mem({start_addr}, {num_words})"), Self::Ext2Inv => write!(f, "ext2_inv"), Self::Ext2INTT => write!(f, "ext2_intt"), + Self::SmtGet => write!(f, "smt_get"), } } } diff --git a/processor/Cargo.toml b/processor/Cargo.toml index c0474748af..de0d3cec54 100644 --- a/processor/Cargo.toml +++ b/processor/Cargo.toml @@ -30,5 +30,6 @@ winter-prover = { package = "winter-prover", version = "0.6", default-features = logtest = { version = "2.0", default-features = false } miden-assembly = { package = "miden-assembly", path = "../assembly", version = "0.6", default-features = false } rand-utils = { package = "winter-rand-utils", version = "0.6" } +test-utils = { package = "miden-test-utils", path = "../test-utils", version = "0.1" } winter-fri = { package = "winter-fri", version = "0.6" } winter-utils = { package = "winter-utils", version = "0.6" } diff --git a/processor/src/advice/mem_provider.rs b/processor/src/advice/mem_provider.rs index c11da7a508..46a6994a9e 100644 --- a/processor/src/advice/mem_provider.rs +++ b/processor/src/advice/mem_provider.rs @@ -1,6 +1,6 @@ use super::{ AdviceInputs, AdviceProvider, AdviceSource, BTreeMap, ExecutionError, Felt, IntoBytes, - MerklePath, MerkleStore, NodeIndex, Vec, Word, + MerklePath, MerkleStore, NodeIndex, StarkField, Vec, Word, }; // MEMORY ADVICE PROVIDER @@ -122,6 +122,19 @@ impl AdviceProvider for MemAdviceProvider { .map_err(ExecutionError::MerkleStoreLookupFailed) } + fn get_leaf_depth( + &self, + root: Word, + tree_depth: &Felt, + index: &Felt, + ) -> Result { + let tree_depth = u8::try_from(tree_depth.as_int()) + .map_err(|_| ExecutionError::InvalidTreeDepth { depth: *tree_depth })?; + self.store + .get_leaf_depth(root, tree_depth, index.as_int()) + .map_err(ExecutionError::MerkleStoreLookupFailed) + } + fn update_merkle_node( &mut self, root: Word, diff --git a/processor/src/advice/mod.rs b/processor/src/advice/mod.rs index af60a53592..5580ed5931 100644 --- a/processor/src/advice/mod.rs +++ b/processor/src/advice/mod.rs @@ -1,4 +1,4 @@ -use super::{ExecutionError, Felt, InputError, Word}; +use super::{ExecutionError, Felt, InputError, StarkField, Word}; use vm_core::{ crypto::merkle::{MerklePath, MerkleStore, NodeIndex}, utils::{ @@ -124,6 +124,21 @@ pub trait AdviceProvider { index: &Felt, ) -> Result; + /// Reconstructs a path from the root until a leaf or empty node and returns its depth. + /// + /// For more information, check [MerkleStore::get_leaf_depth]. + /// + /// # Errors + /// Will return an error if: + /// - The provided `tree_depth` doesn't fit `u8`. + /// - The conditions of [MerkleStore::get_leaf_depth] aren't met. + fn get_leaf_depth( + &self, + root: Word, + tree_depth: &Felt, + index: &Felt, + ) -> Result; + /// Updates a node at the specified depth and index in a Merkle tree with the specified root; /// returns the Merkle path from the updated node to the new root. /// @@ -211,6 +226,15 @@ where T::get_merkle_path(self, root, depth, index) } + fn get_leaf_depth( + &self, + root: Word, + tree_depth: &Felt, + index: &Felt, + ) -> Result { + T::get_leaf_depth(self, root, tree_depth, index) + } + fn update_merkle_node( &mut self, root: Word, diff --git a/processor/src/decorators/mod.rs b/processor/src/decorators/mod.rs index 18bde219f1..75cd08fc66 100644 --- a/processor/src/decorators/mod.rs +++ b/processor/src/decorators/mod.rs @@ -1,10 +1,16 @@ use super::{ AdviceInjector, AdviceProvider, AdviceSource, Decorator, ExecutionError, Felt, Process, - StarkField, + StarkField, Word, +}; +use vm_core::{ + crypto::merkle::EmptySubtreeRoots, utils::collections::Vec, FieldElement, QuadExtension, ONE, + WORD_SIZE, ZERO, }; -use vm_core::{utils::collections::Vec, FieldElement, QuadExtension, WORD_SIZE, ZERO}; use winter_prover::math::fft; +#[cfg(test)] +mod tests; + // TYPE ALIASES // ================================================================================================ type QuadFelt = QuadExtension; @@ -47,6 +53,7 @@ where } AdviceInjector::Ext2Inv => self.inject_ext2_inv_result(), AdviceInjector::Ext2INTT => self.inject_ext2_intt_result(), + AdviceInjector::SmtGet => self.inject_smtget(), } } @@ -278,6 +285,60 @@ where Ok(()) } + + /// Pushes the value and depth flags of a leaf indexed by `key` on a Sparse Merkle tree with + /// the provided `root`. + /// + /// The Sparse Merkle tree is tiered, meaning it will have leaf depths in `{16, 32, 48, 64}`. + /// The depth flags define the tier on which the leaf is located. + /// + /// The operand stack is expected to be arranged as follows: + /// - key, 4 elements. + /// - root of the Sparse Merkle tree, 4 elements. + /// + /// After a successful operation, the advice stack will look as follows: + /// - boolean flag set to `1` if the depth is `16` or `32`. + /// - boolean flag set to `1` if the depth is `16` or `48`. + /// - value word; will be zeroed if the tree don't contain a mapped value for the key. + /// + /// # Errors + /// Will return an error if: + /// - The provided Merkle root doesn't exist on the advice provider + /// - If, at depth `64`, the node value is not zero, and there is no associated leaf value + fn inject_smtget(&mut self) -> Result<(), ExecutionError> { + // fetch the arguments from the operand stack + let key = [self.stack.get(3), self.stack.get(2), self.stack.get(1), self.stack.get(0)]; + let root = [self.stack.get(7), self.stack.get(6), self.stack.get(5), self.stack.get(4)]; + + // fetch the leaf depth from the Merkle store + // an empty tree will return depth `0`, but we always point to the minimum tier (16) + const TREE_DEPTH: Felt = Felt::new(64); + let index = &key[3]; + let depth = self.advice_provider.get_leaf_depth(root, &TREE_DEPTH, index)?.max(16); + + // fetch the node value + let index = index.as_int() >> (64 - depth); + let index = Felt::new(index); + let node = self.advice_provider.get_tree_node(root, &Felt::new(depth as u64), &index)?; + + // set the flags + let is_16_or_32 = if depth == 16 || depth == 32 { ONE } else { ZERO }; + let is_16_or_48 = if depth == 16 || depth == 48 { ONE } else { ZERO }; + self.advice_provider.push_stack(AdviceSource::Value(is_16_or_32))?; + self.advice_provider.push_stack(AdviceSource::Value(is_16_or_48))?; + + // set the node value; zeroed if empty sub-tree + let empty = EmptySubtreeRoots::empty_hashes(64); + if Word::from(empty[depth as usize]) == node { + self.advice_provider.push_stack(AdviceSource::Value(ZERO))?; + self.advice_provider.push_stack(AdviceSource::Value(ZERO))?; + self.advice_provider.push_stack(AdviceSource::Value(ZERO))?; + self.advice_provider.push_stack(AdviceSource::Value(ZERO))?; + } else { + self.advice_provider.push_stack(AdviceSource::Map { key: node })?; + } + Ok(()) + } } // HELPER FUNCTIONS @@ -288,79 +349,3 @@ fn u64_to_u32_elements(value: u64) -> (Felt, Felt) { let lo = Felt::new((value as u32) as u64); (hi, lo) } - -// TESTS -// ================================================================================================ - -#[cfg(test)] -mod tests { - use super::{ - super::{AdviceInputs, Felt, FieldElement, Kernel, Operation, StarkField}, - Process, - }; - use crate::{MemAdviceProvider, StackInputs, Word}; - use vm_core::{ - crypto::merkle::{MerkleStore, MerkleTree}, - AdviceInjector, Decorator, - }; - - #[test] - fn inject_merkle_node() { - let leaves = [init_leaf(1), init_leaf(2), init_leaf(3), init_leaf(4)]; - let tree = MerkleTree::new(leaves.to_vec()).unwrap(); - let store = MerkleStore::default().with_merkle_tree(leaves).unwrap(); - let stack_inputs = [ - tree.root()[0].as_int(), - tree.root()[1].as_int(), - tree.root()[2].as_int(), - tree.root()[3].as_int(), - 1, - tree.depth() as u64, - ]; - - let stack_inputs = StackInputs::try_from_values(stack_inputs).unwrap(); - let advice_inputs = AdviceInputs::default().with_merkle_store(store); - let advice_provider = MemAdviceProvider::from(advice_inputs); - let mut process = Process::new(Kernel::default(), stack_inputs, advice_provider); - process.execute_op(Operation::Noop).unwrap(); - - // push the node onto the advice stack - process - .execute_decorator(&Decorator::Advice(AdviceInjector::MerkleNode)) - .unwrap(); - - // pop the node from the advice stack and push it onto the operand stack - process.execute_op(Operation::AdvPop).unwrap(); - process.execute_op(Operation::AdvPop).unwrap(); - process.execute_op(Operation::AdvPop).unwrap(); - process.execute_op(Operation::AdvPop).unwrap(); - - let expected_stack = build_expected(&[ - leaves[1][3], - leaves[1][2], - leaves[1][1], - leaves[1][0], - Felt::new(2), - Felt::new(1), - tree.root()[3], - tree.root()[2], - tree.root()[1], - tree.root()[0], - ]); - assert_eq!(expected_stack, process.stack.trace_state()); - } - - // HELPER FUNCTIONS - // -------------------------------------------------------------------------------------------- - fn init_leaf(value: u64) -> Word { - [Felt::new(value), Felt::ZERO, Felt::ZERO, Felt::ZERO] - } - - fn build_expected(values: &[Felt]) -> [Felt; 16] { - let mut expected = [Felt::ZERO; 16]; - for (&value, result) in values.iter().zip(expected.iter_mut()) { - *result = value - } - expected - } -} diff --git a/processor/src/decorators/tests.rs b/processor/src/decorators/tests.rs new file mode 100644 index 0000000000..88d46f02a7 --- /dev/null +++ b/processor/src/decorators/tests.rs @@ -0,0 +1,187 @@ +use super::{ + super::{AdviceInputs, Felt, FieldElement, Kernel, Operation, StarkField}, + Process, ZERO, +}; +use crate::{MemAdviceProvider, StackInputs, Word}; +use test_utils::{crypto::get_smt_remaining_key, rand::seeded_word}; +use vm_core::{ + crypto::{ + hash::Rpo256, + merkle::{EmptySubtreeRoots, MerkleStore, MerkleTree, NodeIndex}, + }, + utils::IntoBytes, + AdviceInjector, Decorator, ONE, +}; + +#[test] +fn inject_merkle_node() { + let leaves = [init_leaf(1), init_leaf(2), init_leaf(3), init_leaf(4)]; + let tree = MerkleTree::new(leaves.to_vec()).unwrap(); + let store = MerkleStore::default().with_merkle_tree(leaves).unwrap(); + let stack_inputs = [ + tree.root()[0].as_int(), + tree.root()[1].as_int(), + tree.root()[2].as_int(), + tree.root()[3].as_int(), + 1, + tree.depth() as u64, + ]; + + let stack_inputs = StackInputs::try_from_values(stack_inputs).unwrap(); + let advice_inputs = AdviceInputs::default().with_merkle_store(store); + let advice_provider = MemAdviceProvider::from(advice_inputs); + let mut process = Process::new(Kernel::default(), stack_inputs, advice_provider); + process.execute_op(Operation::Noop).unwrap(); + + // push the node onto the advice stack + process + .execute_decorator(&Decorator::Advice(AdviceInjector::MerkleNode)) + .unwrap(); + + // pop the node from the advice stack and push it onto the operand stack + process.execute_op(Operation::AdvPop).unwrap(); + process.execute_op(Operation::AdvPop).unwrap(); + process.execute_op(Operation::AdvPop).unwrap(); + process.execute_op(Operation::AdvPop).unwrap(); + + let expected_stack = build_expected(&[ + leaves[1][3], + leaves[1][2], + leaves[1][1], + leaves[1][0], + Felt::new(2), + Felt::new(1), + tree.root()[3], + tree.root()[2], + tree.root()[1], + tree.root()[0], + ]); + assert_eq!(expected_stack, process.stack.trace_state()); +} + +#[test] +fn inject_smtget() { + // setup the test + let empty = EmptySubtreeRoots::empty_hashes(64); + let initial_root = Word::from(empty[0]); + let mut seed = 0xfb; + let key = seeded_word(&mut seed); + let value = seeded_word(&mut seed); + + // check leaves on empty trees + for depth in [16, 32, 48, 64] { + // compute the remaining key + let remaining = get_smt_remaining_key(key, depth); + + // compute node value + let depth = Felt::new(depth as u64); + let store = MerkleStore::new(); + let node = Rpo256::merge_in_domain(&[remaining.into(), value.into()], depth).into(); + + // expect absent value with constant depth 16 + let expected = [ONE, ONE, ZERO, ZERO, ZERO, ZERO]; + assert_case_smtget(key, value, node, initial_root, store, &expected); + } + + // check leaves inserted on all tiers + for depth in [16, 32, 48, 64] { + // compute the remaining key + let remaining = get_smt_remaining_key(key, depth); + + // set depth flags + let is_16_or_32 = (depth == 16 || depth == 32).then_some(ONE).unwrap_or(ZERO); + let is_16_or_48 = (depth == 16 || depth == 48).then_some(ONE).unwrap_or(ZERO); + + // compute node value + let index = key[3].as_int() >> 64 - depth; + let index = NodeIndex::new(depth, index).unwrap(); + let depth = Felt::new(depth as u64); + let node = Rpo256::merge_in_domain(&[remaining.into(), value.into()], depth).into(); + + // set tier node value and expect the value from the injector + let mut store = MerkleStore::new(); + let root = store.set_node(initial_root, index, node).unwrap().root; + let expected = [is_16_or_32, is_16_or_48, value[3], value[2], value[1], value[0]]; + assert_case_smtget(key, value, node, root, store, &expected); + } + + // check absent siblings of non-empty trees + for depth in [16, 32, 48, 64] { + // set depth flags + let is_16_or_32 = (depth == 16 || depth == 32).then_some(ONE).unwrap_or(ZERO); + let is_16_or_48 = (depth == 16 || depth == 48).then_some(ONE).unwrap_or(ZERO); + + // compute the index of the absent node + let index = key[3].as_int() >> 64 - depth; + let index = NodeIndex::new(depth, index).unwrap(); + + // compute the sibling index of the target with its remaining key and node + let sibling = index.sibling(); + let mut sibling_key = key; + sibling_key[3] = Felt::new(sibling.value() >> depth.min(63)); + let sibling_node = + Rpo256::merge_in_domain(&[sibling_key.into(), value.into()], Felt::new(depth as u64)) + .into(); + + // run the text, expecting absent target node + let mut store = MerkleStore::new(); + let root = store.set_node(initial_root, sibling, sibling_node).unwrap().root; + let expected = [is_16_or_32, is_16_or_48, ZERO, ZERO, ZERO, ZERO]; + assert_case_smtget(key, value, sibling_node, root, store, &expected); + } +} + +// HELPER FUNCTIONS +// -------------------------------------------------------------------------------------------- + +fn init_leaf(value: u64) -> Word { + [Felt::new(value), Felt::ZERO, Felt::ZERO, Felt::ZERO] +} + +fn build_expected(values: &[Felt]) -> [Felt; 16] { + let mut expected = [Felt::ZERO; 16]; + for (&value, result) in values.iter().zip(expected.iter_mut()) { + *result = value + } + expected +} + +fn assert_case_smtget( + key: Word, + value: Word, + node: Word, + root: Word, + store: MerkleStore, + expected_stack: &[Felt], +) { + // build the process + let stack_inputs = StackInputs::try_from_values([ + root[0].as_int(), + root[1].as_int(), + root[2].as_int(), + root[3].as_int(), + key[0].as_int(), + key[1].as_int(), + key[2].as_int(), + key[3].as_int(), + ]) + .unwrap(); + let advice_inputs = AdviceInputs::default() + .with_merkle_store(store) + .with_map([(node.into_bytes(), value.into_iter().collect())]); + let advice_provider = MemAdviceProvider::from(advice_inputs); + let mut process = Process::new(Kernel::default(), stack_inputs, advice_provider); + + // call the injector and clear the stack + process.execute_op(Operation::Noop).unwrap(); + process.execute_decorator(&Decorator::Advice(AdviceInjector::SmtGet)).unwrap(); + for _ in 0..8 { + process.execute_op(Operation::Drop).unwrap(); + } + + // expect the stack output + for _ in 0..expected_stack.len() { + process.execute_op(Operation::AdvPop).unwrap(); + } + assert_eq!(build_expected(expected_stack), process.stack.trace_state()); +} diff --git a/processor/src/errors.rs b/processor/src/errors.rs index e64ad09f7a..b6b0f29399 100644 --- a/processor/src/errors.rs +++ b/processor/src/errors.rs @@ -18,6 +18,7 @@ pub enum ExecutionError { AdviceKeyNotFound(Word), AdviceStackReadFailed(u32), InvalidNodeIndex { depth: Felt, value: Felt }, + InvalidTreeDepth { depth: Felt }, MerkleUpdateInPlace, MerkleStoreLookupFailed(MerkleError), MerkleStoreUpdateFailed(MerkleError), @@ -55,6 +56,9 @@ impl Display for ExecutionError { fmt, "The provided index {value} is out of bounds for a node at depth {depth}" ), + InvalidTreeDepth { depth } => { + write!(fmt, "The provided {depth} is out of bounds and cannot be represented as an unsigned 8-bits integer") + } MerkleUpdateInPlace => write!(fmt, "Update in place is not supported"), MerkleStoreLookupFailed(reason) => { write!(fmt, "Advice provider Merkle store backend lookup failed: {reason}") diff --git a/stdlib/asm/collections/smt.masm b/stdlib/asm/collections/smt.masm new file mode 100644 index 0000000000..2d5eed87d2 --- /dev/null +++ b/stdlib/asm/collections/smt.masm @@ -0,0 +1,245 @@ +#! Produces a remaining path key and index for depth 16 +#! +#! Input: [v, ...] +#! Output: [(v << 16) >> 16, v >> 48, ...] +#! +#! Cycles: 7 +proc.split_16 + u32split + u32unchecked_divmod.65536 + mul.4294967296 + movup.2 + add +end + +#! Produces a remaining path key and index for depth 48 +#! +#! Input: [v, ...] +#! Output: [(v << 48) >> 48, v >> 16, ...] +#! +#! Cycles: 10 +proc.split_48 + u32split + swap + u32unchecked_divmod.65536 + swap + movup.2 + mul.65536 + add + swap +end + +#! Get the leaf value for depth 16 +#! +#! Input: [V, K, R, ...] +#! Output: [V, R, ...] +#! +#! Cycles: 101 +proc.get_16 + # compute the remaining path and index `i` + # cycles: 15 + swapw exec.split_16 swap movdn.12 + # => [K', V, R, i, ...] + + # prepare the permutation argument for hash in domain + # cycles: 16 + swapw dupw movdnw.2 push.0.16.0.0 movdnw.2 + # => [V, K', D, V, R, i, ...] + + # execute the permutation and preserve the digest + # cycles: 10 + hperm dropw swapw dropw + # => [N, V, R, i, ...] + + # set the flag if the value is zero + # cycles: 19 + dup.7 eq.0 dup.7 eq.0 and dup.6 eq.0 and dup.6 eq.0 and + # => [V == 0, N, V, R, i, ...] + + # conditionally select empty constant or node value + # cycles: 11 + push.17483286922353768131.353378057542380712 + push.1935183237414585408.4820339620987989650 + movup.4 cdropw + # => [N', V, R, i, ...] + + # check the opening value + # cycles: 20 + movupw.2 movup.12 push.16 mtree_get movupw.2 + # => [V', N' R, V, ...] + + # assert the opening matches the computed node value and return + # cycles: 12 + assert_eqw swapw + # => [V, R, ...] +end + +#! Get the leaf value for depth 32 +#! +#! Input: [V, K, R, ...] +#! Output: [V, R, ...] +# +#! Cycles: 94 +proc.get_32 + # compute the remaining path and index `i` + # cycles: 6 + swapw u32split movdn.12 + # => [K', V, R, i, ...] + + # prepare the permutation argument for hash in domain + # cycles: 16 + swapw dupw movdnw.2 push.0.32.0.0 movdnw.2 + # => [V, K', D, V, R, i, ...] + + # execute the permutation and preserve the digest + # cycles: 10 + hperm dropw swapw dropw + # => [N, V, R, i, ...] + + # set the flag if the value is zero + # cycles: 19 + dup.7 eq.0 dup.7 eq.0 and dup.6 eq.0 and dup.6 eq.0 and + # => [V == 0, N, V, R, i, ...] + + # conditionally select empty constant or node value + # cycles: 11 + push.11677748883385181208.15891398395707500576 + push.3790704659934033620.2126099371106695189 + movup.4 cdropw + # => [N', V, R, i, ...] + + # check the opening value + # cycles: 20 + movupw.2 movup.12 push.32 mtree_get movupw.2 + # => [V', N' R, V, ...] + + # assert the opening matches the computed node value and return + # cycles: 12 + assert_eqw swapw + # => [V, R, ...] +end + +#! Get the leaf value for depth 48 +#! +#! Input: [V, K, R, ...] +#! Output: [V, R, ...] +# +#! Cycles: 104 +proc.get_48 + # compute the remaining path and index `i` + # cycles: 18 + swapw exec.split_48 swap movdn.12 + # => [K', V, R, i, ...] + + # prepare the permutation argument for hash in domain + # cycles: 16 + swapw dupw movdnw.2 push.0.48.0.0 movdnw.2 + # => [V, K', D, V, R, i, ...] + + # execute the permutation and preserve the digest + # cycles: 10 + hperm dropw swapw dropw + # => [N, V, R, i, ...] + + # set the flag if the value is zero + # cycles: 19 + dup.7 eq.0 dup.7 eq.0 and dup.6 eq.0 and dup.6 eq.0 and + # => [V == 0, N, V, R, i, ...] + + # conditionally select empty constant or node value + # cycles: 11 + push.10650694022550988030.5634734408638476525 + push.9233115969432897632.1437907447409278328 + movup.4 cdropw + # => [N', V, R, i, ...] + + # check the opening value + # cycles: 20 + movupw.2 movup.12 push.48 mtree_get movupw.2 + # => [V', N' R, V, ...] + + # assert the opening matches the computed node value and return + # cycles: 12 + assert_eqw swapw + # => [V, R, ...] +end + +#! Get the leaf value for depth 64 +#! +#! Input: [V, K, R, ...] +#! Output: [V, R, ...] +# +#! Cycles: 94 +proc.get_64 + # compute the remaining path and index `i` + # cycles: 6 + swapw movdn.11 push.0 + # => [K', V, R, i, ...] + + # prepare the permutation argument for hash in domain + # cycles: 16 + swapw dupw movdnw.2 push.0.64.0.0 movdnw.2 + # => [V, K', D, V, R, i, ...] + + # execute the permutation and preserve the digest + # cycles: 10 + hperm dropw swapw dropw + # => [N, V, R, i, ...] + + # set the flag if the value is zero + # cycles: 19 + dup.7 eq.0 dup.7 eq.0 and dup.6 eq.0 and dup.6 eq.0 and + # => [V == 0, N, V, R, i, ...] + + # conditionally select empty constant or node value + # cycles: 11 + padw movup.4 cdropw + # => [N', V, R, i, ...] + + # check the opening value + # cycles: 20 + movupw.2 movup.12 push.64 mtree_get movupw.2 + # => [V', N' R, V, ...] + + # assert the opening matches the computed node value and return + # cycles: 12 + assert_eqw swapw + # => [V, R, ...] +end + +#! Returns the value stored under the specified key in a Sparse Merkle tree with the specified root. +#! +#! If the value for a given key has not been set, the returned `V` will consist of all zeroes. +#! +#! Input: [K, R, ...] +#! Output: [V, R, ...] +#! +#! Depth 16: 105 cycles +#! Depth 32: 98 cycles +#! Depth 48: 108 cycles +#! Depth 64: 98 cycles +export.get + # invoke adv and fetch target depth flags + adv.smtget adv_push.6 + # => [d in {16, 32}, d in {16, 48}, V, K, R, ...] + + # call the inner procedure depending on the depth + if.true + if.true + # depth 16 + exec.get_16 + else + # depth 32 + exec.get_32 + end + else + if.true + # depth 48 + exec.get_48 + else + # depth 64 + exec.get_64 + end + end + # => [V, R, ...] +end diff --git a/stdlib/docs/smt_collections.md b/stdlib/docs/smt_collections.md new file mode 100644 index 0000000000..d13e8314fe --- /dev/null +++ b/stdlib/docs/smt_collections.md @@ -0,0 +1,5 @@ + +## std::collections::smt +| Procedure | Description | +| ----------- | ------------- | +| get | Returns the value stored under the specified key in a Sparse Merkle tree with the specified root.

If the value for a given key has not been set, the returned `V` will consist of all zeroes.

Input: [K, R, ...]

Output: [V, R, ...]

Depth 16: 105 cycles

Depth 32: 98 cycles

Depth 48: 108 cycles

Depth 64: 98 cycles | diff --git a/stdlib/tests/collections/mod.rs b/stdlib/tests/collections/mod.rs index 488e8fa8e5..a4a55187ba 100644 --- a/stdlib/tests/collections/mod.rs +++ b/stdlib/tests/collections/mod.rs @@ -1 +1,2 @@ mod mmr; +mod smt; diff --git a/stdlib/tests/collections/smt.rs b/stdlib/tests/collections/smt.rs new file mode 100644 index 0000000000..0715fa62ed --- /dev/null +++ b/stdlib/tests/collections/smt.rs @@ -0,0 +1,162 @@ +use crate::build_test; +use test_utils::{ + crypto::get_smt_remaining_key, rand::seeded_word, EmptySubtreeRoots, Felt, IntoBytes, + MerkleStore, NodeIndex, Rpo256, StarkField, Word, ONE, WORD_SIZE, ZERO, +}; + +#[test] +fn smtget_opens_correctly_from_empty_tree() { + // setup the test + let empty = EmptySubtreeRoots::empty_hashes(64); + let initial_root = Word::from(empty[0]); + let mut seed = 1 << 40; + let key = seeded_word(&mut seed); + let value = seeded_word(&mut seed); + + // check empty tree + let store = MerkleStore::new(); + let advice_map = []; + let expected = [ZERO; WORD_SIZE]; + assert_smt_get_opens_correctly(key, expected, initial_root, store, &advice_map); + + // check included leaves for all tiers + for depth in [16, 32, 48, 64] { + // compute the index and node value + let index = key[3].as_int(); + let index = index >> (64 - depth); + let index = NodeIndex::new(depth as u8, index).unwrap(); + let remaining = get_smt_remaining_key(key, depth as u8); + let node = + Rpo256::merge_in_domain(&[remaining.into(), value.into()], Felt::new(depth)).into(); + + // setup the store and run the test + let mut store = MerkleStore::new(); + let root = store.set_node(initial_root, index, node).unwrap().root; + let advice_map = [(node.into_bytes(), value.to_vec())]; + assert_smt_get_opens_correctly(key, value, root, store, &advice_map); + } +} + +#[test] +fn smtget_opens_correctly_from_tree_with_sibling() { + // setup the test + let empty = EmptySubtreeRoots::empty_hashes(64); + let initial_root = Word::from(empty[0]); + let mut seed = 1 << 40; + let key = seeded_word(&mut seed); + let value = seeded_word(&mut seed); + + // check empty tree + let store = MerkleStore::new(); + let advice_map = []; + let expected = [ZERO; WORD_SIZE]; + assert_smt_get_opens_correctly(key, expected, initial_root, store, &advice_map); + + // check included leaves for all tiers + for depth in [16, 32, 48, 64] { + // init a new storage + let mut store = MerkleStore::new(); + + // compute the index and node value + let index = key[3].as_int(); + let index = index >> (64 - depth); + let index = NodeIndex::new(depth as u8, index).unwrap(); + let remaining = get_smt_remaining_key(key, depth as u8); + let node = + Rpo256::merge_in_domain(&[remaining.into(), value.into()], Felt::new(depth)).into(); + + // insert the target node into the storage + let root = store.set_node(initial_root, index, node).unwrap().root; + + // compute the index and sibling value + let sibling = key[3].as_int(); + let sibling = sibling >> (64 - depth); + let sibling = sibling ^ 1; + let sibling = NodeIndex::new(depth as u8, sibling).unwrap(); + let mut sibling_node = Word::from(node); + sibling_node[3] = sibling_node[3] + ONE; + let mut sibling_value = Word::from(value); + sibling_value[3] = sibling_value[3] + ONE; + + // insert the sibling node into the storage + let root = store.set_node(root, sibling, sibling_node).unwrap().root; + + // execute the test and expect the right node to be returned + let advice_map = [ + (node.into_bytes(), value.to_vec()), + (sibling_node.into_bytes(), sibling_value.to_vec()), + ]; + assert_smt_get_opens_correctly(key, value, root, store, &advice_map); + } +} + +#[test] +fn smtget_opens_correctly_for_empty_leaf() { + // setup the test + let empty = EmptySubtreeRoots::empty_hashes(64); + let mut seed = 1 << 50; + + // test different cases where the path diverges on the target + for diverge in [16, 32, 48, 64] { + // setup a new store + let mut store = MerkleStore::new(); + let root = Word::from(empty[0]); + + // generate a new key/value pair + let key = seeded_word(&mut seed); + let value = seeded_word(&mut seed); + + // compute a sibling index and node value + let mut sibling = key[3].as_int(); + sibling ^= 1 << (64 - diverge); + let remaining = get_smt_remaining_key(key, diverge as u8); + let sibling = NodeIndex::new(diverge as u8, sibling >> (64 - diverge)).unwrap(); + let node = + Rpo256::merge_in_domain(&[remaining.into(), value.into()], Felt::new(diverge)).into(); + + // expect ZERO as the leaf isn't included + let root = store.set_node(root, sibling, node).unwrap().root; + let advice_map = [(node.into_bytes(), value.to_vec())]; + let expected = [ZERO; WORD_SIZE]; + assert_smt_get_opens_correctly(key, expected, root, store.clone(), &advice_map); + } +} + +fn assert_smt_get_opens_correctly( + key: Word, + value: Word, + root: Word, + store: MerkleStore, + advice_map: &[([u8; 32], Vec)], +) { + let source = r#" + use.std::collections::smt + + begin + exec.smt::get + end + "#; + let initial_stack = [ + root[0].as_int(), + root[1].as_int(), + root[2].as_int(), + root[3].as_int(), + key[0].as_int(), + key[1].as_int(), + key[2].as_int(), + key[3].as_int(), + ]; + let expected_output = [ + value[3].as_int(), + value[2].as_int(), + value[1].as_int(), + value[0].as_int(), + root[3].as_int(), + root[2].as_int(), + root[1].as_int(), + root[0].as_int(), + ]; + let advice_stack = []; + build_test!(source, &initial_stack, &advice_stack, store, advice_map.iter().cloned()) + .expect_stack(&expected_output); +} diff --git a/test-utils/src/crypto.rs b/test-utils/src/crypto.rs index a0bcb8f7a6..4a4f06815a 100644 --- a/test-utils/src/crypto.rs +++ b/test-utils/src/crypto.rs @@ -1,4 +1,4 @@ -use super::{Felt, FieldElement, Vec, Word}; +use super::{Felt, FieldElement, StarkField, Vec, Word}; // RE-EXPORTS // ================================================================================================ @@ -26,3 +26,13 @@ pub fn init_merkle_leaves(values: &[u64]) -> Vec { pub fn init_merkle_leaf(value: u64) -> Word { [Felt::new(value), Felt::ZERO, Felt::ZERO, Felt::ZERO] } + +/// Returns a remaining path key for a Sparse Merkle tree +pub fn get_smt_remaining_key(mut key: Word, depth: u8) -> Word { + key[3] = Felt::new(match depth { + 16 | 32 | 48 => (key[3].as_int() << depth) >> depth, + 64 => 0, + _ => unreachable!(), + }); + key +} diff --git a/test-utils/src/lib.rs b/test-utils/src/lib.rs index 5133d4556d..55c94b621b 100644 --- a/test-utils/src/lib.rs +++ b/test-utils/src/lib.rs @@ -23,6 +23,10 @@ pub use prover::{MemAdviceProvider, ProofOptions}; pub use test_case::test_case; pub use verifier::ProgramInfo; pub use vm_core::{ + crypto::{ + hash::Rpo256, + merkle::{EmptySubtreeRoots, MerkleStore, NodeIndex}, + }, stack::STACK_TOP_SIZE, utils::{collections, group_slice_elements, group_vector_elements, IntoBytes, ToElements}, Felt, FieldElement, Program, StarkField, Word, ONE, WORD_SIZE, ZERO, @@ -38,15 +42,15 @@ pub mod serde { pub mod crypto; +#[cfg(not(target_family = "wasm"))] +pub mod rand; + mod test_builders; pub use test_builders::*; #[cfg(not(target_family = "wasm"))] pub use proptest; -#[cfg(not(target_family = "wasm"))] -pub use rand_utils as rand; - // TYPE ALIASES // ================================================================================================ diff --git a/test-utils/src/rand.rs b/test-utils/src/rand.rs new file mode 100644 index 0000000000..6d34e524de --- /dev/null +++ b/test-utils/src/rand.rs @@ -0,0 +1,35 @@ +use super::{Felt, Word}; + +pub use rand_utils::*; + +// SEEDED GENERATORS +// ================================================================================================ + +/// Mutates a seed and generates a word deterministically +pub fn seeded_word(seed: &mut u64) -> Word { + let seed = generate_bytes_seed(seed); + prng_array(seed) +} + +/// Mutates a seed and generates an element deterministically +pub fn seeded_element(seed: &mut u64) -> Felt { + let seed = generate_bytes_seed(seed); + let num = prng_array::(seed)[0]; + Felt::new(num) +} + +// HELPERS +// ================================================================================================ + +/// Generate a bytes seed that can be used as input for rand_utils. +/// +/// Increments the argument. +fn generate_bytes_seed(seed: &mut u64) -> [u8; 32] { + // increment the seed + *seed = seed.wrapping_add(1); + + // generate a bytes seed + let mut bytes = [0u8; 32]; + bytes[..8].copy_from_slice(&seed.to_le_bytes()); + bytes +}