cl_interpret/function/
collect_upvars.rs

1//! Collects the "Upvars" of a function at the point of its creation, allowing variable capture
2use crate::env::Environment;
3use cl_ast::{
4    Function, Let, Path, PathPart, Pattern, Sym,
5    ast_visitor::{visit::*, walk::Walk},
6};
7use std::collections::{HashMap, HashSet};
8
9pub fn collect_upvars(f: &Function, env: &Environment) -> super::Upvars {
10    CollectUpvars::new(env).visit(f).finish_copied()
11}
12
13#[derive(Clone, Debug)]
14pub struct CollectUpvars<'env> {
15    env: &'env Environment,
16    upvars: HashMap<Sym, usize>,
17    blacklist: HashSet<Sym>,
18}
19
20impl<'env> CollectUpvars<'env> {
21    pub fn new(env: &'env Environment) -> Self {
22        Self { upvars: HashMap::new(), blacklist: HashSet::new(), env }
23    }
24
25    pub fn finish(&mut self) -> HashMap<Sym, usize> {
26        std::mem::take(&mut self.upvars)
27    }
28
29    pub fn finish_copied(&mut self) -> super::Upvars {
30        let Self { env, upvars, blacklist: _ } = self;
31        std::mem::take(upvars)
32            .into_iter()
33            .filter_map(|(k, v)| env.get_id(v).cloned().map(|v| (k, v)))
34            .collect()
35    }
36
37    pub fn add_upvar(&mut self, name: &Sym) {
38        let Self { env, upvars, blacklist } = self;
39        if blacklist.contains(name) || upvars.contains_key(name) {
40            return;
41        }
42        if let Ok(place) = env.id_of(*name) {
43            upvars.insert(*name, place);
44        }
45    }
46
47    pub fn bind_name(&mut self, name: &Sym) {
48        self.blacklist.insert(*name);
49    }
50
51    pub fn scope(&mut self, f: impl Fn(&mut CollectUpvars<'env>)) {
52        let blacklist = self.blacklist.clone();
53
54        // visit the scope
55        f(self);
56
57        // restore the blacklist
58        self.blacklist = blacklist;
59    }
60}
61
62impl<'a> Visit<'a> for CollectUpvars<'_> {
63    fn visit_block(&mut self, b: &'a cl_ast::Block) {
64        self.scope(|cu| b.children(cu));
65    }
66
67    fn visit_let(&mut self, l: &'a cl_ast::Let) {
68        let Let { mutable, name, ty, init } = l;
69        self.visit_mutability(mutable);
70
71        ty.visit_in(self);
72        // visit the initializer, which may use the bound name
73        init.visit_in(self);
74        // a bound name can never be an upvar
75        self.visit_pattern(name);
76    }
77
78    fn visit_path(&mut self, p: &'a cl_ast::Path) {
79        // TODO: path resolution in environments
80        let Path { absolute: false, parts } = p else {
81            return;
82        };
83        let [PathPart::Ident(name)] = parts.as_slice() else {
84            return;
85        };
86        self.add_upvar(name);
87    }
88
89    fn visit_fielder(&mut self, f: &'a cl_ast::Fielder) {
90        let cl_ast::Fielder { name, init } = f;
91        if let Some(init) = init {
92            self.visit_expr(init);
93        } else {
94            self.add_upvar(name); // fielder without init grabs from env
95        }
96    }
97
98    fn visit_pattern(&mut self, value: &'a cl_ast::Pattern) {
99        match value {
100            Pattern::Name(name) => {
101                self.bind_name(name);
102            }
103            Pattern::RangeExc(_, _) | Pattern::RangeInc(_, _) => {}
104            _ => value.children(self),
105        }
106    }
107
108    fn visit_match_arm(&mut self, value: &'a cl_ast::MatchArm) {
109        // MatchArms bind variables with a very small local scope
110        self.scope(|cu| value.children(cu));
111    }
112}
cl_interpret/function/collect_upvars.rs

cl_interpret/function/
collect_upvars.rs
