miden_assembly/ast/
module.rs

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