Skip to main content

miden_assembly_syntax/sema/
context.rs

1use alloc::{
2    boxed::Box,
3    collections::{BTreeMap, BTreeSet},
4    sync::Arc,
5    vec::Vec,
6};
7
8use miden_debug_types::{SourceFile, SourceManager, SourceSpan, Span, Spanned};
9use miden_utils_diagnostics::{Diagnostic, Severity};
10
11use super::{SemanticAnalysisError, SyntaxError};
12use crate::ast::{
13    constants::{ConstEvalError, eval::CachedConstantValue},
14    *,
15};
16
17/// This maintains the state for semantic analysis of a single [Module].
18pub struct AnalysisContext {
19    constants: BTreeMap<Ident, Constant>,
20    cached_constant_values: BTreeMap<Ident, ConstantValue>,
21    imported: BTreeSet<Ident>,
22    procedures: BTreeSet<ProcedureName>,
23    errors: Vec<SemanticAnalysisError>,
24    source_file: Arc<SourceFile>,
25    source_manager: Arc<dyn SourceManager>,
26    warnings_as_errors: bool,
27}
28
29impl constants::ConstEnvironment for AnalysisContext {
30    type Error = SemanticAnalysisError;
31
32    fn get_source_file_for(&self, span: SourceSpan) -> Option<Arc<SourceFile>> {
33        if span.source_id() == self.source_file.id() {
34            Some(self.source_file.clone())
35        } else {
36            None
37        }
38    }
39    #[inline]
40    fn get(&mut self, name: &Ident) -> Result<Option<CachedConstantValue<'_>>, Self::Error> {
41        if let Some(value) = self.cached_constant_values.get(name) {
42            Ok(Some(CachedConstantValue::Hit(value)))
43        } else if let Some(constant) = self.constants.get(name) {
44            Ok(Some(CachedConstantValue::Miss(&constant.value)))
45        } else if self.imported.contains(name) {
46            // We don't have the definition available yet
47            Ok(None)
48        } else {
49            Err(ConstEvalError::UndefinedSymbol {
50                symbol: name.clone(),
51                source_file: self.get_source_file_for(name.span()),
52            }
53            .into())
54        }
55    }
56    #[inline(always)]
57    fn get_by_path(
58        &mut self,
59        path: Span<&Path>,
60    ) -> Result<Option<CachedConstantValue<'_>>, Self::Error> {
61        if let Some(name) = path.as_ident() {
62            self.get(&name)
63        } else {
64            Ok(None)
65        }
66    }
67
68    #[inline]
69    fn on_eval_completed(&mut self, name: Span<&Path>, value: &ConstantExpr) {
70        let Some(name) = name.as_ident() else {
71            return;
72        };
73        if let Some(value) = value.as_value() {
74            self.cached_constant_values.insert(name, value);
75        } else {
76            self.cached_constant_values.remove(&name);
77        }
78    }
79}
80
81impl AnalysisContext {
82    pub fn new(source_file: Arc<SourceFile>, source_manager: Arc<dyn SourceManager>) -> Self {
83        Self {
84            constants: Default::default(),
85            cached_constant_values: Default::default(),
86            imported: Default::default(),
87            procedures: Default::default(),
88            errors: Default::default(),
89            source_file,
90            source_manager,
91            warnings_as_errors: false,
92        }
93    }
94
95    pub fn set_warnings_as_errors(&mut self, yes: bool) {
96        self.warnings_as_errors = yes;
97    }
98
99    #[inline(always)]
100    pub fn warnings_as_errors(&self) -> bool {
101        self.warnings_as_errors
102    }
103
104    #[inline(always)]
105    pub fn source_manager(&self) -> Arc<dyn SourceManager> {
106        self.source_manager.clone()
107    }
108
109    pub fn register_procedure_name(&mut self, name: ProcedureName) {
110        self.procedures.insert(name);
111    }
112
113    pub fn register_imported_name(&mut self, name: Ident) {
114        self.imported.insert(name);
115    }
116
117    /// Define a new constant `constant`
118    ///
119    /// Returns `Err` if a constant with the same name is already defined
120    pub fn define_constant(&mut self, module: &mut Module, constant: Constant) {
121        if let Err(err) = module.define_constant(constant.clone()) {
122            self.errors.push(err);
123        } else {
124            let name = constant.name.clone();
125            self.constants.insert(name, constant);
126        }
127    }
128
129    /// Register a constant for semantic analysis without defining it in the module.
130    ///
131    /// This is used for enum variants so we can fold discriminants without
132    /// attempting to define the same constant twice.
133    pub fn register_constant(&mut self, constant: Constant) {
134        let name = constant.name.clone();
135        self.cached_constant_values.remove(&name);
136        if let Some(prev) = self.constants.get(&name) {
137            self.errors.push(SemanticAnalysisError::SymbolConflict {
138                span: constant.span,
139                prev_span: prev.span,
140            });
141        } else {
142            self.constants.insert(name, constant);
143        }
144    }
145
146    /// Evaluate constants for validation and cache their values without changing their expressions.
147    pub fn evaluate_constants(&mut self) {
148        self.cached_constant_values.clear();
149        let constants = self.constants.keys().cloned().collect::<Vec<_>>();
150
151        for constant in constants.iter() {
152            let expr = ConstantExpr::Var(Span::new(
153                constant.span(),
154                PathBuf::from(constant.clone()).into(),
155            ));
156            match constants::eval::expr(&expr, self) {
157                Ok(value) => {
158                    if let Some(cached) = value.as_value() {
159                        self.cached_constant_values.insert(constant.clone(), cached);
160                    } else {
161                        self.cached_constant_values.remove(constant);
162                    }
163                },
164                Err(err) => {
165                    self.cached_constant_values.remove(constant);
166                    self.errors.push(err);
167                },
168            }
169        }
170    }
171
172    /// Get the evaluated constant expression bound to `name`.
173    ///
174    /// Returns `Err` if the symbol is undefined
175    pub fn get_evaluated_constant(
176        &self,
177        name: &Ident,
178    ) -> Result<ConstantExpr, SemanticAnalysisError> {
179        if let Some(value) = self.cached_constant_values.get(name) {
180            Ok(value.clone().into())
181        } else if let Some(constant) = self.constants.get(name) {
182            Ok(constant.value.clone())
183        } else {
184            Err(SemanticAnalysisError::SymbolResolutionError(Box::new(
185                SymbolResolutionError::undefined(name.span(), &self.source_manager),
186            )))
187        }
188    }
189
190    pub fn error(&mut self, diagnostic: SemanticAnalysisError) {
191        self.errors.push(diagnostic);
192    }
193
194    pub fn has_errors(&self) -> bool {
195        if self.warnings_as_errors() {
196            return !self.errors.is_empty();
197        }
198        self.errors
199            .iter()
200            .any(|err| matches!(err.severity().unwrap_or(Severity::Error), Severity::Error))
201    }
202
203    pub fn has_failed(&mut self) -> Result<(), SyntaxError> {
204        if self.has_errors() {
205            Err(SyntaxError {
206                source_file: self.source_file.clone(),
207                errors: core::mem::take(&mut self.errors),
208            })
209        } else {
210            Ok(())
211        }
212    }
213
214    pub fn into_result(self) -> Result<(), SyntaxError> {
215        if self.has_errors() {
216            Err(SyntaxError {
217                source_file: self.source_file.clone(),
218                errors: self.errors,
219            })
220        } else {
221            self.emit_warnings();
222            Ok(())
223        }
224    }
225
226    #[cfg(feature = "std")]
227    fn emit_warnings(self) {
228        use crate::diagnostics::Report;
229
230        if !self.errors.is_empty() {
231            // Emit warnings to stderr
232            let warning = Report::from(super::errors::SyntaxWarning {
233                source_file: self.source_file,
234                errors: self.errors,
235            });
236            std::eprintln!("{warning}");
237        }
238    }
239
240    #[cfg(not(feature = "std"))]
241    fn emit_warnings(self) {}
242}
243
244#[cfg(test)]
245mod tests {
246    use alloc::{boxed::Box, string::String, sync::Arc};
247    use core::cell::Cell;
248
249    use super::AnalysisContext;
250    use crate::{
251        Path, PathBuf,
252        ast::{
253            Constant, ConstantExpr, ConstantOp, ConstantValue, Ident, Visibility,
254            constants::{self, eval::CachedConstantValue},
255        },
256        debuginfo::{
257            DefaultSourceManager, SourceContent, SourceLanguage, SourceManager, SourceSpan, Span,
258            Uri,
259        },
260        parser::IntValue,
261    };
262
263    struct CountingEnv<'a> {
264        inner: &'a mut AnalysisContext,
265        hits: Cell<usize>,
266        misses: Cell<usize>,
267    }
268
269    impl<'a> CountingEnv<'a> {
270        fn new(inner: &'a mut AnalysisContext) -> Self {
271            Self {
272                inner,
273                hits: Cell::new(0),
274                misses: Cell::new(0),
275            }
276        }
277
278        fn hits(&self) -> usize {
279            self.hits.get()
280        }
281
282        fn misses(&self) -> usize {
283            self.misses.get()
284        }
285    }
286
287    impl constants::ConstEnvironment for CountingEnv<'_> {
288        type Error = super::SemanticAnalysisError;
289
290        fn get_source_file_for(
291            &self,
292            span: SourceSpan,
293        ) -> Option<Arc<crate::debuginfo::SourceFile>> {
294            <AnalysisContext as constants::ConstEnvironment>::get_source_file_for(self.inner, span)
295        }
296
297        fn get(&mut self, name: &Ident) -> Result<Option<CachedConstantValue<'_>>, Self::Error> {
298            let value = <AnalysisContext as constants::ConstEnvironment>::get(self.inner, name)?;
299            if let Some(ref value) = value {
300                match value {
301                    CachedConstantValue::Hit(_) => self.hits.set(self.hits.get() + 1),
302                    CachedConstantValue::Miss(_) => self.misses.set(self.misses.get() + 1),
303                }
304            }
305            Ok(value)
306        }
307
308        fn get_by_path(
309            &mut self,
310            path: Span<&Path>,
311        ) -> Result<Option<CachedConstantValue<'_>>, Self::Error> {
312            if let Some(name) = path.as_ident() {
313                self.get(&name)
314            } else {
315                <AnalysisContext as constants::ConstEnvironment>::get_by_path(self.inner, path)
316            }
317        }
318
319        fn on_eval_completed(&mut self, name: Span<&Path>, value: &ConstantExpr) {
320            <AnalysisContext as constants::ConstEnvironment>::on_eval_completed(
321                self.inner, name, value,
322            );
323        }
324    }
325
326    fn make_name(i: usize) -> Ident {
327        format!("C{i:05}").parse().expect("generated constant name must be valid")
328    }
329
330    fn make_ref(name: Ident) -> ConstantExpr {
331        let path = Arc::<Path>::from(PathBuf::from(name));
332        ConstantExpr::Var(Span::new(SourceSpan::default(), path))
333    }
334
335    fn make_shared_subexpression_chain(context: &mut AnalysisContext, depth: usize) {
336        for i in 0..depth {
337            let name = make_name(i);
338            let next = make_name(i + 1);
339            context.register_constant(Constant::new(
340                SourceSpan::default(),
341                Visibility::Public,
342                name,
343                ConstantExpr::BinaryOp {
344                    span: SourceSpan::default(),
345                    op: ConstantOp::Add,
346                    lhs: Box::new(make_ref(next.clone())),
347                    rhs: Box::new(make_ref(next)),
348                },
349            ));
350        }
351
352        context.register_constant(Constant::new(
353            SourceSpan::default(),
354            Visibility::Public,
355            make_name(depth),
356            ConstantExpr::Int(Span::new(SourceSpan::default(), IntValue::from(1_u32))),
357        ));
358    }
359
360    #[test]
361    fn semantic_const_eval_memoizes_shared_subexpressions() {
362        let source_manager = Arc::new(DefaultSourceManager::default());
363        let uri =
364            Uri::from(String::from("mem://const-eval-shared-subexpressions").into_boxed_str());
365        let content = SourceContent::new(
366            SourceLanguage::Masm,
367            uri.clone(),
368            String::from("begin\n    nop\nend\n").into_boxed_str(),
369        );
370        let source_file = source_manager.load_from_raw_parts(uri, content);
371        let mut context = AnalysisContext::new(source_file, source_manager);
372
373        // Each Ci references C(i+1) twice, so without memoization the number of misses would
374        // grow exponentially with depth.
375        let depth = 24;
376        make_shared_subexpression_chain(&mut context, depth);
377
378        let root_name = make_name(0);
379        let mut env = CountingEnv::new(&mut context);
380        let root = make_ref(root_name);
381        let result = constants::eval::expr(&root, &mut env)
382            .expect("shared-subexpression constant graph should evaluate");
383
384        assert!(
385            matches!(result.as_value(), Some(ConstantValue::Int(_))),
386            "evaluation should produce a concrete integer constant value"
387        );
388        assert_eq!(env.misses(), depth + 1, "each constant in the chain should miss at most once");
389        assert_eq!(
390            env.hits(),
391            depth,
392            "the second reference to each dependency should be served from cache"
393        );
394    }
395}