miden_assembly/ast/procedure/
resolver.rs

1use alloc::{collections::BTreeMap, vec::Vec};
2
3use super::{ProcedureIndex, ProcedureName, QualifiedProcedureName};
4use crate::{LibraryPath, RpoDigest, SourceSpan, Span, Spanned, ast::Ident};
5
6// RESOLVED PROCEDURE
7// ================================================================================================
8
9/// Represents the result of resolving a [ProcedureName] in the context of a module.
10#[derive(Debug, Clone)]
11pub enum ResolvedProcedure {
12    /// The name was resolved to a procedure definition in the same module at the given index
13    Local(Span<ProcedureIndex>),
14    /// The name was resolved to a procedure exported from another module
15    External(QualifiedProcedureName),
16    /// The name was resolved to a procedure with a known MAST root
17    MastRoot(RpoDigest),
18}
19
20impl Spanned for ResolvedProcedure {
21    fn span(&self) -> SourceSpan {
22        match self {
23            ResolvedProcedure::Local(p) => p.span(),
24            ResolvedProcedure::External(p) => p.span(),
25            ResolvedProcedure::MastRoot(_) => SourceSpan::default(),
26        }
27    }
28}
29
30// LOCAL NAME RESOLVER
31// ================================================================================================
32
33/// A lookup table for procedure names in the context of some module
34pub struct LocalNameResolver {
35    imports: BTreeMap<Ident, Span<LibraryPath>>,
36    resolved: BTreeMap<ProcedureName, ProcedureIndex>,
37    resolutions: Vec<ResolvedProcedure>,
38}
39
40impl LocalNameResolver {
41    /// Try to resolve `name` as a procedure
42    pub fn resolve(&self, name: &ProcedureName) -> Option<ResolvedProcedure> {
43        self.resolved
44            .get(name)
45            .copied()
46            .map(|index| self.resolutions[index.as_usize()].clone())
47    }
48
49    /// Try to resolve `name` to an imported module, returning the [LibraryPath] of that module.
50    pub fn resolve_import(&self, name: &Ident) -> Option<Span<&LibraryPath>> {
51        self.imports.get(name).map(|spanned| spanned.as_ref())
52    }
53
54    /// Get the name of the procedure at `index`
55    ///
56    /// This is guaranteed to resolve if `index` is valid, and will panic if not.
57    pub fn get_name(&self, index: ProcedureIndex) -> &ProcedureName {
58        self.resolved
59            .iter()
60            .find_map(|(k, v)| if v == &index { Some(k) } else { None })
61            .expect("invalid procedure index")
62    }
63
64    /// Extend the set of imports this resolver knows about.
65    pub fn with_imports<I>(mut self, imports: I) -> Self
66    where
67        I: IntoIterator<Item = (Ident, Span<LibraryPath>)>,
68    {
69        self.imports.extend(imports);
70        self
71    }
72}
73
74impl FromIterator<(ProcedureName, ResolvedProcedure)> for LocalNameResolver {
75    /// Construct a [LocalNameResolver] from an iterator of resolved names.
76    fn from_iter<T>(iter: T) -> Self
77    where
78        T: IntoIterator<Item = (ProcedureName, ResolvedProcedure)>,
79    {
80        let mut resolver = Self {
81            imports: Default::default(),
82            resolved: Default::default(),
83            resolutions: Default::default(),
84        };
85        for (name, resolution) in iter {
86            let index = ProcedureIndex::new(resolver.resolutions.len());
87            resolver.resolutions.push(resolution);
88            resolver.resolved.insert(name, index);
89        }
90        resolver
91    }
92}