[3/x] Link Miden packages without cross-context calls lifting/lowering

Cargo.toml

@@ -95,13 +95,13 @@ miden-parsing = "0.1"
 # miden-stdlib = { version = "0.11.0", path = "../vm-aux/stdlib" }
 # miden-package = { version = "0.11.0", path = "../vm-aux/package" }
 
-miden-assembly = { version = "0.11.0", git = "https://github.com/0xPolygonMiden/miden-vm", rev = "fe4afa2ad655772cd1ef1c30d014c383731ab293" }
-miden-core = { version = "0.11.0", git = "https://github.com/0xPolygonMiden/miden-vm", rev = "fe4afa2ad655772cd1ef1c30d014c383731ab293" }
-miden-processor = { version = "0.11.0", git = "https://github.com/0xPolygonMiden/miden-vm", rev = "fe4afa2ad655772cd1ef1c30d014c383731ab293" }
-miden-stdlib = { version = "0.11.0", git = "https://github.com/0xPolygonMiden/miden-vm", rev = "fe4afa2ad655772cd1ef1c30d014c383731ab293", features = [
+miden-assembly = { version = "0.11.0", git = "https://github.com/0xPolygonMiden/miden-vm", rev = "75bd636a6b83fd6cb134b0981663aa28e49a811c" }
+miden-core = { version = "0.11.0", git = "https://github.com/0xPolygonMiden/miden-vm", rev = "75bd636a6b83fd6cb134b0981663aa28e49a811c" }
+miden-processor = { version = "0.11.0", git = "https://github.com/0xPolygonMiden/miden-vm", rev = "75bd636a6b83fd6cb134b0981663aa28e49a811c" }
+miden-stdlib = { version = "0.11.0", git = "https://github.com/0xPolygonMiden/miden-vm", rev = "75bd636a6b83fd6cb134b0981663aa28e49a811c", features = [
     "with-debug-info",
 ] }
-miden-package = { version = "0.11.0", git = "https://github.com/0xPolygonMiden/miden-vm", rev = "fe4afa2ad655772cd1ef1c30d014c383731ab293" }
+miden-package = { version = "0.11.0", git = "https://github.com/0xPolygonMiden/miden-vm", rev = "75bd636a6b83fd6cb134b0981663aa28e49a811c" }
 
 #miden-assembly = { git = "https://github.com/0xPolygonMiden/miden-vm", rev = "828557c28ca1d159bfe42195e7ea73256ce4aa06" }
 #miden-core = { git = "https://github.com/0xPolygonMiden/miden-vm", rev = "828557c28ca1d159bfe42195e7ea73256ce4aa06" }

codegen/masm/src/masm/module.rs

@@ -3,7 +3,7 @@ use std::{collections::BTreeSet, fmt, path::Path, sync::Arc};
 use intrusive_collections::{intrusive_adapter, RBTree, RBTreeAtomicLink};
 use miden_assembly::{
     ast::{self, ModuleKind},
-    LibraryPath,
+    LibraryNamespace, LibraryPath,
 };
 use midenc_hir::{
     diagnostics::{Report, SourceFile, SourceSpan, Span, Spanned},
@@ -166,9 +166,15 @@ impl Module {
         // Create module import table
         for ir_import in self.imports.iter() {
             let span = ir_import.span;
-            let name =
-                ast::Ident::new_with_span(span, ir_import.alias.as_str()).map_err(Report::msg)?;
-            let path = LibraryPath::new(ir_import.name.as_str()).expect("invalid import path");
+            let name = ast::Ident::new_unchecked(Span::new(span, ir_import.alias.as_str().into()));
+            let module_name = ir_import.name.as_str();
+            let path = LibraryPath::new(module_name).unwrap_or_else(|_| {
+                let module_id = ast::Ident::new_unchecked(Span::new(
+                    span,
+                    Arc::from(module_name.to_string().into_boxed_str()),
+                ));
+                LibraryPath::new_from_components(LibraryNamespace::Anon, [module_id])
+            });
             let import = ast::Import {
                 span,
                 name,

codegen/masm/src/masm/program.rs

@@ -1,9 +1,9 @@
-use std::{fmt, path::Path, sync::Arc};
+use std::{collections::BTreeMap, fmt, path::Path, sync::Arc};
 
 use hir::{Signature, Symbol};
 use miden_assembly::{
-    ast::{ModuleKind, ProcedureName},
-    KernelLibrary, Library as CompiledLibrary, LibraryNamespace,
+    ast::{self, ModuleKind, ProcedureName, QualifiedProcedureName},
+    KernelLibrary, Library as CompiledLibrary, LibraryNamespace, LibraryPath, Span,
 };
 use miden_core::{crypto::hash::Rpo256, AdviceMap};
 use midenc_hir::{
@@ -251,18 +251,26 @@ impl Program {
         let mut assembler =
             Assembler::new(session.source_manager.clone()).with_debug_mode(debug_mode);
 
+        let mut lib_modules = Vec::new();
         // Link extra libraries
         for library in self.library.libraries.iter() {
-            if log::log_enabled!(log::Level::Debug) {
-                for module in library.module_infos() {
-                    log::debug!("registering '{}' with assembler", module.path());
-                }
+            for module in library.module_infos() {
+                log::debug!("registering '{}' with assembler", module.path());
+                lib_modules.push(module.path().to_string());
             }
             assembler.add_library(library)?;
         }
 
         // Assemble library
         for module in self.library.modules.iter() {
+            if lib_modules.contains(&module.id.to_string()) {
+                log::warn!(
+                    "module '{}' is already registered with the assembler as library's module, \
+                     skipping",
+                    module.id
+                );
+                continue;
+            }
             log::debug!("adding '{}' to assembler", module.id.as_str());
             let kind = module.kind;
             let module = module.to_ast(debug_mode).map(Box::new)?;
@@ -476,19 +484,28 @@ impl Library {
         let mut assembler =
             Assembler::new(session.source_manager.clone()).with_debug_mode(debug_mode);
 
+        let mut lib_modules = Vec::new();
         // Link extra libraries
         for library in self.libraries.iter() {
-            if log::log_enabled!(log::Level::Debug) {
-                for module in library.module_infos() {
-                    log::debug!("registering '{}' with assembler", module.path());
-                }
+            for module in library.module_infos() {
+                log::debug!("registering '{}' with assembler", module.path());
+                lib_modules.push(module.path().to_string());
             }
             assembler.add_library(library)?;
         }
 
         // Assemble library
         let mut modules = Vec::with_capacity(self.modules.len());
         for module in self.modules.iter() {
+            let module_id = module.id.as_str();
+            if lib_modules.contains(&module_id.to_string()) {
+                log::warn!(
+                    "module '{}' is already registered with the assembler as library's module, \
+                     skipping",
+                    module_id
+                );
+                continue;
+            }
             log::debug!("adding '{}' to assembler", module.id.as_str());
             let module = module.to_ast(debug_mode).map(Box::new)?;
             modules.push(module);
@@ -499,8 +516,54 @@ impl Library {
             .iter()
             .map(|rodata| (rodata.digest, rodata.to_elements()))
             .collect();
-        Ok(Arc::new(lib.with_advice_map(advice_map)))
+        let mut mast_forest = lib.mast_forest().as_ref().clone();
+        mast_forest.advice_map_mut().extend(advice_map);
+        let converted_exports = recover_wasm_cm_interfaces(&lib);
+        let lib = CompiledLibrary::new(Arc::new(mast_forest), converted_exports)?;
+        Ok(Arc::new(lib))
+    }
+}
+
+/// Try to recognize Wasm CM interfaces and transform those exports to have Wasm interface
+/// encoded as module name.
+/// Temporary workaround for:
+/// 1. Temporary exporting multiple interfaces from the same(Wasm core) module (an interface is encoded
+///    in the function name);
+/// 2. Assembler using the current module name to generate exports.
+fn recover_wasm_cm_interfaces(
+    lib: &CompiledLibrary,
+) -> BTreeMap<QualifiedProcedureName, miden_processor::MastNodeId> {
+    let mut exports = BTreeMap::new();
+    for export in lib.exports() {
+        let export_node_id = lib.get_export_node_id(export);
+        if export.module.to_string().starts_with("intrinsics")
+            || export.name.as_str().starts_with("cabi")
+        {
+            // Preserve intrinsics modules and internal Wasm CM `cabi_*` functions
+            exports.insert(export.clone(), export_node_id);
+            continue;
+        }
+
+        if export.name.as_str().contains("/") {
+            // Wasm CM interface
+            let parts = export.name.as_str().split('#').collect::<Vec<_>>();
+            assert_eq!(parts.len(), 2);
+            let module =
+                ast::Ident::new_unchecked(Span::new(SourceSpan::default(), parts[0].into()));
+            let name = parts[1];
+            let path = LibraryPath::new_from_components(LibraryNamespace::Anon, [module]);
+            let name = ast::ProcedureName::new_unchecked(ast::Ident::new_unchecked(Span::new(
+                SourceSpan::default(),
+                Arc::from(name),
+            )));
+            let new_export = QualifiedProcedureName::new(path, name);
+            exports.insert(new_export, export_node_id);
+        } else {
+            // Non-Wasm CM interface, preserve as is
+            exports.insert(export.clone(), export_node_id);
+        }
     }
+    exports
 }
 
 impl fmt::Display for Library {

frontend-wasm/src/component/translator.rs

@@ -26,7 +26,7 @@ use crate::{
         module_env::ParsedModule,
         module_translation_state::ModuleTranslationState,
         types::{EntityIndex, FuncIndex},
-        Module, ModuleImport,
+        Module,
     },
     unsupported_diag, WasmTranslationConfig,
 };
@@ -117,7 +117,7 @@ impl<'a, 'data> ComponentTranslator<'a, 'data> {
             }
         }
         for (name, export) in &wasm_translation.component.exports {
-            self.build_export(export, name, &mut component_builder)?;
+            self.build_export(export, name, None, &mut component_builder)?;
         }
         Ok(component_builder.build())
     }
@@ -183,7 +183,21 @@ impl<'a, 'data> ComponentTranslator<'a, 'data> {
                     self.config,
                     self.session,
                 )?;
-                component_builder.add_module(ir_module.into()).expect("module is already added");
+                // Skip the shim and fixups core modules, they are a workaround for
+                // specify the core instance memory and reallod function for the lowering
+
+                // TODO:
+                // Imported function from the shim has empty module name and "0" as a if name.
+
+                if ir_module.name.as_str() != "wit-component:shim"
+                    && ir_module.name.as_str() != "wit-component:fixups"
+                {
+                    component_builder
+                        .add_module(ir_module.into())
+                        .expect("module is already added");
+                } else {
+                    // dbg!(&ir_module.functions().collect::<Vec<_>>());
+                }
             }
             InstantiateModule::Import(..) => {
                 unsupported_diag!(
@@ -212,8 +226,18 @@ impl<'a, 'data> ComponentTranslator<'a, 'data> {
                 options,
             } => {
                 let module_import = module.imports.get(idx).expect("module import not found");
+                let function_id = module.func_name(module_import.index.unwrap_func()).into();
+                let function_id = FunctionIdent {
+                    module: module.name(),
+                    function: function_id,
+                };
+
+                // TODO:
+                // Find process_list_felt instead empty module name and "0" function name!
+                // Follow module_import.index through the shim modules/imports/exports?
+
                 let runtime_import_idx = self.lower_imports[index];
-                let function_id = function_id_from_import(module_import);
+                // dbg!(&module_import);
                 match self.translate_import(
                     runtime_import_idx,
                     *lower_ty,
@@ -330,19 +354,25 @@ impl<'a, 'data> ComponentTranslator<'a, 'data> {
         &self,
         export: &Export,
         name: &String,
+        interface: Option<String>,
         component_builder: &mut ComponentBuilder,
     ) -> WasmResult<()> {
         match export {
             Export::LiftedFunction { ty, func, options } => {
-                let export_name = Symbol::intern(name).into();
+                dbg!(name);
+                // The inline export does no have an interface name
+                let interface = interface.unwrap_or_default();
+                dbg!(&interface);
+                let export_name = InterfaceFunctionIdent::from_full(interface, name.clone());
                 let export = self.build_export_lifted_function(func, ty, options)?;
                 component_builder.add_export(export_name, export);
                 Ok(())
             }
             Export::Instance(exports) => {
+                let interface = Some(name.clone());
                 // Flatten any(nested) interface instance exports into the IR `Component` exports
-                for (name, export) in exports {
-                    self.build_export(export, name, component_builder)?;
+                for (export_name, export) in exports {
+                    self.build_export(export, export_name, interface.clone(), component_builder)?;
                 }
                 Ok(())
             }
@@ -371,10 +401,10 @@ impl<'a, 'data> ComponentTranslator<'a, 'data> {
         ty: &TypeFuncIndex,
         options: &CanonicalOptions,
     ) -> WasmResult<ComponentExport> {
-        let func_ident = self.func_id_from_core_def(func)?;
+        let core_func_ident = self.func_id_from_core_def(func)?;
         let lifted_func_ty = convert_lifted_func_ty(ty, &self.component_types);
         let export = midenc_hir::ComponentExport {
-            function: func_ident,
+            function: core_func_ident,
             function_ty: lifted_func_ty,
             options: self.translate_canonical_options(options)?,
         };
@@ -444,11 +474,6 @@ impl<'a, 'data> ComponentTranslator<'a, 'data> {
     }
 }
 
-/// Get the function id from the given Wasm core module import
-fn function_id_from_import(module_import: &ModuleImport) -> FunctionIdent {
-    recover_imported_masm_function_id(&module_import.module, module_import.field.as_str())
-}
-
 /// Get the function id from the given Wasm func_idx in the given Wasm core exporting_module
 fn function_id_from_export(exporting_module: &Module, func_idx: FuncIndex) -> FunctionIdent {
     let func_name = exporting_module.func_name(func_idx);

frontend-wasm/src/miden_abi/mod.rs

@@ -92,7 +92,11 @@ pub fn recover_imported_masm_function_id(
              {wasm_function_id}"
         )
     } else {
-        wasm_module_id
+        // Unrecognized module ID, return as is
+        return FunctionIdent {
+            module: Ident::from(wasm_module_id),
+            function: Ident::from(wasm_function_id),
+        };
     };
     // Since `hash-1to1` is an invalid name in Wasm CM (dashed part cannot start with a digit),
     // we need to translate the CM name to the one that is expected at the linking stage

frontend-wasm/src/module/build_ir.rs

@@ -124,7 +124,12 @@ pub fn build_ir_module(
         let func_name = &parsed_module.module.func_name(*func_index);
         let wasm_func_type = module_types[func_type.signature].clone();
         let ir_func_type = ir_func_type(&wasm_func_type, &session.diagnostics)?;
-        let sig = ir_func_sig(&ir_func_type, CallConv::SystemV, Linkage::External);
+        let linkage = if parsed_module.module.is_exported_function(func_index) {
+            Linkage::External
+        } else {
+            Linkage::Internal
+        };
+        let sig = ir_func_sig(&ir_func_type, CallConv::SystemV, linkage);
         let mut module_func_builder = module_builder.function(func_name.as_str(), sig.clone())?;
         let FunctionBodyData { validator, body } = body_data;
         let mut func_validator = validator.into_validator(Default::default());

frontend-wasm/src/module/mod.rs

@@ -199,6 +199,13 @@ impl Module {
         index.index() < self.num_imported_funcs
     }
 
+    pub fn is_exported_function(&self, index: &FuncIndex) -> bool {
+        self.exports.values().any(|export| match export {
+            EntityIndex::Function(func_id) => func_id == index,
+            _ => false,
+        })
+    }
+
     /// Convert a `DefinedTableIndex` into a `TableIndex`.
     #[inline]
     pub fn table_index(&self, defined_table: DefinedTableIndex) -> TableIndex {