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 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 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 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 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 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}