Assembler: enable vendoring of compiled libraries (fixes #1435)

CHANGELOG.md

@@ -5,6 +5,7 @@
 #### Changes
 - Update minimum supported Rust version to 1.84.
 - Change Chiplet Fields to Public (#1629).
+- Added to the `Assembler` the ability to vendor a compiled library.
 
 
 ## 0.12.0 (2025-01-22)

assembly/src/assembler/mast_forest_builder.rs

@@ -1,19 +1,22 @@
 use alloc::{
     collections::{BTreeMap, BTreeSet},
+    sync::Arc,
     vec::Vec,
 };
 use core::ops::{Index, IndexMut};
 
+use miette::{IntoDiagnostic, Report};
 use vm_core::{
     crypto::hash::RpoDigest,
     mast::{
         DecoratorFingerprint, DecoratorId, MastForest, MastNode, MastNodeFingerprint, MastNodeId,
+        Remapping, SubtreeIterator,
     },
     Decorator, DecoratorList, Operation,
 };
 
 use super::{GlobalProcedureIndex, Procedure};
-use crate::AssemblyError;
+use crate::{AssemblyError, Library};
 
 // CONSTANTS
 // ================================================================================================
@@ -59,16 +62,42 @@ pub struct MastForestBuilder {
     /// used as a candidate set of nodes that may be eliminated if the are not referenced by any
     /// other node in the forest and are not a root of any procedure.
     merged_basic_block_ids: BTreeSet<MastNodeId>,
+    /// A MastForest that contains vendored libraries, it's used to find precompiled procedures and
+    /// copy their subtrees instead of inserting external nodes.
+    vendored_mast: Arc<MastForest>,
+    /// Keeps track of the new ids assigned to nodes that are copied from the vendored_mast.
+    vendored_remapping: Remapping,
 }
 
 impl MastForestBuilder {
+    /// Creates a new builder that can access vendored libraries.
+    ///
+    /// When [`Self::vendor_or_ensure_external`] is called, if the root is present in the vendored
+    /// libraries, the body of the function is copied in the MAST Forest being built. Otherwise an
+    /// external node is inserted and the funtion body is expected to be found in a library added at
+    /// runtime.
+    pub fn new<'a>(
+        vendored_libraries: impl IntoIterator<Item = &'a Library>,
+    ) -> Result<Self, Report> {
+        // All vendored library are merged into a single MastForest.
+        let forests = vendored_libraries.into_iter().map(|lib| lib.mast_forest().as_ref());
+        let (vendored_mast, _remapping) = MastForest::merge(forests).into_diagnostic()?;
+        // The adviceMap of the vendored forest is copied to the forest being built.
+        let mut mast_forest = MastForest::default();
+        *mast_forest.advice_map_mut() = vendored_mast.advice_map().clone();
+        Ok(MastForestBuilder {
+            mast_forest,
+            vendored_mast: Arc::new(vendored_mast),
+            ..Self::default()
+        })
+    }
+
     /// Removes the unused nodes that were created as part of the assembly process, and returns the
     /// resulting MAST forest.
     ///
-    /// It also returns the map from old node IDs to new node IDs; or `None` if the `MastForest` was
-    /// unchanged. Any [`MastNodeId`] used in reference to the old [`MastForest`] should be remapped
-    /// using this map.
-    pub fn build(mut self) -> (MastForest, Option<BTreeMap<MastNodeId, MastNodeId>>) {
+    /// It also returns the map from old node IDs to new node IDs. Any [`MastNodeId`] used in
+    /// reference to the old [`MastForest`] should be remapped using this map.
+    pub fn build(mut self) -> (MastForest, BTreeMap<MastNodeId, MastNodeId>) {
         let nodes_to_remove = get_nodes_to_remove(self.merged_basic_block_ids, &self.mast_forest);
         let id_remappings = self.mast_forest.remove_nodes(&nodes_to_remove);
 
@@ -450,9 +479,22 @@ impl MastForestBuilder {
         self.ensure_node(MastNode::new_dyncall())
     }
 
-    /// Adds an external node to the forest, and returns the [`MastNodeId`] associated with it.
-    pub fn ensure_external(&mut self, mast_root: RpoDigest) -> Result<MastNodeId, AssemblyError> {
-        self.ensure_node(MastNode::new_external(mast_root))
+    /// If the root is present in the vendored MAST, its subtree is copied. Otherwise an
+    /// external node is added to the forest.
+    pub fn vendor_or_ensure_external(
+        &mut self,
+        mast_root: RpoDigest,
+    ) -> Result<MastNodeId, AssemblyError> {
+        if let Some(root_id) = self.vendored_mast.find_procedure_root(mast_root) {
+            for old_id in SubtreeIterator::new(&root_id, &self.vendored_mast.clone()) {
+                let node = self.vendored_mast[old_id].remap_children(&self.vendored_remapping);
+                let new_id = self.ensure_node(node)?;
+                self.vendored_remapping.insert(old_id, new_id);
+            }
+            Ok(root_id.remap(&self.vendored_remapping))
+        } else {
+            self.ensure_node(MastNode::new_external(mast_root))
+        }
     }
 
     pub fn set_before_enter(&mut self, node_id: MastNodeId, decorator_ids: Vec<DecoratorId>) {

assembly/src/assembler/mod.rs

@@ -71,6 +71,8 @@ pub struct Assembler {
     warnings_as_errors: bool,
     /// Whether the assembler enables extra debugging information.
     in_debug_mode: bool,
+    /// Collects libraries that can be used during assembly to vendor procedures.
+    vendored_libraries: BTreeMap<RpoDigest, Library>,
 }
 
 impl Default for Assembler {
@@ -82,6 +84,7 @@ impl Default for Assembler {
             module_graph,
             warnings_as_errors: false,
             in_debug_mode: false,
+            vendored_libraries: BTreeMap::new(),
         }
     }
 }
@@ -97,6 +100,7 @@ impl Assembler {
             module_graph,
             warnings_as_errors: false,
             in_debug_mode: false,
+            vendored_libraries: BTreeMap::new(),
         }
     }
 
@@ -170,7 +174,7 @@ impl Assembler {
         module: impl Compile,
         options: CompileOptions,
     ) -> Result<ModuleIndex, Report> {
-        let ids = self.add_modules_with_options(vec![module], options)?;
+        let ids = self.add_modules_with_options([module], options)?;
         Ok(ids[0])
     }
 
@@ -196,7 +200,7 @@ impl Assembler {
                 Ok(module)
             })
             .collect::<Result<Vec<_>, Report>>()?;
-        let ids = self.module_graph.add_ast_modules(modules.into_iter())?;
+        let ids = self.module_graph.add_ast_modules(modules)?;
         Ok(ids)
     }
     /// Adds all modules (defined by ".masm" files) from the specified directory to the module
@@ -223,13 +227,11 @@ impl Assembler {
 
     /// Adds the compiled library to provide modules for the compilation.
     ///
-    /// We only current support adding non-vendored libraries - that is, the source code of exported
-    /// procedures is not included in the program that compiles against the library. The library's
-    /// source code is instead expected to be loaded in the processor at execution time. Hence, all
-    /// calls to library procedures will be compiled down to a [`vm_core::mast::ExternalNode`] (i.e.
-    /// a reference to the procedure's MAST root). This means that when executing a program compiled
-    /// against a library, the processor will not be able to differentiate procedures with the same
-    /// MAST root but different decorators.
+    /// All calls to the library's procedures will be compiled down to a
+    /// [`vm_core::mast::ExternalNode`] (i.e. a reference to the procedure's MAST root).
+    /// The library's source code is expected to be loaded in the processor at execution time.
+    /// This means that when executing a program compiled against a library, the processor will not
+    /// be able to differentiate procedures with the same MAST root but different decorators.
     ///
     /// Hence, it is not recommended to export two procedures that have the same MAST root (i.e. are
     /// identical except for their decorators). Note however that we don't expect this scenario to
@@ -251,6 +253,27 @@ impl Assembler {
         self.add_library(library)?;
         Ok(self)
     }
+
+    /// Adds a compiled library from which procedures will be vendored into the assembled code.
+    ///
+    /// Vendoring in this context means that when a procedure from this library is invoked from the
+    /// assembled code, the entire procedure MAST will be copied into the assembled code. Thus,
+    /// when the resulting code is executed on the VM, the vendored library does not need to be
+    /// provided to the VM to resolve external calls.
+    pub fn add_vendored_library(&mut self, library: impl AsRef<Library>) -> Result<(), Report> {
+        self.add_library(&library)?;
+        self.vendored_libraries
+            .insert(*library.as_ref().digest(), library.as_ref().clone());
+        Ok(())
+    }
+
+    /// Adds a compiled library from which procedures will be vendored into the assembled code.
+    ///
+    /// See [`Self::add_vendored_library`]
+    pub fn with_vendored_library(mut self, library: impl AsRef<Library>) -> Result<Self, Report> {
+        self.add_vendored_library(library)?;
+        Ok(self)
+    }
 }
 
 // ------------------------------------------------------------------------------------------------
@@ -298,9 +321,9 @@ impl Assembler {
         modules: impl IntoIterator<Item = impl Compile>,
         options: CompileOptions,
     ) -> Result<Library, Report> {
-        let ast_module_indices = self.add_modules_with_options(modules, options)?;
+        let mut mast_forest_builder = MastForestBuilder::new(self.vendored_libraries.values())?;
 
-        let mut mast_forest_builder = MastForestBuilder::default();
+        let ast_module_indices = self.add_modules_with_options(modules, options)?;
 
         let mut exports = {
             let mut exports = BTreeMap::new();
@@ -326,11 +349,9 @@ impl Assembler {
         };
 
         let (mast_forest, id_remappings) = mast_forest_builder.build();
-        if let Some(id_remappings) = id_remappings {
-            for (_proc_name, node_id) in exports.iter_mut() {
-                if let Some(&new_node_id) = id_remappings.get(node_id) {
-                    *node_id = new_node_id;
-                }
+        for (_proc_name, node_id) in exports.iter_mut() {
+            if let Some(&new_node_id) = id_remappings.get(node_id) {
+                *node_id = new_node_id;
             }
         }
 
@@ -393,7 +414,8 @@ impl Assembler {
             .ok_or(SemanticAnalysisError::MissingEntrypoint)?;
 
         // Compile the module graph rooted at the entrypoint
-        let mut mast_forest_builder = MastForestBuilder::default();
+        let mut mast_forest_builder = MastForestBuilder::new(self.vendored_libraries.values())?;
+
         self.compile_subgraph(entrypoint, &mut mast_forest_builder)?;
         let entry_node_id = mast_forest_builder
             .get_procedure(entrypoint)
@@ -402,9 +424,7 @@ impl Assembler {
 
         // in case the node IDs changed, update the entrypoint ID to the new value
         let (mast_forest, id_remappings) = mast_forest_builder.build();
-        let entry_node_id = id_remappings
-            .map(|id_remappings| id_remappings[&entry_node_id])
-            .unwrap_or(entry_node_id);
+        let entry_node_id = *id_remappings.get(&entry_node_id).unwrap_or(&entry_node_id);
 
         Ok(Program::with_kernel(
             mast_forest.into(),
@@ -769,8 +789,10 @@ impl Assembler {
         }
     }
 
-    /// Verifies the validity of the MAST root as a procedure root hash, and returns the ID of the
-    /// [`core::mast::ExternalNode`] that wraps it.
+    /// Verifies the validity of the MAST root as a procedure root hash, and adds it to the forest.
+    ///
+    /// If the root is present in the vendored MAST, its subtree is copied. Otherwise an
+    /// external node is added to the forest.
     fn ensure_valid_procedure_mast_root(
         &self,
         kind: InvokeKind,
@@ -822,7 +844,7 @@ impl Assembler {
             Some(_) | None => (),
         }
 
-        mast_forest_builder.ensure_external(mast_root)
+        mast_forest_builder.vendor_or_ensure_external(mast_root)
     }
 }
 

assembly/src/assembler/module_graph/mod.rs

@@ -207,10 +207,6 @@ impl ModuleGraph {
 
     /// Add `module` to the graph.
     ///
-    /// NOTE: This operation only adds a module to the graph, but does not perform the
-    /// important analysis needed for compilation, you must call [recompute] once all modules
-    /// are added to ensure the analysis results reflect the current version of the graph.
-    ///
     /// # Errors
     ///
     /// This operation can fail for the following reasons:
@@ -223,9 +219,8 @@ impl ModuleGraph {
     /// This function will panic if the number of modules exceeds the maximum representable
     /// [ModuleIndex] value, `u16::MAX`.
     pub fn add_ast_module(&mut self, module: Box<Module>) -> Result<ModuleIndex, AssemblyError> {
-        let res = self.add_module(PendingWrappedModule::Ast(module))?;
-        self.recompute()?;
-        Ok(res)
+        let ids = self.add_ast_modules([module])?;
+        Ok(ids[0])
     }
 
     fn add_module(&mut self, module: PendingWrappedModule) -> Result<ModuleIndex, AssemblyError> {
@@ -242,7 +237,7 @@ impl ModuleGraph {
 
     pub fn add_ast_modules(
         &mut self,
-        modules: impl Iterator<Item = Box<Module>>,
+        modules: impl IntoIterator<Item = Box<Module>>,
     ) -> Result<Vec<ModuleIndex>, AssemblyError> {
         let idx = modules
             .into_iter()

assembly/src/tests.rs

@@ -3029,3 +3029,31 @@ fn test_program_serde_with_decorators() {
 
     assert_eq!(original_program, deserialized_program);
 }
+
+#[test]
+fn vendoring() -> TestResult {
+    let context = TestContext::new();
+    let mut mod_parser = ModuleParser::new(ModuleKind::Library);
+    let vendor_lib = {
+        let source = source_file!(&context, "export.bar push.1 end export.prune push.2 end");
+        let mod1 = mod_parser.parse(LibraryPath::new("test::mod1").unwrap(), source).unwrap();
+        Assembler::default().assemble_library([mod1]).unwrap()
+    };
+
+    let lib = {
+        let source = source_file!(&context, "export.foo exec.::test::mod1::bar end");
+        let mod2 = mod_parser.parse(LibraryPath::new("test::mod2").unwrap(), source).unwrap();
+
+        let mut assembler = Assembler::default();
+        assembler.add_vendored_library(vendor_lib)?;
+        assembler.assemble_library([mod2]).unwrap()
+    };
+
+    let expected_lib = {
+        let source = source_file!(&context, "export.foo push.1 end");
+        let mod2 = mod_parser.parse(LibraryPath::new("test::mod2").unwrap(), source).unwrap();
+        Assembler::default().assemble_library([mod2]).unwrap()
+    };
+    assert!(lib == expected_lib);
+    Ok(())
+}