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