From c90513e4758e89f4c56b6a46f5e70454254a8a7a Mon Sep 17 00:00:00 2001 From: Bobbin Threadbare Date: Thu, 17 Aug 2023 01:27:56 -0700 Subject: [PATCH 1/2] feat: implement TSMT delete procedures --- core/Cargo.toml | 2 +- processor/src/advice/mod.rs | 6 +- processor/src/advice/providers.rs | 8 +- .../src/decorators/adv_stack_injectors.rs | 71 +++- processor/src/operations/crypto_ops.rs | 3 +- stdlib/asm/collections/smt.masm | 363 ++++++++++++++++-- stdlib/docs/collections/smt.md | 3 +- stdlib/tests/collections/smt.rs | 105 +++++ 8 files changed, 519 insertions(+), 42 deletions(-) diff --git a/core/Cargo.toml b/core/Cargo.toml index 651d7e4ef2..8a06f05aac 100644 --- a/core/Cargo.toml +++ b/core/Cargo.toml @@ -21,7 +21,7 @@ std = ["math/std", "winter-utils/std"] [dependencies] math = { package = "winter-math", version = "0.6", default-features = false } -crypto = { package = "miden-crypto", git = "https://github.com/0xPolygonMiden/crypto.git", branch = "next", default-features = false } +crypto = { package = "miden-crypto", git = "https://github.com/0xPolygonMiden/crypto.git", branch = "bobbin-leaf-traversal", default-features = false } winter-crypto = { package = "winter-crypto", version = "0.6", default-features = false } winter-utils = { package = "winter-utils", version = "0.6", default-features = false } diff --git a/processor/src/advice/mod.rs b/processor/src/advice/mod.rs index 8382a2075b..250e549d7c 100644 --- a/processor/src/advice/mod.rs +++ b/processor/src/advice/mod.rs @@ -153,7 +153,7 @@ pub trait AdviceProvider { ) -> 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. + /// returns the Merkle path from the updated node to the new root, together with the new root. /// /// The tree is cloned prior to the update. Thus, the advice provider retains the original and /// the updated tree. @@ -171,7 +171,7 @@ pub trait AdviceProvider { depth: &Felt, index: &Felt, value: Word, - ) -> Result; + ) -> Result<(MerklePath, Word), ExecutionError>; /// Creates a new Merkle tree in the advice provider by combining Merkle trees with the /// specified roots. The root of the new tree is defined as `hash(left_root, right_root)`. @@ -267,7 +267,7 @@ where depth: &Felt, index: &Felt, value: Word, - ) -> Result { + ) -> Result<(MerklePath, Word), ExecutionError> { T::update_merkle_node(self, root, depth, index, value) } diff --git a/processor/src/advice/providers.rs b/processor/src/advice/providers.rs index 18c151fa18..b84c488b93 100644 --- a/processor/src/advice/providers.rs +++ b/processor/src/advice/providers.rs @@ -172,7 +172,7 @@ where depth: &Felt, index: &Felt, value: Word, - ) -> Result { + ) -> Result<(MerklePath, Word), ExecutionError> { let node_index = NodeIndex::from_elements(depth, index).map_err(|_| { ExecutionError::InvalidTreeNodeIndex { depth: *depth, @@ -181,7 +181,7 @@ where })?; self.store .set_node(root.into(), node_index, value.into()) - .map(|root| root.path) + .map(|root| (root.path, root.root.into())) .map_err(ExecutionError::MerkleStoreUpdateFailed) } @@ -289,7 +289,7 @@ impl AdviceProvider for MemAdviceProvider { self.provider.get_leaf_depth(root, tree_depth, index) } - fn update_merkle_node(&mut self, root: Word, depth: &Felt, index: &Felt, value: Word) -> Result { + fn update_merkle_node(&mut self, root: Word, depth: &Felt, index: &Felt, value: Word) -> Result<(MerklePath, Word), ExecutionError> { self.provider.update_merkle_node(root, depth, index, value) } @@ -414,7 +414,7 @@ impl AdviceProvider for RecAdviceProvider { self.provider.get_leaf_depth(root, tree_depth, index) } - fn update_merkle_node(&mut self, root: Word, depth: &Felt, index: &Felt, value: Word) -> Result { + fn update_merkle_node(&mut self, root: Word, depth: &Felt, index: &Felt, value: Word) -> Result<(MerklePath, Word), ExecutionError> { self.provider.update_merkle_node(root, depth, index, value) } diff --git a/processor/src/decorators/adv_stack_injectors.rs b/processor/src/decorators/adv_stack_injectors.rs index b22d86b1cb..8b44ecf188 100644 --- a/processor/src/decorators/adv_stack_injectors.rs +++ b/processor/src/decorators/adv_stack_injectors.rs @@ -1,6 +1,9 @@ use super::{super::Ext2InttError, AdviceProvider, AdviceSource, ExecutionError, Process}; use vm_core::{ - crypto::{hash::RpoDigest, merkle::EmptySubtreeRoots}, + crypto::{ + hash::RpoDigest, + merkle::{EmptySubtreeRoots, TieredSmt}, + }, utils::collections::Vec, Felt, FieldElement, QuadExtension, StarkField, Word, ONE, WORD_SIZE, ZERO, }; @@ -368,7 +371,8 @@ where /// # 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 + // get the key, value, and tree root from the stack + let value = self.stack.get_word(0); let key = self.stack.get_word(1); let root = self.stack.get_word(2); @@ -386,7 +390,13 @@ where // 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)?; + let node = self.advice_provider.get_tree_node(root, &Felt::from(depth), &index)?; + + // if the value to be inserted is an empty word, we need to process it as a delete + if value == TieredSmt::EMPTY_VALUE { + self.handle_smt_delete(root, node, depth, index, key)?; + return Ok(()); + } // 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. @@ -411,7 +421,7 @@ where Ok(()) } - // TSMT INSERT HELPER METHODS + // TSMT UPDATE HELPER METHODS // -------------------------------------------------------------------------------------------- /// Retrieves a key-value pair for the specified leaf node from the advice map. @@ -558,6 +568,59 @@ where Ok(()) } + + /// TODO: add comments + fn handle_smt_delete( + &mut self, + root: Word, + node: Word, + depth: u8, + index: Felt, + key: Word, + ) -> Result<(), ExecutionError> { + let empty = EmptySubtreeRoots::empty_hashes(64)[depth as usize]; + + if node == Word::from(empty) { + self.advice_provider.push_stack(AdviceSource::Value(ONE))?; + self.advice_provider.push_stack(AdviceSource::Value(ZERO))?; + + let (is_16_or_32, is_16_or_48) = get_depth_flags(depth); + self.advice_provider.push_stack(AdviceSource::Value(is_16_or_32))?; + self.advice_provider.push_stack(AdviceSource::Value(is_16_or_48))?; + } else { + let (leaf_key, leaf_value) = self.get_smt_upper_leaf_preimage(node).unwrap(); + + if leaf_key != key { + self.advice_provider.push_stack(AdviceSource::Word(leaf_value))?; + self.advice_provider.push_stack(AdviceSource::Word(leaf_key))?; + + self.advice_provider.push_stack(AdviceSource::Value(ONE))?; + self.advice_provider.push_stack(AdviceSource::Value(ONE))?; + + let (is_16_or_32, is_16_or_48) = get_depth_flags(depth); + self.advice_provider.push_stack(AdviceSource::Value(is_16_or_32))?; + self.advice_provider.push_stack(AdviceSource::Value(is_16_or_48))?; + } else { + let (_, new_root) = self.advice_provider.update_merkle_node( + root, + &Felt::from(depth), + &index, + empty.into(), + )?; + + self.advice_provider.push_stack(AdviceSource::Word(leaf_value))?; + self.advice_provider.push_stack(AdviceSource::Word(new_root))?; + + 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))?; + } + } + + Ok(()) + } } // HELPER FUNCTIONS diff --git a/processor/src/operations/crypto_ops.rs b/processor/src/operations/crypto_ops.rs index 7b36720690..75cc998ceb 100644 --- a/processor/src/operations/crypto_ops.rs +++ b/processor/src/operations/crypto_ops.rs @@ -140,7 +140,8 @@ where // get a Merkle path to it. the length of the returned path is expected to match the // specified depth. if the new node is the root of a tree, this instruction will append the // whole sub-tree to this node. - let path = self.advice_provider.update_merkle_node(old_root, &depth, &index, new_node)?; + let (path, _) = + self.advice_provider.update_merkle_node(old_root, &depth, &index, new_node)?; assert_eq!(path.len(), depth.as_int() as usize); let merkle_tree_update = self.chiplets.update_merkle_root(old_node, new_node, &path, index); diff --git a/stdlib/asm/collections/smt.masm b/stdlib/asm/collections/smt.masm index a831a09044..ad90832213 100644 --- a/stdlib/asm/collections/smt.masm +++ b/stdlib/asm/collections/smt.masm @@ -968,41 +968,32 @@ 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 is the actual implementation of the `insert` procedure below, except for two things: +#! - This procedure assumes that value V is not [ZERO; 4], but this is not checked. +#! - This procedure assumes that the relevant flags have already been read from the advice provider. #! -#! This assumes that the value is not [ZERO; 4]. If it is, the procedure fails. -#! -#! Input: [V, K, R, ...] +#! Input: [is_update, f0, f1, f2, V, K, R, ...] #! Output: [V_old, R_new, ...] #! +#! Where: +#! - is_update is a flag specifying whether the insert is just an update of a value under an +#! existing key. +#! - Meaning of the flags f0, f1, and f2 depends on what type of insert is being executed. +#! #! Cycles: #! - Update existing leaf: -#! - Depth 16: 129 -#! - Depth 32: 126 -#! - Depth 48: 131 +#! - Depth 16: 116 +#! - Depth 32: 113 +#! - Depth 48: 118 #! - Insert new leaf: -#! - Depth 16: 100 -#! - Depth 32: 181 -#! - Depth 48: 181 +#! - Depth 16: 81 +#! - Depth 32: 162 +#! - Depth 48: 162 #! - Replace a leaf with a subtree: -#! - Depth 16 -> 32: 242 -#! - Depth 16 -> 48: 263 -#! - Depth 32 -> 48: 253 -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, ...] - +#! - Depth 16 -> 32: 221 +#! - Depth 16 -> 48: 244 +#! - Depth 32 -> 48: 234 +proc.insert_internal # 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, ...] @@ -1124,3 +1115,319 @@ export.insert # => [V_old, 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_new, ...] +#! +#! Cycles: +#! - Update existing leaf: +#! - Depth 16: 137 +#! - Depth 32: 134 +#! - Depth 48: 139 +#! - Insert new leaf: +#! - Depth 16: 102 +#! - Depth 32: 183 +#! - Depth 48: 183 +#! - Replace a leaf with a subtree: +#! - Depth 16 -> 32: 242 +#! - Depth 16 -> 48: 265 +#! - Depth 32 -> 48: 255 +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, ...] + + # execute the actual insert procedure + exec.insert_internal + # => [V_old, R_new, ...] +end + +# DELETE +# ================================================================================================= + +#! Verifies that a node at depth 16 and index defined by the most significant element of K in a +#! tree with root R is a root of an empty subtree. +#! +#! Input: [Z, K, R, ...] +#! Output: [Z, R, ...] +#! +#! Where Z is [ZERO; 4]. +#! +#! Cycles: 26 +proc.verify_empty_node_16 + # (9 cycles) + swapw.2 dup.4 exec.get_top_16_bits push.16 + # => [16, idx, R, K, Z, ...] + + # (4 cycles) + push.EMPTY_16_0.EMPTY_16_1.EMPTY_16_2.EMPTY_16_3 + # => [E16, 16, idx, R, K, Z, ...] + + # (1 cycle) + mtree_verify + # => [E16, 16, idx, R, K, Z, ...] + + # (12 cycles) + dropw drop drop swapw dropw swapw + # => [Z, R, ...] +end + +#! Verifies that a node at depth 32 and index defined by the most significant element of K in a +#! tree with root R is a root of an empty subtree. +#! +#! Input: [Z, K, R, ...] +#! Output: [Z, R, ...] +#! +#! Where Z is [ZERO; 4]. +#! +#! Cycles: 23 +proc.verify_empty_node_32 + # (6 cycles) + swapw.2 dup.4 exec.get_top_32_bits push.32 + # => [32, idx, R, K, Z, ...] + + # (4 cycles) + push.EMPTY_32_0.EMPTY_32_1.EMPTY_32_2.EMPTY_32_3 + # => [E32, 32, idx, R, K, Z, ...] + + # (1 cycle) + mtree_verify + # => [E32, 32, idx, R, K, Z, ...] + + # (12 cycles) + dropw drop drop swapw dropw swapw + # => [Z, K, ...] +end + +#! Verifies that a node at depth 48 and index defined by the most significant element of K in a +#! tree with root R is a root of an empty subtree. +#! +#! Input: [Z, K, R, ...] +#! Output: [Z, R, ...] +#! +#! Where Z is [ZERO; 4]. +#! +#! Cycles: 29 +proc.verify_empty_node_48 + # (12 cycles) + swapw.2 dup.4 exec.get_top_48_bits push.48 + # => [48, idx, R, K, Z, ...] + + # (4 cycles) + push.EMPTY_48_0.EMPTY_48_1.EMPTY_48_2.EMPTY_48_3 + # => [E48, 48, idx, R, K, Z, ...] + + # (1 cycle) + mtree_verify + # => [E48, 48, idx, R, K, Z, ...] + + # (12 cycles) + dropw drop drop swapw dropw swapw + # => [Z, K, ...] +end + +#! Verifies that a leaf node located at depth d and index idx in the tree defined by root R is +#! contains key which is different from K. +#! +#! Input: [d, idx, K, R, Z, ...] +#! Output: [Z, R, ...] +#! +#! Where Z is [ZERO; 4]. +#! +#! Cycles: 50 +proc.verify_leaf_with_another_key + # load the leaf key K_e from the advice provider (11 cycles) + movdn.5 movdn.5 push.0 dup.5 push.0.0 swapw adv_push.4 + # => [K_e, K, 0, 0, d, 0, d, idx, R, Z, ...] + + # make sure K_e and K are not the same (17 cycles) + repeat.4 + dup.3 movup.8 eq + end + and and and assertz + # => [K_e, 0, 0, d, 0, d, idx, R, Z, ...] + + # load leaf value V_e from the advice provider (4 cycles) + adv_push.4 + # => [V_e, K_e, 0, 0, d, 0, d, idx, R, Z, ...] + + # compute N = hash([K_e, K_v], domain=d) - (10 cycles) + hperm dropw swapw dropw + # => [N, d, idx, R, Z, ...] + + # verify that node N exists in the tree with root R at the specified index and depth + mtree_verify + # => [N, d, idx, R, Z, ...] + + # clean up the stack and return (7 cycles) + dropw drop drop swapw + # => [Z, R, ...] +end + +#! Removes a key-value under key K from the Tiered Sparse Merkle tree defined by root R, and +#! returns the new tree root together with the value previously associated with key K. +#! +#! If the key-value pair is not in the tree, this proves that the key-value pair is not in the +#! tree but does not modify the tree itself. +#! +#! This assumes that the value Z is [ZERO; 4], but this is not checked. +#! +#! Input: [key_not_in_tree, f0, f1, f2, Z, K, R, ...] +#! Output: [V_old, R_new, ...] +#! +#! Cycles: +#! - Key not in the tree (key with common prefix in the tree): +#! - Depth 16: 68 +#! - Depth 32: 65 +#! - Depth 48: 71 +#! - Key not in the tree (no key with common prefix): +#! - Depth 16: 34 +#! - Depth 32: 31 +#! - Depth 48: 37 +#! - Remove a leaf from the tree: +#! - 121 - 284 cycles +proc.delete + if.true # --- key is not in the tree ---------------------------------------------------------- + if.true # --- key with common prefix in the tree --- + if.true + if.true + # (10 cycles) + swapw.2 swapw dup exec.get_top_16_bits push.16 + # => [16, idx, K, R, Z, ...] + + exec.verify_leaf_with_another_key + else + # (7 cycles) + swapw.2 swapw dup exec.get_top_32_bits push.32 + # => [32, idx, K, R, Z, ...] + + exec.verify_leaf_with_another_key + end + else + if.true + # (13 cycles) + swapw.2 swapw dup exec.get_top_48_bits push.48 + # => [48, idx, K, R, Z, ...] + + exec.verify_leaf_with_another_key + else + # depth 64 - currently not implemented + push.0 assert + end + end + else # --- no key with common prefix in the tree --- + if.true + if.true + exec.verify_empty_node_16 + else + exec.verify_empty_node_32 + end + else + if.true + exec.verify_empty_node_48 + else + # depth 64 - currently not implemented + push.0 assert + end + end + end + else # --- key is in the tree --------------------------------------------------------------- + # load the root of the tree with the leaf removed from the advice provider (2 cycles) + push.0 adv_loadw + # => [R_new, Z, K, R, ...] + + # load the value to be removed from the advice provider (6 cycles) + dupw swapw.2 adv_loadw + # => [V_old, R_new, R_new, K, R, ...] + + # prepare the stack for TSMT insert operation (9 cycles) + swapw.3 dupw.3 adv.push_smtinsert adv_push.4 + # => [f0, f1, f2, f3, V_old, K, R_new, R_new, V_old, R, ...] + + # insert key-pair (K, V_old) into TSMT with root R_new - i.e., insert the key-value pair + # to be removed into the TSMT with value already removed. if all values were provided + # correctly, we should get TSMT with the root prior to the key-pair's removal. + # (81 - 244 cycles) + exec.insert_internal + # => [Z, R_old, R_new, V_old, R, ...] + + # make sure the value we got back is an empty node (8 cycles) + assertz assertz assertz assertz + # => [R_old, R_new, V_old, R, ...] + + # Make sure R and R_old are the same (13 cycles) + swapw swapw.3 assert_eqw + # => [V_old, R_new, ...] + end +end + +# SET +# ================================================================================================= + +#! Sets the value associated with key K to V in a Sparse Merkle tree with root R. Returns the new +#! root of the tree together with the value previously associated with key K. +#! +#! If no value was previously associated with K, [ZERO; 4] is returned. +#! +#! Unlike the `insert` procedure defined above, this procedure allows for values to be set to +#! [ZERO; 4]. +#! +#! Input: [V, K, R, ...] +#! Output: [V_old, R_new, ...] +#! +#! Cycles: +#! - Update existing leaf: +#! - Depth 16: 137 +#! - Depth 32: 133 +#! - Depth 48: 139 +#! - Insert new leaf: +#! - Depth 16: 102 +#! - Depth 32: 183 +#! - Depth 48: 183 +#! - Replace a leaf with a subtree: +#! - Depth 16 -> 32: 242 +#! - Depth 16 -> 48: 265 +#! - Depth 32 -> 48: 255 +#! - Remove a key-value pair: +#! - Key-value pair not in tree: 52 - 93 +#! - Key-value pair is in tree: 142 - 305 +export.set + # arrange the data needed for the update procedure on the advice stack and move the first + # 4 flags onto the operand stack; meaning of the flags f0, f1, f2, and f3 depends on what + # type of update is being executed (4 cycles) + adv.push_smtinsert adv_push.4 + # => [f0, f1, f2, f3, V, K, R, ...] + + # determine if the value is an empty word (15 cycles) + repeat.4 + dup.7 push.0 eq + end + and and and + # => [is_empty_value, f0, f1, f2, f3, V, K, R, ...] + + # the value is an empty word execute the delete procedure; otherwise execute the internal + # insert procedure because we already have all required data on the advice provider and know + # that the value being inserted is not an empty word. + if.true + exec.delete + else + exec.insert_internal + end +end diff --git a/stdlib/docs/collections/smt.md b/stdlib/docs/collections/smt.md index 2ff3771647..bed60441bb 100644 --- a/stdlib/docs/collections/smt.md +++ b/stdlib/docs/collections/smt.md @@ -3,4 +3,5 @@ | 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

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_new, ...]

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 16 -> 32: 242

- Depth 16 -> 48: 263

- Depth 32 -> 48: 253 | +| 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_new, ...]

Cycles:

- Update existing leaf:

- Depth 16: 137

- Depth 32: 134

- Depth 48: 139

- Insert new leaf:

- Depth 16: 102

- Depth 32: 183

- Depth 48: 183

- Replace a leaf with a subtree:

- Depth 16 -> 32: 242

- Depth 16 -> 48: 265

- Depth 32 -> 48: 255 | +| set | Sets the value associated with key K to V in a Sparse Merkle tree with root R. Returns the new

root of the tree together with the value previously associated with key K.

If no value was previously associated with K, [ZERO; 4] is returned.

Unlike the `insert` procedure defined above, this procedure allows for values to be set to

[ZERO; 4].

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

Output: [V_old, R_new, ...]

Cycles:

- Update existing leaf:

- Depth 16: 137

- Depth 32: 133

- Depth 48: 139

- Insert new leaf:

- Depth 16: 102

- Depth 32: 183

- Depth 48: 183

- Replace a leaf with a subtree:

- Depth 16 -> 32: 242

- Depth 16 -> 48: 265

- Depth 32 -> 48: 255

- Remove a key-value pair:

- Key-value pair not in tree: 52 - 93

- Key-value pair is in tree: 142 - 305 | diff --git a/stdlib/tests/collections/smt.rs b/stdlib/tests/collections/smt.rs index d5b10993a6..448de05c62 100644 --- a/stdlib/tests/collections/smt.rs +++ b/stdlib/tests/collections/smt.rs @@ -362,6 +362,111 @@ fn assert_insert( } } +// SET TESTS +// ================================================================================================ + +#[test] +fn tsmt_set_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]; + + // set a value under key_a into an empty tree; this inserts one entry into the advice map + let init_smt = smt.clone(); + smt.insert(key_a.into(), val_a1); + let new_map_entries = [build_node_entry(key_a, val_a1, 16)]; + assert_set(&init_smt, key_a, EMPTY_VALUE, val_a1, smt.root().into(), &new_map_entries); + + // update a value under key_a; this inserts one entry into the advice map + let init_smt = smt.clone(); + smt.insert(key_a.into(), val_a2); + let new_map_entries = [build_node_entry(key_a, val_a2, 16)]; + assert_set(&init_smt, key_a, val_a1, val_a2, smt.root().into(), &new_map_entries); + + // set an empty value for a previously un-set key; this should not change the tree + let raw_b = 0b00000000_10000000_11111111_11111111_11111111_11111111_11111111_11111111_u64; + let key_b = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_b)]); + assert_set(&smt, key_b, EMPTY_VALUE, EMPTY_VALUE, smt.root().into(), &[]); + + // set an empty value for a previously un-set key which shares 16-bit prefix with A; + // this should not change the tree + let raw_c = 0b00000000_00000000_01111111_11111111_11111111_11111111_11111111_11111111_u64; + let key_c = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_c)]); + assert_set(&smt, key_c, EMPTY_VALUE, EMPTY_VALUE, smt.root().into(), &[]); + + // set the value at key A to an empty word + let init_smt = smt.clone(); + smt.insert(key_a.into(), EMPTY_VALUE); + assert_set(&init_smt, key_a, val_a2, EMPTY_VALUE, smt.root().into(), &[]); +} + +fn assert_set( + init_smt: &TieredSmt, + key: RpoDigest, + old_value: Word, + new_value: Word, + new_root: RpoDigest, + new_map_entries: &[AdvMapEntry], +) { + let old_root = init_smt.root(); + let source = r#" + use.std::collections::smt + + begin + exec.smt::set + 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 = stack_top_to_ints(&[ + 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); + let process = build_test!(source, &initial_stack, &[], store, adv_map.clone()) + .execute_process() + .unwrap(); + + // check the returned values + let stack = stack_to_ints(&process.stack.trace_state()); + assert_eq!(stack, expected_output); + + // remove the initial key-value pairs from the advice map + let mut new_adv_map = process.advice_provider.map().clone(); + for (key, value) in adv_map.iter() { + let init_value = new_adv_map.remove(key).unwrap(); + assert_eq!(value, &init_value); + } + + // make sure the remaining values in the advice map are the same as expected new entries + assert_eq!(new_adv_map.len(), new_map_entries.len()); + for (key, val) in new_map_entries { + let old_val = new_adv_map.get(key).unwrap(); + assert_eq!(old_val, val); + } +} + // HELPER FUNCTIONS // ================================================================================================ From bdc7109231ba5db7324139f330fd56a8f0cbb616 Mon Sep 17 00:00:00 2001 From: Bobbin Threadbare Date: Sat, 19 Aug 2023 02:52:23 -0700 Subject: [PATCH 2/2] feat: finish implemented basic TSMT functionality --- CHANGELOG.md | 3 + core/Cargo.toml | 2 +- core/src/operations/decorators/advice.rs | 3 + processor/src/advice/mod.rs | 38 +++- processor/src/advice/providers.rs | 33 ++- .../src/decorators/adv_stack_injectors.rs | 170 ++++++++++++++-- stdlib/asm/collections/smt.masm | 12 +- stdlib/tests/collections/smt.rs | 188 +++++++++++++++++- 8 files changed, 415 insertions(+), 34 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 3467a431ee..d6724da9a7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -19,6 +19,9 @@ - Allowed the assembler to produce programs with "phantom" calls (#1019). - Added `TraceLenSummary` struct which holds information about traces lengths to the `ExecutionTrace` (#1029). +#### Stdlib +- Completed `std::collections::smt` module by implementing `insert` and `set` procedures (#1036, #1038, #1046). + ## 0.6.1 (2023-06-29) - Fixed `no-std` compilation for `miden-core`, `miden-assembly`, and `miden-processor` crates. diff --git a/core/Cargo.toml b/core/Cargo.toml index 8a06f05aac..651d7e4ef2 100644 --- a/core/Cargo.toml +++ b/core/Cargo.toml @@ -21,7 +21,7 @@ std = ["math/std", "winter-utils/std"] [dependencies] math = { package = "winter-math", version = "0.6", default-features = false } -crypto = { package = "miden-crypto", git = "https://github.com/0xPolygonMiden/crypto.git", branch = "bobbin-leaf-traversal", default-features = false } +crypto = { package = "miden-crypto", git = "https://github.com/0xPolygonMiden/crypto.git", branch = "next", default-features = false } winter-crypto = { package = "winter-crypto", version = "0.6", default-features = false } winter-utils = { package = "winter-utils", version = "0.6", default-features = false } diff --git a/core/src/operations/decorators/advice.rs b/core/src/operations/decorators/advice.rs index 854a4c7469..ca2b12d6c5 100644 --- a/core/src/operations/decorators/advice.rs +++ b/core/src/operations/decorators/advice.rs @@ -167,6 +167,9 @@ pub enum AdviceInjector { /// - 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] /// - Complex insert: [f0, f1, ZERO (is_simple_insert), ZERO (is_update), E_KEY, E_VALUE] + /// - Delete against an empty subtree: [d0, d1, ZERO (is_leaf), ONE (key_not_set)] + /// - Delete against another leaf: [d0, d1, ONE (is_leaf), ONE (key_not_set), KEY, VALUE] + /// - Delete against own leaf: [ZERO, ZERO, ZERO, ZERO (key_not_set), NEW_ROOT, OLD_VALUE] /// /// Where: /// - ROOT and NEW_ROOT are the roots of the TSMT before and after the insert respectively. diff --git a/processor/src/advice/mod.rs b/processor/src/advice/mod.rs index 250e549d7c..e21d0bbc57 100644 --- a/processor/src/advice/mod.rs +++ b/processor/src/advice/mod.rs @@ -92,6 +92,12 @@ pub trait AdviceProvider { /// Returns an error if the value specified by the advice source cannot be obtained. fn push_stack(&mut self, source: AdviceSource) -> Result<(), ExecutionError>; + // ADVICE MAP + // -------------------------------------------------------------------------------------------- + + /// Returns a reference to the value(s) associated with the specified key in the advice map. + fn get_mapped_values(&self, key: &[u8; 32]) -> Option<&[Felt]>; + /// Inserts the provided value into the advice map under the specified key. /// /// The values in the advice map can be moved onto the advice stack by invoking @@ -101,13 +107,7 @@ pub trait AdviceProvider { /// are replaced with the specified values. fn insert_into_map(&mut self, key: Word, values: Vec) -> Result<(), ExecutionError>; - // ADVICE MAP - // -------------------------------------------------------------------------------------------- - - /// Returns a reference to the value(s) associated with the specified key in the advice map. - fn get_mapped_values(&self, key: &[u8; 32]) -> Option<&[Felt]>; - - // ADVISE SETS + // MERKLE STORE // -------------------------------------------------------------------------------------------- /// Returns a node at the specified depth and index in a Merkle tree with the given root. @@ -152,6 +152,21 @@ pub trait AdviceProvider { index: &Felt, ) -> Result; + /// Returns node value and index of a leaf node in the subtree of the specified root, if and + /// only if this is the only leaf in the entire subtree. Otherwise, None is returned. + /// + /// The root itself is assumed to be located at the specified index in a tree with the provided + /// depth. + /// + /// # Errors + /// Returns an error if a three for the specified root does not exist in the advice provider. + fn find_lone_leaf( + &self, + root: Word, + root_index: NodeIndex, + tree_depth: u8, + ) -> Result, ExecutionError>; + /// 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, together with the new root. /// @@ -261,6 +276,15 @@ where T::get_leaf_depth(self, root, tree_depth, index) } + fn find_lone_leaf( + &self, + root: Word, + root_index: NodeIndex, + tree_depth: u8, + ) -> Result, ExecutionError> { + T::find_lone_leaf(self, root, root_index, tree_depth) + } + fn update_merkle_node( &mut self, root: Word, diff --git a/processor/src/advice/providers.rs b/processor/src/advice/providers.rs index b84c488b93..73fc05021b 100644 --- a/processor/src/advice/providers.rs +++ b/processor/src/advice/providers.rs @@ -103,18 +103,19 @@ where Ok(()) } - fn insert_into_map(&mut self, key: Word, values: Vec) -> Result<(), ExecutionError> { - self.map.insert(key.into_bytes(), values); - Ok(()) - } - // ADVICE MAP // -------------------------------------------------------------------------------------------- + fn get_mapped_values(&self, key: &[u8; 32]) -> Option<&[Felt]> { self.map.get(key).map(|v| v.as_slice()) } - // ADVISE SETS + fn insert_into_map(&mut self, key: Word, values: Vec) -> Result<(), ExecutionError> { + self.map.insert(key.into_bytes(), values); + Ok(()) + } + + // MERKLE STORE // -------------------------------------------------------------------------------------------- fn get_tree_node( @@ -166,6 +167,18 @@ where .map_err(ExecutionError::MerkleStoreLookupFailed) } + fn find_lone_leaf( + &self, + root: Word, + root_index: NodeIndex, + tree_depth: u8, + ) -> Result, ExecutionError> { + self.store + .find_lone_leaf(root.into(), root_index, tree_depth) + .map(|leaf| leaf.map(|(index, leaf)| (index, leaf.into()))) + .map_err(ExecutionError::MerkleStoreLookupFailed) + } + fn update_merkle_node( &mut self, root: Word, @@ -289,6 +302,10 @@ impl AdviceProvider for MemAdviceProvider { self.provider.get_leaf_depth(root, tree_depth, index) } + fn find_lone_leaf(&self, root: Word, root_index: NodeIndex, tree_depth: u8) -> Result, ExecutionError> { + self.provider.find_lone_leaf(root, root_index, tree_depth) + } + fn update_merkle_node(&mut self, root: Word, depth: &Felt, index: &Felt, value: Word) -> Result<(MerklePath, Word), ExecutionError> { self.provider.update_merkle_node(root, depth, index, value) } @@ -414,6 +431,10 @@ impl AdviceProvider for RecAdviceProvider { self.provider.get_leaf_depth(root, tree_depth, index) } + fn find_lone_leaf(&self, root: Word, root_index: NodeIndex, tree_depth: u8) -> Result, ExecutionError> { + self.provider.find_lone_leaf(root, root_index, tree_depth) + } + fn update_merkle_node(&mut self, root: Word, depth: &Felt, index: &Felt, value: Word) -> Result<(MerklePath, Word), ExecutionError> { self.provider.update_merkle_node(root, depth, index, value) } diff --git a/processor/src/decorators/adv_stack_injectors.rs b/processor/src/decorators/adv_stack_injectors.rs index 8b44ecf188..f11195d00e 100644 --- a/processor/src/decorators/adv_stack_injectors.rs +++ b/processor/src/decorators/adv_stack_injectors.rs @@ -1,10 +1,10 @@ use super::{super::Ext2InttError, AdviceProvider, AdviceSource, ExecutionError, Process}; use vm_core::{ crypto::{ - hash::RpoDigest, - merkle::{EmptySubtreeRoots, TieredSmt}, + hash::{Rpo256, RpoDigest}, + merkle::{EmptySubtreeRoots, NodeIndex, TieredSmt}, }, - utils::collections::Vec, + utils::collections::{btree_map::Entry, BTreeMap, Vec}, Felt, FieldElement, QuadExtension, StarkField, Word, ONE, WORD_SIZE, ZERO, }; use winter_prover::math::fft; @@ -387,15 +387,14 @@ where unimplemented!("handling of depth=64 tier hasn't been implemented yet"); } - // get the value of the node a this index/depth + // get the value of the node at 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::from(depth), &index)?; // if the value to be inserted is an empty word, we need to process it as a delete if value == TieredSmt::EMPTY_VALUE { - self.handle_smt_delete(root, node, depth, index, key)?; - return Ok(()); + return self.handle_smt_delete(root, node, depth, index, key); } // figure out what kind of insert we are doing; possible options are: @@ -569,7 +568,22 @@ where Ok(()) } - /// TODO: add comments + /// Prepares the advice stack for a TSMT deletion operation. Specifically, the advice stack + /// will be arranged as follows (depending on the type of the node which occupies the location + /// at which the node for the specified key should be present): + /// + /// - Root of empty subtree: [d0, d1, ZERO (is_leaf), ONE (key_not_set)] + /// - Leaf for another key: [d0, d1, ONE (is_leaf), ONE (key_not_set), KEY, VALUE] + /// - Leaf for the provided key: [ZERO, ZERO, ZERO, ZERO (key_not_set), NEW_ROOT, 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`. + /// - KEY and VALUE is the key-value pair of a leaf node occupying the location of the node + /// for the specified key. Note that KEY may be the same as the specified key or different + /// from the specified key if the location is occupied by a different key-value pair. + /// - NEW_ROOT is the new root of the TSMT post deletion. + /// - OLD_VALUE is the value which is to be replaced with [ZERO; 4]. fn handle_smt_delete( &mut self, root: Word, @@ -578,35 +592,80 @@ where index: Felt, key: Word, ) -> Result<(), ExecutionError> { - let empty = EmptySubtreeRoots::empty_hashes(64)[depth as usize]; + let empty = EmptySubtreeRoots::empty_hashes(TieredSmt::MAX_DEPTH)[depth as usize]; if node == Word::from(empty) { + // if the node to be replaced is already an empty node, we set key_not_set = ONE, + // and is_leaf = ZERO self.advice_provider.push_stack(AdviceSource::Value(ONE))?; self.advice_provider.push_stack(AdviceSource::Value(ZERO))?; + // set depth flags based on node's depth let (is_16_or_32, is_16_or_48) = get_depth_flags(depth); self.advice_provider.push_stack(AdviceSource::Value(is_16_or_32))?; self.advice_provider.push_stack(AdviceSource::Value(is_16_or_48))?; } else { - let (leaf_key, leaf_value) = self.get_smt_upper_leaf_preimage(node).unwrap(); + // if the node is not a root of an empty subtree, it must be a leaf; thus we can get + // the key and the value stored in the leaf. + let (leaf_key, leaf_value) = self.get_smt_upper_leaf_preimage(node)?; if leaf_key != key { + // if the node to be replaced is a leaf for different key, we push that key-value + // pair onto the advice stack and set key_not_set = ONE and is_leaf = ONE + self.advice_provider.push_stack(AdviceSource::Word(leaf_value))?; self.advice_provider.push_stack(AdviceSource::Word(leaf_key))?; self.advice_provider.push_stack(AdviceSource::Value(ONE))?; self.advice_provider.push_stack(AdviceSource::Value(ONE))?; + // set depth flags based on node's depth let (is_16_or_32, is_16_or_48) = get_depth_flags(depth); self.advice_provider.push_stack(AdviceSource::Value(is_16_or_32))?; self.advice_provider.push_stack(AdviceSource::Value(is_16_or_48))?; } else { - let (_, new_root) = self.advice_provider.update_merkle_node( - root, - &Felt::from(depth), - &index, - empty.into(), - )?; + // if the key which we want to set to [ZERO; 4] does have an associated value, + // we update the tree in the advice provider to get the new root, then push the root + // and the old value onto the advice stack, key_not_set = ZERO, and also push 3 + // ZERO values for padding + let new_root = match self.find_lone_sibling(root, depth, &index)? { + Some((sibling, new_index)) => { + // if the node to be deleted has a lone sibling, we need to move it to a + // higher tier. + + // first, we compute the value of the new node on the higher tier + let (leaf_key, leaf_val) = self.get_smt_upper_leaf_preimage(*sibling)?; + let new_node = Rpo256::merge_in_domain( + &[leaf_key.into(), leaf_val.into()], + new_index.depth().into(), + ); + + // then we insert the node and its pre-image into the advice provider + let mut elements = leaf_key.to_vec(); + elements.extend_from_slice(&leaf_val); + self.advice_provider.insert_into_map(new_node.into(), elements)?; + + // and finally we update the tree in the advice provider + let (_, new_root) = self.advice_provider.update_merkle_node( + root, + &new_index.depth().into(), + &new_index.value().into(), + new_node.into(), + )?; + new_root + } + None => { + // if the node does not have a lone sibling, we just replace it with an + // empty node + let (_, new_root) = self.advice_provider.update_merkle_node( + root, + &Felt::from(depth), + &index, + empty.into(), + )?; + new_root + } + }; self.advice_provider.push_stack(AdviceSource::Word(leaf_value))?; self.advice_provider.push_stack(AdviceSource::Word(new_root))?; @@ -621,6 +680,87 @@ where Ok(()) } + + /// Returns info about a lone sibling of a leaf specified by depth and index parameters in the + /// Tiered Sparse Merkle tree defined by the specified root. If no lone siblings exist for the + /// specified parameters, None is returned. + /// + /// A lone sibling is defined as a leaf which has a common root with the specified leaf at a + /// higher tier such that the subtree starting at this root contains only these two leaves. + /// + /// In addition to the leaf node itself, this also returns the index of the common root at a + /// higher tier. + fn find_lone_sibling( + &self, + root: Word, + depth: u8, + index: &Felt, + ) -> Result, ExecutionError> { + debug_assert!(matches!(depth, 16 | 32 | 48)); + + // if the leaf is on the first tier (depth=16), we don't care about lone siblings as they + // cannot be moved to a higher tier. + if depth == TieredSmt::TIER_SIZE { + return Ok(None); + } + + let empty = &EmptySubtreeRoots::empty_hashes(TieredSmt::MAX_DEPTH)[..=depth as usize]; + + // get the path to the leaf node + let path: Vec<_> = self.advice_provider.get_merkle_path(root, &depth.into(), index)?.into(); + + // traverse the path from the leaf up to the root, keeping track of all non-empty nodes; + // here we ignore the top 16 depths because lone siblings cannot be moved to a higher tier + // from tier at depth 16. + let mut non_empty_nodes = BTreeMap::new(); + for (depth, sibling) in (TieredSmt::TIER_SIZE..=depth).rev().zip(path.iter()) { + // map the depth of each node to the tier it would "round up" to. For example, 17 maps + // to tier 1, 32 also maps to tier 1, but 33 maps to tier 2. + let tier = (depth - 1) / TieredSmt::TIER_SIZE; + + // if the node is non-empty, insert it into the map, but if a node for the same tier + // is already in the map, stop traversing the tree. we do this because if two nodes in + // a given tier are non-empty a lone sibling cannot exist at this tier or any higher + // tier. to indicate the the tier cannot contain a lone sibling, we set the value in + // the map to None. + if sibling != &empty[depth as usize] { + match non_empty_nodes.entry(tier) { + Entry::Vacant(entry) => { + entry.insert(Some((depth, *sibling))); + } + Entry::Occupied(mut entry) => { + entry.insert(None); + break; + } + } + } + } + + // take the deepest non-empty node and check if its subtree contains just a single leaf + if let Some((_, Some((node_depth, node)))) = non_empty_nodes.pop_last() { + let mut node_index = NodeIndex::new(depth, index.as_int()).expect("invalid node index"); + node_index.move_up_to(node_depth); + let node_index = node_index.sibling(); + + if let Some((mut leaf_index, leaf)) = self.advice_provider.find_lone_leaf( + node.into(), + node_index, + TieredSmt::MAX_DEPTH, + )? { + // if the node's subtree does contain a single leaf, figure out to which depth + // we can move it up to. we do this by taking the next tier down from the tier + // which contained at least one non-empty node on the path from the original leaf + // up to the root. if there were no non-empty nodes on this path, we default to + // the first tier (i.e., depth 16). + let target_tier = non_empty_nodes.keys().last().map(|&t| t + 1).unwrap_or(1); + leaf_index.move_up_to(target_tier * TieredSmt::TIER_SIZE); + + return Ok(Some((leaf.into(), leaf_index))); + } + } + + Ok(None) + } } // HELPER FUNCTIONS diff --git a/stdlib/asm/collections/smt.masm b/stdlib/asm/collections/smt.masm index ad90832213..72b95b156b 100644 --- a/stdlib/asm/collections/smt.masm +++ b/stdlib/asm/collections/smt.masm @@ -1244,7 +1244,7 @@ proc.verify_empty_node_48 end #! Verifies that a leaf node located at depth d and index idx in the tree defined by root R is -#! contains key which is different from K. +#! contained key which is different from K. #! #! Input: [d, idx, K, R, Z, ...] #! Output: [Z, R, ...] @@ -1287,11 +1287,15 @@ end #! If the key-value pair is not in the tree, this proves that the key-value pair is not in the #! tree but does not modify the tree itself. #! -#! This assumes that the value Z is [ZERO; 4], but this is not checked. -#! #! Input: [key_not_in_tree, f0, f1, f2, Z, K, R, ...] #! Output: [V_old, R_new, ...] #! +#! Where: +#! - key_not_in_tree is a flag specifying whether the specified key is present in the tree defined +#! by R. +#! - Meaning of the flags f0, f1, and f2 depends on what type of delete is being executed. +#! - Z is [ZERO; 4]. This is assumed but not checked. +#! #! Cycles: #! - Key not in the tree (key with common prefix in the tree): #! - Depth 16: 68 @@ -1349,7 +1353,7 @@ proc.delete end end else # --- key is in the tree --------------------------------------------------------------- - # load the root of the tree with the leaf removed from the advice provider (2 cycles) + # load from the advice provider the root of the tree with the leaf removed (2 cycles) push.0 adv_loadw # => [R_new, Z, K, R, ...] diff --git a/stdlib/tests/collections/smt.rs b/stdlib/tests/collections/smt.rs index 448de05c62..60f21f2f05 100644 --- a/stdlib/tests/collections/smt.rs +++ b/stdlib/tests/collections/smt.rs @@ -397,12 +397,198 @@ fn tsmt_set_16() { let key_c = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_c)]); assert_set(&smt, key_c, EMPTY_VALUE, EMPTY_VALUE, smt.root().into(), &[]); - // set the value at key A to an empty word + // set the value at key A to an empty word; this removes node A from the tree let init_smt = smt.clone(); smt.insert(key_a.into(), EMPTY_VALUE); assert_set(&init_smt, key_a, val_a2, EMPTY_VALUE, smt.root().into(), &[]); } +#[test] +fn tsmt_set_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]; + + // this tests a complex insertion when a leaf node moves from depth 16 to depth 32; this + // moves the original node to depth 32, and thus two new entries are added to the advice map + let init_smt = smt.clone(); + smt.insert(key_b.into(), val_b); + let new_map_entries = [build_node_entry(key_a, val_a, 32), build_node_entry(key_b, val_b, 32)]; + assert_set(&init_smt, key_b, EMPTY_VALUE, val_b, smt.root().into(), &new_map_entries); + + // update a value under key_a; this adds one new entry to the advice map + let init_smt = smt.clone(); + let val_a2 = [ONE, ZERO, ZERO, ONE]; + smt.insert(key_a.into(), val_a2); + let new_map_entries = [build_node_entry(key_a, val_a2, 32)]; + assert_set(&init_smt, key_a, val_a, val_a2, smt.root().into(), &new_map_entries); + + // insert a value under key_c which has the same 16-bit prefix as A and B; this inserts a new + // node at depth 32, and thus adds one entry to the advice map + 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); + let new_map_entries = [build_node_entry(key_c, val_c, 32)]; + assert_set(&init_smt, key_c, EMPTY_VALUE, val_c, smt.root().into(), &new_map_entries); + + // remove C from the tree + let init_smt = smt.clone(); + smt.insert(key_c.into(), EMPTY_VALUE); + assert_set(&init_smt, key_c, val_c, EMPTY_VALUE, smt.root().into(), &[]); + + // remove A from the tree; this should move B to depth 16, and thus inserts a new entry into + // the advice map for node B at depth 16 + let init_smt = smt.clone(); + smt.insert(key_a.into(), EMPTY_VALUE); + let new_map_entries = [build_node_entry(key_b, val_b, 16)]; + assert_set(&init_smt, key_a, val_a2, EMPTY_VALUE, smt.root().into(), &new_map_entries); +} + +#[test] +fn tsmt_set_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_00111111_11111111_11111111_11111111_u64; + let key_b = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_b)]); + let val_b = [ONE, ONE, ZERO, ZERO]; + + // this tests a complex insertion when a leaf moves from depth 16 to depth 48; this moves + // node at depth 16 to depth 48 and inserts a new node at depth 48 + let init_smt = smt.clone(); + smt.insert(key_b.into(), val_b); + let new_map_entries = [build_node_entry(key_a, val_a, 48), build_node_entry(key_b, val_b, 48)]; + assert_set(&init_smt, key_b, EMPTY_VALUE, val_b, smt.root().into(), &new_map_entries); + + // update a value under key_a; this inserts one entry into the advice map + let init_smt = smt.clone(); + let val_a2 = [ONE, ZERO, ZERO, ONE]; + smt.insert(key_a.into(), val_a2); + let new_map_entries = [build_node_entry(key_a, val_a2, 48)]; + assert_set(&init_smt, key_a, val_a, val_a2, smt.root().into(), &new_map_entries); + + // insert a value under key_c which has the same 32-bit prefix as A and B; this inserts + // one entry into the advice map + let raw_c = 0b00000000_00000000_11111111_11111111_00111111_01111111_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); + let new_map_entries = [build_node_entry(key_c, val_c, 48)]; + assert_set(&init_smt, key_c, EMPTY_VALUE, val_c, smt.root().into(), &new_map_entries); + + // at this point the tree has 3 nodes A, B, C, all 3 are a depth 48 and share the same 32-bit + // prefix; also B and C share the same 34-bit prefix. + + // remove node A from the tree; since B and C share the same 34-bit prefix, they remain at + // depth 48 + let init_smt = smt.clone(); + smt.insert(key_a.into(), EMPTY_VALUE); + assert_set(&init_smt, key_a, val_a2, EMPTY_VALUE, smt.root().into(), &[]); + + // remove node B from the tree; this will move node C to depth 16, and thus inserts a new + // entry into the advice map for node C at depth 16 + let init_smt = smt.clone(); + smt.insert(key_b.into(), EMPTY_VALUE); + let new_map_entries = [build_node_entry(key_c, val_c, 16)]; + assert_set(&init_smt, key_b, val_b, EMPTY_VALUE, smt.root().into(), &new_map_entries); +} + +#[test] +fn tsmt_set_48_from_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_01111111_11111111_11111111_11111111_u64; + let key_b = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_b)]); + let val_b = [ONE, ONE, ZERO, ZERO]; + smt.insert(key_b.into(), val_b); + + // insert a value under key_c which has the same 32-bit prefix as A + 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]; + + // this tests a complex insertion when a leaf moves from depth 32 to depth 48; two new + // entries are added to the advice map + let init_smt = smt.clone(); + smt.insert(key_c.into(), val_c); + let new_map_entries = [build_node_entry(key_a, val_a, 48), build_node_entry(key_c, val_c, 48)]; + assert_set(&init_smt, key_c, EMPTY_VALUE, val_c, smt.root().into(), &new_map_entries); + + // remove C from the tree; this should cause leaf A to move to depth 32 + let init_smt = smt.clone(); + smt.insert(key_c.into(), EMPTY_VALUE); + let new_map_entries = [build_node_entry(key_a, val_a, 32)]; + assert_set(&init_smt, key_c, val_c, EMPTY_VALUE, smt.root().into(), &new_map_entries); + + // remove B from the tree; this should cause leaf A to move to depth 16 + let init_smt = smt.clone(); + smt.insert(key_b.into(), EMPTY_VALUE); + let new_map_entries = [build_node_entry(key_a, val_a, 16)]; + assert_set(&init_smt, key_b, val_b, EMPTY_VALUE, smt.root().into(), &new_map_entries); +} + +#[test] +fn tsmt_set_48_lone_sibling_move_to_32() { + let mut smt = TieredSmt::default(); + + // depth 48 leaf + 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); + + // depth 48 leaf + let raw_b = 0b00000000_00000000_11111111_11111111_11111111_00111111_11111111_11111111_u64; + let key_b = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_b)]); + let val_b = [ONE, ONE, ZERO, ZERO]; + smt.insert(key_b.into(), val_b); + + // depth 32 leaf + let raw_c = 0b00000000_00000000_11111111_11111110_11111111_11111111_11111111_11111111_u64; + let key_c = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_c)]); + let val_c = [ONE, ZERO, ZERO, ZERO]; + smt.insert(key_c.into(), val_c); + + // depth 32 leaf + let raw_d = 0b00000000_00000000_11111111_11111101_11111111_11111111_11111111_11111111_u64; + let key_d = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_d)]); + let val_d = [ONE, ONE, ZERO, ZERO]; + smt.insert(key_d, val_d); + + // remove leaf a such that it key_b, val_b should move to depth 32 + let init_smt = smt.clone(); + smt.insert(key_a.into(), EMPTY_VALUE); + let new_map_entries = [build_node_entry(key_b, val_b, 32)]; + assert_set(&init_smt, key_a, val_a, EMPTY_VALUE, smt.root().into(), &new_map_entries); +} + fn assert_set( init_smt: &TieredSmt, key: RpoDigest,