lutra_compiler/resolver/
module.rs1use 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 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 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 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 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}