1
//! Free-variable analysis for lambda bodies.
2
//!
3
//! Tier 1.5 of the error-processing model rewrites defuns and lambdas
4
//! to emit as real wasm functions instead of inline templates. A real
5
//! wasm function only sees its own params + globals + table imports,
6
//! so any binding the body references that lives in the surrounding
7
//! lexical scope must travel in via a closure environment struct.
8
//! [`compute_captures`] picks out exactly those names; subsequent
9
//! slices use the result to shape the env-struct (#34) and to emit the
10
//! body fn (#29).
11
//!
12
//! What counts as a capture: a name `n` referenced in the body where
13
//! the in-scope `SymbolTable` resolves `n` to a runtime binding
14
//! (`Expr::WasmRuntime` / `Expr::WasmLocal`) and `n` is **not** rebound
15
//! by an enclosing form inside the body (the lambda's own params, or
16
//! a nested LET / LET* / LAMBDA / LABELS / DOLIST / DO / DO*).
17
//! Constants, builtins, special forms, and globally-defined defuns
18
//! don't need capture — they remain reachable from the wasm function
19
//! via the same global symbol table the call-site sees.
20
//!
21
//! The pass is purely structural: it walks `Expr` and never emits
22
//! wasm. Body code is not modified — captures are just enumerated for
23
//! later phases.
24

            
25
#[cfg(test)]
26
mod tests;
27

            
28
use crate::ast::{Expr, LambdaParams};
29
use crate::runtime::SymbolTable;
30

            
31
/// Ordered set of captured names. Insertion-order is preserved so the
32
/// env-struct field layout is deterministic — the body emitter and
33
/// the call-site env-construction phase walk the same order.
34
#[derive(Debug, Default, Clone, PartialEq, Eq)]
35
pub struct CaptureSet {
36
    names: Vec<String>,
37
}
38

            
39
impl CaptureSet {
40
348
    pub fn insert(&mut self, name: &str) {
41
348
        if !self.names.iter().any(|n| n == name) {
42
347
            self.names.push(name.to_string());
43
347
        }
44
348
    }
45

            
46
3945
    pub fn iter(&self) -> impl Iterator<Item = &str> {
47
3945
        self.names.iter().map(String::as_str)
48
3945
    }
49

            
50
    #[cfg(test)]
51
    #[must_use]
52
4
    pub fn len(&self) -> usize {
53
4
        self.names.len()
54
4
    }
55

            
56
    #[cfg(test)]
57
    #[must_use]
58
5
    pub fn contains(&self, name: &str) -> bool {
59
5
        self.names.iter().any(|n| n == name)
60
5
    }
61
}
62

            
63
/// Compute the set of free variables in `body` that aren't bound by
64
/// `params` and that resolve to a runtime binding in `symbols`.
65
#[must_use]
66
3956
pub fn compute_captures(symbols: &SymbolTable, params: &LambdaParams, body: &Expr) -> CaptureSet {
67
3956
    let mut captures = CaptureSet::default();
68
3956
    let bound = collect_param_names(params);
69
3956
    walk(symbols, &bound, body, &mut captures);
70
3956
    captures
71
3956
}
72

            
73
3956
fn collect_param_names(params: &LambdaParams) -> Vec<String> {
74
3956
    let mut names = Vec::with_capacity(
75
3956
        params.required.len()
76
3956
            + params.optional.len()
77
3956
            + usize::from(params.rest.is_some())
78
3956
            + params.key.len()
79
3956
            + params.aux.len(),
80
    );
81
3956
    names.extend(params.required.iter().cloned());
82
3956
    names.extend(params.optional.iter().map(|(n, _)| n.clone()));
83
3956
    if let Some(rest) = &params.rest {
84
        names.push(rest.clone());
85
3956
    }
86
3956
    names.extend(params.key.iter().map(|(n, _)| n.clone()));
87
3956
    names.extend(params.aux.iter().map(|(n, _)| n.clone()));
88
3956
    names
89
3956
}
90

            
91
16032
fn walk(symbols: &SymbolTable, bound: &[String], expr: &Expr, captures: &mut CaptureSet) {
92
16032
    match expr {
93
9276
        Expr::Symbol(name) => visit_symbol(symbols, bound, name, captures),
94
4163
        Expr::List(elems) => walk_list(symbols, bound, elems, captures),
95
        Expr::Cons(car, cdr) => {
96
            walk(symbols, bound, car, captures);
97
            walk(symbols, bound, cdr, captures);
98
        }
99
205
        Expr::Quote(_) | Expr::Quasiquote(_) => {
100
205
            // Quoted forms don't reference runtime bindings — symbols
101
205
            // inside a quote are data, not variable references.
102
205
        }
103
        Expr::Unquote(inner) | Expr::UnquoteSplicing(inner) => {
104
            walk(symbols, bound, inner, captures);
105
        }
106
        Expr::Lambda(inner_params, inner_body) => {
107
            // A nested lambda introduces its own params; anything the
108
            // outer body sees as captured by the *outer* lambda must
109
            // also be captured if it slips through here without
110
            // rebinding. We recurse with `bound` extended by the inner
111
            // params so the inner-body walk distinguishes its own
112
            // params from the outer scope.
113
            let mut extended = bound.to_vec();
114
            extended.extend(collect_param_names(inner_params));
115
            walk(symbols, &extended, inner_body, captures);
116
        }
117
2388
        _ => {
118
2388
            // Atoms (Nil, Bool, Number, String, Keyword, Bytes,
119
2388
            // RuntimeValue, WasmRuntime, WasmLocal) carry no symbol
120
2388
            // references that could be captures.
121
2388
        }
122
    }
123
16032
}
124

            
125
9276
fn visit_symbol(symbols: &SymbolTable, bound: &[String], name: &str, captures: &mut CaptureSet) {
126
11180
    if bound.iter().any(|n| n == name) {
127
4697
        return;
128
4579
    }
129
4579
    let Some(sym) = symbols.lookup(name) else {
130
14
        return;
131
    };
132
4217
    if matches!(
133
4565
        sym.value(),
134
        Some(Expr::WasmRuntime(_) | Expr::WasmLocal(_, _))
135
348
    ) {
136
348
        captures.insert(name);
137
4224
    }
138
9276
}
139

            
140
4163
fn walk_list(symbols: &SymbolTable, bound: &[String], elems: &[Expr], captures: &mut CaptureSet) {
141
4163
    let head = elems
142
4163
        .first()
143
4163
        .and_then(Expr::as_symbol)
144
4163
        .map(str::to_uppercase);
145
4163
    match head.as_deref() {
146
4163
        Some("LET") => walk_let(symbols, bound, elems, captures, false),
147
4161
        Some("LET*") => walk_let(symbols, bound, elems, captures, true),
148
4161
        Some("LAMBDA") => walk_lambda_form(symbols, bound, elems, captures),
149
4091
        Some("LABELS") => walk_labels(symbols, bound, elems, captures),
150
4091
        Some("DOLIST" | "DOTIMES") => walk_iteration(symbols, bound, elems, captures),
151
4090
        Some("DO" | "DO*") => walk_do(
152
            symbols,
153
            bound,
154
            elems,
155
            captures,
156
            head.as_deref() == Some("DO*"),
157
        ),
158
4090
        Some("QUOTE") => {}
159
        _ => {
160
12000
            for elem in elems {
161
12000
                walk(symbols, bound, elem, captures);
162
12000
            }
163
        }
164
    }
165
4163
}
166

            
167
2
fn walk_let(
168
2
    symbols: &SymbolTable,
169
2
    bound: &[String],
170
2
    elems: &[Expr],
171
2
    captures: &mut CaptureSet,
172
2
    sequential: bool,
173
2
) {
174
2
    let Some(bindings_expr) = elems.get(1) else {
175
        for elem in &elems[1..] {
176
            walk(symbols, bound, elem, captures);
177
        }
178
        return;
179
    };
180
2
    let bindings = match bindings_expr {
181
2
        Expr::List(b) => b,
182
        _ => {
183
            for elem in &elems[1..] {
184
                walk(symbols, bound, elem, captures);
185
            }
186
            return;
187
        }
188
    };
189

            
190
2
    let mut scope = bound.to_vec();
191
2
    for binding in bindings {
192
2
        let (name, init) = parse_let_binding(binding);
193
2
        if let Some(init_expr) = init {
194
2
            let init_scope = if sequential { &scope[..] } else { bound };
195
2
            walk(symbols, init_scope, init_expr, captures);
196
        }
197
2
        if let Some(n) = name {
198
2
            scope.push(n);
199
2
        }
200
    }
201
2
    for body_form in &elems[2..] {
202
2
        walk(symbols, &scope, body_form, captures);
203
2
    }
204
2
}
205

            
206
2
fn parse_let_binding(binding: &Expr) -> (Option<String>, Option<&Expr>) {
207
2
    match binding {
208
        Expr::Symbol(name) => (Some(name.clone()), None),
209
2
        Expr::List(parts) => match parts.as_slice() {
210
            [Expr::Symbol(name)] => (Some(name.clone()), None),
211
2
            [Expr::Symbol(name), init] => (Some(name.clone()), Some(init)),
212
            _ => (None, None),
213
        },
214
        _ => (None, None),
215
    }
216
2
}
217

            
218
70
fn walk_lambda_form(
219
70
    symbols: &SymbolTable,
220
70
    bound: &[String],
221
70
    elems: &[Expr],
222
70
    captures: &mut CaptureSet,
223
70
) {
224
    // (LAMBDA (params...) body...) — params follow head; everything
225
    // after is the body. We only capture names visible in `bound` here
226
    // for this nested lambda's body, so once we know the params we
227
    // recurse on each body form with the extended bound set.
228
70
    let Some(params_expr) = elems.get(1) else {
229
        return;
230
    };
231
70
    let mut scope = bound.to_vec();
232
70
    if let Expr::List(params) = params_expr {
233
70
        for p in params {
234
70
            if let Expr::Symbol(name) = p
235
70
                && !is_lambda_keyword(name)
236
70
            {
237
70
                scope.push(name.clone());
238
70
            } else if let Expr::List(spec) = p
239
                && let Some(Expr::Symbol(name)) = spec.first()
240
            {
241
                scope.push(name.clone());
242
            }
243
        }
244
    }
245
70
    for body_form in &elems[2..] {
246
70
        walk(symbols, &scope, body_form, captures);
247
70
    }
248
70
}
249

            
250
70
fn is_lambda_keyword(name: &str) -> bool {
251
    matches!(
252
70
        name.to_ascii_uppercase().as_str(),
253
70
        "&OPTIONAL" | "&REST" | "&KEY" | "&AUX"
254
    )
255
70
}
256

            
257
fn walk_labels(symbols: &SymbolTable, bound: &[String], elems: &[Expr], captures: &mut CaptureSet) {
258
    let Some(defs_expr) = elems.get(1) else {
259
        return;
260
    };
261
    let mut scope = bound.to_vec();
262
    if let Expr::List(defs) = defs_expr {
263
        for def in defs {
264
            if let Expr::List(d) = def
265
                && let Some(Expr::Symbol(name)) = d.first()
266
            {
267
                scope.push(name.clone());
268
            }
269
        }
270
        // Now walk each def's params + body with all label names in
271
        // scope so mutual references resolve.
272
        for def in defs {
273
            if let Expr::List(d) = def
274
                && d.len() >= 3
275
            {
276
                let mut def_scope = scope.clone();
277
                if let Expr::List(params) = &d[1] {
278
                    for p in params {
279
                        if let Expr::Symbol(pn) = p {
280
                            def_scope.push(pn.clone());
281
                        }
282
                    }
283
                }
284
                for body_form in &d[2..] {
285
                    walk(symbols, &def_scope, body_form, captures);
286
                }
287
            }
288
        }
289
    }
290
    for body_form in &elems[2..] {
291
        walk(symbols, &scope, body_form, captures);
292
    }
293
}
294

            
295
1
fn walk_iteration(
296
1
    symbols: &SymbolTable,
297
1
    bound: &[String],
298
1
    elems: &[Expr],
299
1
    captures: &mut CaptureSet,
300
1
) {
301
    // (DOLIST (var list-expr [result]) body...) and (DOTIMES (var
302
    // count-expr [result]) body...) — same shape, var is in scope
303
    // only inside body.
304
1
    let Some(spec) = elems.get(1) else {
305
        return;
306
    };
307
1
    let (var, exprs): (Option<String>, Vec<&Expr>) = match spec {
308
1
        Expr::List(parts) => {
309
1
            let var = parts.first().and_then(|p| match p {
310
1
                Expr::Symbol(n) => Some(n.clone()),
311
                _ => None,
312
1
            });
313
1
            (var, parts.iter().skip(1).collect())
314
        }
315
        _ => (None, vec![]),
316
    };
317
1
    for e in exprs {
318
1
        walk(symbols, bound, e, captures);
319
1
    }
320
1
    let mut scope = bound.to_vec();
321
1
    if let Some(v) = var {
322
1
        scope.push(v);
323
1
    }
324
1
    for body_form in &elems[2..] {
325
1
        walk(symbols, &scope, body_form, captures);
326
1
    }
327
1
}
328

            
329
fn walk_do(
330
    symbols: &SymbolTable,
331
    bound: &[String],
332
    elems: &[Expr],
333
    captures: &mut CaptureSet,
334
    sequential: bool,
335
) {
336
    // (DO ((var init step)...) (test result...) body...)
337
    let Some(bindings_expr) = elems.get(1) else {
338
        return;
339
    };
340
    let bindings = match bindings_expr {
341
        Expr::List(b) => b,
342
        _ => return,
343
    };
344

            
345
    let mut scope = bound.to_vec();
346
    for binding in bindings {
347
        if let Expr::List(parts) = binding {
348
            if let Some(Expr::Symbol(name)) = parts.first() {
349
                if let Some(init) = parts.get(1) {
350
                    let init_scope = if sequential { &scope[..] } else { bound };
351
                    walk(symbols, init_scope, init, captures);
352
                }
353
                scope.push(name.clone());
354
                if let Some(step) = parts.get(2) {
355
                    walk(symbols, &scope, step, captures);
356
                }
357
            }
358
        } else if let Expr::Symbol(name) = binding {
359
            scope.push(name.clone());
360
        }
361
    }
362
    if let Some(test_clause) = elems.get(2)
363
        && let Expr::List(parts) = test_clause
364
    {
365
        for p in parts {
366
            walk(symbols, &scope, p, captures);
367
        }
368
    }
369
    for body_form in elems.iter().skip(3) {
370
        walk(symbols, &scope, body_form, captures);
371
    }
372
}