Skip to main content

lutra_compiler/resolver/
module.rs

1use indexmap::IndexMap;
2use itertools::Itertools;
3
4use crate::Span;
5use crate::diagnostic::Diagnostic;
6use crate::pr;
7
8pub fn init_root(root_module_def: pr::ModuleDef) -> Result<pr::ModuleDef, Vec<Diagnostic>> {
9    let mut root = pr::ModuleDef {
10        defs: IndexMap::new(),
11    };
12
13    let diagnostics = root.populate_module(root_module_def.defs);
14    if !diagnostics.is_empty() {
15        return Err(diagnostics);
16    }
17
18    Ok(root)
19}
20
21impl pr::ModuleDef {
22    /// Get definition by fully qualified ident.
23    pub fn get(&self, fq_ident: &pr::Path) -> Option<&pr::Def> {
24        let sub_module = self.get_submodule(fq_ident.parent())?;
25        sub_module.defs.get(fq_ident.last())
26    }
27
28    /// Get definition by fully qualified ident and return remaining steps into the def.
29    pub fn try_get<'a, 's>(&'a self, steps: &'s [String]) -> Option<(&'a pr::Def, &'s [String])> {
30        let mut curr_mod = self;
31        for (index, step) in steps.iter().enumerate() {
32            let def = curr_mod.defs.get(step)?;
33            if let pr::DefKind::Module(sub_module) = &def.kind {
34                curr_mod = sub_module;
35            } else {
36                return Some((def, &steps[(index + 1)..]));
37            }
38        }
39        None
40    }
41
42    /// Get an exclusive reference to definition by fully qualified ident.
43    pub fn get_mut(&mut self, ident: &pr::Path) -> Option<&mut pr::Def> {
44        let module = self.get_module_mut(ident.parent())?;
45
46        module.defs.get_mut(ident.last())
47    }
48
49    pub fn get_submodule(&self, path: &[String]) -> Option<&pr::ModuleDef> {
50        let mut curr_mod = self;
51        for step in path {
52            let def = curr_mod.defs.get(step)?;
53            curr_mod = def.kind.as_module()?;
54        }
55        Some(curr_mod)
56    }
57
58    pub fn get_module_mut(&mut self, path: &[String]) -> Option<&mut pr::ModuleDef> {
59        let mut curr_mod = self;
60        for step in path {
61            let def = curr_mod.defs.get_mut(step)?;
62            curr_mod = def.kind.as_module_mut()?;
63        }
64        Some(curr_mod)
65    }
66
67    pub fn iter_defs(&self) -> impl Iterator<Item = (&String, &pr::Def)> {
68        self.defs.iter()
69    }
70
71    pub fn iter_defs_re(&self) -> impl Iterator<Item = (pr::Path, &pr::Def)> {
72        let non_modules = (self.defs.iter())
73            .filter(|(_, d)| !d.kind.is_module())
74            .map(|(name, d)| (pr::Path::from_name(name), d));
75
76        let sub_defs = (self.defs.iter())
77            .filter(|(_, d)| d.kind.is_module())
78            .flat_map(|(name, d)| {
79                let sub_module = d.kind.as_module().unwrap();
80                sub_module
81                    .iter_defs_re()
82                    .map(|(p, d)| (p.prepend(pr::Path::from_name(name)), d))
83                    .collect_vec()
84            });
85
86        non_modules.chain(sub_defs)
87    }
88
89    pub(super) fn take_unresolved(&mut self, ident: &pr::Path) -> (pr::DefKind, Option<Span>) {
90        let def = self.get_mut(ident).unwrap();
91        let unresolved = def.kind.as_unresolved_mut().unwrap();
92        (*unresolved.take().unwrap(), def.span)
93    }
94
95    pub(super) fn insert_unresolved(&mut self, ident: &pr::Path, def_kind: pr::DefKind) {
96        let def = self.get_mut(ident).unwrap();
97        *def.kind.as_unresolved_mut().unwrap() = Some(Box::new(def_kind));
98    }
99
100    pub(super) fn populate_module(&mut self, defs: IndexMap<String, pr::Def>) -> Vec<Diagnostic> {
101        let mut diagnostics = Vec::new();
102
103        for (name, def) in defs {
104            let kind = match def.kind {
105                pr::DefKind::Module(module_def) => {
106                    // init new module and recurse
107                    let mut new_mod = pr::ModuleDef::default();
108                    diagnostics.extend(new_mod.populate_module(module_def.defs));
109
110                    pr::DefKind::Module(new_mod)
111                }
112                kind => pr::DefKind::Unresolved(Some(Box::new(kind))),
113            };
114            self.defs.insert(name, pr::Def { kind, ..def });
115        }
116        diagnostics
117    }
118}