Skip to content

Commit

Permalink
addressed comments
Browse files Browse the repository at this point in the history
  • Loading branch information
bobbinth committed Aug 10, 2023
1 parent a7c22fc commit 6a5e2f9
Show file tree
Hide file tree
Showing 4 changed files with 73 additions and 112 deletions.
51 changes: 12 additions & 39 deletions processor/src/decorators/adv_stack_injectors.rs
Original file line number Diff line number Diff line change
Expand Up @@ -352,14 +352,13 @@ where
/// Operand stack: [OLD_VALUE, NEW_ROOT, ...]
/// Advice stack, depends on the type of insert:
/// - Simple insert at depth 16: [d0, d1, ONE, ZERO]
/// - Simple insert at depth 32 or 48: [d0, d1, ONE, ZERO, P_NODE, CHILD_1, CHILD_2]
/// - Update of an existing leaf: [ONE, d0, d1, ONE, VALUE]
/// - Simple insert at depth 32 or 48: [d0, d1, ONE, ZERO, P_NODE]
/// - Update of an existing leaf: [ONE, d0, d1, ONE, 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, with CHILD1 and
/// CHILD2 being its immediate children.
/// - 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.
Expand Down Expand Up @@ -390,7 +389,7 @@ where
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 an insert we are doing; possible options are:
// 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).
Expand All @@ -399,14 +398,11 @@ where
// 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; we also provide immediate children of the node so that
// in the VM, we can prove that this is an internal node rather than a leaf node
let (p_node, child1, child2) =
self.get_node_and_children(root, index.as_int() >> 16, depth - 16)?;

// the order of the nodes on the advice stack is arranged so that we read the node
// first, followed by the left child, and then the right child.
for &element in p_node.iter().chain(child1.iter()).chain(child2.iter()).rev() {
// 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))?;
}
}
Expand All @@ -429,9 +425,11 @@ where
// 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 {
(ZERO, ONE)
// return is_update = ZERO, is_simple_insert = ZERO
(ZERO, ZERO)
}
};

Expand All @@ -454,31 +452,6 @@ where
}
Ok(())
}

// HELPER METHODS
// --------------------------------------------------------------------------------------------

/// Returns a node located at the specified index and depth in tree defined by the specified
/// root. In addition to the node itself, immediate children of the node are also returned.
fn get_node_and_children(
&self,
root: Word,
index: u64,
depth: u8,
) -> Result<(Word, Word, Word), ExecutionError> {
let node_depth = Felt::from(depth);
let node_index = Felt::from(index);
let node = self.advice_provider.get_tree_node(root, &node_depth, &node_index)?;

let child_depth = Felt::from(depth + 1);
let child1_index = Felt::from(index << 1);
let child1 = self.advice_provider.get_tree_node(root, &child_depth, &child1_index)?;

let child2_index = Felt::from((index << 1) + 1);
let child2 = self.advice_provider.get_tree_node(root, &child_depth, &child2_index)?;

Ok((node, child1, child2))
}
}

// HELPER FUNCTIONS
Expand Down
2 changes: 1 addition & 1 deletion processor/src/decorators/tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -154,7 +154,7 @@ fn inject_smtinsert() {
let key_a = build_key(raw_a);
let val_a = [Felt::new(3), Felt::new(5), Felt::new(7), Felt::new(9)];

// this is be a simple insertion at depth 16, and thus the flags should look as follows:
// 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;
Expand Down
130 changes: 59 additions & 71 deletions stdlib/asm/collections/smt.masm
Original file line number Diff line number Diff line change
Expand Up @@ -31,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, ...]
Expand All @@ -55,7 +65,7 @@ end
#! 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, ...]
Expand Down Expand Up @@ -129,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, ...]
Expand Down Expand Up @@ -203,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, ...]
Expand Down Expand Up @@ -291,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
Expand All @@ -319,8 +329,8 @@ end
#!
#! Where:
#! - R is the initial root of the TSMT, and R_new is the new root of the TSMT.
#! - d and idx is the depth and index (at that depth) of the leaf node to be updated.
#! - K and V is the key-value pair for the leaf node where V is a new value for key K.
#! - 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:
Expand Down Expand Up @@ -420,19 +430,20 @@ proc.insert_16
# => [0, 0, 0, 0, R_new, ...]
end

#! Inserts a new leaf node at depths 32.
#! 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.
#! - P is the internal node at depth 16 on the path to the new leaf node.
#!
#! This procedure consists of two high-level steps:
#! - First, insert N = hash([K, V], domain=32) into a subtree with root P.
#! - Then, insert the new root P into the TSMT with root R.
#! - 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.
Expand All @@ -441,7 +452,7 @@ end
#! - Node at depth 16 is an internal node.
#! - Node at depth 32 is a root of an empty subtree.
#!
#! Cycles: 183
#! Cycles: 154
proc.insert_32.2
# load the value of P from the advice provider (5 cycles)
adv_push.4 swapw.2
Expand Down Expand Up @@ -473,14 +484,15 @@ proc.insert_32.2

# 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 u32split swap drop u32unchecked_divmod.65536
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 (30 cycles)
# 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, ...]

Expand All @@ -492,35 +504,19 @@ proc.insert_32.2
# => [P_new, X, X, R, ...]

# prepare the stack for mtree_set operation against R; here we load idx_hi from loc[0][0]
# (7 cycles)
swapw swapw.3 loc_load.0 push.16
# => [16, idx_hi, R, P_new, X, X, ...]
# (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, X, ...]

# load previously saved P with P_old to make sure they are the same (6 cycles)
swapw swapw.3 loc_loadw.1
# => [P, P_old, X, R_new, ...]

# make sure P and P_old are the same (12 cycles)
repeat.4
dup.4 assert_eq
end
# => [P, X, R_new, ...]
# => [P_old, R_new, X, ...]

# load values A and B from the advice provider and prepare to compute H = hash([A, B]); we'll
# use this to prove that P is an internal node (11 cycles)
swapw adv_loadw padw swapw adv_push.4
# => [B, A, 0, 0, 0, 0, P, R_new, ...]
# load previously saved P to compare it with P_old (6 cycles)
swapw swapw.2 loc_loadw.1
# => [P, P_old, R_new, ...]

# compute H = hash([A, B]) (10 cycles)
hperm dropw swapw dropw
# => [H, P, R_new, ...]

# make sure P and H are the same; this proves that the node is an internal node because leaf
# nodes are always hashed in a specific domain (11 cycles)
# make sure P and P_old are the same (11 cycles)
assert_eqw
# => [R_new, ...]

Expand All @@ -529,7 +525,7 @@ proc.insert_32.2
# => [0, 0, 0, 0, R_new, ...]
end

#! Inserts a new leaf node at depths 48.
#! Inserts a new leaf node at depth 48.
#!
#! Input: [V, K, R, ...];
#! Output:[0, 0, 0, 0, R_new, ...]
Expand Down Expand Up @@ -575,47 +571,32 @@ proc.insert_48.2
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 (30 cycles)
# 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)
# 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]
# (7 cycles)
swapw swapw.3 loc_load.0 push.32
# => [32, idx_hi, R, P_new, X, X, ...]
# (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, X, ...]
# => [P_old, R_new, X, ...]

# load previously saved P with P_old to make sure they are the same (6 cycles)
swapw swapw.3 loc_loadw.1
# => [P, P_old, X, R_new, ...]

# make sure P and P_old are the same (12 cycles)
repeat.4
dup.4 assert_eq
end
# => [P, X, R_new, ...]

# load values A and B from the advice provider and prepare to compute H = hash([A, B]); we'll
# use this to prove that P is an internal node (11 cycles)
swapw adv_loadw padw swapw adv_push.4
# => [B, A, 0, 0, 0, 0, P, R_new, ...]

# compute H = hash([A, B]) (10 cycles)
hperm dropw swapw dropw
# => [H, P, R_new, ...]
swapw swapw.2 loc_loadw.1
# => [P, P_old, R_new, ...]

# make sure P and H are the same; this proves that the node is an internal node because leaf
# nodes are always hashed in a specific domain (11 cycles)
# make sure P and P_old are the same (11 cycles)
assert_eqw
# => [R_new, ...]

Expand All @@ -642,8 +623,8 @@ end
#! - Depth 48: 131
#! - Insert new leaf:
#! - Depth 16: 100
#! - Depth 32: 210
#! - Depth 48: 210
#! - Depth 32: 181
#! - Depth 48: 181
#! - Replace a leaf with a subtree:
#! - Depth 32: TODO
#! - Depth 48: TODO
Expand All @@ -655,15 +636,19 @@ export.insert
and and and assertz
# => [V, K, R, ...]

# invoke adv and fetch target depth flags (4 cycles)
# 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
# => [TODO]
# => [is_update, f0, f1, f2, V, K, R, ...]

# call the inner procedure depending on the type of insert and depth
if.true # --- update leaf -----------------------------------------------------
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
Expand All @@ -673,9 +658,10 @@ export.insert
exec.update_16_32_48
else # --- update a leaf node at depth 32 ---
drop
# => [V, K, R, ...]

#(5 cycles)
dup.4 u32split swap drop
dup.4 exec.get_top_32_bits
push.32
# => [32, idx, V, K, R, ...]

Expand All @@ -684,6 +670,7 @@ export.insert
else
if.true # --- update a leaf node at depth 48 ---
drop
# => [V, K, R, ...]

# (10 cycles)
dup.4 exec.get_top_48_bits
Expand All @@ -697,6 +684,7 @@ export.insert
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
Expand Down
Loading

0 comments on commit 6a5e2f9

Please sign in to comment.