Merge pull request #912 from 0xPolygonMiden/bobbin-smt-insert

Implement insert functionality for SMT in stdlib

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<W: ByteWriter>(&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()?;

assembly/src/ast/parsers/adv_ops.rs

@@ -33,6 +33,10 @@ pub fn parse_adv_inject(op: &Token) -> Result<Node, ParsingError> {
             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 => {

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}"),
         }

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

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),
         }