Skip to main content

knave_expand/core/
references.rs

1use rustc_hash::FxHashMap;
2
3use super::{BindingArena, BindingId, BindingKind, Control, Expr, ExprKind, Leaf, Pattern};
4
5#[derive(Default)]
6pub struct References {
7    globals: Vec<BindingId>,
8    uses_compile_time_code: bool,
9    arm_outers: FxHashMap<BindingId, BindingId>,
10}
11
12impl References {
13    pub fn of(value: &Expr, bindings: &BindingArena) -> Self {
14        let mut references = Self::default();
15        references.collect(value, bindings);
16        references
17    }
18
19    pub fn globals(&self) -> &[BindingId] {
20        &self.globals
21    }
22
23    pub fn into_globals(self) -> Vec<BindingId> {
24        self.globals
25    }
26
27    pub fn uses_compile_time_code(&self) -> bool {
28        self.uses_compile_time_code
29    }
30
31    fn collect(&mut self, expr: &Expr, bindings: &BindingArena) {
32        match expr.kind() {
33            ExprKind::Leaf(Leaf::Var(binding)) => self.note_read(*binding, bindings),
34            ExprKind::Leaf(Leaf::Quote(_)) => self.uses_compile_time_code = true,
35            ExprKind::Control(Control::Match(_, arms)) => {
36                for arm in arms {
37                    self.collect_arm_outers(&arm.pattern);
38                }
39            }
40            _ => {}
41        }
42        for child in expr.children() {
43            self.collect(child, bindings);
44        }
45    }
46
47    fn collect_arm_outers(&mut self, pattern: &Pattern) {
48        self.arm_outers.extend(pattern.arm_outer());
49        for nested in pattern.nested() {
50            self.collect_arm_outers(nested);
51        }
52    }
53
54    fn note_read(&mut self, binding: BindingId, bindings: &BindingArena) {
55        self.note_reference(binding, bindings);
56        if let Some(outer) = self.arm_outers.get(&binding).copied() {
57            self.note_reference(outer, bindings);
58        }
59    }
60
61    fn note_reference(&mut self, binding: BindingId, bindings: &BindingArena) {
62        match bindings.get(binding).kind() {
63            BindingKind::Global(_) => self.globals.push(binding),
64            BindingKind::CodeBuiltin(_) => self.uses_compile_time_code = true,
65            _ => {}
66        }
67    }
68}