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