feat: implement kernel compilation in the assembler (#1418)

CHANGELOG.md

@@ -15,15 +15,17 @@
 - Relaxed the parser to allow one branch of an `if.(true|false)` to be empty
 - Added support for immediate values for `u32and`, `u32or`, `u32xor` and `u32not` bitwise instructions (#1362).
 - Optimized `std::sys::truncate_stuck` procedure (#1384).
-- Updated CI and Makefile to standardise it accross Miden repositories (#1342).
+- Updated CI and Makefile to standardize it across Miden repositories (#1342).
 - Add serialization/deserialization for `MastForest` (#1370)
 - Assembler: add the ability to compile MAST libraries, and to assemble a program using compiled libraries (#1401)
 - Updated CI to support `CHANGELOG.md` modification checking and `no changelog` label (#1406)
 - Introduced `MastForestError` to enforce `MastForest` node count invariant (#1394)
-- Added functions to `MastForest` and `MastForestBuilder` to add and ensure nodes with fewer LOC (#1404, #1412)
-- Made `Assembler` single-use (#1409)
+- Added functions to `MastForestBuilder` to allow ensuring of nodes with fewer LOC (#1404)
+- [BREAKING] Made `Assembler` single-use (#1409).
 - Removed `ProcedureCache` from the assembler (#1411).
-- Added `Assembler::assemble_library()` (#1413)
+- Added functions to `MastForest` and `MastForestBuilder` to add and ensure nodes with fewer LOC (#1404, #1412)
+- Added `Assembler::assemble_library()` (#1413).
+- Added `Assembler::assemble_kernel()` (#1418).
 
 #### Changed
 

assembly/src/assembler/id.rs

@@ -54,3 +54,14 @@ impl ModuleIndex {
         self.0 as usize
     }
 }
+
+impl core::ops::Add<ProcedureIndex> for ModuleIndex {
+    type Output = GlobalProcedureIndex;
+
+    fn add(self, rhs: ProcedureIndex) -> Self::Output {
+        GlobalProcedureIndex {
+            module: self,
+            index: rhs,
+        }
+    }
+}

assembly/src/assembler/mod.rs

@@ -1,9 +1,7 @@
 use crate::{
-    ast::{
-        self, FullyQualifiedProcedureName, InvocationTarget, InvokeKind, ModuleKind, ProcedureIndex,
-    },
+    ast::{self, FullyQualifiedProcedureName, InvocationTarget, InvokeKind, ModuleKind},
     diagnostics::Report,
-    library::CompiledLibrary,
+    library::{CompiledLibrary, KernelLibrary},
     sema::SemanticAnalysisError,
     AssemblyError, Compile, CompileOptions, Library, LibraryNamespace, LibraryPath, RpoDigest,
     Spanned,
@@ -62,18 +60,13 @@ pub struct Assembler {
 // ------------------------------------------------------------------------------------------------
 /// Constructors
 impl Assembler {
-    /// Start building an [Assembler]
-    pub fn new() -> Self {
-        Self::default()
-    }
-
-    /// Start building an [`Assembler`] with the given [`Kernel`].
-    pub fn with_kernel(kernel: Kernel) -> Self {
-        let mut assembler = Self::new();
-
-        assembler.module_graph.set_kernel(None, kernel);
-
-        assembler
+    /// Start building an [`Assembler`] with a kernel defined by the provided [KernelLibrary].
+    pub fn with_kernel(kernel_lib: KernelLibrary) -> Self {
+        let (kernel, kernel_module, _) = kernel_lib.into_parts();
+        Self {
+            module_graph: ModuleGraph::with_kernel(kernel, kernel_module),
+            ..Default::default()
+        }
     }
 
     /// Sets the default behavior of this assembler with regard to warning diagnostics.
@@ -149,7 +142,8 @@ impl Assembler {
     pub fn add_compiled_library(&mut self, library: CompiledLibrary) -> Result<(), Report> {
         self.module_graph
             .add_compiled_modules(library.into_module_infos())
-            .map_err(Report::from)
+            .map_err(Report::from)?;
+        Ok(())
     }
 
     /// Adds the library to provide modules for the compilation.
@@ -235,7 +229,11 @@ impl Assembler {
 // ------------------------------------------------------------------------------------------------
 /// Compilation/Assembly
 impl Assembler {
-    /// Assembles a set of modules into a library.
+    /// Assembles a set of modules into a [CompiledLibrary].
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if parsing or compilation of the specified modules fails.
     pub fn assemble_library(
         mut self,
         modules: impl Iterator<Item = impl Compile>,
@@ -254,29 +252,14 @@ impl Assembler {
         let exports = {
             let mut exports = Vec::new();
 
-            for ast_module_idx in ast_module_indices {
+            for module_idx in ast_module_indices {
                 // Note: it is safe to use `unwrap_ast()` here, since all modules looped over are
                 // AST (we just added them to the module graph)
-                let ast_module = self.module_graph[ast_module_idx].unwrap_ast().clone();
-
-                for (proc_idx, procedure) in ast_module.procedures().enumerate() {
-                    // Only add exports; locals will be added if they are in the call graph rooted
-                    // at those procedures
-                    if !procedure.visibility().is_exported() {
-                        continue;
-                    }
-
-                    let gid = GlobalProcedureIndex {
-                        module: ast_module_idx,
-                        index: ProcedureIndex::new(proc_idx),
-                    };
+                let ast_module = self.module_graph[module_idx].unwrap_ast().clone();
 
-                    self.compile_subgraph(gid, false, &mut mast_forest_builder)?;
-
-                    exports.push(FullyQualifiedProcedureName::new(
-                        ast_module.path().clone(),
-                        procedure.name().clone(),
-                    ));
+                for (proc_idx, fqn) in ast_module.exported_procedures() {
+                    self.compile_subgraph(module_idx + proc_idx, false, &mut mast_forest_builder)?;
+                    exports.push(fqn);
                 }
             }
 
@@ -286,6 +269,41 @@ impl Assembler {
         Ok(CompiledLibrary::new(mast_forest_builder.build(), exports)?)
     }
 
+    /// Assembles the provided module into a [KernelLibrary] intended to be used as a Kernel.
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if parsing or compilation of the specified modules fails.
+    pub fn assemble_kernel(mut self, module: impl Compile) -> Result<KernelLibrary, Report> {
+        let options = CompileOptions {
+            kind: ModuleKind::Kernel,
+            warnings_as_errors: self.warnings_as_errors,
+            path: Some(LibraryPath::from(LibraryNamespace::Kernel)),
+        };
+
+        let module = module.compile_with_options(options)?;
+        let module_idx = self.module_graph.add_ast_module(module)?;
+
+        self.module_graph.recompute()?;
+
+        let mut mast_forest_builder = MastForestBuilder::default();
+
+        // Note: it is safe to use `unwrap_ast()` here, since all modules looped over are
+        // AST (we just added them to the module graph)
+        let ast_module = self.module_graph[module_idx].unwrap_ast().clone();
+
+        let exports = ast_module
+            .exported_procedures()
+            .map(|(proc_idx, fqn)| {
+                self.compile_subgraph(module_idx + proc_idx, false, &mut mast_forest_builder)?;
+                Ok(fqn)
+            })
+            .collect::<Result<Vec<FullyQualifiedProcedureName>, Report>>()?;
+
+        let library = CompiledLibrary::new(mast_forest_builder.build(), exports)?;
+        Ok(library.try_into()?)
+    }
+
     /// Compiles the provided module into a [`Program`]. The resulting program can be executed on
     /// Miden VM.
     ///

assembly/src/assembler/module_graph/mod.rs

@@ -21,9 +21,12 @@ use crate::{
         ProcedureName, ResolvedProcedure,
     },
     library::{ModuleInfo, ProcedureInfo},
-    AssemblyError, LibraryPath, RpoDigest, Spanned,
+    AssemblyError, LibraryNamespace, LibraryPath, RpoDigest, Spanned,
 };
 
+// WRAPPER STRUCTS
+// ================================================================================================
+
 /// Wraps all supported representations of a procedure in the module graph.
 ///
 /// Currently, there are two supported representations:
@@ -168,15 +171,15 @@ impl ModuleGraph {
     pub fn add_compiled_modules(
         &mut self,
         module_infos: impl Iterator<Item = ModuleInfo>,
-    ) -> Result<(), AssemblyError> {
+    ) -> Result<Vec<ModuleIndex>, AssemblyError> {
         let module_indices: Vec<ModuleIndex> = module_infos
             .map(|module| self.add_module(PendingWrappedModule::Info(module)))
             .collect::<Result<_, _>>()?;
 
         self.recompute()?;
 
         // Register all procedures as roots
-        for module_index in module_indices {
+        for &module_index in module_indices.iter() {
             for (proc_index, proc) in self[module_index].unwrap_info().clone().procedure_infos() {
                 let gid = GlobalProcedureIndex {
                     module: module_index,
@@ -187,7 +190,7 @@ impl ModuleGraph {
             }
         }
 
-        Ok(())
+        Ok(module_indices)
     }
 
     /// Add `module` to the graph.
@@ -238,9 +241,27 @@ impl ModuleGraph {
 // ------------------------------------------------------------------------------------------------
 /// Kernels
 impl ModuleGraph {
-    pub(super) fn set_kernel(&mut self, kernel_index: Option<ModuleIndex>, kernel: Kernel) {
-        self.kernel_index = kernel_index;
-        self.kernel = kernel;
+    /// Returns a new [ModuleGraph] instantiated from the provided kernel and kernel info module.
+    ///
+    /// Note: it is assumed that kernel and kernel_module are consistent, but this is not checked.
+    ///
+    /// TODO: consider passing `KerneLibrary` into this constructor as a parameter instead.
+    pub(super) fn with_kernel(kernel: Kernel, kernel_module: ModuleInfo) -> Self {
+        assert!(!kernel.is_empty());
+        assert_eq!(kernel_module.path(), &LibraryPath::from(LibraryNamespace::Kernel));
+
+        // add kernel module to the graph
+        // TODO: simplify this to avoid using Self::add_compiled_modules()
+        let mut graph = Self::default();
+        let module_indexes = graph
+            .add_compiled_modules([kernel_module].into_iter())
+            .expect("failed to add kernel module to the module graph");
+        assert_eq!(module_indexes[0], ModuleIndex::new(0), "kernel should be the first module");
+
+        graph.kernel_index = Some(module_indexes[0]);
+        graph.kernel = kernel;
+
+        graph
     }
 
     pub fn kernel(&self) -> &Kernel {

assembly/src/assembler/tests.rs

@@ -66,7 +66,7 @@ fn nested_blocks() {
         }
     }
 
-    let assembler = Assembler::new().with_library(&DummyLibrary::default()).unwrap();
+    let assembler = Assembler::default().with_library(&DummyLibrary::default()).unwrap();
 
     // The expected `MastForest` for the program (that we will build by hand)
     let mut expected_mast_forest_builder = MastForestBuilder::default();
@@ -200,7 +200,7 @@ fn nested_blocks() {
 /// forest.
 #[test]
 fn duplicate_procedure() {
-    let assembler = Assembler::new();
+    let assembler = Assembler::default();
 
     let program_source = r#"
         proc.foo
@@ -227,7 +227,7 @@ fn duplicate_procedure() {
 /// Ensures that equal MAST nodes don't get added twice to a MAST forest
 #[test]
 fn duplicate_nodes() {
-    let assembler = Assembler::new();
+    let assembler = Assembler::default();
 
     let program_source = r#"
     begin
@@ -317,7 +317,7 @@ fn explicit_fully_qualified_procedure_references() {
         }
     }
 
-    let assembler = Assembler::new().with_library(&DummyLibrary::default()).unwrap();
+    let assembler = Assembler::default().with_library(&DummyLibrary::default()).unwrap();
 
     let program = r#"
     begin

assembly/src/ast/module.rs

@@ -6,7 +6,10 @@ use alloc::{
 };
 use core::fmt;
 
-use super::{Export, Import, LocalNameResolver, ProcedureIndex, ProcedureName, ResolvedProcedure};
+use super::{
+    Export, FullyQualifiedProcedureName, Import, LocalNameResolver, ProcedureIndex, ProcedureName,
+    ResolvedProcedure,
+};
 use crate::{
     ast::{AliasTarget, AstSerdeOptions, Ident},
     diagnostics::{Report, SourceFile},
@@ -365,6 +368,26 @@ impl Module {
         self.procedures.iter_mut()
     }
 
+    /// Returns procedures exported from this module.
+    ///
+    /// Each exported procedure is represented by its local procedure index and a fully qualified
+    /// name.
+    pub fn exported_procedures(
+        &self,
+    ) -> impl Iterator<Item = (ProcedureIndex, FullyQualifiedProcedureName)> + '_ {
+        self.procedures.iter().enumerate().filter_map(|(proc_idx, p)| {
+            // skip un-exported procedures
+            if !p.visibility().is_exported() {
+                return None;
+            }
+
+            let proc_idx = ProcedureIndex::new(proc_idx);
+            let fqn = FullyQualifiedProcedureName::new(self.path().clone(), p.name().clone());
+
+            Some((proc_idx, fqn))
+        })
+    }
+
     /// Get an iterator over the imports declared in this module.
     ///
     /// See [Import] for details on what information is available for imports.