Merge pull request #627 from 0xMiden/greenhat/i626-advice-map-sdk

Advice map loading in Miden SDK

Author: bitwalker

codegen/masm/intrinsics/advice.masm

@@ -0,0 +1,25 @@
+#! ADV_PUSH_MAPVALN: Pushes a list of field elements onto the advice stack. 
+#! The list is looked up in the advice map using `key` as the key.
+#!
+#! Stack Input:  [key3, key2, key1, key0, ...]
+#! Stack Output: [num_elements, ...] (returns the number of elements pushed)
+#!
+#! The key is a Word (4 field elements) used to lookup values in the advice map
+export.adv_push_mapvaln
+    # Stack: [key3, key2, key1, key0, ...]
+    
+    # Call the VM instruction to push values from advice map
+    adv.push_mapvaln
+    # Stack: [key3, key2, key1, key0, ...]
+    # Side effect: [n, ele1, ele2, ...] pushed to advice stack where n is the count
+    
+    # Drop the key word from the stack
+    dropw
+    # Stack: [...]
+    
+    # The adv.push_mapvaln instruction pushes the count on top of the advice stack
+    # We need to pop it from the advice stack to return it
+    # Note: The advice stack now has [ele1, ele2, ...] with count removed
+    adv_push.1
+    # Stack: [num_elements, ...]
+end

codegen/masm/src/intrinsics.rs

@@ -8,12 +8,14 @@ pub const I32_INTRINSICS_MODULE_NAME: &str = "intrinsics::i32";
 pub const I64_INTRINSICS_MODULE_NAME: &str = "intrinsics::i64";
 pub const MEM_INTRINSICS_MODULE_NAME: &str = "intrinsics::mem";
 pub const CRYPTO_INTRINSICS_MODULE_NAME: &str = "intrinsics::crypto";
+pub const ADVICE_INTRINSICS_MODULE_NAME: &str = "intrinsics::advice";
 
-pub const INTRINSICS_MODULE_NAMES: [&str; 4] = [
+pub const INTRINSICS_MODULE_NAMES: [&str; 5] = [
     I32_INTRINSICS_MODULE_NAME,
     I64_INTRINSICS_MODULE_NAME,
     MEM_INTRINSICS_MODULE_NAME,
     CRYPTO_INTRINSICS_MODULE_NAME,
+    ADVICE_INTRINSICS_MODULE_NAME,
 ];
 
 const I32_INTRINSICS: &str =
@@ -24,9 +26,11 @@ const MEM_INTRINSICS: &str =
     include_str!(concat!(env!("CARGO_MANIFEST_DIR"), "/intrinsics/mem.masm"));
 const CRYPTO_INTRINSICS: &str =
     include_str!(concat!(env!("CARGO_MANIFEST_DIR"), "/intrinsics/crypto.masm"));
+const ADVICE_INTRINSICS: &str =
+    include_str!(concat!(env!("CARGO_MANIFEST_DIR"), "/intrinsics/advice.masm"));
 
 /// This is a mapping of intrinsics module name to the raw MASM source for that module
-const INTRINSICS: [(&str, &str, &str); 4] = [
+const INTRINSICS: [(&str, &str, &str); 5] = [
     (
         I32_INTRINSICS_MODULE_NAME,
         I32_INTRINSICS,
@@ -47,6 +51,11 @@ const INTRINSICS: [(&str, &str, &str); 4] = [
         CRYPTO_INTRINSICS,
         concat!(env!("CARGO_MANIFEST_DIR"), "/intrinsics/crypto.masm"),
     ),
+    (
+        ADVICE_INTRINSICS_MODULE_NAME,
+        ADVICE_INTRINSICS,
+        concat!(env!("CARGO_MANIFEST_DIR"), "/intrinsics/advice.masm"),
+    ),
 ];
 
 /// This helper loads the named module from the set of intrinsics modules defined in this crate.

frontend/wasm/src/intrinsics/advice.rs

@@ -0,0 +1,78 @@
+use midenc_dialect_hir::HirOpBuilder;
+use midenc_hir::{
+    dialects::builtin::FunctionRef,
+    interner::{symbols, Symbol},
+    Builder, FunctionType, SmallVec, SourceSpan, SymbolNameComponent, Type, ValueRef,
+};
+
+use crate::{error::WasmResult, module::function_builder_ext::FunctionBuilderExt};
+
+pub(crate) const MODULE_ID: &str = "intrinsics::advice";
+/// The module path prefix for advice intrinsics, not including the function name
+pub const MODULE_PREFIX: &[SymbolNameComponent] = &[
+    SymbolNameComponent::Root,
+    SymbolNameComponent::Component(symbols::Intrinsics),
+    SymbolNameComponent::Component(symbols::Advice),
+];
+
+/// Get the [FunctionType] of an advice intrinsic, if it is implemented as a function.
+///
+/// Returns `None` for intrinsics which are unknown, or correspond to native instructions.
+pub fn function_type(function: Symbol) -> Option<FunctionType> {
+    match function.as_str() {
+        "adv_push_mapvaln" => {
+            // The WASM import signature: takes 4 f32 values (Word) and returns 1 f32
+            let sig = FunctionType::new(
+                midenc_hir::CallConv::Wasm,
+                vec![
+                    Type::I32, // key0
+                    Type::I32, // key1
+                    Type::I32, // key2
+                    Type::I32, // key3
+                ],
+                vec![Type::I32], // Returns number of elements pushed
+            );
+            Some(sig)
+        }
+        _ => None,
+    }
+}
+
+/// Convert a call to an advice intrinsic function into instruction(s)
+pub fn convert_advice_intrinsics<B: ?Sized + Builder>(
+    function: Symbol,
+    function_ref: Option<FunctionRef>,
+    args: &[ValueRef],
+    builder: &mut FunctionBuilderExt<'_, B>,
+    span: SourceSpan,
+) -> WasmResult<SmallVec<[ValueRef; 1]>> {
+    let function_ref =
+        function_ref.unwrap_or_else(|| panic!("expected '{function}' to have been declared"));
+
+    match function.as_str() {
+        "adv_push_mapvaln" => {
+            // The WASM import has 4 parameters (key0-3) and returns 1 f32
+            assert_eq!(args.len(), 4, "{function} takes exactly four arguments (key0-3)");
+
+            let func = function_ref.borrow();
+            let signature = func.signature().clone();
+            drop(func);
+
+            // Call the function with all arguments
+            // The intrinsics::io::adv_push_mapvaln function will be mapped to the MASM adv_push_mapvaln
+            let exec = builder.exec(function_ref, signature, args.iter().copied(), span)?;
+
+            // Extract the return value from the exec operation
+            let borrow = exec.borrow();
+            let results = borrow.as_ref().results();
+            let result_vals: SmallVec<[ValueRef; 1]> =
+                results.iter().map(|op_res| op_res.borrow().as_value_ref()).collect();
+
+            // The function returns the number of elements pushed as i32
+            Ok(result_vals)
+        }
+        _ => {
+            panic!("unsupported io intrinsic: '{function}'")
+        }
+    }
+}

frontend/wasm/src/intrinsics/felt.rs

@@ -12,7 +12,7 @@ pub(crate) const MODULE_ID: &str = "intrinsics::felt";
 pub(crate) const MODULE_PREFIX: &[SymbolNameComponent] = &[
     SymbolNameComponent::Root,
     SymbolNameComponent::Component(symbols::Intrinsics),
-    SymbolNameComponent::Component(symbols::Felt),
+    SymbolNameComponent::Component(symbols::FeltModule),
 ];
 
 /// Convert a call to a felt op intrinsic function into instruction(s)

frontend/wasm/src/intrinsics/intrinsic.rs

@@ -4,7 +4,7 @@ use midenc_hir::{
     FunctionType, SymbolNameComponent, SymbolPath, Type,
 };
 
-use super::{crypto, debug, felt, mem};
+use super::{advice, crypto, debug, felt, mem};
 
 /// Error raised when an attempt is made to use or load an unrecognized intrinsic
 #[derive(Debug, thiserror::Error, Diagnostic)]
@@ -26,6 +26,8 @@ pub enum Intrinsic {
     Felt(Symbol),
     /// A cryptographic intrinsic
     Crypto(Symbol),
+    /// An advice intrinsic
+    Advice(Symbol),
 }
 
 /// Attempt to recognize an intrinsic function from the given [SymbolPath].
@@ -62,8 +64,9 @@ impl TryFrom<&SymbolPath> for Intrinsic {
         match kind {
             symbols::Debug => Ok(Self::Debug(function)),
             symbols::Mem => Ok(Self::Mem(function)),
-            symbols::Felt => Ok(Self::Felt(function)),
+            symbols::FeltModule => Ok(Self::Felt(function)),
             symbols::Crypto => Ok(Self::Crypto(function)),
+            symbols::Advice => Ok(Self::Advice(function)),
             _ => Err(UnknownIntrinsicError(path.clone())),
         }
     }
@@ -83,8 +86,9 @@ impl Intrinsic {
         match self {
             Self::Debug(_) => symbols::Debug,
             Self::Mem(_) => symbols::Mem,
-            Self::Felt(_) => symbols::Felt,
+            Self::Felt(_) => symbols::FeltModule,
             Self::Crypto(_) => symbols::Crypto,
+            Self::Advice(_) => symbols::Advice,
         }
     }
 
@@ -95,6 +99,7 @@ impl Intrinsic {
             Self::Mem(_) => SymbolPath::from_iter(mem::MODULE_PREFIX.iter().copied()),
             Self::Felt(_) => SymbolPath::from_iter(felt::MODULE_PREFIX.iter().copied()),
             Self::Crypto(_) => SymbolPath::from_iter(crypto::MODULE_PREFIX.iter().copied()),
+            Self::Advice(_) => SymbolPath::from_iter(advice::MODULE_PREFIX.iter().copied()),
         }
     }
 
@@ -104,7 +109,8 @@ impl Intrinsic {
             Self::Debug(function)
             | Self::Mem(function)
             | Self::Felt(function)
-            | Self::Crypto(function) => *function,
+            | Self::Crypto(function)
+            | Self::Advice(function) => *function,
         }
     }
 
@@ -138,6 +144,7 @@ impl Intrinsic {
                     _ => None,
                 }
             }
+            Self::Advice(function) => advice::function_type(*function),
         }
     }
 
@@ -150,6 +157,10 @@ impl Intrinsic {
                 mem::function_type(*function).map(IntrinsicsConversionResult::FunctionType)
             }
             Self::Debug(_) | Self::Felt(_) => Some(IntrinsicsConversionResult::MidenVmOp),
+            Self::Advice(_function) => self
+                .function_type()
+                .map(IntrinsicsConversionResult::FunctionType)
+                .or(Some(IntrinsicsConversionResult::MidenVmOp)),
             // Crypto intrinsics are converted to function calls
             Self::Crypto(_function) => self
                 .function_type()

frontend/wasm/src/intrinsics/mod.rs

@@ -2,6 +2,7 @@ mod intrinsic;
 
 pub use self::intrinsic::*;
 
+pub mod advice;
 pub mod crypto;
 pub mod debug;
 pub mod felt;
@@ -25,6 +26,7 @@ fn modules() -> &'static FxHashSet<SymbolPath> {
         s.insert(SymbolPath::from_iter(felt::MODULE_PREFIX.iter().copied()));
         s.insert(SymbolPath::from_iter(debug::MODULE_PREFIX.iter().copied()));
         s.insert(SymbolPath::from_iter(crypto::MODULE_PREFIX.iter().copied()));
+        s.insert(SymbolPath::from_iter(advice::MODULE_PREFIX.iter().copied()));
         s
     });
     &MODULES
@@ -51,6 +53,9 @@ pub fn convert_intrinsics_call<B: ?Sized + Builder>(
         Intrinsic::Crypto(function) => {
             crypto::convert_crypto_intrinsics(function, function_ref, args, builder, span)
         }
+        Intrinsic::Advice(function) => {
+            advice::convert_advice_intrinsics(function, function_ref, args, builder, span)
+        }
     }
 }
 

frontend/wasm/src/miden_abi/mod.rs

@@ -119,6 +119,8 @@ pub fn recover_imported_masm_module(wasm_module_id: &str) -> Result<SymbolPath,
         Ok(SymbolPath::from_masm_module_id(intrinsics::debug::MODULE_ID))
     } else if wasm_module_id.starts_with("miden:core-intrinsics/intrinsics-crypto") {
         Ok(SymbolPath::from_masm_module_id(intrinsics::crypto::MODULE_ID))
+    } else if wasm_module_id.starts_with("miden:core-intrinsics/intrinsics-advice") {
+        Ok(SymbolPath::from_masm_module_id(intrinsics::advice::MODULE_ID))
     } else if wasm_module_id.starts_with("miden:core-base/account") {
         Ok(SymbolPath::from_masm_module_id(tx_kernel::account::MODULE_ID))
     } else if wasm_module_id.starts_with("miden:core-base/note") {

frontend/wasm/src/miden_abi/stdlib/mem.rs

@@ -15,6 +15,7 @@ pub(crate) const MODULE_PREFIX: &[SymbolNameComponent] = &[
 
 pub(crate) const PIPE_WORDS_TO_MEMORY: &str = "pipe_words_to_memory";
 pub(crate) const PIPE_DOUBLE_WORDS_TO_MEMORY: &str = "pipe_double_words_to_memory";
+pub(crate) const PIPE_PREIMAGE_TO_MEMORY: &str = "pipe_preimage_to_memory";
 
 pub(crate) fn signatures() -> ModuleFunctionTypeMap {
     let mut m: ModuleFunctionTypeMap = Default::default();
@@ -52,6 +53,20 @@ pub(crate) fn signatures() -> ModuleFunctionTypeMap {
             ],
         ),
     );
+    funcs.insert(
+        Symbol::from(PIPE_PREIMAGE_TO_MEMORY),
+        FunctionType::new(
+            CallConv::Wasm,
+            [
+                Felt, // num_words
+                I32,  // write_ptr
+                Felt, Felt, Felt, Felt, // COM (commitment)
+            ],
+            [
+                I32, // write_ptr'
+            ],
+        ),
+    );
     m.insert(SymbolPath::from_iter(MODULE_PREFIX.iter().copied()), funcs);
     m
 }

frontend/wasm/src/miden_abi/transform.rs

@@ -30,6 +30,7 @@ fn get_transform_strategy(path: &SymbolPath) -> Option<TransformStrategy> {
                 stdlib::mem::PIPE_WORDS_TO_MEMORY | stdlib::mem::PIPE_DOUBLE_WORDS_TO_MEMORY => {
                     Some(TransformStrategy::ReturnViaPointer)
                 }
+                stdlib::mem::PIPE_PREIMAGE_TO_MEMORY => Some(TransformStrategy::NoTransform),
                 _ => None,
             },
             symbols::Crypto => match components.next()?.as_symbol_name() {

hir-symbol/src/symbols.toml

@@ -106,8 +106,10 @@ blake3 = {}
 crypto = {}
 debug = {}
 dsa = {}
+felt_module = { value = "felt" }
 hashes = {}
 intrinsics = {}
+advice = {}
 miden = {}
 mem = {}
 note = {}