Plafer mast library compilation (#1401)

* Introduce `ExternalNode` * Replace `Assembler.node_id_by_digest` map * add TODOP * Add `Host::get_mast_forest` * Move kernel and entrypoint out of `MastForest` * Remove ProgramError * docs * cleanup Program constructors * fix docs * Make `Program.kernel` an `Arc` * fix executable * invoke_mast_root: fix external node creation logic * add failing test * don't make root in `combine_mast_node_ids` and `compile_body` * fix External docs * fmt * fix `entrypoint` doc * Rename `Program::new_with_kernel()` * Document `MastForestStore` and `MemMastForestStore` * fix syscall * execute_* functions: use `MastForest` * `Program`: Remove `Arc` around kernel * remove `Arc` around `MastForest` in `Program` * Return error on malformed host * Simplify `DefaultHost` * `MastForest::add_node()`: add docs * fmt * add failing `duplicate_procedure()` test * Introduce `MastForestBuilder` * Rename `mod tests` -> `testing` * add `duplicate_node()` test * changelog * Program: use `assert!()` instead of `debug_assert!()` * `MastForest::make_root()`: add assert * fmt * Serialization for `MastNodeId` * serialization for MastNode variants except basic block * MastForest serialization scaffolding * define `MastNodeType` constructor from `MastNode` * test join serialization of MastNodeType * `MastNodeType` serialization of split * Revert "serialization for MastNode variants except basic block" This reverts commit efc24fd. * add TODOP * impl Deserializable for `MastForest` (scaffold) * mast_node_to_info() scaffold * try_info_to_mast_node scaffold * Rename `EncodedMastNodeType` * add info module * encode operations into `data` field * decode operations * implement `BasicBlockNode::num_operations_and_decorators()` * OperationOrDecoratorIterator * basic block node: move tests in new file * operation_or_decorator_iterator test * Implement `Operation::with_opcode_and_data()` * encode decorators * implement `decode_decorator()` * fix exec invocation * no else blk special case * add procedure roots comment * implement forgotten `todo!()` * `serialize_deserialize_all_nodes` test * `decode_operations_and_decorators`: fix bit check * confirm_assumptions test scaffold * minor adjustments * Introduce `StringTableBuilder` * naming * test confirm_operation_and_decorator_structure * remove TODOP * remove unused `MastNode::new_dyncall()` * Remove `Error` type * add TODOP * complete test `serialize_deserialize_all_nodes` * check digest on deserialization * remove TODOP * safely decode mast node ids * use method syntax in `MastNodeType` decoding * TODOPs * rewrite <= expression * new `MastNodeType` * implement `Deserializable` for `MastNodeType` * migrate tests to new * Use new MastNodeType * rename string_table_builder_ module * implement `BasicBlockDataBuilder` * add TODOP * BasicBlockDataDecoder * use `BasicBlockDataDecoder` * add headers * add `MastNodeInfo` method * return `Result` instead of `Option` * Remove TODOP * docs * chore: add section separators and fix typos * refactor: change type of the error code of u32assert2 from Felt to u32 (#1382) * impl `Serializable` for `Operation` * impl Deserializable for `Operation` * `StringTableBuilder`: switch to using blake 3 * `EncodedDecoratorVariant`: moved discriminant bit logic to `discriminant()` method * Remove basic block offset * Cargo: don't specify patch versions * make deserialization more efficient * num-traits and num-derive: set default-features false * Remove `OperationData` * `StringRef`: move string length to data buffer * store offset in block * Use `source.read_u32/u64()` * Update `MastNodeInfo` docstring * rename arguments in `encode_u32_pair` * Use basic block offset in deserialization * `BasicBlockDataDecoder`: use `ByteReader::read_u16/32()` methods * `StringTableBuilder`: fix comment * Remove `StringRef` in favor of `DataOffset` * cleanup `MastNodeType` serialization * derive `Copy` for `MastNodeType` * `MastNodeType` tests * add `MastNodeType` tests * use assert * fix asserts * `ModuleGraph::recompute()` reverse edge caller/callee * Implement `Assembler::assemble_library()` * changelog * fix docs * Introduce `CompiledFQDN` * Introduce `WrapperModule` to module graph * split `ModuleGraph::add_module()` * fix compile errors from API changes * fix debug structs * fix `Assembler::get_module_exports()` * fix `process_graph_worklist` * fix procedure * fix `NameResolver` * move `CompiledModule` * `CompiledLibrary::into_compiled_modules` * `Assembler::add_compiled_library()` * changelog * fix `assemble_library()` signature * test `compiled_library()` * nits * register mast roots in `Assembler::add_compiled_library()` * fix resolve * `ModuleGraph::topological_sort_from_root`: only include AST procedures * `Assembler::resolve_target()`: look for digest in module graph first * remove `AssemblyContext::allow_phantom_calls` flag * remove TODOP * `ResolvedProcedure` is no longer `Spanned` * improve test * remove TODOP * `CompiledProcedure` -> `ProcedureInfo` * Document `CompiledLibrary` * Rename `CompiledModule` -> `ModuleInfo` * Refactor `ModuleInfo` * `ModuleWrapper` -> `WrappedModule` * Document `PendingModuleWrapper` * document `Assembler::assemble_library()` * fix TODOP * rename * fix test * cleanup `ModuleGraph::topological_sort_from_root` * fix CI * re-implement `Spanned` for `ResolvedProcedure` * reintroduce proper error message * remove unused methods * Remove all `allow(unused)` methods * Document `unwrap_ast()` call * `NameResolver`: remove use of `unwrap_ast()` * Document or remove all calls to `WrappedModule.unwrap_ast()` * rename `PendingWrappedModule` * Add `ModuleGraph::add_compiled_modules()` * Remove `ModuleGraph::add_module_info()` * refactor: remove Assembler::compile_program() internal method * refactor: remove Assembler::assemble_with_options() internal method --------- Co-authored-by: Bobbin Threadbare <[email protected]> Co-authored-by: Andrey Khmuro <[email protected]>
0xPolygonMiden · Jul 29, 2024 · 78b55a9 · 78b55a9
1 parent 8a5497e
commit 78b55a9
Show file tree

Hide file tree

Showing 14 changed files with 575 additions and 296 deletions.
diff --git a/CHANGELOG.md b/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 `MastForest` and `MastForestBuilder` to add and ensure nodes with fewer LOC (#1404, #1412)

diff --git a/assembly/src/assembler/basic_block_builder.rs b/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};

diff --git a/assembly/src/assembler/instruction/mod.rs b/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;

diff --git a/assembly/src/assembler/instruction/procedures.rs b/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 {

diff --git a/assembly/src/assembler/mod.rs b/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(),
                     ));
                 }
@@ -284,74 +293,33 @@ impl Assembler {
     ///
     /// Returns an error if parsing or compilation of the specified program fails, or if the source
     /// doesn't have an entrypoint.
-    pub fn assemble_program(self, source: impl Compile) -> Result<Program, Report> {
-        let opts = CompileOptions {
+    pub fn assemble_program(mut self, source: impl Compile) -> Result<Program, Report> {
+        let options = CompileOptions {
+            kind: ModuleKind::Executable,
             warnings_as_errors: self.warnings_as_errors,
-            ..CompileOptions::default()
+            path: Some(LibraryPath::from(LibraryNamespace::Exec)),
         };
 
-        self.assemble_with_options(source, opts)
-    }
-
-    /// Compiles the provided module into a [Program] using the provided options.
-    ///
-    /// The resulting program can be executed on Miden VM.
-    ///
-    /// # Errors
-    ///
-    /// Returns an error if parsing or compilation of the specified program fails, or the options
-    /// are invalid.
-    fn assemble_with_options(
-        mut self,
-        source: impl Compile,
-        options: CompileOptions,
-    ) -> Result<Program, Report> {
-        if options.kind != ModuleKind::Executable {
-            return Err(Report::msg(
-                "invalid compile options: assemble_with_opts_in_context requires that the kind be 'executable'",
-            ));
-        }
-
-        let mast_forest_builder = MastForestBuilder::default();
-
-        let program = source.compile_with_options(CompileOptions {
-            // Override the module name so that we always compile the executable
-            // module as #exe
-            path: Some(LibraryPath::from(LibraryNamespace::Exec)),
-            ..options
-        })?;
+        let program = source.compile_with_options(options)?;
         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)?;
 
-        self.compile_program(entrypoint, mast_forest_builder)
-    }
-
-    /// Compile the provided [Module] into a [Program].
-    ///
-    /// Ensures that the [`MastForest`] entrypoint is set to the entrypoint of the program.
-    ///
-    /// Returns an error if the provided Miden Assembly is invalid.
-    fn compile_program(
-        mut self,
-        entrypoint: GlobalProcedureIndex,
-        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());
-
         // Compile the module graph rooted at the entrypoint
+        let mut mast_forest_builder = MastForestBuilder::default();
         let entry_procedure = self.compile_subgraph(entrypoint, true, &mut mast_forest_builder)?;
 
         Ok(Program::with_kernel(
@@ -371,16 +339,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 +368,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 +387,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 +432,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 +568,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"),
+                    },
+                }
             }
         }
     }