Implement TSMT leaf insertion into advice map

CHANGELOG.md

@@ -7,6 +7,7 @@
 - Added support for nested modules (#992).
 - Added support for the arithmetic expressions in constant values (#1026).
 - Added support for module aliases (#1037).
+- Added `adv.insert_hperm` decorator (#1042).
 
 #### VM Internals
 - Simplified range checker and removed 1 main and 1 auxiliary trace column (#949).

assembly/src/ast/nodes/advice.rs

@@ -27,6 +27,7 @@ pub enum AdviceInjectorNode {
     InsertMem,
     InsertHdword,
     InsertHdwordImm { domain: u8 },
+    InsertHperm,
 }
 
 impl From<&AdviceInjectorNode> for AdviceInjector {
@@ -59,6 +60,7 @@ impl From<&AdviceInjectorNode> for AdviceInjector {
             InsertHdwordImm { domain } => Self::HdwordToMap {
                 domain: Felt::from(*domain),
             },
+            InsertHperm => Self::HpermToMap,
         }
     }
 }
@@ -79,6 +81,7 @@ impl fmt::Display for AdviceInjectorNode {
             InsertMem => write!(f, "insert_mem"),
             InsertHdword => write!(f, "insert_hdword"),
             InsertHdwordImm { domain } => write!(f, "insert_hdword.{domain}"),
+            InsertHperm => writeln!(f, "insert_hperm"),
         }
     }
 }
@@ -98,6 +101,7 @@ const PUSH_MTNODE: u8 = 8;
 const INSERT_MEM: u8 = 9;
 const INSERT_HDWORD: u8 = 10;
 const INSERT_HDWORD_IMM: u8 = 11;
+const INSERT_HPERM: u8 = 12;
 
 impl Serializable for AdviceInjectorNode {
     fn write_into<W: ByteWriter>(&self, target: &mut W) {
@@ -124,6 +128,7 @@ impl Serializable for AdviceInjectorNode {
                 target.write_u8(INSERT_HDWORD_IMM);
                 target.write_u8(*domain);
             }
+            InsertHperm => target.write_u8(INSERT_HPERM),
         }
     }
 }
@@ -158,6 +163,7 @@ impl Deserializable for AdviceInjectorNode {
                 let domain = source.read_u8()?;
                 Ok(AdviceInjectorNode::InsertHdwordImm { domain })
             }
+            INSERT_HPERM => Ok(AdviceInjectorNode::InsertHperm),
             val => Err(DeserializationError::InvalidValue(val.to_string())),
         }
     }

assembly/src/ast/parsers/adv_ops.rs

@@ -81,6 +81,10 @@ pub fn parse_adv_inject(op: &Token) -> Result<Node, ParsingError> {
             }
             _ => return Err(ParsingError::extra_param(op)),
         },
+        "insert_hperm" => match op.num_parts() {
+            2 => AdvInject(InsertHperm),
+            _ => return Err(ParsingError::extra_param(op)),
+        },
         _ => return Err(ParsingError::invalid_op(op)),
     };
 

core/src/operations/decorators/advice.rs

@@ -214,6 +214,21 @@ pub enum AdviceInjector {
     /// Where KEY is computed as hash(A || B, domain), where domain is provided via the immediate
     /// value.
     HdwordToMap { domain: Felt },
+
+    /// Reads three words from the operand stack and inserts the top two words into the advice map
+    /// under the key defined by applying an RPO permutation to all three words.
+    ///
+    /// Inputs:
+    ///   Operand stack: [B, A, C, ...]
+    ///   Advice map: {...}
+    ///
+    /// Outputs:
+    ///   Operand stack: [B, A, C, ...]
+    ///   Advice map: {KEY: [a0, a1, a2, a3, b0, b1, b2, b3]}
+    ///
+    /// Where KEY is computed by extracting the digest elements from hperm([C, A, B]). For example,
+    /// if C is [0, d, 0, 0], KEY will be set as hash(A || B, d).
+    HpermToMap,
 }
 
 impl fmt::Display for AdviceInjector {
@@ -238,6 +253,7 @@ impl fmt::Display for AdviceInjector {
             Self::SmtInsert => write!(f, "smt_insert"),
             Self::MemToMap => write!(f, "mem_to_map"),
             Self::HdwordToMap { domain } => write!(f, "hdword_to_map.{domain}"),
+            Self::HpermToMap => write!(f, "hperm_to_map"),
         }
     }
 }

docs/src/user_docs/assembly/io_operations.md

@@ -47,14 +47,16 @@ Advice injectors fall into two categories: (1) injectors which push new data ont
 
 | Instruction                                  | Stack_input                | Stack_output               | Notes                                                                                                                                                                                                                                           |
 | -------------------------------------------- | -------------------------- | -------------------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
-| adv.push_mapval <br> adv.push_mapval.*s*     | [K, ... ]                  | [K, ... ]                  | Pushes a list of field elements onto the advice stack. The list is looked up in the advice map using word $K$ as the key. If offset $s$ is provided, the key is taken starting from item $s$ on the stack.                                      |
+| adv.push_mapval <br> adv.push_mapval.*s*     | [K, ... ]                  | [K, ... ]                  | Pushes a list of field elements onto the advice stack. The list is looked up in the advice map using word $K$ as the key. If offset $s$ is provided, the key is taken starting from item $s$ on the stack. |
 | adv.push_mapvaln <br> adv.push_mapvaln.*s*   | [K, ... ]                  | [K, ... ]                  | Pushes a list of field elements together with the number of elements onto the advice stack. The list is looked up in the advice map using word $K$ as the key. If offset $s$ is provided, the key is taken starting from item $s$ on the stack. |
-| adv.push_mtnode                              | [d, i, R, ... ]            | [d, i, R, ... ]            | Pushes a node of a Merkle tree with root $R$ at depth $d$ and index $i$ from Merkle store onto the advice stack.                                                                                                                                |
-| adv.push_u64div                              | [b1, b0, a1, a0, ...]      | [b1, b0, a1, a0, ...]      | Pushes the result of `u64` division $a / b$ onto the advice stack. Both $a$ and $b$ are represented using 32-bit limbs. The result consists of both the quotient and the remainder.                                                             |
-| adv.push_ext2intt                            | [osize, isize, iptr, ... ] | [osize, isize, iptr, ... ] | Given evaluations of a polynomial over some specified domain, interpolates the evaluations into a polynomial in coefficient form and pushes the result into the advice stack.                                                                   |
-| adv.smt_get                                  | [K, R, ... ]               | [K, R, ... ]               | Pushes values onto the advice stack which are required for successful retrieval of a  value under the key $K$ from a Sparse Merkle Tree with root $R$.                                                                                          |
-| adv.insert_mem                               | [K, a, b, ... ]            | [K, a, b, ... ]            | Reads words $data \leftarrow mem[a] .. mem[b]$ from memory, and save the data into $advice\_map[K] \leftarrow data$.                                                                                                                            |
-| adv.insert_hdword <br> adv.insert_hdword.*d* | [B, A, ... ]               | [B, A, ... ]               | Reads top two words from the stack, computes a key as $K \leftarrow hash(A || b, d)$, and saves the data into $advice\_map[K] \leftarrow [A, B]$. $d$ is an optional domain value which can be between $0$ and $255$, default value $0$.        |
+| adv.push_mtnode                              | [d, i, R, ... ]            | [d, i, R, ... ]            | Pushes a node of a Merkle tree with root $R$ at depth $d$ and index $i$ from Merkle store onto the advice stack. |
+| adv.push_u64div                              | [b1, b0, a1, a0, ...]      | [b1, b0, a1, a0, ...]      | Pushes the result of `u64` division $a / b$ onto the advice stack. Both $a$ and $b$ are represented using 32-bit limbs. The result consists of both the quotient and the remainder. |
+| adv.push_ext2intt                            | [osize, isize, iptr, ... ] | [osize, isize, iptr, ... ] | Given evaluations of a polynomial over some specified domain, interpolates the evaluations into a polynomial in coefficient form and pushes the result into the advice stack. |
+| adv.smt_get                                  | [K, R, ... ]               | [K, R, ... ]               | Pushes values onto the advice stack which are required for successful retrieval of a  value under the key $K$ from a Sparse Merkle Tree with root $R$. |
+| adv.smt_insert                               | [V, K, R, ...]             | [V, K, R, ...]             | Pushes values onto the advice stack which are required for successful insertion of a key-value pair $(K, V)$ into a Sparse Merkle Tree with root $R$. |
+| adv.insert_mem                               | [K, a, b, ... ]            | [K, a, b, ... ]            | Reads words $data \leftarrow mem[a] .. mem[b]$ from memory, and save the data into $advice\_map[K] \leftarrow data$. |
+| adv.insert_hdword <br> adv.insert_hdword.*d* | [B, A, ... ]               | [B, A, ... ]               | Reads top two words from the stack, computes a key as $K \leftarrow hash(A || b, d)$, and saves the data into $advice\_map[K] \leftarrow [A, B]$. $d$ is an optional domain value which can be between $0$ and $255$, default value $0$. |
+| adv.insert_hperm                             | [B, A, C, ...]             | [B, A, C, ...]             | Reads top three words from the stack, computes a key as $K \leftarrow permute(C, A, B).digest$, and saves data into $advice\_mpa[K] \leftarrow [A, B]$. |
 
 ### Random access memory
 

processor/src/advice/mod.rs

@@ -103,6 +103,7 @@ pub trait AdviceProvider {
 
     // 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]>;
 

processor/src/advice/source.rs

@@ -9,13 +9,13 @@ pub enum AdviceSource {
     /// Puts a single value onto the advice stack.
     Value(Felt),
 
-    /// Puts a word (4 elements) ont the the stack.
+    /// Puts a word (4 elements) onto the stack.
     Word(Word),
 
     /// Fetches a list of elements under the specified key from the advice map and pushes them onto
     /// the advice stack.
     ///
-    /// If `include_len` is set to true, this also pushes the number of elements ont the advice
+    /// If `include_len` is set to true, this also pushes the number of elements onto the advice
     /// stack.
     ///
     /// Note: this operation doesn't consume the map element so it can be called multiple times

processor/src/decorators/adv_map_injectors.rs

@@ -1,5 +1,9 @@
 use super::{AdviceProvider, ExecutionError, Process};
-use vm_core::{crypto::hash::Rpo256, utils::collections::Vec, Felt, StarkField, WORD_SIZE, ZERO};
+use vm_core::{
+    crypto::hash::{Rpo256, RpoDigest},
+    utils::collections::Vec,
+    Felt, StarkField, WORD_SIZE, ZERO,
+};
 
 // ADVICE INJECTORS
 // ================================================================================================
@@ -72,6 +76,50 @@ where
         self.advice_provider.insert_into_map(key.into(), values)
     }
 
+    /// Reads three words from the operand stack and inserts the top two words into the advice map
+    /// under the key defined by applying an RPO permutation to all three words.
+    ///
+    /// Inputs:
+    ///   Operand stack: [B, A, C, ...]
+    ///   Advice map: {...}
+    ///
+    /// Outputs:
+    ///   Operand stack: [B, A, C, ...]
+    ///   Advice map: {KEY: [a0, a1, a2, a3, b0, b1, b2, b3]}
+    ///
+    /// Where KEY is computed by extracting the digest elements from hperm([C, A, B]). For example,
+    /// if C is [0, d, 0, 0], KEY will be set as hash(A || B, d).
+    pub(super) fn insert_hperm_into_adv_map(&mut self) -> Result<(), ExecutionError> {
+        // read the state from the stack
+        let mut state = [
+            self.stack.get(11),
+            self.stack.get(10),
+            self.stack.get(9),
+            self.stack.get(8),
+            self.stack.get(7),
+            self.stack.get(6),
+            self.stack.get(5),
+            self.stack.get(4),
+            self.stack.get(3),
+            self.stack.get(2),
+            self.stack.get(1),
+            self.stack.get(0),
+        ];
+
+        // get the values to be inserted into the advice map from the state
+        let values = state[Rpo256::RATE_RANGE].to_vec();
+
+        // apply the permutation to the state and extract the key from it
+        Rpo256::apply_permutation(&mut state);
+        let key = RpoDigest::new(
+            state[Rpo256::DIGEST_RANGE]
+                .try_into()
+                .expect("failed to extract digest from state"),
+        );
+
+        self.advice_provider.insert_into_map(key.into(), values)
+    }
+
     // HELPER METHODS
     // --------------------------------------------------------------------------------------------
 

processor/src/decorators/mod.rs

@@ -48,6 +48,7 @@ where
             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),
+            AdviceInjector::HpermToMap => self.insert_hperm_into_adv_map(),
         }
     }
 

stdlib/asm/collections/smt.masm

@@ -127,7 +127,7 @@ end
 #!
 #! Cycles: 20
 proc.extract_index_32_16
-    # slit teh most significant elements of U and V into 32-bit chunks (4 cycles)
+    # split the most significant elements of U and V into 32-bit chunks (4 cycles)
     dup.4 u32split dup.2 u32split
     # => [u3_hi, u3_lo, v3_hi, v3_lo, U, V, ...]
 
@@ -435,6 +435,9 @@ proc.update_16_32_48.2
     movdn.3 movup.2 drop push.0 swapw.2 loc_storew.1 swapw
     # => [V, K, 0, 0, d, 0, R, ...]
 
+    # insert N |-> [K, V] into the advice map (0 cycles)
+    adv.insert_hperm
+
     # compute the hash of the node N = hash([K, V], domain=d) - (1 cycle)
     hperm
     # => [X, N, X, R, ...]
@@ -490,10 +493,13 @@ proc.insert_16
     swapw dup exec.get_top_16_bits
     # => [idx, K, V, R, ...]
 
-    # prepare the stack for computing leaf node value (6 cycles)
+    # prepare the stack for computing N = hash([K, V], domain=16) (6 cycles)
     movdn.8 push.0.16.0.0 swapw.2
     # => [V, K, 0, 0, 16, 0, idx, R, ...]
 
+    # insert N |-> [K, V] into the advice map (0 cycles)
+    adv.insert_hperm
+
     # compute leaf node value as N = hash([K, V], domain=16) (10 cycles)
     hperm dropw swapw dropw
     # => [N, idx, R, ...]
@@ -554,6 +560,9 @@ proc.insert_32.2
     push.0.32.0.0 swapw.2
     # => [V, K, 0, 0, 32, 0, P, R, ...]
 
+    # insert N |-> [K, V] into the advice map (0 cycles)
+    adv.insert_hperm
+
     # compute N = hash([K, V], domain=32) (1 cycle)
     hperm
     # => [X, N, X, P, R, ...]
@@ -634,6 +643,9 @@ proc.insert_48.2
     push.0.48.0.0 swapw.2
     # => [V, K, 0, 0, 48, 0, P, R, ...]
 
+    # insert N |-> [K, V] into the advice map (0 cycles)
+    adv.insert_hperm
+
     # compute N = hash([K, V], domain=48) (1 cycle)
     hperm
     # => [X, N, X, P, R, ...]
@@ -695,7 +707,7 @@ end
 
 #! Replaces a leaf node at depth 16 with a subtree containing two leaf nodes at depth 32 such that
 #! one of the leaf nodes commits to a key-value pair equal to the leaf node at depth 16, and the
-#! other leaf node comments to the key-value pair being inserted.
+#! other leaf node commits to the key-value pair being inserted.
 #!
 #! Input:  [idx_lo_e, idx_lo_n, idx_hi, K_e, K, V, R, ...]
 #! Output: [0, 0, 0, 0, R_new, ...]
@@ -738,6 +750,9 @@ proc.replace_32.3
     loc_loadw.1 push.0.32.0.0 swapw.2 swapw.3 loc_loadw.2
     # => [V_e, K_e, 0, 0, 32, 0, P, K, V, R, ...]
 
+    # insert M |-> [K_e, V_e] into the advice map (0 cycles)
+    adv.insert_hperm
+
     # compute M = hash([K_e, V_e], domain=32) - 1 cycle
     hperm
     # => [X, M, X, P, K, V, R, ...]
@@ -759,6 +774,9 @@ proc.replace_32.3
     dropw swapw dropw swapdw push.0.32.0.0 swapw.2
     # => [V, K, 0, 0, 32, 0, T, P, R, ...]
 
+    # insert N |-> [K, V] into the advice map (0 cycles)
+    adv.insert_hperm
+
     # compute N = hash([K, V], domain=32) - 1 cycle
     hperm
     # => [X, N, X, T, P, R, ...]
@@ -767,7 +785,7 @@ proc.replace_32.3
     dropw swapw.2 swapw loc_loadw.0 swap drop movup.2 drop
     # => [16, idx_lo_n, T, N, P, R, ...]
 
-    # insert M into an empty subtree rooted at T; this leaves a root of empty subtree at depth 32
+    # insert N into an empty subtree rooted at T; this leaves a root of empty subtree at depth 32
     # on the stack - though, we don't need to verify this (29 cycles)
     mtree_set
     # => [E32, P_new, P, R, ...]
@@ -792,7 +810,7 @@ end
 
 #! Replaces a leaf node at depth 16 or 32 with a subtree containing two leaf nodes at depth 48
 #! such that one of the leaf nodes commits to a key-value pair equal to the leaf node at the
-#! original depth, and the other leaf node comments to the key-value pair being inserted.
+#! original depth, and the other leaf node commits to the key-value pair being inserted.
 #!
 #! Input:  [E, idx_lo_e, idx_lo_n, idx_hi, d, K_e, K, V, R, ...]
 #! Output: [0, 0, 0, 0, R_new, ...]
@@ -808,7 +826,7 @@ end
 #!
 #! This procedure consists of three high-level steps:
 #! - First, insert M = hash([K_e, V_e], domain=48) into an empty subtree at depth 48 - d, where
-#!   K_e and V_e are the key-value pair for the existing leaf node at depth d. This outputs the
+#!   K_e and V_e is the key-value pair for the existing leaf node at depth d. This outputs the
 #!   new root of the subtree T.
 #! - Then, insert N = hash([K, V], domain=48) into a subtree with root T. This outputs the new
 #!   root of the subtree P_new.
@@ -857,6 +875,9 @@ proc.replace_48.4
     loc_loadw.2 push.0.48.0.0 swapw.2 swapw.3 loc_loadw.1
     # => [V_e, K_e, 0, 0, 48, 0, P, K, V, R, ...]
 
+    # insert M |-> [K_e, V_e] into the advice map (0 cycles)
+    adv.insert_hperm
+
     # compute M = hash([K_e, V_e], domain=48) (1 cycle)
     hperm
     # => [X, M, X, P, K, V, R, ...]
@@ -890,6 +911,9 @@ proc.replace_48.4
     push.0.48.0.0 swapw.2
     # => [V, K, 0, 0, 48, 0, T, P, R, ...]
 
+    # insert N |-> [K, V] into the advice map (0 cycles)
+    adv.insert_hperm
+
     # compute N = hash([K, V], domain=48) - (1 cycles)
     hperm
     # => [X, N, X, T, P, R, ...]