diff --git a/assembly/src/ast/nodes/advice.rs b/assembly/src/ast/nodes/advice.rs index 413fc31bd5..00d054da9e 100644 --- a/assembly/src/ast/nodes/advice.rs +++ b/assembly/src/ast/nodes/advice.rs @@ -18,6 +18,7 @@ pub enum AdviceInjectorNode { PushU64div, PushExt2intt, PushSmtGet, + PushSmtInsert, PushMapVal, PushMapValImm { offset: u8 }, PushMapValN, @@ -35,6 +36,7 @@ impl From<&AdviceInjectorNode> for AdviceInjector { PushU64div => Self::DivU64, PushExt2intt => Self::Ext2Intt, PushSmtGet => Self::SmtGet, + PushSmtInsert => Self::SmtInsert, PushMapVal => Self::MapValueToStack { include_len: false, key_offset: 0, @@ -68,6 +70,7 @@ impl fmt::Display for AdviceInjectorNode { PushU64div => write!(f, "push_u64div"), PushExt2intt => write!(f, "push_ext2intt"), PushSmtGet => write!(f, "push_smtget"), + PushSmtInsert => write!(f, "push_smtinsert"), PushMapVal => write!(f, "push_mapval"), PushMapValImm { offset } => write!(f, "push_mapval.{offset}"), PushMapValN => write!(f, "push_mapvaln"), @@ -86,14 +89,15 @@ impl fmt::Display for AdviceInjectorNode { const PUSH_U64DIV: u8 = 0; const PUSH_EXT2INTT: u8 = 1; const PUSH_SMTGET: u8 = 2; -const PUSH_MAPVAL: u8 = 3; -const PUSH_MAPVAL_IMM: u8 = 4; -const PUSH_MAPVALN: u8 = 5; -const PUSH_MAPVALN_IMM: u8 = 6; -const PUSH_MTNODE: u8 = 7; -const INSERT_MEM: u8 = 8; -const INSERT_HDWORD: u8 = 9; -const INSERT_HDWORD_IMM: u8 = 10; +const PUSH_SMTINSERT: u8 = 3; +const PUSH_MAPVAL: u8 = 4; +const PUSH_MAPVAL_IMM: u8 = 5; +const PUSH_MAPVALN: u8 = 6; +const PUSH_MAPVALN_IMM: u8 = 7; +const PUSH_MTNODE: u8 = 8; +const INSERT_MEM: u8 = 9; +const INSERT_HDWORD: u8 = 10; +const INSERT_HDWORD_IMM: u8 = 11; impl Serializable for AdviceInjectorNode { fn write_into(&self, target: &mut W) { @@ -102,6 +106,7 @@ impl Serializable for AdviceInjectorNode { PushU64div => target.write_u8(PUSH_U64DIV), PushExt2intt => target.write_u8(PUSH_EXT2INTT), PushSmtGet => target.write_u8(PUSH_SMTGET), + PushSmtInsert => target.write_u8(PUSH_SMTINSERT), PushMapVal => target.write_u8(PUSH_MAPVAL), PushMapValImm { offset } => { target.write_u8(PUSH_MAPVAL_IMM); @@ -129,6 +134,7 @@ impl Deserializable for AdviceInjectorNode { PUSH_U64DIV => Ok(AdviceInjectorNode::PushU64div), PUSH_EXT2INTT => Ok(AdviceInjectorNode::PushExt2intt), PUSH_SMTGET => Ok(AdviceInjectorNode::PushSmtGet), + PUSH_SMTINSERT => Ok(AdviceInjectorNode::PushSmtInsert), PUSH_MAPVAL => Ok(AdviceInjectorNode::PushMapVal), PUSH_MAPVAL_IMM => { let offset = source.read_u8()?; diff --git a/assembly/src/ast/parsers/adv_ops.rs b/assembly/src/ast/parsers/adv_ops.rs index 4c7042484d..1393796a05 100644 --- a/assembly/src/ast/parsers/adv_ops.rs +++ b/assembly/src/ast/parsers/adv_ops.rs @@ -33,6 +33,10 @@ pub fn parse_adv_inject(op: &Token) -> Result { 2 => AdvInject(PushSmtGet), _ => return Err(ParsingError::extra_param(op)), }, + "push_smtinsert" => match op.num_parts() { + 2 => AdvInject(PushSmtInsert), + _ => return Err(ParsingError::extra_param(op)), + }, "push_mapval" => match op.num_parts() { 2 => AdvInject(PushMapVal), 3 => { diff --git a/core/src/operations/decorators/advice.rs b/core/src/operations/decorators/advice.rs index d583c22f57..d87ec83219 100644 --- a/core/src/operations/decorators/advice.rs +++ b/core/src/operations/decorators/advice.rs @@ -181,6 +181,9 @@ pub enum AdviceInjector { /// Where KEY is computed as hash(A || B, domain), where domain is provided via the immediate /// value. HdwordToMap { domain: Felt }, + + /// TODO: add docs + SmtInsert, } impl fmt::Display for AdviceInjector { @@ -202,6 +205,7 @@ impl fmt::Display for AdviceInjector { Self::Ext2Inv => write!(f, "ext2_inv"), Self::Ext2Intt => write!(f, "ext2_intt"), Self::SmtGet => write!(f, "smt_get"), + Self::SmtInsert => write!(f, "smt_insert"), Self::MemToMap => write!(f, "mem_to_map"), Self::HdwordToMap { domain } => write!(f, "hdword_to_map.{domain}"), } diff --git a/processor/src/decorators/adv_stack_injectors.rs b/processor/src/decorators/adv_stack_injectors.rs index 1d9a2e6f2d..ec0f0ae8ca 100644 --- a/processor/src/decorators/adv_stack_injectors.rs +++ b/processor/src/decorators/adv_stack_injectors.rs @@ -338,6 +338,120 @@ where Ok(()) } + + /// Pushes values onto the advice stack which are required for successful insertion of a + /// key-value pair into a Sparse Merkle Tree data structure. + /// + /// The Sparse Merkle Tree is tiered, meaning it will have leaf depths in `{16, 32, 48, 64}`. + /// + /// Inputs: + /// Operand stack: [VALUE, KEY, ROOT, ...] + /// Advice stack: [...] + /// + /// Outputs: + /// Operand stack: [OLD_VALUE, NEW_ROOT, ...] + /// Advice stack, depends on the type of insert: + /// - Simple insert at depth 16: [d0, d1, ONE (is_simple_insert), ZERO (is_update)] + /// - Simple insert at depth 32 or 48: [d0, d1, ONE (is_simple_insert), ZERO (is_update), P_NODE] + /// - Update of an existing leaf: [ZERO (padding), d0, d1, ONE (is_update), OLD_VALUE] + /// + /// Where: + /// - d0 is a boolean flag set to `1` if the depth is `16` or `48`. + /// - d1 is a boolean flag set to `1` if the depth is `16` or `32`. + /// - P_NODE is an internal node located at the tier above the insert tier. + /// - VALUE is the value to be inserted. + /// - OLD_VALUE is the value previously associated with the specified KEY. + /// - ROOT and NEW_ROOT are the roots of the TSMT prior and post the insert respectively. + /// + /// # Errors + /// Will return an error if the provided Merkle root doesn't exist on the advice provider. + /// + /// # Panics + /// Will panic as unimplemented if the target depth is `64`. + pub(super) fn push_smtinsert_inputs(&mut self) -> Result<(), ExecutionError> { + // get the key and tree root from the stack + let key = [self.stack.get(7), self.stack.get(6), self.stack.get(5), self.stack.get(4)]; + let root = [self.stack.get(11), self.stack.get(10), self.stack.get(9), self.stack.get(8)]; + + // determine the depth of the first leaf or an empty tree node + let index = &key[3]; + let depth = self.advice_provider.get_leaf_depth(root, &SMT_MAX_TREE_DEPTH, index)?; + debug_assert!(depth < 65); + + // map the depth value to its tier; this rounds up depth to 16, 32, 48, or 64 + let depth = SMT_NORMALIZED_DEPTHS[depth as usize]; + if depth == 64 { + unimplemented!("handling of depth=64 tier hasn't been implemented yet"); + } + + // get the value of the node a this index/depth + 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)?; + + // figure out what kind of insert we are doing; possible options are: + // - if the node is a root of an empty subtree, this is a simple insert. + // - if the node is a leaf, this could be either an update (for the same key), or a + // complex insert (i.e., the existing leaf needs to be moved to a lower tier). + let empty = EmptySubtreeRoots::empty_hashes(64)[depth as usize]; + let (is_update, is_simple_insert) = if node == Word::from(empty) { + // handle simple insert case + if depth == 32 || depth == 48 { + // for depth 32 and 48, we need to provide the internal node located on the tier + // above the insert tier + let p_index = Felt::from(index.as_int() >> 16); + let p_depth = Felt::from(depth - 16); + let p_node = self.advice_provider.get_tree_node(root, &p_depth, &p_index)?; + for &element in p_node.iter().rev() { + self.advice_provider.push_stack(AdviceSource::Value(element))?; + } + } + + // return is_update = ZERO, is_simple_insert = ONE + (ZERO, ONE) + } else { + // if the node is a leaf node, push the elements mapped to this node onto the advice + // stack; the elements should be [KEY, VALUE], with key located at the top of the + // advice stack. + self.advice_provider.push_stack(AdviceSource::Map { + key: node, + include_len: false, + })?; + + // remove the KEY from the advice stack, leaving only the VALUE on the stack + let leaf_key = self.advice_provider.pop_stack_word()?; + + // if the key for the value to be inserted is the same as the leaf's key, we are + // dealing with a simple update. otherwise, we are dealing with a complex insert + // (i.e., the leaf needs to be moved to a lower tier). + if leaf_key == key { + // return is_update = ONE, is_simple_insert = ZERO + (ONE, ZERO) + } else { + // return is_update = ZERO, is_simple_insert = ZERO + (ZERO, ZERO) + } + }; + + // set the flags used to determine which tier the insert is happening at + 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_update))?; + if is_update == ONE { + // for update we don't need to specify whether we are dealing with an insert; but we + // insert an extra ONE at the end so that we can read 4 values from the advice stack + // regardless of which branch is taken. + self.advice_provider.push_stack(AdviceSource::Value(is_16_or_32))?; + self.advice_provider.push_stack(AdviceSource::Value(is_16_or_48))?; + self.advice_provider.push_stack(AdviceSource::Value(ZERO))?; + } else { + self.advice_provider.push_stack(AdviceSource::Value(is_simple_insert))?; + self.advice_provider.push_stack(AdviceSource::Value(is_16_or_32))?; + self.advice_provider.push_stack(AdviceSource::Value(is_16_or_48))?; + } + Ok(()) + } } // HELPER FUNCTIONS diff --git a/processor/src/decorators/mod.rs b/processor/src/decorators/mod.rs index 27b92166fa..811f53b712 100644 --- a/processor/src/decorators/mod.rs +++ b/processor/src/decorators/mod.rs @@ -45,6 +45,7 @@ where AdviceInjector::Ext2Inv => self.push_ext2_inv_result(), AdviceInjector::Ext2Intt => self.push_ext2_intt_result(), AdviceInjector::SmtGet => self.push_smtget_inputs(), + AdviceInjector::SmtInsert => self.push_smtinsert_inputs(), AdviceInjector::MemToMap => self.insert_mem_values_into_adv_map(), AdviceInjector::HdwordToMap { domain } => self.insert_hdword_into_adv_map(*domain), } diff --git a/processor/src/decorators/tests.rs b/processor/src/decorators/tests.rs index 17aee7542c..6cccdd8eb5 100644 --- a/processor/src/decorators/tests.rs +++ b/processor/src/decorators/tests.rs @@ -3,11 +3,11 @@ use super::{ Process, }; use crate::{MemAdviceProvider, StackInputs, Word}; -use test_utils::{crypto::get_smt_remaining_key, rand::seeded_word}; +use test_utils::rand::seeded_word; use vm_core::{ crypto::{ hash::{Rpo256, RpoDigest}, - merkle::{EmptySubtreeRoots, MerkleStore, MerkleTree, NodeIndex}, + merkle::{EmptySubtreeRoots, MerkleStore, MerkleTree, NodeIndex, TieredSmt}, }, utils::IntoBytes, AdviceInjector, Decorator, ONE, ZERO, @@ -60,6 +60,9 @@ fn push_merkle_node() { assert_eq!(expected_stack, process.stack.trace_state()); } +// SMTGET TESTS +// ================================================================================================ + #[test] fn push_smtget() { // setup the test @@ -71,13 +74,10 @@ fn push_smtget() { // check leaves on empty trees for depth in [16, 32, 48] { - // compute the remaining key - let remaining = get_smt_remaining_key(key, depth); - // compute node value let depth_element = Felt::from(depth); let store = MerkleStore::new(); - let node = Rpo256::merge_in_domain(&[remaining.into(), value.into()], depth_element); + let node = Rpo256::merge_in_domain(&[key.into(), value.into()], depth_element); // expect absent value with constant depth 16 let expected = [ZERO, ZERO, ZERO, ZERO, ZERO, ZERO, ZERO, ZERO, ZERO, ONE, ONE]; @@ -86,9 +86,6 @@ fn push_smtget() { // check leaves inserted on all tiers for depth in [16, 32, 48] { - // 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); @@ -97,7 +94,7 @@ fn push_smtget() { let index = key[3].as_int() >> 64 - depth; let index = NodeIndex::new(depth, index).unwrap(); let depth_element = Felt::from(depth); - let node = Rpo256::merge_in_domain(&[remaining.into(), value.into()], depth_element); + let node = Rpo256::merge_in_domain(&[key.into(), value.into()], depth_element); // set tier node value and expect the value from the injector let mut store = MerkleStore::new(); @@ -108,10 +105,10 @@ fn push_smtget() { value[2], value[1], value[0], - remaining[3], - remaining[2], - remaining[1], - remaining[0], + key[3], + key[2], + key[1], + key[0], is_16_or_32, is_16_or_48, ]; @@ -144,8 +141,81 @@ fn push_smtget() { } } +// SMTINSERT TESTS +// ================================================================================================ + +#[test] +fn inject_smtinsert() { + let mut smt = TieredSmt::default(); + + // --- insert into empty tree --------------------------------------------- + + let raw_a = 0b_01101001_01101100_00011111_11111111_10010110_10010011_11100000_00000000_u64; + let key_a = build_key(raw_a); + let val_a = [Felt::new(3), Felt::new(5), Felt::new(7), Felt::new(9)]; + + // this is a simple insertion at depth 16, and thus the flags should look as follows: + let is_update = ZERO; + let is_simple_insert = ONE; + let is_16_or_32 = ONE; + let is_16_or_48 = ONE; + let expected_stack = [is_update, is_simple_insert, is_16_or_32, is_16_or_48]; + let process = prepare_smt_insert(key_a, val_a, &smt, expected_stack.len(), Vec::new()); + assert_eq!(build_expected(&expected_stack), process.stack.trace_state()); + + // --- update same key with different value ------------------------------- + + // insert val_a into the tree so that val_b overwrites it + smt.insert(key_a.into(), val_a); + let val_b = [ONE, ONE, ZERO, ZERO]; + + // this is a simple update, and thus the flags should look as follows: + let is_update = ONE; + let is_16_or_32 = ONE; + let is_16_or_48 = ONE; + + // also, the old value should be present in the advice stack: + let expected_stack = [ + val_a[3], + val_a[2], + val_a[1], + val_a[0], + is_update, + is_16_or_32, + is_16_or_48, + ZERO, + ]; + let adv_map = vec![build_adv_map_entry(key_a, val_a, 16)]; + let process = prepare_smt_insert(key_a, val_b, &smt, expected_stack.len(), adv_map); + assert_eq!(build_expected(&expected_stack), process.stack.trace_state()); +} + +fn prepare_smt_insert( + key: Word, + value: Word, + smt: &TieredSmt, + adv_stack_depth: usize, + adv_map: Vec<([u8; 32], Vec)>, +) -> Process { + let root: Word = smt.root().into(); + let store = MerkleStore::from(smt); + + let stack_inputs = build_stack_inputs(value, key, root); + let advice_inputs = AdviceInputs::default().with_merkle_store(store).with_map(adv_map); + let mut process = build_process(stack_inputs, advice_inputs); + + process.execute_op(Operation::Noop).unwrap(); + process + .execute_decorator(&Decorator::Advice(AdviceInjector::SmtInsert)) + .unwrap(); + + move_adv_to_stack(&mut process, adv_stack_depth); + + process +} + // HELPER FUNCTIONS -// -------------------------------------------------------------------------------------------- +// ================================================================================================ fn init_leaf(value: u64) -> Word { [Felt::new(value), Felt::ZERO, Felt::ZERO, Felt::ZERO] @@ -160,7 +230,7 @@ fn build_expected(values: &[Felt]) -> [Felt; 16] { } fn assert_case_smtget( - depth: u8, + _depth: u8, key: Word, value: Word, node: RpoDigest, @@ -169,19 +239,8 @@ fn assert_case_smtget( 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 remaining = get_smt_remaining_key(key, depth); - let mapped = remaining.into_iter().chain(value.into_iter()).collect(); + let stack_inputs = build_stack_inputs(key, root.into(), Word::default()); + let mapped = key.into_iter().chain(value.into_iter()).collect(); let advice_inputs = AdviceInputs::default() .with_merkle_store(store) .with_map([(node.into_bytes(), mapped)]); @@ -192,13 +251,60 @@ fn assert_case_smtget( // 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 { + + // replace operand stack contents with the data on the advice stack + move_adv_to_stack(&mut process, expected_stack.len()); + + assert_eq!(build_expected(expected_stack), process.stack.trace_state()); +} + +fn build_process( + stack_inputs: StackInputs, + adv_inputs: AdviceInputs, +) -> Process { + let advice_provider = MemAdviceProvider::from(adv_inputs); + Process::new(Kernel::default(), stack_inputs, advice_provider, ExecutionOptions::default()) +} + +fn build_stack_inputs(w0: Word, w1: Word, w2: Word) -> StackInputs { + StackInputs::try_from_values([ + w2[0].as_int(), + w2[1].as_int(), + w2[2].as_int(), + w2[3].as_int(), + w1[0].as_int(), + w1[1].as_int(), + w1[2].as_int(), + w1[3].as_int(), + w0[0].as_int(), + w0[1].as_int(), + w0[2].as_int(), + w0[3].as_int(), + ]) + .unwrap() +} + +fn build_key(prefix: u64) -> Word { + [ONE, ONE, ONE, Felt::new(prefix)] +} + +/// Removes all items from the operand stack and pushes the specified number of values from +/// the advice tack onto it. +fn move_adv_to_stack(process: &mut Process, adv_stack_depth: usize) { + let stack_depth = process.stack.depth(); + for _ in 0..stack_depth { process.execute_op(Operation::Drop).unwrap(); } - // expect the stack output - for _ in 0..expected_stack.len() { + for _ in 0..adv_stack_depth { process.execute_op(Operation::AdvPop).unwrap(); } - assert_eq!(build_expected(expected_stack), process.stack.trace_state()); +} + +fn build_adv_map_entry(key: Word, val: Word, depth: u8) -> ([u8; 32], Vec) { + let node = Rpo256::merge_in_domain(&[key.into(), val.into()], Felt::from(depth)); + let mut elements = Vec::new(); + elements.extend_from_slice(&key); + elements.extend_from_slice(&val); + (node.into(), elements) } diff --git a/stdlib/asm/collections/smt.masm b/stdlib/asm/collections/smt.masm index b7f93a4112..96faca1fbe 100644 --- a/stdlib/asm/collections/smt.masm +++ b/stdlib/asm/collections/smt.masm @@ -16,6 +16,9 @@ const.EMPTY_48_1=5634734408638476525 const.EMPTY_48_2=9233115969432897632 const.EMPTY_48_3=1437907447409278328 +# HELPER METHODS +# ================================================================================================= + #! Extracts 16 most significant bits from the passed-in value. #! #! Input: [v, ...] @@ -28,6 +31,16 @@ proc.get_top_16_bits u32unchecked_shr.16 end +#! Extracts 32 most significant bits from the passed-in value. +#! +#! Input: [v, ...] +#! Output: [v >> 32, ...] +#! +#! Cycles: 3 +proc.get_top_32_bits + u32split swap drop +end + #! Extracts 48 most significant bits from the passed-in value. #! #! Input: [v, ...] @@ -43,13 +56,16 @@ proc.get_top_48_bits add end +# GET +# ================================================================================================= + #! Get the leaf value for depth 16. #! #! Input: [K, R, ...] #! Output: [V, R, ...] #! #! Cycles: 85 -proc.get16.2 +proc.get_16.2 # compute index of the node by extracting top 16 bits from the key (8 cycles) dup exec.get_top_16_bits movdn.4 # => [K, i, R, ...] @@ -123,7 +139,7 @@ end #! Output: [V, R, ...] #! #! Cycles: 81 -proc.get32.2 +proc.get_32.2 # compute index of the node by extracting top 16 bits from the key (4 cycles) dup u32split movdn.5 drop # => [K, i, R, ...] @@ -197,7 +213,7 @@ end #! Output: [V, R, ...] #! #! Cycles: 88 -proc.get48.2 +proc.get_48.2 # compute index of the node by extracting top 48 bits from the key (11 cycles) dup exec.get_top_48_bits movdn.4 # => [K, i, R, ...] @@ -275,6 +291,7 @@ end #! Depth 16: 91 cycles #! Depth 32: 87 cycles #! Depth 48: 94 cycles +#! Depth 64: unimplemented export.get # invoke adv and fetch target depth flags adv.push_smtget adv_push.2 @@ -284,15 +301,15 @@ export.get if.true if.true # depth 16 - exec.get16 + exec.get_16 else # depth 32 - exec.get32 + exec.get_32 end else if.true # depth 48 - exec.get48 + exec.get_48 else # depth 64 # currently not implemented @@ -301,3 +318,393 @@ export.get end # => [V, R, ...] end + +# INSERT +# ================================================================================================= + +#! Updates a leaf node at depths 16, 32, or 48. +#! +#! Input: [d, idx, V, K, R, ...]; +#! Output: [V_old, R_new, ...] +#! +#! Where: +#! - R is the initial root of the TSMT, and R_new is the new root of the TSMT. +#! - d, idx are the depth and index (at that depth) of the leaf node to be updated. +#! - K, V are the key-value pair for the leaf node where V is a new value for key K. +#! - V_old is the value previously stored under key K. +#! +#! This procedure succeeds only if: +#! - Node to be replaced at (d, idx) is a leaf node for the same key K. +#! +#! Cycles: 101 +proc.update_16_32_48.2 + # save [idx, d, 0, 0] in loc[0] (5 cycles) + push.0.0 loc_storew.0 + # => [0, 0, d, idx, V, K, R, ...] + + # prepare the stack for computing N = hash([K, V], domain=d), and also save K into loc[1] + # (10 cycles) + movdn.3 movup.2 drop push.0 swapw.2 loc_storew.1 swapw + # => [V, K, 0, 0, d, 0, R, ...] + + # compute the hash of the node N = hash([K, V], domain=d) - (1 cycle) + hperm + # => [X, N, X, R, ...] + + # prepare the stack for the mtree_set operation (8 cycles) + swapw.3 swapw swapw.2 loc_loadw.0 drop drop + # => [d, idx, R, N, X, ...] + + # insert the new leaf node into the tree at the specified index/depth; this also leaves the + # previous value of the node on the stack (29 cycle) + mtree_set + # => [N_old, R_new, X, ...] + + # verify that N_old is a leaf node for the same key K + + # prepare the stack for computing E = hash([K, V_old], domain=d); value of V_old is read + # from the advice provider and is saved into loc[0] (21 cycles) + swapw.2 loc_loadw.0 movdn.3 push.0 movup.3 push.0.0.0 loc_loadw.1 adv_push.4 loc_storew.0 + # => [V_old, K, 0, 0, d, 0, R_new, N_old, ...] + + # compute E = hash([K, V_old], domain=d) + # (10 cycle) + hperm dropw swapw dropw + # => [E, R_new, N_old, ...] + + # make sure E and N_old are the same (14 cycles) + swapw swapw.2 + repeat.4 + dup.4 assert_eq + end + # => [E, R_new, ...] + + # load the old value (which we saved previously) onto the stack (3 cycles) + loc_loadw.0 + # => [V_old, R_new, ...] +end + +#! Inserts a new leaf node at depth 16. +#! +#! Input: [V, K, R, ...]; +#! Output:[0, 0, 0, 0, R_new, ...] +#! +#! Where: +#! - R is the initial root of the TSMT, and R_new is the new root of the TSMT. +#! - K and V is the key-value pair for the leaf node to be inserted. +#! +#! This procedure succeeds only if: +#! - Node to be replaced at depth 16 is a root of an empty subtree. +#! +#! Cycles: 73 +proc.insert_16 + # extract 16-bit index from the key (8 cycles) + swapw dup exec.get_top_16_bits + # => [idx, K, V, R, ...] + + # prepare the stack for computing leaf node value (6 cycles) + movdn.8 push.0.16.0.0 swapw.2 + # => [V, K, 0, 0, 16, 0, idx, R, ...] + + # compute leaf node value as N = hash([K, V], domain=16) (10 cycles) + hperm dropw swapw dropw + # => [N, idx, R, ...] + + # prepare the stack for mtree_set operation (4 cycles) + swapw movup.8 movdn.4 push.16 + # => [16, idx, R, N, ...] + + # insert the node into the tree at depth 16; this also leaves the old value of the node on the + # stack (29 cycle) + mtree_set + # => [N_old, R_new, ...] + + # verify that the old value of the node was a root of an empty subtree for depth 16 (12 cycles) + push.EMPTY_16_3 assert_eq + push.EMPTY_16_2 assert_eq + push.EMPTY_16_1 assert_eq + push.EMPTY_16_0 assert_eq + + # put the return value onto the stack and return (4 cycles) + padw + # => [0, 0, 0, 0, R_new, ...] +end + +#! Inserts a new leaf node at depth 32. +#! +#! Input: [V, K, R, ...]; +#! Output:[0, 0, 0, 0, R_new, ...] +#! +#! Where: +#! - R is the initial root of the TSMT, and R_new is the new root of the TSMT. +#! - K, V is the key-value pair for the leaf node to be inserted into the TSMT. +#! +#! This procedure consists of two high-level steps: +#! - First, insert N = hash([K, V], domain=32) into a subtree with root P, where P is the +#! internal node at depth 16 on the path to the new leaf node. This outputs the new root +#! of the subtree P_new. +#! - Then, insert P_new into the TSMT with root R. +#! +#! We do this to minimize the number of hashes consumed by the procedure for Merkle path +#! verification. Specifically, Merkle path verification will require exactly 64 hashes. +#! +#! This procedure succeeds only if: +#! - Node at depth 16 is an internal node. +#! - Node at depth 32 is a root of an empty subtree. +#! +#! Cycles: 154 +proc.insert_32.2 + # load the value of P from the advice provider (5 cycles) + adv_push.4 swapw.2 + # => [K, V, P, R, ...] + + # save k3 into loc[0][0] (4 cycles) + dup loc_store.0 + # => [K, V, P, R, ...] + + # prepare the stack for computing N = hash([K, V], domain=32) - (5 cycles) + push.0.32.0.0 swapw.2 + # => [V, K, 0, 0, 32, 0, P, R, ...] + + # compute N = hash([K, V], domain=32) (1 cycle) + hperm + # => [X, N, X, P, R, ...] + + # save P into loc[1] to be used later (5 cycles) + swapw.3 loc_storew.1 + # => [P, N, X, X, R, ...] + + # make sure P is not a root of an empty subtree at depth 16 (17 cycles) + dup push.EMPTY_16_3 eq + dup.2 push.EMPTY_16_2 eq + dup.4 push.EMPTY_16_1 eq + dup.6 push.EMPTY_16_0 eq + and and and assertz + # => [P, N, X, X, R, ...] + + # load k3 from memory, extract upper 32 bits from it and split them into two 16-bit values + # such that the top 16-bits are in idx_hi and the next 16 bits are in idx_lo (9 cycles) + loc_load.0 exec.get_top_32_bits u32unchecked_divmod.65536 + # => [idx_lo, idx_hi, P, N, X, X, R, ...] + + # save idx_hi into loc[0][0] to be used later (5 cycles) + swap loc_store.0 + # => [idx_lo, P, N, X, X, R, ...] + + # replace node at idx_lo in P with N, the old value of the node is left on the stack; this also + # proves that P is a leaf node because a leaf node cannot have children at depth 16 (30 cycles) + push.16 mtree_set + # => [N_old, P_new, X, X, R, ...] + + # make sure that N_old is a root of an empty subtree at depth 32 (12 cycles) + push.EMPTY_32_3 assert_eq + push.EMPTY_32_2 assert_eq + push.EMPTY_32_1 assert_eq + push.EMPTY_32_0 assert_eq + # => [P_new, X, X, R, ...] + + # prepare the stack for mtree_set operation against R; here we load idx_hi from loc[0][0] + # (11 cycles) + swapw.2 dropw swapw.2 loc_load.0 push.16 + # => [16, idx_hi, R, P_new, X, ...] + + # insert P_new into tree with root R at depth 16 and idx_hi index (29 cycles) + mtree_set + # => [P_old, R_new, X, ...] + + # load previously saved P to compare it with P_old (6 cycles) + swapw swapw.2 loc_loadw.1 + # => [P, P_old, R_new, ...] + + # make sure P and P_old are the same (11 cycles) + assert_eqw + # => [R_new, ...] + + # put the return value onto the stack and return (4 cycles) + padw + # => [0, 0, 0, 0, R_new, ...] +end + +#! Inserts a new leaf node at depth 48. +#! +#! Input: [V, K, R, ...]; +#! Output:[0, 0, 0, 0, R_new, ...] +#! +#! This procedure is nearly identical to the insert_32 procedure above, adjusted for the use of +#! constants and idx_hi/idx_lo computation. It may be possible to combine the two at the expense +#! of extra 10 - 20 cycles. +proc.insert_48.2 + # load the value of P from the advice provider (5 cycles) + adv_push.4 swapw.2 + # => [K, V, P, R, ...] + + # save k3 into loc[0][0] (4 cycles) + dup loc_store.0 + # => [K, V, P, R, ...] + + # prepare the stack for computing N = hash([K, V], domain=48) - (5 cycles) + push.0.48.0.0 swapw.2 + # => [V, K, 0, 0, 48, 0, P, R, ...] + + # compute N = hash([K, V], domain=48) (1 cycle) + hperm + # => [X, N, X, P, R, ...] + + # save P into loc[1] to be used later (5 cycles) + swapw.3 loc_storew.1 + # => [P, N, X, X, R, ...] + + # make sure P is not a root of an empty subtree at depth 32 (17 cycles) + dup push.EMPTY_32_3 eq + dup.2 push.EMPTY_32_2 eq + dup.4 push.EMPTY_32_1 eq + dup.6 push.EMPTY_32_0 eq + and and and assertz + # => [P, N, X, X, R, ...] + + # load k3 from memory, extract upper 48 bits from it and split them into two values such that + # the top 32-bits are in idx_hi and the next 16 bits are in idx_lo (9 cycles) + loc_load.0 u32split swap u32unchecked_divmod.65536 drop + # => [idx_lo, idx_hi, P, N, X, X, R, ...] + + # save idx_hi into loc[0][0] to be used later (5 cycles) + swap loc_store.0 + # => [idx_lo, P, N, X, X, R, ...] + + # replace node at idx_lo in P with N, the old value of the node is left on the stack; this also + # proves that P is a leaf node because a leaf node cannot have children at depth 16 (30 cycles) + push.16 mtree_set + # => [N_old, P_new, X, X, R, ...] + + # make sure that N_old is a root of an empty subtree at depth 48 (12 cycles) + push.EMPTY_48_3 assert_eq + push.EMPTY_48_2 assert_eq + push.EMPTY_48_1 assert_eq + push.EMPTY_48_0 assert_eq + # => [P_new, X, X, R, ...] + + # prepare the stack for mtree_set operation against R; here we load idx_hi from loc[0][0] + # (11 cycles) + swapw.2 dropw swapw.2 loc_load.0 push.32 + # => [32, idx_hi, R, P_new, X, ...] + + # insert P_new into tree with root R at depth 32 and idx_hi index (29 cycles) + mtree_set + # => [P_old, R_new, X, ...] + + # load previously saved P with P_old to make sure they are the same (6 cycles) + swapw swapw.2 loc_loadw.1 + # => [P, P_old, R_new, ...] + + # make sure P and P_old are the same (11 cycles) + assert_eqw + # => [R_new, ...] + + # put the return value onto the stack and return (4 cycles) + padw + # => [0, 0, 0, 0, R_new, ...] +end + +#! Inserts the specified value into a Sparse Merkle Tree with the specified root under the +#! specified key. +#! +#! The value previously stored in the SMT under this key is left on the stack together with +#! the updated tree root. +#! +#! This assumes that the value is not [ZERO; 4]. If it is, the procedure fails. +#! +#! Input: [V, K, R, ...]; +#! Output:[V_old, R', ...] +#! +#! Cycles: +#! - Update existing leaf: +#! - Depth 16: 129 +#! - Depth 32: 126 +#! - Depth 48: 131 +#! - Insert new leaf: +#! - Depth 16: 100 +#! - Depth 32: 181 +#! - Depth 48: 181 +#! - Replace a leaf with a subtree: +#! - Depth 32: TODO +#! - Depth 48: TODO +export.insert + # make sure the value is not [ZERO; 4] (17 cycles) + repeat.4 + dup.3 eq.0 + end + and and and assertz + # => [V, K, R, ...] + + # arrange the data needed for the insert procedure on the advice stack and move the + # first 4 flags onto the operand stack; meaning of the flags f0, f1, and f2 depends + # on what type of insert is being executed (4 cycles) + adv.push_smtinsert adv_push.4 + # => [is_update, f0, f1, f2, V, K, R, ...] + + # call the inner procedure depending on the type of insert and depth + if.true # --- update leaf ------------------------------------------------- + # => [is_16_or_32, is_16_or_48, ZERO, V, K, R, ...] + if.true + if.true # --- update a leaf node at depth 16 --- + drop + # => [V, K, R, ...] + + # (cycles 8) + dup.4 exec.get_top_16_bits + push.16 + # => [16, idx, V, K, R, ...] + + exec.update_16_32_48 + else # --- update a leaf node at depth 32 --- + drop + # => [V, K, R, ...] + + #(5 cycles) + dup.4 exec.get_top_32_bits + push.32 + # => [32, idx, V, K, R, ...] + + exec.update_16_32_48 + end + else + if.true # --- update a leaf node at depth 48 --- + drop + # => [V, K, R, ...] + + # (10 cycles) + dup.4 exec.get_top_48_bits + push.48 + # => [48, idx, V, K, R, ...] + + exec.update_16_32_48 + else + # depth 64 - currently not implemented + push.0 assert + end + end + else + # => [is_simple_insert, is_16_or_32, is_16_or_48, V, K, R, ...] + if.true # --- inset new leaf ---------------------------------------------- + if.true + if.true + exec.insert_16 + else + exec.insert_32 + end + else + if.true + exec.insert_48 + else + # depth 64 - currently not implemented + push.0 assert + end + end + else # --- replace leaf with subtree ---------------------------------- + # TODO: implement replace leaf with subtree + push.0 assert + end + end + + # => [V, R, ...] +end diff --git a/stdlib/docs/collections/smt.md b/stdlib/docs/collections/smt.md index 85ee7a889d..edb31797fe 100644 --- a/stdlib/docs/collections/smt.md +++ b/stdlib/docs/collections/smt.md @@ -2,4 +2,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: 91 cycles

Depth 32: 87 cycles

Depth 48: 94 cycles | +| 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: 91 cycles

Depth 32: 87 cycles

Depth 48: 94 cycles

Depth 64: unimplemented | +| insert | Inserts the specified value into a Sparse Merkle Tree with the specified root under the

specified key.

The value previously stored in the SMT under this key is left on the stack together with

the updated tree root.

This assumes that the value is not [ZERO; 4]. If it is, the procedure fails.

Input: [V, K, R, ...];

Output:[V_old, R', ...]

Cycles:

- Update existing leaf:

- Depth 16: 129

- Depth 32: 126

- Depth 48: 131

- Insert new leaf:

- Depth 16: 100

- Depth 32: 181

- Depth 48: 181

- Replace a leaf with a subtree:

- Depth 32: TODO

- Depth 48: TODO | diff --git a/stdlib/tests/collections/smt.rs b/stdlib/tests/collections/smt.rs index 9618368cd9..10f6c4ad60 100644 --- a/stdlib/tests/collections/smt.rs +++ b/stdlib/tests/collections/smt.rs @@ -13,7 +13,7 @@ const EMPTY_VALUE: Word = TieredSmt::EMPTY_VALUE; // ================================================================================================ #[test] -fn smtget_depth_16() { +fn tsmt_get_16() { let mut smt = TieredSmt::default(); // create a key @@ -40,7 +40,7 @@ fn smtget_depth_16() { } #[test] -fn smtget_depth_32() { +fn tsmt_get_32() { let mut smt = TieredSmt::default(); // populate the tree with two key-value pairs sharing the same 16-bit prefix for the keys @@ -75,7 +75,7 @@ fn smtget_depth_32() { } #[test] -fn smtget_depth_48() { +fn tsmt_get_48() { let mut smt = TieredSmt::default(); // populate the tree with two key-value pairs sharing the same 32-bit prefix for the keys @@ -109,6 +109,140 @@ fn smtget_depth_48() { assert_smt_get_opens_correctly(&smt, key_e, EMPTY_VALUE); } +// INSERTS +// ================================================================================================ + +#[test] +fn tsmt_insert_16() { + let mut smt = TieredSmt::default(); + + let raw_a = 0b00000000_00000000_11111111_11111111_11111111_11111111_11111111_11111111_u64; + let key_a = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_a)]); + let val_a1 = [ONE, ZERO, ZERO, ZERO]; + let val_a2 = [ONE, ONE, ZERO, ZERO]; + + // insert a value under key_a into an empty tree + let init_smt = smt.clone(); + smt.insert(key_a.into(), val_a1); + assert_insert(&init_smt, key_a, EMPTY_VALUE, val_a1, smt.root().into()); + + // update a value under key_a + let init_smt = smt.clone(); + smt.insert(key_a.into(), val_a2); + assert_insert(&init_smt, key_a, val_a1, val_a2, smt.root().into()); +} + +#[test] +fn tsmt_insert_32() { + let mut smt = TieredSmt::default(); + + // insert a value under key_a into an empty tree + let raw_a = 0b00000000_00000000_11111111_11111111_11111111_11111111_11111111_11111111_u64; + let key_a = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_a)]); + let val_a = [ONE, ZERO, ZERO, ZERO]; + smt.insert(key_a.into(), val_a); + + // insert a value under key_b which has the same 16-bit prefix as A + let raw_b = 0b00000000_00000000_01111111_11111111_11111111_11111111_11111111_11111111_u64; + let key_b = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_b)]); + let val_b = [ONE, ONE, ZERO, ZERO]; + + // TODO: test this insertion once complex inserts are working + smt.insert(key_b.into(), val_b); + + // update a value under key_a + let init_smt = smt.clone(); + let val_a2 = [ONE, ZERO, ZERO, ONE]; + smt.insert(key_a.into(), val_a2); + assert_insert(&init_smt, key_a, val_a, val_a2, smt.root().into()); + + // insert a value under key_c which has the same 16-bit prefix as A and B + let raw_c = 0b00000000_00000000_00111111_11111111_11111111_11111111_11111111_11111111_u64; + let key_c = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_c)]); + let val_c = [ONE, ONE, ONE, ZERO]; + + let init_smt = smt.clone(); + smt.insert(key_c.into(), val_c); + assert_insert(&init_smt, key_c, EMPTY_VALUE, val_c, smt.root().into()); +} + +#[test] +fn tsmt_insert_48() { + let mut smt = TieredSmt::default(); + + // insert a value under key_a into an empty tree + let raw_a = 0b00000000_00000000_11111111_11111111_11111111_11111111_11111111_11111111_u64; + let key_a = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_a)]); + let val_a = [ONE, ZERO, ZERO, ZERO]; + smt.insert(key_a.into(), val_a); + + // insert a value under key_b which has the same 32-bit prefix as A + let raw_b = 0b00000000_00000000_11111111_11111111_01111111_11111111_11111111_11111111_u64; + let key_b = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_b)]); + let val_b = [ONE, ONE, ZERO, ZERO]; + + // TODO: test this insertion once complex inserts are working + smt.insert(key_b.into(), val_b); + + // update a value under key_a + let init_smt = smt.clone(); + let val_a2 = [ONE, ZERO, ZERO, ONE]; + smt.insert(key_a.into(), val_a2); + assert_insert(&init_smt, key_a, val_a, val_a2, smt.root().into()); + + // insert a value under key_c which has the same 32-bit prefix as A and B + let raw_c = 0b00000000_00000000_11111111_11111111_00111111_11111111_11111111_11111111_u64; + let key_c = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_c)]); + let val_c = [ONE, ONE, ONE, ZERO]; + + let init_smt = smt.clone(); + smt.insert(key_c.into(), val_c); + assert_insert(&init_smt, key_c, EMPTY_VALUE, val_c, smt.root().into()); +} + +fn assert_insert( + init_smt: &TieredSmt, + key: RpoDigest, + old_value: Word, + new_value: Word, + new_root: RpoDigest, +) { + let old_root = init_smt.root(); + let source = r#" + use.std::collections::smt + + begin + exec.smt::insert + end + "#; + let initial_stack = [ + old_root[0].as_int(), + old_root[1].as_int(), + old_root[2].as_int(), + old_root[3].as_int(), + key[0].as_int(), + key[1].as_int(), + key[2].as_int(), + key[3].as_int(), + new_value[0].as_int(), + new_value[1].as_int(), + new_value[2].as_int(), + new_value[3].as_int(), + ]; + let expected_output = [ + old_value[3].as_int(), + old_value[2].as_int(), + old_value[1].as_int(), + old_value[0].as_int(), + new_root[3].as_int(), + new_root[2].as_int(), + new_root[1].as_int(), + new_root[0].as_int(), + ]; + let (store, adv_map) = build_advice_inputs(init_smt); + build_test!(source, &initial_stack, &[], store, adv_map).expect_stack(&expected_output); +} + // TEST HELPERS // ================================================================================================ @@ -143,6 +277,13 @@ fn assert_smt_get_opens_correctly(smt: &TieredSmt, key: RpoDigest, value: Word) root[0].as_int(), ]; + let (store, advice_map) = build_advice_inputs(smt); + let advice_stack = []; + build_test!(source, &initial_stack, &advice_stack, store, advice_map.into_iter()) + .expect_stack(&expected_output); +} + +fn build_advice_inputs(smt: &TieredSmt) -> (MerkleStore, Vec<([u8; 32], Vec)>) { let store = MerkleStore::from(smt); let advice_map = smt .upper_leaves() @@ -153,7 +294,5 @@ fn assert_smt_get_opens_correctly(smt: &TieredSmt, key: RpoDigest, value: Word) }) .collect::>(); - let advice_stack = []; - build_test!(source, &initial_stack, &advice_stack, store, advice_map.into_iter()) - .expect_stack(&expected_output); + (store, advice_map) } diff --git a/test-utils/src/crypto.rs b/test-utils/src/crypto.rs index e06c516098..5a807f32b8 100644 --- a/test-utils/src/crypto.rs +++ b/test-utils/src/crypto.rs @@ -1,4 +1,4 @@ -use super::{Felt, FieldElement, StarkField, Vec, Word}; +use super::{Felt, FieldElement, Vec, Word}; // RE-EXPORTS // ================================================================================================ @@ -32,13 +32,3 @@ 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 -}