lutra_compiler/pr/
def.rs

1use enum_as_inner::EnumAsInner;
2use indexmap::IndexMap;
3use itertools::Itertools;
4
5use super::{Call, CallArg, Expr, ExprKind, Literal, Path, Ref, Ty, TyFuncParam};
6use crate::Span;
7
8/// Definition.
9#[derive(Debug, Clone, PartialEq)]
10pub struct Def {
11    pub kind: DefKind,
12
13    pub annotations: Vec<Anno>,
14
15    /// Code span of the whole definition (including doc comments and annotations)
16    pub span: Option<Span>,
17
18    /// Code span of definition name
19    pub span_name: Option<Span>,
20}
21
22#[derive(Debug, EnumAsInner, PartialEq, Clone)]
23pub enum DefKind {
24    Module(ModuleDef),
25    Expr(ExprDef),
26    /// External function definition. Always contains a `TyKind::Func`.
27    External(Ty),
28    Ty(TyDef),
29    Import(ImportDef),
30    Anno(AnnoDef),
31}
32
33/// Annotation type definition. Introduced with `anno name(const params...)`.
34#[derive(Debug, PartialEq, Clone)]
35pub struct AnnoDef {
36    pub params: Vec<TyFuncParam>,
37}
38
39#[derive(PartialEq, Clone, Default)]
40pub struct ModuleDef {
41    pub imports: Vec<Def>,
42
43    pub defs: IndexMap<String, Def>,
44
45    // Self-annotations, defined within the module
46    pub annotations: Vec<Anno>,
47
48    /// Span covering the content of this module, i.e. the region inside the
49    /// braces for inline modules, or the whole file body for file-based
50    /// submodules. This excludes inner doc comments or annotations.
51    pub span_content: Option<Span>,
52}
53
54impl ModuleDef {
55    /// Get definition by fully qualified ident.
56    pub fn get(&self, fq_ident: &Path) -> Option<&Def> {
57        let sub_module = self.get_module(fq_ident.parent()?)?;
58        sub_module.defs.get(fq_ident.last())
59    }
60
61    /// Get definition by fully qualified ident and return remaining steps into the def.
62    pub fn try_get<'a, 's>(&'a self, steps: &'s [String]) -> Option<(&'a Def, &'s [String])> {
63        let mut curr_mod = self;
64        for (index, step) in steps.iter().enumerate() {
65            let def = curr_mod.defs.get(step)?;
66            if let DefKind::Module(sub_module) = &def.kind {
67                curr_mod = sub_module;
68            } else {
69                return Some((def, &steps[(index + 1)..]));
70            }
71        }
72        None
73    }
74
75    /// Get an exclusive reference to definition by fully qualified ident.
76    pub fn get_mut(&mut self, ident: &Path) -> Option<&mut Def> {
77        let module = self.get_module_mut(ident.parent()?)?;
78
79        module.defs.get_mut(ident.last())
80    }
81
82    pub fn get_module<'a>(&'a self, path: &[String]) -> Option<&'a ModuleDef> {
83        let mut curr_mod = self;
84        for step in path {
85            let def = curr_mod.defs.get(step)?;
86            curr_mod = def.kind.as_module()?;
87        }
88        Some(curr_mod)
89    }
90
91    pub fn get_module_mut(&mut self, path: &[String]) -> Option<&mut ModuleDef> {
92        let mut curr_mod = self;
93        for step in path {
94            let def = curr_mod.defs.get_mut(step)?;
95            curr_mod = def.kind.as_module_mut()?;
96        }
97        Some(curr_mod)
98    }
99
100    pub fn iter_defs(&self) -> impl Iterator<Item = (&String, &Def)> {
101        self.defs.iter()
102    }
103
104    pub fn iter_defs_re(&self) -> impl Iterator<Item = (Path, &Def)> {
105        let non_modules = (self.defs.iter())
106            .filter(|(_, d)| !d.kind.is_module())
107            .map(|(name, d)| (Path::from_name(name), d));
108
109        let sub_defs = (self.defs.iter())
110            .filter(|(_, d)| d.kind.is_module())
111            .flat_map(|(name, d)| {
112                let sub_module = d.kind.as_module().unwrap();
113                sub_module
114                    .iter_defs_re()
115                    .map(|(p, d)| (Path::from_name(name).append(p), d))
116                    .collect_vec()
117            });
118
119        non_modules.chain(sub_defs)
120    }
121
122    /// Collect all non-module def paths under this module into a set.
123    pub fn collect_unresolved(&self) -> std::collections::HashSet<Path> {
124        let p = Path::empty();
125        self.collect_unresolved_re(&p)
126    }
127
128    fn collect_unresolved_re(&self, prefix: &Path) -> std::collections::HashSet<Path> {
129        let mut paths = std::collections::HashSet::new();
130        for (name, def) in &self.defs {
131            let mut path = prefix.clone();
132            path.push(name.clone());
133            if let DefKind::Module(sub) = &def.kind {
134                paths.extend(sub.collect_unresolved_re(&path));
135            } else {
136                paths.insert(path);
137            }
138        }
139        paths
140    }
141
142    pub fn get_doc_at<'a>(&'a self, fq: &Path) -> Option<&'a str> {
143        self.get_anno_at(fq, Anno::as_std_doc)
144    }
145
146    pub fn get_anno_at<'a, R: 'a>(
147        &'a self,
148        fq: &Path,
149        matcher: impl Fn(&'a Anno) -> Option<R>,
150    ) -> Option<R> {
151        if fq.is_empty() {
152            self.annotations.iter().find_map(matcher)
153        } else {
154            self.get(fq)?.annotations.iter().find_map(matcher)
155        }
156    }
157}
158
159#[derive(Debug, PartialEq, Clone)]
160pub struct ExprDef {
161    pub value: Box<Expr>,
162
163    pub constant: bool,
164    pub ty: Option<Ty>,
165}
166
167#[derive(Debug, PartialEq, Clone)]
168pub struct TyDef {
169    pub ty: Ty,
170    pub is_framed: bool,
171    pub framed_label: Option<String>,
172}
173
174#[derive(Debug, PartialEq, Clone)]
175pub struct ImportDef {
176    pub kind: ImportKind,
177    pub span: Option<Span>,
178}
179
180#[derive(Debug, PartialEq, Clone)]
181pub enum ImportKind {
182    /// Represents `std::sql::from as read_table`
183    Single(Path, Option<String>),
184
185    /// Represents `std::sql::{...}`
186    Many(Path, Vec<ImportDef>),
187
188    /// Represents `std::sql::*`
189    Star(Path),
190}
191
192impl ImportDef {
193    pub fn new_simple(path: Path, span: Option<Span>) -> Self {
194        Self {
195            kind: ImportKind::Single(path, None),
196            span,
197        }
198    }
199}
200
201impl ImportKind {
202    pub fn as_simple(&self) -> Option<&Path> {
203        let ImportKind::Single(path, _) = self else {
204            return None;
205        };
206        Some(path)
207    }
208
209    pub fn path(&self) -> &Path {
210        match self {
211            ImportKind::Single(path, _) => path,
212            ImportKind::Many(path, _) => path,
213            ImportKind::Star(path) => path,
214        }
215    }
216
217    pub fn path_mut(&mut self) -> &mut Path {
218        match self {
219            ImportKind::Single(p, _) => p,
220            ImportKind::Many(p, _) => p,
221            ImportKind::Star(p) => p,
222        }
223    }
224}
225
226#[derive(Debug, Clone, PartialEq)]
227pub struct Anno {
228    pub expr: Box<Expr>,
229}
230
231impl Anno {
232    pub fn new(name: Path, args: Vec<CallArg>, span: Span) -> Self {
233        let subject = Box::new(Expr::new_with_span(name, span));
234        let expr = Box::new(Expr::new_with_span(Call { subject, args }, span));
235        Anno { expr }
236    }
237
238    pub fn get_args_mut(&mut self) -> Option<&mut Vec<CallArg>> {
239        self.expr.kind.as_call_mut().map(|c| &mut c.args)
240    }
241
242    pub fn as_named<'a>(&'a self, name: &[&str]) -> Option<&'a [CallArg]> {
243        if is_named_fq(&self.expr, name) {
244            return Some(&[]);
245        }
246        if let ExprKind::Call(call) = &self.expr.kind
247            && is_named_fq(&call.subject, name)
248        {
249            return Some(&call.args);
250        }
251        None
252    }
253
254    // --- std annotations ---
255
256    /// Construct a `@std::doc("...")` annotation.
257    pub fn new_std_doc(content: String, span: Span) -> Self {
258        let content = Expr::new_with_span(Literal::Text(content), span);
259        let args = vec![CallArg::simple(content)];
260        Anno::new(Path::new(["std", "doc"].to_vec()), args, span)
261    }
262
263    /// Matches as `@std::doc("...")` annotation.
264    pub fn as_std_doc(&self) -> Option<&str> {
265        let args = (self.as_named(&["std", "doc"])).or_else(|| self.as_named(&["doc"]))?;
266        // fallback is needed for docs on defs in std
267        Some(args.first()?.expr.kind.as_literal()?.as_text()?)
268    }
269
270    /// Matches as `@std::hidden` annotation.
271    pub fn as_hidden(&self) -> Option<()> {
272        let _ = self
273            .as_named(&["std", "hidden"])
274            .or_else(|| self.as_named(&["hidden"]))?;
275        // fallback is needed for docs on defs in std
276        Some(())
277    }
278
279    /// Matches as `@std::metadata(name = "...")` annotation.
280    pub fn as_std_metadata(&self) -> Option<&str> {
281        let args = self
282            .as_named(&["std", "metadata"])
283            .or_else(|| self.as_named(&["metadata"]))?;
284        // fallback is needed for the @metadata on std itself
285        Some(args.first()?.expr.kind.as_literal()?.as_text()?)
286    }
287
288    /// Matches as `@std::runner("...")` annotation.
289    pub fn as_std_runner(&self) -> Option<&str> {
290        let args = self
291            .as_named(&["std", "runner"])
292            .or_else(|| self.as_named(&["runner"]))?;
293        // fallback is needed for the @runner on std itself
294        Some(args.first()?.expr.kind.as_literal()?.as_text()?)
295    }
296
297    /// Matches as `@std::schema` annotation.
298    pub fn as_std_schema(&self) -> Option<()> {
299        let _ = self.as_named(&["std", "schema"])?;
300        Some(())
301    }
302
303    /// Matches as `@std::rust_derive([...])` annotation.
304    pub fn as_std_rust_derive(&self) -> Option<Vec<&str>> {
305        let args = self.as_named(&["std", "rust_derive"])?;
306        let items = args.first()?.expr.kind.as_array()?;
307        items
308            .iter()
309            .map(|i| Some(i.kind.as_literal()?.as_text()?.as_str()))
310            .collect()
311    }
312}
313
314fn is_named_fq(expr: &Expr, fq: &[&str]) -> bool {
315    if let Some(Ref::Global(target)) = &expr.target {
316        return target.as_steps() == fq;
317    }
318    expr.kind.as_ident().is_some_and(|i| i.as_steps() == fq)
319}
320
321impl Def {
322    pub fn new<K: Into<DefKind>>(kind: K) -> Def {
323        Def {
324            kind: kind.into(),
325            annotations: Vec::new(),
326
327            span: None,
328            span_name: None,
329        }
330    }
331
332    pub fn dummy() -> Self {
333        Def::new(DefKind::dummy())
334    }
335
336    pub fn get_doc(&self) -> Option<&str> {
337        self.get_anno(Anno::as_std_doc)
338    }
339
340    pub fn get_anno<'a, R: 'a>(&'a self, matcher: impl Fn(&'a Anno) -> Option<R>) -> Option<R> {
341        self.annotations.iter().find_map(matcher)
342    }
343}
344
345impl DefKind {
346    pub fn dummy() -> Self {
347        DefKind::Module(ModuleDef::default())
348    }
349}
350
351impl From<ModuleDef> for DefKind {
352    fn from(value: ModuleDef) -> Self {
353        DefKind::Module(value)
354    }
355}
356impl From<ImportDef> for DefKind {
357    fn from(value: ImportDef) -> Self {
358        DefKind::Import(value)
359    }
360}
361impl From<Expr> for DefKind {
362    fn from(expr: Expr) -> Self {
363        DefKind::Expr(ExprDef {
364            value: Box::new(expr),
365            ty: None,
366            constant: false,
367        })
368    }
369}
370
371impl std::fmt::Debug for ModuleDef {
372    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
373        let mut ds = f.debug_struct("ModuleDef");
374
375        if !self.annotations.is_empty() {
376            ds.field("annotations", &self.annotations);
377        }
378
379        if self.defs.len() < 15 {
380            ds.field("defs", &DebugNames(&self.defs));
381        } else {
382            ds.field("defs", &format!("... {} entries ...", self.defs.len()));
383        }
384        ds.finish()
385    }
386}
387
388struct DebugNames<'a>(&'a IndexMap<String, Def>);
389
390impl<'a> std::fmt::Debug for DebugNames<'a> {
391    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
392        let mut dm = f.debug_map();
393        for (n, def) in self.0.iter().sorted_by_key(|x| x.0) {
394            dm.entry(n, def);
395        }
396        dm.finish()
397    }
398}
lutra_compiler/pr/def.rs

lutra_compiler/pr/
def.rs
