Merge pull request #155 from 0xPolygonMiden/greenhat/i144-stdlib

[3/x] Miden stdlib functions transformation

frontend-wasm/src/code_translator/mod.rs

@@ -26,7 +26,7 @@ use wasmparser::{MemArg, Operator};
 use crate::{
     error::{WasmError, WasmResult},
     intrinsics::{convert_intrinsics_call, is_miden_intrinsics_module},
-    miden_abi::{is_miden_sdk_module, transform::transform_miden_abi_call},
+    miden_abi::{is_miden_abi_module, transform::transform_miden_abi_call},
     module::{
         func_translation_state::{ControlStackFrame, ElseData, FuncTranslationState},
         function_builder_ext::FunctionBuilderExt,
@@ -161,6 +161,25 @@ pub fn translate_operator(
                 unsupported_diag!(diagnostics, "MemoryCopy: only single memory is supported");
             }
         }
+        Operator::MemoryFill { mem } => {
+            // See semantics at https://webassembly.github.io/spec/core/exec/instructions.html#exec-memory-fill
+
+            // This is a temporary workaround until we have this instruction
+            // properly implemented in IR. I encountered the Wasm `memory.fill`
+            // instruction in the `GlobalAlloc::alloc_zeroed` function, which is
+            // used to zero out the memory allocated by `GlobalAlloc::alloc`, but
+            // since the memory in Miden VM is guaranteed to be initialized to
+            // zeros we can ignore this instruction in this case for now
+            // see https://github.com/0xPolygonMiden/compiler/issues/156
+            if *mem != 0 {
+                unsupported_diag!(diagnostics, "MemoryFill: only single memory is supported");
+            }
+            let _num_bytes = state.pop1();
+            let val = state.pop1();
+            let _dst_i32 = state.pop1();
+            // Fail if the value is not zero, i.e. the memory is not zeroed
+            builder.ins().assert_eq_imm(Immediate::I32(0), val, span);
+        }
         /******************************* Load instructions ***********************************/
         Operator::I32Load8U { memarg } => {
             translate_load_zext(U8, I32, memarg, state, builder, span)
@@ -658,7 +677,7 @@ fn translate_call(
         let results = convert_intrinsics_call(func_id, args, builder, span);
         func_state.popn(num_wasm_args);
         func_state.pushn(&results);
-    } else if is_miden_sdk_module(func_id.module.as_symbol()) {
+    } else if is_miden_abi_module(func_id.module.as_symbol()) {
         // Miden SDK function call, transform the call to the Miden ABI if needed
         let results = transform_miden_abi_call(func_id, args, builder, span, diagnostics);
         assert_eq!(

frontend-wasm/src/miden_abi/mod.rs

@@ -1,8 +1,13 @@
+pub(crate) mod stdlib;
 pub(crate) mod transform;
 pub(crate) mod tx_kernel;
 
 use miden_core::crypto::hash::RpoDigest;
 use miden_hir::{FunctionType, Symbol};
+use rustc_hash::FxHashMap;
+
+pub(crate) type FunctionTypeMap = FxHashMap<&'static str, FunctionType>;
+pub(crate) type ModuleFunctionTypeMap = FxHashMap<&'static str, FunctionTypeMap>;
 
 /// Parse the stable import function name and the hex encoded digest from the function name
 pub fn parse_import_function_digest(import_name: &str) -> Result<(String, RpoDigest), String> {
@@ -21,13 +26,44 @@ pub fn parse_import_function_digest(import_name: &str) -> Result<(String, RpoDig
     ))
 }
 
-pub fn is_miden_sdk_module(module_id: Symbol) -> bool {
-    tx_kernel::types().contains_key(module_id.as_str())
+pub fn is_miden_abi_module(module_id: Symbol) -> bool {
+    is_miden_stdlib_module(module_id) || is_miden_sdk_module(module_id)
+}
+
+pub fn miden_abi_function_type(module_id: Symbol, function_id: Symbol) -> FunctionType {
+    if is_miden_stdlib_module(module_id) {
+        miden_stdlib_function_type(module_id, function_id)
+    } else {
+        miden_sdk_function_type(module_id, function_id)
+    }
+}
+
+fn is_miden_sdk_module(module_id: Symbol) -> bool {
+    tx_kernel::signatures().contains_key(module_id.as_str())
 }
 
 /// Get the target Miden ABI tx kernel function type for the given module and function id
 pub fn miden_sdk_function_type(module_id: Symbol, function_id: Symbol) -> FunctionType {
-    let funcs = tx_kernel::types()
+    let funcs = tx_kernel::signatures()
+        .get(module_id.as_str())
+        .expect(format!("No Miden ABI function types found for module {}", module_id).as_str());
+    funcs.get(function_id.as_str()).cloned().expect(
+        format!(
+            "No Miden ABI function type found for function {} in module {}",
+            function_id, module_id
+        )
+        .as_str(),
+    )
+}
+
+fn is_miden_stdlib_module(module_id: Symbol) -> bool {
+    stdlib::signatures().contains_key(module_id.as_str())
+}
+
+/// Get the target Miden ABI stdlib function type for the given module and function id
+#[inline(always)]
+fn miden_stdlib_function_type(module_id: Symbol, function_id: Symbol) -> FunctionType {
+    let funcs = stdlib::signatures()
         .get(module_id.as_str())
         .expect(format!("No Miden ABI function types found for module {}", module_id).as_str());
     funcs.get(function_id.as_str()).cloned().expect(

frontend-wasm/src/miden_abi/stdlib.rs

@@ -0,0 +1,19 @@
+//! Function types and lowered signatures for the Miden stdlib API functions
+
+use std::sync::OnceLock;
+
+use super::ModuleFunctionTypeMap;
+
+pub(crate) mod crypto;
+pub(crate) mod mem;
+
+pub(crate) fn signatures() -> &'static ModuleFunctionTypeMap {
+    static TYPES: OnceLock<ModuleFunctionTypeMap> = OnceLock::new();
+    TYPES.get_or_init(|| {
+        let mut m: ModuleFunctionTypeMap = Default::default();
+        m.extend(crypto::hashes::signatures());
+        m.extend(crypto::dsa::signatures());
+        m.extend(mem::signatures());
+        m
+    })
+}

frontend-wasm/src/miden_abi/stdlib/crypto.rs

@@ -0,0 +1,2 @@
+pub(crate) mod dsa;
+pub(crate) mod hashes;

frontend-wasm/src/miden_abi/stdlib/crypto/dsa.rs

@@ -0,0 +1,17 @@
+use miden_hir::FunctionType;
+use miden_hir_type::Type::*;
+
+use crate::miden_abi::{FunctionTypeMap, ModuleFunctionTypeMap};
+
+pub(crate) const RPO_FALCON512_VERIFY: &str = "rpo_falcon512_verify";
+
+pub(crate) fn signatures() -> ModuleFunctionTypeMap {
+    let mut m: ModuleFunctionTypeMap = Default::default();
+    let mut funcs: FunctionTypeMap = Default::default();
+    funcs.insert(
+        RPO_FALCON512_VERIFY,
+        FunctionType::new([Felt, Felt, Felt, Felt, Felt, Felt, Felt, Felt], []),
+    );
+    m.insert("miden:prelude/std_crypto_dsa", funcs);
+    m
+}

frontend-wasm/src/miden_abi/stdlib/crypto/hashes.rs

@@ -0,0 +1,33 @@
+use miden_hir::FunctionType;
+use miden_hir_type::Type::*;
+
+use crate::miden_abi::{FunctionTypeMap, ModuleFunctionTypeMap};
+
+pub(crate) const BLAKE3_HASH_1TO1: &str = "blake3_hash_1to1";
+pub(crate) const BLAKE3_HASH_2TO1: &str = "blake3_hash_2to1";
+
+pub(crate) fn signatures() -> ModuleFunctionTypeMap {
+    let mut m: ModuleFunctionTypeMap = Default::default();
+    let mut crypto: FunctionTypeMap = Default::default();
+    crypto.insert(
+        BLAKE3_HASH_1TO1,
+        //Accepts and returns a 8 Felt elements
+        FunctionType::new(
+            [Felt, Felt, Felt, Felt, Felt, Felt, Felt, Felt],
+            [Felt, Felt, Felt, Felt, Felt, Felt, Felt, Felt],
+        ),
+    );
+    crypto.insert(
+        BLAKE3_HASH_2TO1,
+        // Accepts 16 and returns a 8 Felt elements
+        FunctionType::new(
+            [
+                Felt, Felt, Felt, Felt, Felt, Felt, Felt, Felt, Felt, Felt, Felt, Felt, Felt, Felt,
+                Felt, Felt,
+            ],
+            [Felt, Felt, Felt, Felt, Felt, Felt, Felt, Felt],
+        ),
+    );
+    m.insert("miden:prelude/std_crypto_hashes", crypto);
+    m
+}

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

@@ -0,0 +1,45 @@
+use miden_hir::FunctionType;
+use miden_hir_type::Type::*;
+
+use crate::miden_abi::{FunctionTypeMap, ModuleFunctionTypeMap};
+
+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) fn signatures() -> ModuleFunctionTypeMap {
+    let mut m: ModuleFunctionTypeMap = Default::default();
+    let mut funcs: FunctionTypeMap = Default::default();
+    funcs.insert(
+        PIPE_WORDS_TO_MEMORY,
+        FunctionType::new(
+            [
+                Felt, // num_words
+                Felt, // write_ptr
+            ],
+            [
+                Felt, Felt, Felt, Felt, // HASH
+                Felt, // write_ptr'
+            ],
+        ),
+    );
+    funcs.insert(
+        PIPE_DOUBLE_WORDS_TO_MEMORY,
+        FunctionType::new(
+            [
+                Felt, Felt, Felt, Felt, // C
+                Felt, Felt, Felt, Felt, // B
+                Felt, Felt, Felt, Felt, // A
+                Felt, // write_ptr
+                Felt, // end_ptr
+            ],
+            [
+                Felt, Felt, Felt, Felt, // C
+                Felt, Felt, Felt, Felt, // B
+                Felt, Felt, Felt, Felt, // A
+                Felt, // write_ptr
+            ],
+        ),
+    );
+    m.insert("miden:prelude/std_mem", funcs);
+    m
+}

frontend-wasm/src/miden_abi/transform.rs

@@ -1,7 +1,7 @@
 use miden_diagnostics::DiagnosticsHandler;
-use miden_hir::{FunctionIdent, InstBuilder, SourceSpan, Type::*, Value};
+use miden_hir::{FunctionIdent, Immediate, InstBuilder, SourceSpan, Type::*, Value};
 
-use super::tx_kernel;
+use super::{stdlib, tx_kernel};
 use crate::module::function_builder_ext::FunctionBuilderExt;
 
 /// The strategy to use for transforming a function call
@@ -20,6 +20,11 @@ fn get_transform_strategy(function_id: &str) -> TransformStrategy {
         tx_kernel::NOTE_GET_INPUTS => TransformStrategy::ListReturn,
         tx_kernel::ACCOUNT_ADD_ASSET => TransformStrategy::ReturnViaPointer,
         tx_kernel::ACCOUNT_GET_ID => TransformStrategy::NoTransform,
+        stdlib::crypto::hashes::BLAKE3_HASH_1TO1 => TransformStrategy::ReturnViaPointer,
+        stdlib::crypto::hashes::BLAKE3_HASH_2TO1 => TransformStrategy::ReturnViaPointer,
+        stdlib::crypto::dsa::RPO_FALCON512_VERIFY => TransformStrategy::NoTransform,
+        stdlib::mem::PIPE_WORDS_TO_MEMORY => TransformStrategy::ReturnViaPointer,
+        stdlib::mem::PIPE_DOUBLE_WORDS_TO_MEMORY => TransformStrategy::ReturnViaPointer,
         _ => panic!("No transform strategy found for function {}", function_id),
     }
 }
@@ -69,7 +74,7 @@ pub fn list_return(
     results[0..1].to_vec()
 }
 
-/// The Miden ABI function returns on the stack and we want to return via a pointer argument
+/// The Miden ABI function returns felts on the stack and we want to return via a pointer argument
 pub fn return_via_pointer(
     func_id: FunctionIdent,
     args: &[Value],
@@ -86,12 +91,19 @@ pub fn return_via_pointer(
     assert_eq!(ptr_arg_ty, I32);
     let ptr_u32 = builder.ins().cast(ptr_arg, U32, span);
     for (idx, value) in results.iter().enumerate() {
+        let value_ty = builder.data_flow_graph().value_type(*value);
+        assert_eq!(value_ty, &Felt, "In return_via_pointer, expected only Felt value type returns");
         let eff_ptr = if idx == 0 {
             ptr_u32
         } else {
-            builder
-                .ins()
-                .add_imm_checked(ptr_u32, miden_hir::Immediate::I32(idx as i32), span)
+            // We're storing the stack-returned felts(64-bit) values in the
+            // memory that from the Rust "side" point of view is byte-addressed,
+            // meaning that for example in array of felts the second felt is
+            // expected to be +8 from the first and so on. So we need to
+            // multiply the index by 8 so that the subsequent Rust code finds
+            // the values in the expected locations.
+            let imm = Immediate::I32(idx as i32 * 8);
+            builder.ins().add_imm_checked(ptr_u32, imm, span)
         };
         let value_ty = builder.data_flow_graph().value_type(*value).clone();
         let addr = builder.ins().inttoptr(eff_ptr, Ptr(value_ty.into()), span);

frontend-wasm/src/miden_abi/tx_kernel.rs

@@ -5,7 +5,9 @@ use std::sync::OnceLock;
 
 use miden_hir::FunctionType;
 use miden_hir_type::Type::*;
-use rustc_hash::FxHashMap;
+
+use super::ModuleFunctionTypeMap;
+use crate::miden_abi::FunctionTypeMap;
 
 pub const NOTE_MODULE_NAME: &str = "miden:tx_kernel/note";
 pub const NOTE_GET_INPUTS: &str = "get_inputs";
@@ -14,25 +16,22 @@ pub const ACCOUNT_MODULE_NAME: &str = "miden:tx_kernel/account";
 pub const ACCOUNT_ADD_ASSET: &str = "add_asset";
 pub const ACCOUNT_GET_ID: &str = "get_id";
 
-type FunctionTypeMap = FxHashMap<&'static str, FunctionType>;
-type ModuleFunctionTypeMap = FxHashMap<&'static str, FunctionTypeMap>;
-
-pub(crate) fn types() -> &'static ModuleFunctionTypeMap {
+pub(crate) fn signatures() -> &'static ModuleFunctionTypeMap {
     static TYPES: OnceLock<ModuleFunctionTypeMap> = OnceLock::new();
     TYPES.get_or_init(|| {
         let mut m: ModuleFunctionTypeMap = Default::default();
 
         let mut note: FunctionTypeMap = Default::default();
-        note.insert(NOTE_GET_INPUTS, FunctionType::new_miden([Felt], [I32, Felt]));
+        note.insert(NOTE_GET_INPUTS, FunctionType::new([Felt], [I32, Felt]));
         m.insert(NOTE_MODULE_NAME, note);
 
         let mut account: FunctionTypeMap = Default::default();
         account.insert(
             ACCOUNT_ADD_ASSET,
             // Accepts and returns word
-            FunctionType::new_miden([Felt, Felt, Felt, Felt], [Felt, Felt, Felt, Felt]),
+            FunctionType::new([Felt, Felt, Felt, Felt], [Felt, Felt, Felt, Felt]),
         );
-        account.insert(ACCOUNT_GET_ID, FunctionType::new_miden([], [Felt]));
+        account.insert(ACCOUNT_GET_ID, FunctionType::new([], [Felt]));
         m.insert(ACCOUNT_MODULE_NAME, account);
 
         m

frontend-wasm/src/module/build_ir.rs

@@ -8,7 +8,7 @@ use super::{module_translation_state::ModuleTranslationState, Module};
 use crate::{
     error::WasmResult,
     intrinsics::is_miden_intrinsics_module,
-    miden_abi::miden_sdk_function_type,
+    miden_abi::miden_abi_function_type,
     module::{
         func_translator::FuncTranslator,
         module_env::{FunctionBodyData, ModuleEnvironment, ParsedModule},
@@ -84,7 +84,7 @@ pub fn translate_module_as_component(
                     // ignore intrinsics imports
                     continue;
                 }
-                let function_ty = miden_sdk_function_type(
+                let function_ty = miden_abi_function_type(
                     ext_func.module.as_symbol(),
                     ext_func.function.as_symbol(),
                 );

frontend-wasm/src/module/module_translation_state.rs

@@ -7,7 +7,7 @@ use super::{instance::ModuleArgument, ir_func_type, EntityIndex, FuncIndex, Modu
 use crate::{
     error::WasmResult,
     intrinsics::is_miden_intrinsics_module,
-    miden_abi::{is_miden_sdk_module, miden_sdk_function_type, parse_import_function_digest},
+    miden_abi::{is_miden_abi_module, miden_abi_function_type, parse_import_function_digest},
     translation_utils::sig_from_funct_type,
     WasmError,
 };
@@ -115,9 +115,9 @@ impl ModuleTranslationState {
         diagnostics: &DiagnosticsHandler,
     ) -> WasmResult<FunctionIdent> {
         let (func_id, wasm_sig) = self.functions[&index].clone();
-        let sig: Signature = if is_miden_sdk_module(func_id.module.as_symbol()) {
+        let sig: Signature = if is_miden_abi_module(func_id.module.as_symbol()) {
             let ft =
-                miden_sdk_function_type(func_id.module.as_symbol(), func_id.function.as_symbol());
+                miden_abi_function_type(func_id.module.as_symbol(), func_id.function.as_symbol());
             Signature::new(
                 ft.params.into_iter().map(AbiParam::new),
                 ft.results.into_iter().map(AbiParam::new),