add Assembler::assemble_library()

assembly/src/assembler/mod.rs

@@ -1,8 +1,10 @@
 use crate::{
     ast::{
-        self, FullyQualifiedProcedureName, Instruction, InvocationTarget, InvokeKind, ModuleKind,
+        self, AliasTarget, Export, FullyQualifiedProcedureName, Instruction, InvocationTarget,
+        InvokeKind, ModuleKind, ProcedureIndex,
     },
     diagnostics::Report,
+    library::{CompiledLibrary, ProcedureInfo},
     sema::SemanticAnalysisError,
     AssemblyError, Compile, CompileOptions, Felt, Library, LibraryNamespace, LibraryPath,
     RpoDigest, Spanned, ONE, ZERO,
@@ -225,6 +227,114 @@ impl Assembler {
 // ------------------------------------------------------------------------------------------------
 /// Compilation/Assembly
 impl Assembler {
+    /// Assembles a set of modules into a library.
+    pub fn assemble_library(
+        mut self,
+        modules: impl Iterator<Item = impl Compile>,
+    ) -> Result<CompiledLibrary, Report> {
+        let module_ids: Vec<ModuleIndex> = modules
+            .map(|module| {
+                let module = module.compile_with_options(CompileOptions::for_library())?;
+
+                Ok(self.module_graph.add_module(module)?)
+            })
+            .collect::<Result<_, Report>>()?;
+        self.module_graph.recompute()?;
+
+        let mut mast_forest_builder = MastForestBuilder::default();
+
+        self.assemble_graph(&mut mast_forest_builder)?;
+
+        let exports = {
+            let mut exports = Vec::new();
+            for module_id in module_ids {
+                let module = self.module_graph.get_module(module_id).unwrap();
+                let module_path = module.path();
+
+                let exports_in_module: Vec<FullyQualifiedProcedureName> =
+                    self.get_module_exports(module_id, &mast_forest_builder).map(|procedures| {
+                        procedures
+                            .into_iter()
+                            .map(|proc| {
+                                FullyQualifiedProcedureName::new(module_path.clone(), proc.name)
+                            })
+                            .collect()
+                    })?;
+
+                exports.extend(exports_in_module);
+            }
+
+            exports
+        };
+
+        Ok(CompiledLibrary::new(mast_forest_builder.build(), exports)?)
+    }
+
+    fn get_module_exports(
+        &self,
+        module_index: ModuleIndex,
+        mast_forest_builder: &MastForestBuilder,
+    ) -> Result<Vec<ProcedureInfo>, Report> {
+        let mut exports = Vec::new();
+        for (index, procedure) in self.module_graph[module_index].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 = match procedure {
+                        Export::Procedure(_) => GlobalProcedureIndex {
+                            module: module_index,
+                            index: ProcedureIndex::new(index),
+                        },
+                        Export::Alias(ref alias) => {
+                            match alias.target() {
+                                AliasTarget::MastRoot(digest) => {
+                                    self.module_graph.get_procedure_index_by_digest(digest)
+                                        .unwrap_or_else(|| {
+                                            panic!(
+                                                "compilation apparently succeeded, but did not find a \
+                                                        entry in the procedure cache for alias '{}', i.e. '{}'",
+                                                alias.name(),
+                                                digest
+                                            );
+                                        })
+                                }
+                                AliasTarget::Path(ref name)=> {
+                                    self.module_graph.find(alias.source_file(), name)?
+                                }
+                            }
+                        }
+                    };
+            let proc = mast_forest_builder.get_procedure(gid).unwrap_or_else(|| match procedure {
+                Export::Procedure(ref proc) => {
+                    panic!(
+                        "compilation apparently succeeded, but did not find a \
+                                entry in the procedure cache for '{}'",
+                        proc.name()
+                    )
+                }
+                Export::Alias(ref alias) => {
+                    panic!(
+                        "compilation apparently succeeded, but did not find a \
+                                entry in the procedure cache for alias '{}', i.e. '{}'",
+                        alias.name(),
+                        alias.target()
+                    );
+                }
+            });
+
+            let compiled_proc = ProcedureInfo {
+                name: proc.name().clone(),
+                digest: mast_forest_builder[proc.body_node_id()].digest(),
+            };
+
+            exports.push(compiled_proc);
+        }
+
+        Ok(exports)
+    }
+
     /// Compiles the provided module into a [`Program`]. The resulting program can be executed on
     /// Miden VM.
     ///
@@ -309,6 +419,20 @@ impl Assembler {
         ))
     }
 
+    /// Compile all of the uncompiled procedures in the module graph, placing them
+    /// in the procedure cache once compiled.
+    ///
+    /// Returns an error if any of the provided Miden Assembly is invalid.
+    fn assemble_graph(
+        &mut self,
+        mast_forest_builder: &mut MastForestBuilder,
+    ) -> Result<(), Report> {
+        let mut worklist = self.module_graph.topological_sort().to_vec();
+        assert!(!worklist.is_empty());
+        self.process_graph_worklist(&mut worklist, None, mast_forest_builder)
+            .map(|_| ())
+    }
+
     /// Compile the uncompiled procedure in the module graph which are members of the subgraph
     /// rooted at `root`, placing them in the MAST forest builder once compiled.
     ///

assembly/src/assembler/module_graph/mod.rs

@@ -9,12 +9,15 @@ pub use self::name_resolver::{CallerInfo, ResolvedTarget};
 
 use alloc::{boxed::Box, collections::BTreeMap, sync::Arc, vec::Vec};
 use core::ops::Index;
+use std::borrow::Cow;
 use vm_core::Kernel;
 
 use smallvec::{smallvec, SmallVec};
 
 use self::{analysis::MaybeRewriteCheck, name_resolver::NameResolver, rewrites::ModuleRewriter};
 use super::{GlobalProcedureIndex, ModuleIndex};
+use crate::ast::ResolvedProcedure;
+use crate::diagnostics::{RelatedLabel, SourceFile};
 use crate::{
     ast::{
         Export, FullyQualifiedProcedureName, InvocationTarget, Module, Procedure, ProcedureIndex,
@@ -140,6 +143,16 @@ impl ModuleGraph {
 // ------------------------------------------------------------------------------------------------
 /// Analysis
 impl ModuleGraph {
+    /// Get a slice representing the topological ordering of this graph.
+    ///
+    /// The slice is ordered such that when a node is encountered, all of its dependencies come
+    /// after it in the slice. Thus, by walking the slice in reverse, we visit the leaves of the
+    /// graph before any of the dependents of those leaves. We use this property to resolve MAST
+    /// roots for the entire program, bottom-up.
+    pub fn topological_sort(&self) -> &[GlobalProcedureIndex] {
+        self.topo.as_slice()
+    }
+
     /// Recompute the module graph.
     ///
     /// This should be called any time `add_module`, `add_library`, etc., are called, when all such
@@ -416,6 +429,66 @@ impl ModuleGraph {
         Ok(())
     }
 
+    /// Resolves a [FullyQualifiedProcedureName] to its defining [Procedure].
+    pub fn find(
+        &self,
+        source_file: Option<Arc<SourceFile>>,
+        name: &FullyQualifiedProcedureName,
+    ) -> Result<GlobalProcedureIndex, AssemblyError> {
+        let mut next = Cow::Borrowed(name);
+        let mut caller = source_file.clone();
+        loop {
+            let module_index = self.find_module_index(&next.module).ok_or_else(|| {
+                AssemblyError::UndefinedModule {
+                    span: next.span(),
+                    source_file: caller.clone(),
+                    path: name.module.clone(),
+                }
+            })?;
+            let module = &self.modules[module_index.as_usize()];
+
+            match module.resolve(&next.name) {
+                Some(ResolvedProcedure::Local(index)) => {
+                    let id = GlobalProcedureIndex {
+                        module: module_index,
+                        index: index.into_inner(),
+                    };
+                    break Ok(id);
+                }
+                Some(ResolvedProcedure::External(fqn)) => {
+                    // If we see that we're about to enter an infinite resolver loop because
+                    // of a recursive alias, return an error
+                    if name == &fqn {
+                        break Err(AssemblyError::RecursiveAlias {
+                            source_file: caller.clone(),
+                            name: name.clone(),
+                        });
+                    }
+                    next = Cow::Owned(fqn);
+                    caller = module.source_file();
+                }
+                Some(ResolvedProcedure::MastRoot(ref digest)) => {
+                    if let Some(id) = self.get_procedure_index_by_digest(digest) {
+                        break Ok(id);
+                    }
+                    break Err(AssemblyError::Failed {
+                        labels: vec![RelatedLabel::error("undefined procedure")
+                            .with_source_file(source_file)
+                            .with_labeled_span(next.span(), "unable to resolve this reference")],
+                    });
+                }
+                None => {
+                    // No such procedure known to `module`
+                    break Err(AssemblyError::Failed {
+                        labels: vec![RelatedLabel::error("undefined procedure")
+                            .with_source_file(source_file)
+                            .with_labeled_span(next.span(), "unable to resolve this reference")],
+                    });
+                }
+            }
+        }
+    }
+
     /// Resolve a [LibraryPath] to a [ModuleIndex] in this graph
     pub fn find_module_index(&self, name: &LibraryPath) -> Option<ModuleIndex> {
         self.modules.iter().position(|m| m.path() == name).map(ModuleIndex::new)

assembly/src/library/error.rs

@@ -59,3 +59,13 @@ pub enum LibraryError {
     #[cfg(feature = "std")]
     Io(#[from] std::io::Error),
 }
+
+#[derive(Debug, thiserror::Error, Diagnostic)]
+pub enum CompiledLibraryError {
+    #[error("Invalid exports: MAST forest has {roots_len} procedure roots, but exports have {exports_len}")]
+    #[diagnostic()]
+    InvalidExports {
+        exports_len: usize,
+        roots_len: usize,
+    },
+}