Plafer mast library compilation

CHANGELOG.md

@@ -16,7 +16,8 @@
 - 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).
-- Added serialization/deserialization for `MastForest` (#1370)
+- 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 `MastForestBuilder` to allow ensuring of nodes with fewer LOC (#1404)

assembly/src/assembler/basic_block_builder.rs

@@ -1,8 +1,7 @@
-use crate::AssemblyError;
+use crate::{ast::Instruction, AssemblyError};
 
 use super::{
-    mast_forest_builder::MastForestBuilder, BodyWrapper, Decorator, DecoratorList, Instruction,
-    ProcedureContext,
+    mast_forest_builder::MastForestBuilder, BodyWrapper, Decorator, DecoratorList, ProcedureContext,
 };
 use alloc::{borrow::Borrow, string::ToString, vec::Vec};
 use vm_core::{mast::MastNodeId, AdviceInjector, AssemblyOp, Operation};

assembly/src/assembler/instruction/mod.rs

@@ -1,10 +1,10 @@
 use super::{
     ast::InvokeKind, mast_forest_builder::MastForestBuilder, Assembler, BasicBlockBuilder, Felt,
-    Instruction, Operation, ProcedureContext, ONE, ZERO,
+    Operation, ProcedureContext,
 };
-use crate::{diagnostics::Report, utils::bound_into_included_u64, AssemblyError};
+use crate::{ast::Instruction, diagnostics::Report, utils::bound_into_included_u64, AssemblyError};
 use core::ops::RangeBounds;
-use vm_core::{mast::MastNodeId, Decorator};
+use vm_core::{mast::MastNodeId, Decorator, ONE, ZERO};
 
 mod adv_ops;
 mod crypto_ops;

assembly/src/assembler/instruction/procedures.rs

@@ -58,7 +58,10 @@ impl Assembler {
                         callee: proc.fully_qualified_name().clone(),
                     })
                     .and_then(|module| {
-                        if module.is_kernel() {
+                        // Note: this module is guaranteed to be of AST variant, since we have the
+                        // AST of a procedure contained in it (i.e. `proc`). Hence, it must be that
+                        // the entire module is in AST representation as well.
+                        if module.unwrap_ast().is_kernel() {
                             Ok(())
                         } else {
                             Err(AssemblyError::InvalidSysCallTarget {

assembly/src/assembler/mod.rs

@@ -1,17 +1,17 @@
 use crate::{
     ast::{
-        self, FullyQualifiedProcedureName, Instruction, InvocationTarget, InvokeKind, ModuleKind,
-        ProcedureIndex,
+        self, FullyQualifiedProcedureName, InvocationTarget, InvokeKind, ModuleKind, ProcedureIndex,
     },
     diagnostics::Report,
     library::CompiledLibrary,
     sema::SemanticAnalysisError,
-    AssemblyError, Compile, CompileOptions, Felt, Library, LibraryNamespace, LibraryPath,
-    RpoDigest, Spanned, ONE, ZERO,
+    AssemblyError, Compile, CompileOptions, Library, LibraryNamespace, LibraryPath, RpoDigest,
+    Spanned,
 };
 use alloc::{sync::Arc, vec::Vec};
 use mast_forest_builder::MastForestBuilder;
-use vm_core::{mast::MastNodeId, Decorator, DecoratorList, Kernel, Operation, Program};
+use module_graph::{ProcedureWrapper, WrappedModule};
+use vm_core::{mast::MastNodeId, Decorator, DecoratorList, Felt, Kernel, Operation, Program};
 
 mod basic_block_builder;
 mod id;
@@ -140,11 +140,18 @@ impl Assembler {
         let module = module.compile_with_options(options)?;
         assert_eq!(module.kind(), kind, "expected module kind to match compilation options");
 
-        self.module_graph.add_module(module)?;
+        self.module_graph.add_ast_module(module)?;
 
         Ok(())
     }
 
+    /// Adds the compiled library to provide modules for the compilation.
+    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)
+    }
+
     /// Adds the library to provide modules for the compilation.
     pub fn with_library<L>(mut self, library: &L) -> Result<Self, Report>
     where
@@ -233,11 +240,11 @@ impl Assembler {
         mut self,
         modules: impl Iterator<Item = impl Compile>,
     ) -> Result<CompiledLibrary, Report> {
-        let module_indices: Vec<ModuleIndex> = modules
+        let ast_module_indices: Vec<ModuleIndex> = modules
             .map(|module| {
                 let module = module.compile_with_options(CompileOptions::for_library())?;
 
-                Ok(self.module_graph.add_module(module)?)
+                Ok(self.module_graph.add_ast_module(module)?)
             })
             .collect::<Result<_, Report>>()?;
         self.module_graph.recompute()?;
@@ -247,25 +254,27 @@ impl Assembler {
         let exports = {
             let mut exports = Vec::new();
 
-            for module_idx in module_indices {
-                let module = self.module_graph.get_module(module_idx).unwrap();
+            for ast_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 module.procedures().enumerate() {
+                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: module_idx,
+                        module: ast_module_idx,
                         index: ProcedureIndex::new(proc_idx),
                     };
 
                     self.compile_subgraph(gid, false, &mut mast_forest_builder)?;
 
                     exports.push(FullyQualifiedProcedureName::new(
-                        module.path().clone(),
+                        ast_module.path().clone(),
                         procedure.name().clone(),
                     ));
                 }
@@ -323,14 +332,16 @@ impl Assembler {
         assert!(program.is_executable());
 
         // Recompute graph with executable module, and start compiling
-        let module_index = self.module_graph.add_module(program)?;
+        let ast_module_index = self.module_graph.add_ast_module(program)?;
         self.module_graph.recompute()?;
 
-        // Find the executable entrypoint
-        let entrypoint = self.module_graph[module_index]
+        // Find the executable entrypoint Note: it is safe to use `unwrap_ast()` here, since this is
+        // the module we just added, which is in AST representation.
+        let entrypoint = self.module_graph[ast_module_index]
+            .unwrap_ast()
             .index_of(|p| p.is_main())
             .map(|index| GlobalProcedureIndex {
-                module: module_index,
+                module: ast_module_index,
                 index,
             })
             .ok_or(SemanticAnalysisError::MissingEntrypoint)?;
@@ -349,7 +360,7 @@ impl Assembler {
         mut mast_forest_builder: MastForestBuilder,
     ) -> Result<Program, Report> {
         // Raise an error if we are called with an invalid entrypoint
-        assert!(self.module_graph[entrypoint].name().is_main());
+        assert!(self.module_graph.get_procedure_unsafe(entrypoint).name().is_main());
 
         // Compile the module graph rooted at the entrypoint
         let entry_procedure = self.compile_subgraph(entrypoint, true, &mut mast_forest_builder)?;
@@ -371,16 +382,22 @@ impl Assembler {
         is_entrypoint: bool,
         mast_forest_builder: &mut MastForestBuilder,
     ) -> Result<Arc<Procedure>, Report> {
-        let mut worklist = self.module_graph.topological_sort_from_root(root).map_err(|cycle| {
-            let iter = cycle.into_node_ids();
-            let mut nodes = Vec::with_capacity(iter.len());
-            for node in iter {
-                let module = self.module_graph[node.module].path();
-                let proc = self.module_graph[node].name();
-                nodes.push(format!("{}::{}", module, proc));
-            }
-            AssemblyError::Cycle { nodes }
-        })?;
+        let mut worklist: Vec<GlobalProcedureIndex> = self
+            .module_graph
+            .topological_sort_from_root(root)
+            .map_err(|cycle| {
+                let iter = cycle.into_node_ids();
+                let mut nodes = Vec::with_capacity(iter.len());
+                for node in iter {
+                    let module = self.module_graph[node.module].path();
+                    let proc = self.module_graph.get_procedure_unsafe(node);
+                    nodes.push(format!("{}::{}", module, proc.name()));
+                }
+                AssemblyError::Cycle { nodes }
+            })?
+            .into_iter()
+            .filter(|&gid| self.module_graph.get_procedure_unsafe(gid).is_ast())
+            .collect();
 
         assert!(!worklist.is_empty());
 
@@ -394,6 +411,7 @@ impl Assembler {
         Ok(compiled.expect("compilation succeeded but root not found in cache"))
     }
 
+    /// Compiles all procedures in the `worklist`.
     fn process_graph_worklist(
         &mut self,
         worklist: &mut Vec<GlobalProcedureIndex>,
@@ -412,7 +430,12 @@ impl Assembler {
             let is_entry = entrypoint == Some(procedure_gid);
 
             // Fetch procedure metadata from the graph
-            let module = &self.module_graph[procedure_gid.module];
+            let module = match &self.module_graph[procedure_gid.module] {
+                WrappedModule::Ast(ast_module) => ast_module,
+                // Note: if the containing module is in `Info` representation, there is nothing to
+                // compile.
+                WrappedModule::Info(_) => continue,
+            };
             let ast = &module[procedure_gid.index];
             let num_locals = ast.num_locals();
             let name = FullyQualifiedProcedureName {
@@ -452,7 +475,8 @@ impl Assembler {
         let gid = proc_ctx.id();
         let num_locals = proc_ctx.num_locals();
 
-        let proc = self.module_graph[gid].unwrap_procedure();
+        let wrapper_proc = self.module_graph.get_procedure_unsafe(gid);
+        let proc = wrapper_proc.unwrap_ast().unwrap_procedure();
         let proc_body_id = if num_locals > 0 {
             // for procedures with locals, we need to update fmp register before and after the
             // procedure body is executed. specifically:
@@ -587,10 +611,13 @@ impl Assembler {
         match resolved {
             ResolvedTarget::Phantom(digest) => Ok(digest),
             ResolvedTarget::Exact { gid } | ResolvedTarget::Resolved { gid, .. } => {
-                Ok(mast_forest_builder
-                    .get_procedure(gid)
-                    .map(|p| p.mast_root())
-                    .expect("expected callee to have been compiled already"))
+                match mast_forest_builder.get_procedure(gid) {
+                    Some(p) => Ok(p.mast_root()),
+                    None => match self.module_graph.get_procedure_unsafe(gid) {
+                        ProcedureWrapper::Info(p) => Ok(p.digest),
+                        ProcedureWrapper::Ast(_) => panic!("Did not find procedure {gid:?} neither in module graph nor procedure cache"),
+                    },
+                }
             }
         }
     }

assembly/src/assembler/module_graph/debug.rs

@@ -17,62 +17,107 @@ struct DisplayModuleGraph<'a>(&'a ModuleGraph);
 impl<'a> fmt::Debug for DisplayModuleGraph<'a> {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         f.debug_set()
-            .entries(self.0.modules.iter().enumerate().flat_map(|(index, m)| {
-                m.procedures().enumerate().filter_map(move |(i, export)| {
-                    if matches!(export, Export::Alias(_)) {
-                        None
-                    } else {
-                        let gid = GlobalProcedureIndex {
-                            module: ModuleIndex::new(index),
-                            index: ProcedureIndex::new(i),
-                        };
-                        let out_edges = self.0.callgraph.out_edges(gid);
-                        Some(DisplayModuleGraphNodeWithEdges { gid, out_edges })
-                    }
-                })
+            .entries(self.0.modules.iter().enumerate().flat_map(|(module_index, m)| {
+                match m {
+                    WrappedModule::Ast(m) => m
+                        .procedures()
+                        .enumerate()
+                        .filter_map(move |(i, export)| {
+                            if matches!(export, Export::Alias(_)) {
+                                None
+                            } else {
+                                let gid = GlobalProcedureIndex {
+                                    module: ModuleIndex::new(module_index),
+                                    index: ProcedureIndex::new(i),
+                                };
+                                let out_edges = self.0.callgraph.out_edges(gid);
+                                Some(DisplayModuleGraphNodeWithEdges { gid, out_edges })
+                            }
+                        })
+                        .collect::<Vec<_>>(),
+                    WrappedModule::Info(m) => m
+                        .procedure_infos()
+                        .map(|(proc_index, _proc)| {
+                            let gid = GlobalProcedureIndex {
+                                module: ModuleIndex::new(module_index),
+                                index: proc_index,
+                            };
+
+                            let out_edges = self.0.callgraph.out_edges(gid);
+                            DisplayModuleGraphNodeWithEdges { gid, out_edges }
+                        })
+                        .collect::<Vec<_>>(),
+                }
             }))
             .finish()
     }
 }
 
 #[doc(hidden)]
-struct DisplayModuleGraphNodes<'a>(&'a Vec<Arc<Module>>);
+struct DisplayModuleGraphNodes<'a>(&'a Vec<WrappedModule>);
 
 impl<'a> fmt::Debug for DisplayModuleGraphNodes<'a> {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         f.debug_list()
-            .entries(self.0.iter().enumerate().flat_map(|(index, m)| {
-                m.procedures().enumerate().filter_map(move |(i, export)| {
-                    if matches!(export, Export::Alias(_)) {
-                        None
-                    } else {
-                        Some(DisplayModuleGraphNode {
-                            module: ModuleIndex::new(index),
-                            index: ProcedureIndex::new(i),
+            .entries(self.0.iter().enumerate().flat_map(|(module_index, m)| {
+                let module_index = ModuleIndex::new(module_index);
+
+                match m {
+                    WrappedModule::Ast(m) => m
+                        .procedures()
+                        .enumerate()
+                        .filter_map(move |(proc_index, export)| {
+                            if matches!(export, Export::Alias(_)) {
+                                None
+                            } else {
+                                Some(DisplayModuleGraphNode {
+                                    module: module_index,
+                                    index: ProcedureIndex::new(proc_index),
+                                    path: m.path(),
+                                    proc_name: export.name(),
+                                    ty: GraphNodeType::Ast,
+                                })
+                            }
+                        })
+                        .collect::<Vec<_>>(),
+                    WrappedModule::Info(m) => m
+                        .procedure_infos()
+                        .map(|(proc_index, proc)| DisplayModuleGraphNode {
+                            module: module_index,
+                            index: proc_index,
                             path: m.path(),
-                            proc: export,
+                            proc_name: &proc.name,
+                            ty: GraphNodeType::Compiled,
                         })
-                    }
-                })
+                        .collect::<Vec<_>>(),
+                }
             }))
             .finish()
     }
 }
 
+#[derive(Debug)]
+enum GraphNodeType {
+    Ast,
+    Compiled,
+}
+
 #[doc(hidden)]
 struct DisplayModuleGraphNode<'a> {
     module: ModuleIndex,
     index: ProcedureIndex,
     path: &'a LibraryPath,
-    proc: &'a Export,
+    proc_name: &'a ProcedureName,
+    ty: GraphNodeType,
 }
 
 impl<'a> fmt::Debug for DisplayModuleGraphNode<'a> {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         f.debug_struct("Node")
             .field("id", &format_args!("{}:{}", &self.module.as_usize(), &self.index.as_usize()))
             .field("module", &self.path)
-            .field("name", &self.proc.name())
+            .field("name", &self.proc_name)
+            .field("type", &self.ty)
             .finish()
     }
 }