miden_assembly_syntax/ast/
module.rs

1use alloc::{boxed::Box, string::String, sync::Arc, vec::Vec};
2use core::fmt;
3
4use miden_core::{
5    AdviceMap,
6    utils::{ByteReader, ByteWriter, Deserializable, DeserializationError, Serializable},
7};
8use miden_debug_types::{SourceFile, SourceSpan, Span, Spanned};
9use miden_utils_diagnostics::Report;
10
11use super::{
12    DocString, Export, Import, LocalNameResolver, ProcedureIndex, ProcedureName,
13    QualifiedProcedureName, ResolvedProcedure,
14};
15use crate::{
16    LibraryNamespace, LibraryPath,
17    ast::{AliasTarget, Ident},
18    parser::ModuleParser,
19    sema::SemanticAnalysisError,
20};
21
22// MODULE KIND
23// ================================================================================================
24
25/// Represents the kind of a [Module].
26///
27/// The three different kinds have slightly different rules on what syntax is allowed, as well as
28/// what operations can be performed in the body of procedures defined in the module. See the
29/// documentation for each variant for a summary of these differences.
30#[derive(Debug, Default, Copy, Clone, PartialEq, Eq)]
31#[repr(u8)]
32pub enum ModuleKind {
33    /// A library is a simple container of code that must be included into an executable module to
34    /// form a complete program.
35    ///
36    /// Library modules cannot use the `begin`..`end` syntax, which is used to define the
37    /// entrypoint procedure for an executable. Aside from this, they are free to use all other
38    /// MASM syntax.
39    #[default]
40    Library = 0,
41    /// An executable is the root module of a program, and provides the entrypoint for executing
42    /// that program.
43    ///
44    /// As the executable module is the root module, it may not export procedures for other modules
45    /// to depend on, it may only import and call externally-defined procedures, or private
46    /// locally-defined procedures.
47    ///
48    /// An executable module must contain a `begin`..`end` block.
49    Executable = 1,
50    /// A kernel is like a library module, but is special in a few ways:
51    ///
52    /// * Its code always executes in the root context, so it is stateful in a way that normal
53    ///   libraries cannot replicate. This can be used to provide core services that would otherwise
54    ///   not be possible to implement.
55    ///
56    /// * The procedures exported from the kernel may be the target of the `syscall` instruction,
57    ///   and in fact _must_ be called that way.
58    ///
59    /// * Kernels may not use `syscall` or `call` instructions internally.
60    Kernel = 2,
61}
62
63impl ModuleKind {
64    pub fn is_executable(&self) -> bool {
65        matches!(self, Self::Executable)
66    }
67
68    pub fn is_kernel(&self) -> bool {
69        matches!(self, Self::Kernel)
70    }
71
72    pub fn is_library(&self) -> bool {
73        matches!(self, Self::Library)
74    }
75}
76
77impl fmt::Display for ModuleKind {
78    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
79        match self {
80            Self::Library => f.write_str("library"),
81            Self::Executable => f.write_str("executable"),
82            Self::Kernel => f.write_str("kernel"),
83        }
84    }
85}
86
87impl Serializable for ModuleKind {
88    fn write_into<W: ByteWriter>(&self, target: &mut W) {
89        target.write_u8(*self as u8)
90    }
91}
92
93impl Deserializable for ModuleKind {
94    fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
95        match source.read_u8()? {
96            0 => Ok(Self::Library),
97            1 => Ok(Self::Executable),
98            2 => Ok(Self::Kernel),
99            n => Err(DeserializationError::InvalidValue(format!("invalid module kind tag: {n}"))),
100        }
101    }
102}
103
104// MODULE
105// ================================================================================================
106
107/// The abstract syntax tree for a single Miden Assembly module.
108///
109/// All module kinds share this AST representation, as they are largely identical. However, the
110/// [ModuleKind] dictates how the parsed module is semantically analyzed and validated.
111#[derive(Clone)]
112pub struct Module {
113    /// The span covering the entire definition of this module.
114    span: SourceSpan,
115    /// The documentation associated with this module.
116    ///
117    /// Module documentation is provided in Miden Assembly as a documentation comment starting on
118    /// the first line of the module. All other documentation comments are attached to the item the
119    /// precede in the module body.
120    docs: Option<DocString>,
121    /// The fully-qualified path representing the name of this module.
122    path: LibraryPath,
123    /// The kind of module this represents.
124    kind: ModuleKind,
125    /// The imports defined in the module body.
126    pub(crate) imports: Vec<Import>,
127    /// The procedures (defined or re-exported) in the module body.
128    ///
129    /// NOTE: Despite the name, the procedures in this set are not necessarily exported, the
130    /// individual procedure item must be checked to determine visibility.
131    pub(crate) procedures: Vec<Export>,
132    /// AdviceMap that this module expects to be loaded in the host before executing.
133    pub(crate) advice_map: AdviceMap,
134}
135
136/// Constants
137impl Module {
138    /// File extension for a Assembly Module.
139    pub const FILE_EXTENSION: &'static str = "masm";
140
141    /// Name of the root module.
142    pub const ROOT: &'static str = "mod";
143
144    /// File name of the root module.
145    pub const ROOT_FILENAME: &'static str = "mod.masm";
146}
147
148/// Construction
149impl Module {
150    /// Creates a new [Module] with the specified `kind` and fully-qualified path, e.g.
151    /// `std::math::u64`.
152    pub fn new(kind: ModuleKind, path: LibraryPath) -> Self {
153        Self {
154            span: Default::default(),
155            docs: None,
156            path,
157            kind,
158            imports: Default::default(),
159            procedures: Default::default(),
160            advice_map: Default::default(),
161        }
162    }
163
164    /// An alias for creating the default, but empty, `#kernel` [Module].
165    pub fn new_kernel() -> Self {
166        Self::new(ModuleKind::Kernel, LibraryNamespace::Kernel.into())
167    }
168
169    /// An alias for creating the default, but empty, `$exec` [Module].
170    pub fn new_executable() -> Self {
171        Self::new(ModuleKind::Executable, LibraryNamespace::Exec.into())
172    }
173
174    /// Specifies the source span in the source file in which this module was defined, that covers
175    /// the full definition of this module.
176    pub fn with_span(mut self, span: SourceSpan) -> Self {
177        self.span = span;
178        self
179    }
180
181    /// Sets the [LibraryPath] for this module
182    pub fn set_path(&mut self, path: LibraryPath) {
183        self.path = path;
184    }
185
186    /// Sets the [LibraryNamespace] for this module
187    pub fn set_namespace(&mut self, ns: LibraryNamespace) {
188        self.path.set_namespace(ns);
189    }
190
191    /// Sets the documentation for this module
192    pub fn set_docs(&mut self, docs: Option<Span<String>>) {
193        self.docs = docs.map(DocString::new);
194    }
195
196    /// Like [Module::with_span], but does not require ownership of the [Module].
197    pub fn set_span(&mut self, span: SourceSpan) {
198        self.span = span;
199    }
200
201    /// Defines a procedure, raising an error if the procedure is invalid, or conflicts with a
202    /// previous definition
203    pub fn define_procedure(&mut self, export: Export) -> Result<(), SemanticAnalysisError> {
204        if self.is_kernel() && matches!(export, Export::Alias(_)) {
205            return Err(SemanticAnalysisError::ReexportFromKernel { span: export.span() });
206        }
207        if let Some(prev) = self.resolve(export.name()) {
208            let prev_span = prev.span();
209            Err(SemanticAnalysisError::SymbolConflict { span: export.span(), prev_span })
210        } else {
211            self.procedures.push(export);
212            Ok(())
213        }
214    }
215
216    /// Defines an import, raising an error if the import is invalid, or conflicts with a previous
217    /// definition.
218    pub fn define_import(&mut self, import: Import) -> Result<(), SemanticAnalysisError> {
219        if let Some(prev_import) = self.resolve_import(&import.name) {
220            let prev_span = prev_import.span;
221            return Err(SemanticAnalysisError::ImportConflict { span: import.span, prev_span });
222        }
223
224        if let Some(prev_defined) = self.procedures.iter().find(|e| e.name().eq(&import.name)) {
225            let prev_span = prev_defined.span();
226            return Err(SemanticAnalysisError::SymbolConflict { span: import.span, prev_span });
227        }
228
229        self.imports.push(import);
230
231        Ok(())
232    }
233}
234
235/// Parsing
236impl Module {
237    /// Parse a [Module], `name`, of the given [ModuleKind], from `source_file`.
238    pub fn parse(
239        name: LibraryPath,
240        kind: ModuleKind,
241        source_file: Arc<SourceFile>,
242    ) -> Result<Box<Self>, Report> {
243        let mut parser = Self::parser(kind);
244        parser.parse(name, source_file)
245    }
246
247    /// Get a [ModuleParser] for parsing modules of the provided [ModuleKind]
248    pub fn parser(kind: ModuleKind) -> ModuleParser {
249        ModuleParser::new(kind)
250    }
251}
252
253/// Metadata
254impl Module {
255    /// Get the name of this specific module, i.e. the last component of the [LibraryPath] that
256    /// represents the fully-qualified name of the module, e.g. `u64` in `std::math::u64`
257    pub fn name(&self) -> &str {
258        self.path.last()
259    }
260
261    /// Get the fully-qualified name of this module, e.g. `std::math::u64`
262    pub fn path(&self) -> &LibraryPath {
263        &self.path
264    }
265
266    /// Get the namespace of this module, e.g. `std` in `std::math::u64`
267    pub fn namespace(&self) -> &LibraryNamespace {
268        self.path.namespace()
269    }
270
271    /// Returns true if this module belongs to the provided namespace.
272    pub fn is_in_namespace(&self, namespace: &LibraryNamespace) -> bool {
273        self.path.namespace() == namespace
274    }
275
276    /// Get the module documentation for this module, if it was present in the source code the
277    /// module was parsed from
278    pub fn docs(&self) -> Option<Span<&str>> {
279        self.docs.as_ref().map(|spanned| spanned.as_spanned_str())
280    }
281
282    /// Get the type of module this represents:
283    ///
284    /// See [ModuleKind] for details on the different types of modules.
285    pub fn kind(&self) -> ModuleKind {
286        self.kind
287    }
288
289    /// Override the type of module this represents.
290    ///
291    /// See [ModuleKind] for details on what the different types are.
292    pub fn set_kind(&mut self, kind: ModuleKind) {
293        self.kind = kind;
294    }
295
296    /// Returns true if this module is an executable module.
297    #[inline(always)]
298    pub fn is_executable(&self) -> bool {
299        self.kind.is_executable()
300    }
301
302    /// Returns true if this module is the top-level kernel module.
303    #[inline(always)]
304    pub fn is_kernel(&self) -> bool {
305        self.kind.is_kernel() && self.path.is_kernel_path()
306    }
307
308    /// Returns true if this module is a kernel module.
309    #[inline(always)]
310    pub fn is_in_kernel(&self) -> bool {
311        self.kind.is_kernel()
312    }
313
314    /// Returns true if this module has an entrypoint procedure defined,
315    /// i.e. a `begin`..`end` block.
316    pub fn has_entrypoint(&self) -> bool {
317        self.index_of(|p| p.is_main()).is_some()
318    }
319
320    /// Returns a reference to the advice map derived from this module
321    pub fn advice_map(&self) -> &AdviceMap {
322        &self.advice_map
323    }
324
325    /// Get an iterator over the procedures defined in this module.
326    ///
327    /// The entity returned is an [Export], which abstracts over locally-defined procedures and
328    /// re-exported procedures from imported modules.
329    pub fn procedures(&self) -> core::slice::Iter<'_, Export> {
330        self.procedures.iter()
331    }
332
333    /// Same as [Module::procedures], but returns mutable references.
334    pub fn procedures_mut(&mut self) -> core::slice::IterMut<'_, Export> {
335        self.procedures.iter_mut()
336    }
337
338    /// Returns procedures exported from this module.
339    ///
340    /// Each exported procedure is represented by its local procedure index and a fully qualified
341    /// name.
342    pub fn exported_procedures(
343        &self,
344    ) -> impl Iterator<Item = (ProcedureIndex, QualifiedProcedureName)> + '_ {
345        self.procedures.iter().enumerate().filter_map(|(proc_idx, p)| {
346            // skip un-exported procedures
347            if !p.visibility().is_exported() {
348                return None;
349            }
350
351            let proc_idx = ProcedureIndex::new(proc_idx);
352            let fqn = QualifiedProcedureName::new(self.path().clone(), p.name().clone());
353
354            Some((proc_idx, fqn))
355        })
356    }
357
358    /// Get an iterator over the imports declared in this module.
359    ///
360    /// See [Import] for details on what information is available for imports.
361    pub fn imports(&self) -> core::slice::Iter<'_, Import> {
362        self.imports.iter()
363    }
364
365    /// Same as [Self::imports], but returns mutable references to each import.
366    pub fn imports_mut(&mut self) -> core::slice::IterMut<'_, Import> {
367        self.imports.iter_mut()
368    }
369
370    /// Get an iterator over the "dependencies" of a module, i.e. what library namespaces we expect
371    /// to find imported procedures in.
372    ///
373    /// For example, if we have imported `std::math::u64`, then we would expect to find a library
374    /// on disk named `std.masl`, although that isn't a strict requirement. This notion of
375    /// dependencies may go away with future packaging-related changed.
376    pub fn dependencies(&self) -> impl Iterator<Item = &LibraryNamespace> {
377        self.import_paths().map(|import| import.namespace())
378    }
379
380    /// Get the procedure at `index` in this module's procedure table.
381    ///
382    /// The procedure returned may be either a locally-defined procedure, or a re-exported
383    /// procedure. See [Export] for details.
384    pub fn get(&self, index: ProcedureIndex) -> Option<&Export> {
385        self.procedures.get(index.as_usize())
386    }
387
388    /// Get the [ProcedureIndex] for the first procedure in this module's procedure table which
389    /// returns true for `predicate`.
390    pub fn index_of<F>(&self, predicate: F) -> Option<ProcedureIndex>
391    where
392        F: FnMut(&Export) -> bool,
393    {
394        self.procedures.iter().position(predicate).map(ProcedureIndex::new)
395    }
396
397    /// Get the [ProcedureIndex] for the procedure whose name is `name` in this module's procedure
398    /// table, _if_ that procedure is exported.
399    ///
400    /// Non-exported procedures can be retrieved by using [Module::index_of].
401    pub fn index_of_name(&self, name: &ProcedureName) -> Option<ProcedureIndex> {
402        self.index_of(|p| p.name() == name && p.visibility().is_exported())
403    }
404
405    /// Resolves `name` to a procedure within the local scope of this module
406    pub fn resolve(&self, name: &ProcedureName) -> Option<ResolvedProcedure> {
407        let index =
408            self.procedures.iter().position(|p| p.name() == name).map(ProcedureIndex::new)?;
409        match &self.procedures[index.as_usize()] {
410            Export::Procedure(proc) => {
411                Some(ResolvedProcedure::Local(Span::new(proc.name().span(), index)))
412            },
413            Export::Alias(alias) => match alias.target() {
414                AliasTarget::MastRoot(digest) => Some(ResolvedProcedure::MastRoot(**digest)),
415                AliasTarget::ProcedurePath(path) | AliasTarget::AbsoluteProcedurePath(path) => {
416                    Some(ResolvedProcedure::External(path.clone()))
417                },
418            },
419        }
420    }
421
422    /// Construct a search structure that can resolve procedure names local to this module
423    pub fn resolver(&self) -> LocalNameResolver {
424        LocalNameResolver::from_iter(self.procedures.iter().enumerate().map(|(i, p)| match p {
425            Export::Procedure(p) => (
426                p.name().clone(),
427                ResolvedProcedure::Local(Span::new(p.name().span(), ProcedureIndex::new(i))),
428            ),
429            Export::Alias(p) => {
430                let target = match p.target() {
431                    AliasTarget::MastRoot(digest) => ResolvedProcedure::MastRoot(**digest),
432                    AliasTarget::ProcedurePath(path) | AliasTarget::AbsoluteProcedurePath(path) => {
433                        ResolvedProcedure::External(path.clone())
434                    },
435                };
436                (p.name().clone(), target)
437            },
438        }))
439        .with_imports(
440            self.imports
441                .iter()
442                .map(|import| (import.name.clone(), Span::new(import.span(), import.path.clone()))),
443        )
444    }
445
446    /// Resolves `module_name` to an [Import] within the context of this module
447    pub fn resolve_import(&self, module_name: &Ident) -> Option<&Import> {
448        self.imports.iter().find(|import| &import.name == module_name)
449    }
450
451    /// Same as [Module::resolve_import], but returns a mutable reference to the [Import]
452    pub fn resolve_import_mut(&mut self, module_name: &Ident) -> Option<&mut Import> {
453        self.imports.iter_mut().find(|import| &import.name == module_name)
454    }
455
456    /// Return an iterator over the paths of all imports in this module
457    pub fn import_paths(&self) -> impl Iterator<Item = &LibraryPath> + '_ {
458        self.imports.iter().map(|import| &import.path)
459    }
460}
461
462impl core::ops::Index<ProcedureIndex> for Module {
463    type Output = Export;
464
465    #[inline]
466    fn index(&self, index: ProcedureIndex) -> &Self::Output {
467        &self.procedures[index.as_usize()]
468    }
469}
470
471impl core::ops::IndexMut<ProcedureIndex> for Module {
472    #[inline]
473    fn index_mut(&mut self, index: ProcedureIndex) -> &mut Self::Output {
474        &mut self.procedures[index.as_usize()]
475    }
476}
477
478impl Spanned for Module {
479    fn span(&self) -> SourceSpan {
480        self.span
481    }
482}
483
484impl Eq for Module {}
485
486impl PartialEq for Module {
487    fn eq(&self, other: &Self) -> bool {
488        self.kind == other.kind
489            && self.path == other.path
490            && self.docs == other.docs
491            && self.imports == other.imports
492            && self.procedures == other.procedures
493    }
494}
495
496/// Debug representation of this module
497impl fmt::Debug for Module {
498    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
499        f.debug_struct("Module")
500            .field("docs", &self.docs)
501            .field("path", &self.path)
502            .field("kind", &self.kind)
503            .field("imports", &self.imports)
504            .field("procedures", &self.procedures)
505            .finish()
506    }
507}
508
509/// Pretty-printed representation of this module as Miden Assembly text format
510///
511/// NOTE: Delegates to the [crate::prettier::PrettyPrint] implementation internally
512impl fmt::Display for Module {
513    /// Writes this [Module] as formatted MASM code into the formatter.
514    ///
515    /// The formatted code puts each instruction on a separate line and preserves correct
516    /// indentation for instruction blocks.
517    #[inline]
518    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
519        use crate::prettier::PrettyPrint;
520
521        self.pretty_print(f)
522    }
523}
524
525/// The pretty-printer for [Module]
526impl crate::prettier::PrettyPrint for Module {
527    fn render(&self) -> crate::prettier::Document {
528        use crate::prettier::*;
529
530        let mut doc = self
531            .docs
532            .as_ref()
533            .map(|docstring| docstring.render() + nl())
534            .unwrap_or(Document::Empty);
535
536        for (i, import) in self.imports.iter().enumerate() {
537            if i > 0 {
538                doc += nl();
539            }
540            doc += import.render();
541        }
542
543        if !self.imports.is_empty() {
544            doc += nl();
545        }
546
547        let mut export_index = 0;
548        for export in self.procedures.iter() {
549            if export.is_main() {
550                continue;
551            }
552            if export_index > 0 {
553                doc += nl();
554            }
555            doc += export.render();
556            export_index += 1;
557        }
558
559        if let Some(main) = self.procedures().find(|p| p.is_main()) {
560            if export_index > 0 {
561                doc += nl();
562            }
563            doc += main.render();
564        }
565
566        doc
567    }
568}