feat(Assembler): add support for vendoring compiled libraries

assembly/src/assembler/mod.rs

@@ -6,7 +6,7 @@ use module_graph::{ProcedureWrapper, WrappedModule};
 use vm_core::{
     crypto::hash::RpoDigest,
     debuginfo::SourceSpan,
-    mast::{DecoratorId, MastNodeId},
+    mast::{DecoratorId, MastForest, MastNodeId},
     DecoratorList, Felt, Kernel, Operation, Program, WORD_SIZE,
 };
 
@@ -71,6 +71,7 @@ pub struct Assembler {
     warnings_as_errors: bool,
     /// Whether the assembler enables extra debugging information.
     in_debug_mode: bool,
+    vendored_libraries: Vec<Library>,
 }
 
 impl Default for Assembler {
@@ -82,6 +83,7 @@ impl Default for Assembler {
             module_graph,
             warnings_as_errors: false,
             in_debug_mode: false,
+            vendored_libraries: vec![],
         }
     }
 }
@@ -97,6 +99,7 @@ impl Assembler {
             module_graph,
             warnings_as_errors: false,
             in_debug_mode: false,
+            vendored_libraries: vec![],
         }
     }
 
@@ -251,6 +254,36 @@ impl Assembler {
         self.add_library(library)?;
         Ok(self)
     }
+
+    pub fn vendor_library(&mut self, library: impl AsRef<Library>) -> Result<(), Report> {
+        self.add_library(&library)?;
+        self.vendored_libraries.push(library.as_ref().clone());
+        Ok(())
+    }
+}
+
+fn vendor_mast(
+    mast_forest: MastForest,
+    vendored_libraries: Vec<Library>,
+    exports: BTreeMap<QualifiedProcedureName, MastNodeId>,
+) -> Result<(MastForest, BTreeMap<QualifiedProcedureName, MastNodeId>), Report> {
+    use miette::IntoDiagnostic;
+
+    let mut forests: Vec<MastForest> = vendored_libraries
+        .into_iter()
+        .map(|lib| (**lib.mast_forest()).clone())
+        .collect::<Vec<MastForest>>();
+    forests.push(mast_forest);
+    let (mast_forest, remapping) = vm_core::mast::MastForest::merge(&forests).into_diagnostic()?;
+    let forest_idx = forests.len() - 1;
+    let ids = exports
+        .into_iter()
+        .map(|(name, id)| {
+            let id = remapping.map_root(forest_idx, &id).unwrap_or(id);
+            (name, id)
+        })
+        .collect();
+    Ok((mast_forest, ids))
 }
 
 // ------------------------------------------------------------------------------------------------
@@ -334,6 +367,8 @@ impl Assembler {
             }
         }
 
+        let (mast_forest, exports) = vendor_mast(mast_forest, self.vendored_libraries, exports)?;
+
         Ok(Library::new(mast_forest.into(), exports)?)
     }
 
@@ -406,6 +441,16 @@ impl Assembler {
             .map(|id_remappings| id_remappings[&entry_node_id])
             .unwrap_or(entry_node_id);
 
+        use std::str::FromStr;
+
+        let dummy_name = QualifiedProcedureName::from_str("bla::main")?;
+        let (mast_forest, exports) = vendor_mast(
+            mast_forest,
+            self.vendored_libraries,
+            [(dummy_name.clone(), entry_node_id)].into(),
+        )?;
+        let entry_node_id = exports[&dummy_name];
+
         Ok(Program::with_kernel(
             mast_forest.into(),
             entry_node_id,

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 nop end export.prune nop 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.vendor_library(vendor_lib)?;
+        assembler.assemble_library([mod2]).unwrap()
+    };
+
+    let expected_lib = {
+        let source = source_file!(&context, "export.foo nop end");
+        let mod2 = mod_parser.parse(LibraryPath::new("test::mod2").unwrap(), source).unwrap();
+        Assembler::default().assemble_library([mod2]).unwrap()
+    };
+    assert!(lib == expected_lib);
+    Ok(())
+}