air_parser/ast/
module.rs

1use std::collections::{BTreeMap, HashSet};
2
3use miden_diagnostics::{DiagnosticsHandler, Severity, SourceSpan, Span, Spanned};
4
5use crate::{ast::*, sema::SemanticAnalysisError};
6
7/// This is a type alias used to clarify that an identifier refers to a module
8pub type ModuleId = Identifier;
9
10#[derive(Debug, Copy, Clone, PartialEq, Eq)]
11pub enum ModuleType {
12    /// Only one root module may be defined in an AirScript program, using `def`.
13    ///
14    /// The root module has no restrictions on what sections it can contain, and in a
15    /// sense "provides" restricted sections to other modules in the program, e.g. the trace columns.
16    Root,
17    /// Any number of library modules are permitted in an AirScript program, using `module`.
18    ///
19    /// Library modules are restricted from declaring the following sections:
20    ///
21    /// * public_inputs
22    /// * trace_columns
23    /// * boundary_constraints
24    /// * integrity_constraints
25    ///
26    /// However, they are allowed to define constants, functions, and the periodic_columns section.
27    Library,
28}
29
30/// This represents the parsed contents of a single AirScript module
31///
32/// When parsing successfully produces a [Module], it is guaranteed that:
33///
34/// * Fields which are only allowed in root modules are empty/unset in library modules
35/// * Fields which must be present in root modules are guaranteed to be present in a root module
36/// * It is guaranteed that at least one boundary constraint and one integrity constraint are
37///   present in a root module
38/// * No duplicate module-level declarations were present
39/// * All globally-visible declarations are unique
40///
41/// However, most validation is not run at this stage, but rather later when producing
42/// a [Program]. In particular, variable declarations are not checked, and imports are only
43/// partially validated here, in that we check for obviously overlapping imports, but cannot
44/// fully validate them until later. Likewise we do not validate constraints, look for invalid
45/// variable usages, etc.
46#[derive(Debug, Spanned)]
47pub struct Module {
48    #[span]
49    pub span: SourceSpan,
50    pub name: ModuleId,
51    pub ty: ModuleType,
52    pub imports: BTreeMap<ModuleId, Import>,
53    pub constants: BTreeMap<Identifier, Constant>,
54    pub evaluators: BTreeMap<Identifier, EvaluatorFunction>,
55    pub functions: BTreeMap<Identifier, Function>,
56    pub periodic_columns: BTreeMap<Identifier, PeriodicColumn>,
57    pub public_inputs: BTreeMap<Identifier, PublicInput>,
58    pub trace_columns: Vec<TraceSegment>,
59    pub buses: BTreeMap<Identifier, Bus>,
60    pub boundary_constraints: Option<Span<Vec<Statement>>>,
61    pub integrity_constraints: Option<Span<Vec<Statement>>>,
62}
63impl Module {
64    /// Constructs an empty module of the specified type, with the given span and name.
65    ///
66    /// # SAFETY
67    ///
68    /// Similar to [Program::new], this function simply constructs an empty [Module], and
69    /// so it does not uphold any guarantees described in it's documentation. It is up to
70    /// the caller to guarantee that they construct a valid module that upholds those
71    /// guarantees, otherwise it is expected that compilation will panic at some point down
72    /// the line.
73    pub fn new(ty: ModuleType, span: SourceSpan, name: ModuleId) -> Self {
74        Self {
75            span,
76            name,
77            ty,
78            imports: Default::default(),
79            constants: Default::default(),
80            evaluators: Default::default(),
81            functions: Default::default(),
82            buses: Default::default(),
83            periodic_columns: Default::default(),
84            public_inputs: Default::default(),
85            trace_columns: vec![],
86            boundary_constraints: None,
87            integrity_constraints: None,
88        }
89    }
90
91    /// Constructs a module of the specified type, with the given span and name, using the
92    /// provided declarations.
93    ///
94    /// The resulting module has had some initial semantic analysis performed as described
95    /// in the module docs. It is expected that this module will be added to a [Library] and
96    /// used to construct a [Program] before it is considered fully validated.
97    pub fn from_declarations(
98        diagnostics: &DiagnosticsHandler,
99        ty: ModuleType,
100        span: SourceSpan,
101        name: Identifier,
102        mut declarations: Vec<Declaration>,
103    ) -> Result<Self, SemanticAnalysisError> {
104        let mut module = Self::new(ty, span, name);
105
106        // Keep track of named items in this module while building it from
107        // the set of declarations we received. We want to produce modules
108        // which are known to have no name conflicts in their declarations,
109        // including explicitly imported names. Wildcard imports will be
110        // checked in later analysis.
111        let mut names = HashSet::<NamespacedIdentifier>::default();
112
113        for declaration in declarations.drain(..) {
114            match declaration {
115                Declaration::Import(import) => {
116                    module.declare_import(diagnostics, &mut names, import)?;
117                }
118                Declaration::Constant(constant) => {
119                    module.declare_constant(diagnostics, &mut names, constant)?;
120                }
121                Declaration::EvaluatorFunction(evaluator) => {
122                    module.declare_evaluator(diagnostics, &mut names, evaluator)?;
123                }
124                Declaration::Function(function) => {
125                    module.declare_function(diagnostics, &mut names, function)?;
126                }
127                Declaration::PeriodicColumns(mut columns) => {
128                    for column in columns.drain(..) {
129                        module.declare_periodic_column(diagnostics, &mut names, column)?;
130                    }
131                }
132                Declaration::PublicInputs(mut inputs) => {
133                    if module.is_library() {
134                        invalid_section_in_library(diagnostics, "public_inputs", span);
135                        return Err(SemanticAnalysisError::RootSectionInLibrary(span));
136                    }
137                    for input in inputs.item.drain(..) {
138                        module.declare_public_input(diagnostics, &mut names, input)?;
139                    }
140                }
141                Declaration::Trace(segments) => {
142                    module.declare_trace_segments(diagnostics, &mut names, segments)?;
143                }
144                Declaration::BoundaryConstraints(statements) => {
145                    module.declare_boundary_constraints(diagnostics, statements)?;
146                }
147                Declaration::IntegrityConstraints(statements) => {
148                    module.declare_integrity_constraints(diagnostics, statements)?;
149                }
150                Declaration::Buses(mut buses) => {
151                    for bus in buses.drain(..) {
152                        module.declare_bus(diagnostics, &mut names, bus)?;
153                    }
154                }
155            }
156        }
157
158        if module.is_root() {
159            if module.trace_columns.is_empty() {
160                diagnostics.diagnostic(Severity::Error)
161                    .with_message("missing trace_columns section")
162                    .with_note("Root modules must contain a trace_columns section with at least a `main` trace declared")
163                    .emit();
164                return Err(SemanticAnalysisError::Invalid);
165            }
166
167            if !module.trace_columns.iter().any(|ts| ts.name == "$main") {
168                diagnostics.diagnostic(Severity::Error)
169                    .with_message("missing main trace declaration")
170                    .with_note("Root modules must contain a trace_columns section with at least a `main` trace declared")
171                    .emit();
172                return Err(SemanticAnalysisError::Invalid);
173            }
174
175            if module.boundary_constraints.is_none() || module.integrity_constraints.is_none() {
176                return Err(SemanticAnalysisError::MissingConstraints);
177            }
178
179            if module.public_inputs.is_empty() {
180                return Err(SemanticAnalysisError::MissingPublicInputs);
181            }
182        }
183
184        Ok(module)
185    }
186
187    fn declare_import(
188        &mut self,
189        diagnostics: &DiagnosticsHandler,
190        names: &mut HashSet<NamespacedIdentifier>,
191        import: Span<Import>,
192    ) -> Result<(), SemanticAnalysisError> {
193        use std::collections::btree_map::Entry;
194
195        let span = import.span();
196        match import.item {
197            Import::All { module: name } => {
198                if name == self.name {
199                    return Err(SemanticAnalysisError::ImportSelf(name.span()));
200                }
201                match self.imports.entry(name) {
202                    Entry::Occupied(mut entry) => {
203                        let first = entry.key().span();
204                        match entry.get_mut() {
205                            Import::All { .. } => {
206                                diagnostics
207                                    .diagnostic(Severity::Warning)
208                                    .with_message("duplicate module import")
209                                    .with_primary_label(span, "duplicate import occurs here")
210                                    .with_secondary_label(first, "original import was here")
211                                    .emit();
212                            }
213                            Import::Partial { items, .. } => {
214                                for item in items.iter() {
215                                    diagnostics
216                                        .diagnostic(Severity::Warning)
217                                        .with_message("redundant item import")
218                                        .with_primary_label(item.span(), "this import is redundant")
219                                        .with_secondary_label(
220                                            name.span(),
221                                            "because this import imports all items already",
222                                        )
223                                        .emit();
224                                }
225                                entry.insert(import.item);
226                            }
227                        }
228                    }
229                    Entry::Vacant(entry) => {
230                        entry.insert(import.item);
231                    }
232                }
233
234                Ok(())
235            }
236            Import::Partial {
237                module: name,
238                mut items,
239            } => {
240                if name == self.name {
241                    return Err(SemanticAnalysisError::ImportSelf(name.span()));
242                }
243                match self.imports.entry(name) {
244                    Entry::Occupied(mut entry) => match entry.get_mut() {
245                        Import::All { module: prev } => {
246                            diagnostics
247                                .diagnostic(Severity::Warning)
248                                .with_message("redundant module import")
249                                .with_primary_label(name.span(), "this import is redundant")
250                                .with_secondary_label(
251                                    prev.span(),
252                                    "because this import includes all items already",
253                                )
254                                .emit();
255                        }
256                        Import::Partial {
257                            items: prev_items, ..
258                        } => {
259                            for item in items.drain() {
260                                if let Some(prev) = prev_items.get(&item) {
261                                    diagnostics
262                                        .diagnostic(Severity::Warning)
263                                        .with_message("redundant item import")
264                                        .with_primary_label(item.span(), "this import is redundant")
265                                        .with_secondary_label(
266                                            prev.span(),
267                                            "because it was already imported here",
268                                        )
269                                        .emit();
270                                    continue;
271                                }
272                                prev_items.insert(item);
273                                let name = if item.is_uppercase() {
274                                    NamespacedIdentifier::Binding(item)
275                                } else {
276                                    NamespacedIdentifier::Function(item)
277                                };
278                                if let Some(prev) = names.replace(name) {
279                                    conflicting_declaration(
280                                        diagnostics,
281                                        "import",
282                                        prev.span(),
283                                        item.span(),
284                                    );
285                                    return Err(SemanticAnalysisError::NameConflict(item.span()));
286                                }
287                            }
288                        }
289                    },
290                    Entry::Vacant(entry) => {
291                        for item in items.iter().copied() {
292                            let name = if item.is_uppercase() {
293                                NamespacedIdentifier::Binding(item)
294                            } else {
295                                NamespacedIdentifier::Function(item)
296                            };
297                            if let Some(prev) = names.replace(name) {
298                                conflicting_declaration(
299                                    diagnostics,
300                                    "import",
301                                    prev.span(),
302                                    item.span(),
303                                );
304                                return Err(SemanticAnalysisError::NameConflict(item.span()));
305                            }
306                        }
307                        entry.insert(Import::Partial {
308                            module: name,
309                            items,
310                        });
311                    }
312                }
313
314                Ok(())
315            }
316        }
317    }
318
319    fn declare_constant(
320        &mut self,
321        diagnostics: &DiagnosticsHandler,
322        names: &mut HashSet<NamespacedIdentifier>,
323        constant: Constant,
324    ) -> Result<(), SemanticAnalysisError> {
325        if !constant.name.is_uppercase() {
326            diagnostics
327                .diagnostic(Severity::Error)
328                .with_message("constant identifiers must be uppercase ASCII characters, e.g. FOO")
329                .with_primary_label(
330                    constant.name.span(),
331                    "this is an invalid constant identifier",
332                )
333                .emit();
334            return Err(SemanticAnalysisError::Invalid);
335        }
336
337        if let Some(prev) = names.replace(NamespacedIdentifier::Binding(constant.name)) {
338            conflicting_declaration(diagnostics, "constant", prev.span(), constant.name.span());
339            return Err(SemanticAnalysisError::NameConflict(constant.name.span()));
340        }
341
342        // Validate constant expression
343        if let ConstantExpr::Matrix(matrix) = &constant.value {
344            let expected_len = matrix
345                .first()
346                .expect("expected matrix to have at least one row")
347                .len();
348            for vector in matrix.iter().skip(1) {
349                if expected_len != vector.len() {
350                    diagnostics
351                        .diagnostic(Severity::Error)
352                        .with_message("invalid constant")
353                        .with_primary_label(
354                            constant.span(),
355                            "invalid matrix literal: mismatched dimensions",
356                        )
357                        .with_note(
358                            "Matrix constants must have the same number of columns in each row",
359                        )
360                        .emit();
361                    return Err(SemanticAnalysisError::Invalid);
362                }
363            }
364        }
365        assert_eq!(self.constants.insert(constant.name, constant), None);
366
367        Ok(())
368    }
369
370    fn declare_evaluator(
371        &mut self,
372        diagnostics: &DiagnosticsHandler,
373        names: &mut HashSet<NamespacedIdentifier>,
374        evaluator: EvaluatorFunction,
375    ) -> Result<(), SemanticAnalysisError> {
376        if let Some(prev) = names.replace(NamespacedIdentifier::Function(evaluator.name)) {
377            conflicting_declaration(diagnostics, "evaluator", prev.span(), evaluator.name.span());
378            return Err(SemanticAnalysisError::NameConflict(evaluator.name.span()));
379        }
380
381        self.evaluators.insert(evaluator.name, evaluator);
382
383        Ok(())
384    }
385
386    fn declare_function(
387        &mut self,
388        diagnostics: &DiagnosticsHandler,
389        names: &mut HashSet<NamespacedIdentifier>,
390        function: Function,
391    ) -> Result<(), SemanticAnalysisError> {
392        if let Some(prev) = names.replace(NamespacedIdentifier::Function(function.name)) {
393            conflicting_declaration(diagnostics, "function", prev.span(), function.name.span());
394            return Err(SemanticAnalysisError::NameConflict(function.name.span()));
395        }
396
397        self.functions.insert(function.name, function);
398
399        Ok(())
400    }
401
402    fn declare_bus(
403        &mut self,
404        diagnostics: &DiagnosticsHandler,
405        names: &mut HashSet<NamespacedIdentifier>,
406        bus: Bus,
407    ) -> Result<(), SemanticAnalysisError> {
408        if let Some(prev) = names.replace(NamespacedIdentifier::Binding(bus.name)) {
409            conflicting_declaration(diagnostics, "bus", prev.span(), bus.name.span());
410            return Err(SemanticAnalysisError::NameConflict(bus.name.span()));
411        }
412
413        self.buses.insert(bus.name, bus);
414
415        Ok(())
416    }
417
418    fn declare_periodic_column(
419        &mut self,
420        diagnostics: &DiagnosticsHandler,
421        names: &mut HashSet<NamespacedIdentifier>,
422        column: PeriodicColumn,
423    ) -> Result<(), SemanticAnalysisError> {
424        if let Some(prev) = names.replace(NamespacedIdentifier::Binding(column.name)) {
425            conflicting_declaration(
426                diagnostics,
427                "periodic column",
428                prev.span(),
429                column.name.span(),
430            );
431            return Err(SemanticAnalysisError::NameConflict(column.name.span()));
432        }
433
434        match column.period() {
435            n if n > 0 && n.is_power_of_two() => {
436                assert_eq!(self.periodic_columns.insert(column.name, column), None);
437
438                Ok(())
439            }
440            _ => {
441                diagnostics.diagnostic(Severity::Error)
442                    .with_message("invalid periodic column declaration")
443                    .with_primary_label(column.span(), "periodic columns must have a non-zero cycle length which is a power of two")
444                    .emit();
445                Err(SemanticAnalysisError::Invalid)
446            }
447        }
448    }
449
450    fn declare_public_input(
451        &mut self,
452        diagnostics: &DiagnosticsHandler,
453        names: &mut HashSet<NamespacedIdentifier>,
454        input: PublicInput,
455    ) -> Result<(), SemanticAnalysisError> {
456        if self.is_library() {
457            return Err(SemanticAnalysisError::RootSectionInLibrary(input.span()));
458        }
459
460        if let Some(prev) = names.replace(NamespacedIdentifier::Binding(input.name())) {
461            conflicting_declaration(
462                diagnostics,
463                "public input",
464                prev.span(),
465                input.name().span(),
466            );
467            Err(SemanticAnalysisError::NameConflict(input.name().span()))
468        } else {
469            assert_eq!(self.public_inputs.insert(input.name(), input), None);
470            Ok(())
471        }
472    }
473
474    fn declare_trace_segments(
475        &mut self,
476        diagnostics: &DiagnosticsHandler,
477        names: &mut HashSet<NamespacedIdentifier>,
478        mut segments: Span<Vec<TraceSegment>>,
479    ) -> Result<(), SemanticAnalysisError> {
480        let span = segments.span();
481        if self.is_library() {
482            invalid_section_in_library(diagnostics, "trace_columns", span);
483            return Err(SemanticAnalysisError::RootSectionInLibrary(span));
484        }
485
486        for segment in segments.iter() {
487            if let Some(prev) = names.replace(NamespacedIdentifier::Binding(segment.name)) {
488                conflicting_declaration(
489                    diagnostics,
490                    "trace segment",
491                    prev.span(),
492                    segment.name.span(),
493                );
494                return Err(SemanticAnalysisError::NameConflict(segment.name.span()));
495            }
496            for binding in segment.bindings.iter() {
497                let binding_name = binding.name.expect("expected binding name");
498                if let Some(prev) = names.replace(NamespacedIdentifier::Binding(binding_name)) {
499                    conflicting_declaration(
500                        diagnostics,
501                        "trace binding",
502                        prev.span(),
503                        binding_name.span(),
504                    );
505                    return Err(SemanticAnalysisError::NameConflict(binding_name.span()));
506                }
507            }
508        }
509
510        self.trace_columns.append(&mut segments.item);
511
512        Ok(())
513    }
514
515    fn declare_boundary_constraints(
516        &mut self,
517        diagnostics: &DiagnosticsHandler,
518        statements: Span<Vec<Statement>>,
519    ) -> Result<(), SemanticAnalysisError> {
520        let span = statements.span();
521        if self.is_library() {
522            invalid_section_in_library(diagnostics, "boundary_constraints", span);
523            return Err(SemanticAnalysisError::RootSectionInLibrary(span));
524        }
525
526        if let Some(prev) = self.boundary_constraints.as_ref() {
527            conflicting_declaration(diagnostics, "boundary_constraints", prev.span(), span);
528            return Err(SemanticAnalysisError::Invalid);
529        }
530
531        if !statements.iter().any(|s| s.has_constraints()) {
532            diagnostics
533                .diagnostic(Severity::Error)
534                .with_message("at least one boundary constraint must be declared")
535                .with_primary_label(span, "missing constraint declaration in this section")
536                .emit();
537            return Err(SemanticAnalysisError::Invalid);
538        }
539
540        self.boundary_constraints = Some(statements);
541
542        Ok(())
543    }
544
545    fn declare_integrity_constraints(
546        &mut self,
547        diagnostics: &DiagnosticsHandler,
548        statements: Span<Vec<Statement>>,
549    ) -> Result<(), SemanticAnalysisError> {
550        let span = statements.span();
551        if self.is_library() {
552            invalid_section_in_library(diagnostics, "integrity_constraints", span);
553            return Err(SemanticAnalysisError::RootSectionInLibrary(span));
554        }
555
556        if let Some(prev) = self.integrity_constraints.as_ref() {
557            conflicting_declaration(diagnostics, "integrity_constraints", prev.span(), span);
558            return Err(SemanticAnalysisError::Invalid);
559        }
560
561        if !statements.iter().any(|s| s.has_constraints()) {
562            diagnostics
563                .diagnostic(Severity::Error)
564                .with_message("at least one integrity constraint must be declared")
565                .with_primary_label(span, "missing constraint declaration in this section")
566                .emit();
567            return Err(SemanticAnalysisError::Invalid);
568        }
569
570        self.integrity_constraints = Some(statements);
571
572        Ok(())
573    }
574
575    #[inline(always)]
576    pub fn is_root(&self) -> bool {
577        !self.is_library()
578    }
579
580    #[inline(always)]
581    pub fn is_library(&self) -> bool {
582        self.ty == ModuleType::Library
583    }
584
585    /// Traverse all of the items exported from this module
586    pub fn exports(&self) -> impl Iterator<Item = Export<'_>> + '_ {
587        self.constants
588            .values()
589            .map(Export::Constant)
590            .chain(self.evaluators.values().map(Export::Evaluator))
591    }
592
593    /// Get the export with the given identifier, if it can be found
594    pub fn get(&self, id: &Identifier) -> Option<Export<'_>> {
595        if id.is_uppercase() {
596            self.constants.get(id).map(Export::Constant)
597        } else {
598            self.evaluators.get(id).map(Export::Evaluator)
599        }
600    }
601}
602impl Eq for Module {}
603impl PartialEq for Module {
604    fn eq(&self, other: &Self) -> bool {
605        self.name == other.name
606            && self.ty == other.ty
607            && self.imports == other.imports
608            && self.constants == other.constants
609            && self.evaluators == other.evaluators
610            && self.functions == other.functions
611            && self.periodic_columns == other.periodic_columns
612            && self.public_inputs == other.public_inputs
613            && self.trace_columns == other.trace_columns
614            && self.boundary_constraints == other.boundary_constraints
615            && self.integrity_constraints == other.integrity_constraints
616    }
617}
618
619fn invalid_section_in_library(diagnostics: &DiagnosticsHandler, ty: &str, span: SourceSpan) {
620    diagnostics
621        .diagnostic(Severity::Error)
622        .with_message(format!("invalid {ty} declaration"))
623        .with_primary_label(span, "this section is not permitted in a library module")
624        .emit();
625}
626
627fn conflicting_declaration(
628    diagnostics: &DiagnosticsHandler,
629    ty: &str,
630    prev: SourceSpan,
631    current: SourceSpan,
632) {
633    diagnostics
634        .diagnostic(Severity::Error)
635        .with_message(format!("invalid {ty} declaration"))
636        .with_primary_label(current, "this conflicts with a previous declaration")
637        .with_secondary_label(prev, "previously defined here")
638        .emit();
639}
air_parser/ast/module.rs

air_parser/ast/
module.rs
