1
//! Runtime-type inference for the let-body's tail value.
2
//!
3
//! `infer_runtime_type_with_body` walks the body for an accumulator
4
//! `(setf <name> (cons V <name>))` and reads the cons car's
5
//! `PairElement` so the promoted runtime type matches what the codegen
6
//! path emits. The walker tracks `DOLIST` / `DO` / `LET` loop-var
7
//! bindings so nested setfs resolve their car symbol against the
8
//! right type. Without this, a wrong-shape fallback surfaces as a
9
//! `ref.cast` mismatch in downstream consumers.
10

            
11
use crate::ast::{Expr, PairElement, WasmType};
12
use crate::compiler::expr::eval_value;
13
use crate::runtime::{Symbol, SymbolKind, SymbolTable};
14

            
15
/// `elems[from..]`, never panicking on a malformed/short binder form: a `(do)`
16
/// / `(dolist)` / `(let)` with fewer than `from` elements yields an empty slice
17
/// rather than an out-of-range slice index. This walk is a best-effort
18
/// accumulator-type inference hint, so a too-short form simply has no body to
19
/// scan; the malformed form is rejected by the real binding-compile path later.
20
4240
fn tail(elems: &[Expr], from: usize) -> &[Expr] {
21
4240
    elems.get(from..).unwrap_or(&[])
22
4240
}
23

            
24
7240
pub(super) fn infer_runtime_type(expr: &Expr) -> WasmType {
25
    // Bytes share StringRef's `i8_array` representation but aren't a
26
    // `compile_for_stack` literal, so they're classified here, not in the
27
    // shared `classify_stack_type`.
28
7240
    if let Expr::Bytes(_) = expr {
29
        return WasmType::StringRef;
30
7240
    }
31
7240
    match crate::compiler::expr::classify_stack_type(expr) {
32
7104
        Some(ty) => ty,
33
        // For any remaining shape (List, Symbol, Quote, etc.) fall
34
        // back to a generic PairRef. PairElement::I32 is the
35
        // least-restrictive guess; downstream consumers refuse it
36
        // explicitly when wrong, so a disagreement surfaces as a
37
        // compile error rather than a runtime trap.
38
136
        None => WasmType::PairRef(PairElement::I32),
39
    }
40
7240
}
41

            
42
/// Runtime local type for a let-var the body MUTATES via `setf`/`set!`, or
43
/// `None` if the var is never assigned. A const-initialized accumulator
44
/// (`(let ((acc 0)) … (setf acc …))` / `(let ((r nil)) … (setf r (cons …)))`)
45
/// must become a real wasm local so `setf` emits `local.set` — otherwise the
46
/// assignment takes the const eval-rebind path (no runtime wasm) and is lost
47
/// across loop scopes (a DO/dolist accumulator silently stays at its init). The
48
/// element type of a cons accumulator comes from the same setf walk the eval
49
/// path uses; any other assigned var sizes from its init value.
50
9748
pub(super) fn assigned_var_runtime_type(
51
9748
    body: &[Expr],
52
9748
    name: &str,
53
9748
    init_val: &Expr,
54
9748
    symbols: &SymbolTable,
55
9748
) -> Option<WasmType> {
56
10428
    if !body.iter().any(|form| form_assigns_var(form, name)) {
57
2856
        return None;
58
6892
    }
59
6892
    let mut env = symbols.clone();
60
6892
    if let Some(elem) = walk_for_setf_pair_element(body, name, &mut env) {
61
1988
        return Some(WasmType::PairRef(elem));
62
4904
    }
63
    // The rhs isn't a CONS accumulator. Size the local from the FIRST `setf`
64
    // rhs that resolves to a concrete RUNTIME type (e.g. `(setf out
65
    // (get-input-entities))` or `(setf out (a-user-fn))` → PairRef) so it
66
    // matches what the setf stores — a `nil` init / a `(setf out nil)` first
67
    // store would otherwise size it as the falsy `Bool` (i32) and the later ref
68
    // store would fail wasm validation. Each rhs is resolved on a CLONE — the
69
    // same depth-guarded const-fold/inline the compiler uses for type inference
70
    // everywhere, so a user-fn rhs is typed by its body's result type (not
71
    // skipped) and a non-terminating recursion surfaces as an `Err` (skipped via
72
    // `.ok()`), never a hang. Const / nil rhs (no runtime type) are skipped; if
73
    // NO store yields a runtime type we fall back to the init's type. A
74
    // genuinely conflicting second runtime store is left for codegen to reject
75
    // (it owns the single local's type).
76
4904
    collect_setf_rhs(body, name)
77
4904
        .into_iter()
78
5108
        .find_map(|rhs| {
79
5108
            eval_value(&mut env.clone(), rhs)
80
5108
                .ok()
81
5108
                .and_then(|v| v.wasm_type())
82
5108
        })
83
4904
        .or_else(|| Some(infer_runtime_type(init_val)))
84
9748
}
85

            
86
/// Every `(setf <name> rhs)` / `(set! <name> rhs)` rhs anywhere in `body`, in
87
/// source order (recursing into nested forms). Used by
88
/// [`assigned_var_runtime_type`] to size a setf-mutated local from its first
89
/// runtime-typed store when it is not a CONS accumulator.
90
4904
fn collect_setf_rhs<'a>(body: &'a [Expr], name: &str) -> Vec<&'a Expr> {
91
4904
    let mut out = Vec::new();
92
4904
    collect_setf_rhs_into(body, name, &mut out);
93
4904
    out
94
4904
}
95

            
96
39390
fn collect_setf_rhs_into<'a>(body: &'a [Expr], name: &str, out: &mut Vec<&'a Expr>) {
97
106098
    for form in body {
98
106098
        let Expr::List(elems) = form else { continue };
99
34486
        if let Some(Expr::Symbol(head)) = elems.first()
100
32850
            && (head == "SETF" || head == "SET!")
101
        {
102
6332
            for pair in elems[1..].chunks(2) {
103
6332
                if let (Some(Expr::Symbol(place)), Some(rhs)) = (pair.first(), pair.get(1))
104
6332
                    && place == name
105
5244
                {
106
5244
                    out.push(rhs);
107
5244
                }
108
            }
109
28222
        }
110
        // Don't descend into a binder that rebinds `name` — its body's setfs
111
        // target a DIFFERENT (shadowing) variable.
112
34486
        if !binder_rebinds(elems, name) {
113
34486
            collect_setf_rhs_into(elems, name, out);
114
34486
        }
115
    }
116
39390
}
117

            
118
/// Recursively detects a `(setf <name> …)` / `(set! <name> …)` targeting
119
/// `name` anywhere in `form`, NOT descending into a sub-form that rebinds
120
/// `name` (its inner setfs target a shadowing variable).
121
125928
fn form_assigns_var(form: &Expr, name: &str) -> bool {
122
125928
    let Expr::List(elems) = form else {
123
74582
        return false;
124
    };
125
51346
    if let Some(Expr::Symbol(head)) = elems.first()
126
46758
        && (head == "SETF" || head == "SET!")
127
    {
128
7368
        for pair in elems[1..].chunks(2) {
129
7368
            if let Some(Expr::Symbol(place)) = pair.first()
130
7368
                && place == name
131
            {
132
6892
                return true;
133
476
            }
134
        }
135
43978
    }
136
44454
    if binder_rebinds(elems, name) {
137
204
        return false;
138
44250
    }
139
115500
    elems.iter().any(|e| form_assigns_var(e, name))
140
125928
}
141

            
142
/// Whether the list form headed by `elems` binds `name` (a lambda param, or a
143
/// `let`/`let*`/`do`/`do*`/`dolist` variable), shadowing an enclosing binding.
144
/// Used to keep the setf-promotion scans from attributing a shadowed inner
145
/// assignment to the outer binding. Mirrors `lambda::param_infer`'s sibling.
146
120192
fn binder_rebinds(elems: &[Expr], name: &str) -> bool {
147
120192
    let Some(Expr::Symbol(head)) = elems.first() else {
148
6428
        return false;
149
    };
150
113764
    let in_specs = |spec: Option<&Expr>| {
151
6222
        matches!(spec, Some(Expr::List(items)) if items.iter().any(|b| match b {
152
476
            Expr::Symbol(s) => s == name,
153
5746
            Expr::List(parts) => matches!(parts.first(), Some(Expr::Symbol(s)) if s == name),
154
            _ => false,
155
6222
        }))
156
5950
    };
157
113764
    match head.to_ascii_uppercase().as_str() {
158
113764
        "LAMBDA" | "LET" | "LET*" | "DO" | "DO*" => in_specs(elems.get(1)),
159
107814
        "DOLIST" => matches!(
160
5530
            elems.get(1),
161
5530
            Some(Expr::List(spec)) if matches!(spec.first(), Some(Expr::Symbol(s)) if s == name)
162
        ),
163
102284
        _ => false,
164
    }
165
120192
}
166

            
167
/// When the body's tail returns a Symbol whose value flows from a
168
/// `(setf <name> (cons V <name>))` accumulator earlier in the body,
169
/// resolve the inferred runtime type by walking the body for that
170
/// setf and reading the cons car's PairElement. The walker tracks
171
/// dolist/do loop-var bindings so a setf nested inside a dolist
172
/// resolves the loop var's element type (e.g. `tag` inside `(dolist
173
/// (tag (entities-for ...)) ...)` resolves via the dolist's list
174
/// expression's PairElement). Without this, the cons car's symbol
175
/// would be unbound in the eval-time symbol table.
176
612
pub(super) fn infer_runtime_type_with_body(
177
612
    val: &Expr,
178
612
    result_expr: &Expr,
179
612
    body: &[Expr],
180
612
    symbols: &SymbolTable,
181
612
) -> WasmType {
182
612
    if let Expr::Symbol(name) = result_expr {
183
136
        let mut env = symbols.clone();
184
136
        if let Some(elem) = walk_for_setf_pair_element(body, name, &mut env) {
185
            return WasmType::PairRef(elem);
186
136
        }
187
476
    }
188
612
    infer_runtime_type(val)
189
612
}
190

            
191
/// Walks `body` for `(setf <name> <rhs>)` and infers the matching
192
/// PairElement from the cons car of the rhs. Tracks DOLIST/DO/LET
193
/// loop-var bindings into `env` so nested setfs resolve their car
194
/// symbol against the right type.
195
7028
fn walk_for_setf_pair_element(
196
7028
    body: &[Expr],
197
7028
    target: &str,
198
7028
    env: &mut SymbolTable,
199
7028
) -> Option<PairElement> {
200
11386
    for form in body {
201
11386
        if let Some(elem) = walk_form(form, target, env) {
202
1988
            return Some(elem);
203
9398
        }
204
    }
205
5040
    None
206
7028
}
207

            
208
109156
fn walk_form(form: &Expr, target: &str, env: &mut SymbolTable) -> Option<PairElement> {
209
109156
    let Expr::List(elems) = form else {
210
67904
        return None;
211
    };
212
    // A binder that rebinds `target` shadows it: any `(setf target …)` inside
213
    // refers to the inner variable, not the accumulator we're typing.
214
41252
    if binder_rebinds(elems, target) {
215
        return None;
216
41252
    }
217
41252
    let head = elems.first().and_then(|e| match e {
218
41048
        Expr::Symbol(s) => Some(s.as_str()),
219
204
        _ => None,
220
41252
    });
221
41252
    match head {
222
41048
        Some("SETF") => {
223
8320
            for pair in elems[1..].chunks(2) {
224
8320
                if let Some(Expr::Symbol(var)) = pair.first()
225
8320
                    && var == target
226
7232
                    && let Some(rhs) = pair.get(1)
227
7232
                    && let Some(elem) = step_pair_element(rhs, env)
228
                {
229
1988
                    return Some(elem);
230
6332
                }
231
            }
232
12664
            for e in &elems[1..] {
233
12664
                if let Some(found) = walk_form(e, target, env) {
234
                    return Some(found);
235
12664
                }
236
            }
237
6264
            None
238
        }
239
32796
        Some("DOLIST") => {
240
1912
            if let Some(Expr::List(var_spec)) = elems.get(1)
241
1912
                && let Some(Expr::Symbol(var_name)) = var_spec.first()
242
1912
                && let Some(list_expr) = var_spec.get(1)
243
            {
244
1912
                let mut probe = env.clone();
245
1912
                if let Ok(list_val) = eval_value(&mut probe, list_expr)
246
1912
                    && let Some(WasmType::PairRef(elem)) = list_val.wasm_type()
247
1368
                {
248
1368
                    env.define(
249
1368
                        Symbol::new(var_name, SymbolKind::Variable)
250
1368
                            .with_value(Expr::WasmRuntime(elem.as_wasm_type())),
251
1368
                    );
252
1368
                }
253
            }
254
2048
            for e in tail(elems, 2) {
255
2048
                if let Some(found) = walk_form(e, target, env) {
256
206
                    return Some(found);
257
1842
                }
258
            }
259
1706
            None
260
        }
261
30884
        Some("DO") | Some("DO*") => {
262
2124
            if let Some(Expr::List(vars)) = elems.get(1) {
263
2124
                for v in vars {
264
2124
                    if let Expr::List(vparts) = v
265
2124
                        && let Some(Expr::Symbol(vname)) = vparts.first()
266
2124
                        && let Some(init) = vparts.get(1)
267
                    {
268
2124
                        let mut probe = env.clone();
269
2124
                        if let Ok(init_val) = eval_value(&mut probe, init) {
270
2124
                            // A DO var is a RUNTIME value of the init's stack
271
2124
                            // type, not the const init itself: an integer init
272
2124
                            // (`(i 0 …)`) makes `i` a runtime Index (I32), so
273
2124
                            // `(cons i …)` is an I32 cell. Binding the const `0`
274
2124
                            // would resolve its pair element as Ratio (the
275
2124
                            // numeric-literal cell slot) and mis-size the
276
2124
                            // accumulator → a CAR ref.cast trap.
277
2124
                            env.define(
278
2124
                                Symbol::new(vname, SymbolKind::Variable)
279
2124
                                    .with_value(Expr::WasmRuntime(infer_runtime_type(&init_val))),
280
2124
                            );
281
2124
                        }
282
                    }
283
                }
284
            }
285
2124
            for e in tail(elems, 3) {
286
2124
                if let Some(found) = walk_form(e, target, env) {
287
1510
                    return Some(found);
288
614
                }
289
            }
290
614
            None
291
        }
292
28760
        Some("LET") | Some("LET*") => {
293
204
            if let Some(Expr::List(bindings)) = elems.get(1) {
294
204
                for b in bindings {
295
204
                    if let Expr::List(parts) = b
296
204
                        && let Some(Expr::Symbol(bname)) = parts.first()
297
                    {
298
204
                        let init_val = match parts.get(1) {
299
204
                            Some(init) => {
300
204
                                let mut probe = env.clone();
301
204
                                eval_value(&mut probe, init).unwrap_or(Expr::Nil)
302
                            }
303
                            None => Expr::Nil,
304
                        };
305
204
                        env.define(Symbol::new(bname, SymbolKind::Variable).with_value(init_val));
306
                    }
307
                }
308
            }
309
204
            for e in tail(elems, 2) {
310
204
                if let Some(found) = walk_form(e, target, env) {
311
                    return Some(found);
312
204
                }
313
            }
314
204
            None
315
        }
316
        _ => {
317
80730
            for e in elems {
318
80730
                if let Some(found) = walk_form(e, target, env) {
319
1716
                    return Some(found);
320
79014
                }
321
            }
322
27044
            None
323
        }
324
    }
325
109156
}
326

            
327
/// Inline copy of `iteration::common::step_pair_element` to avoid a
328
/// cyclic `pub(super)` exposure. Same semantics: peeks at a CONS
329
/// expression's car and resolves the corresponding PairElement.
330
18612
fn step_pair_element(expr: &Expr, symbols: &SymbolTable) -> Option<PairElement> {
331
18612
    let Expr::List(elems) = expr else {
332
12808
        return None;
333
    };
334
5804
    if let Some(Expr::Symbol(name)) = elems.first()
335
5804
        && name == "CONS"
336
1988
        && elems.len() == 3
337
1988
        && let Some(elem) = resolve_pair_element(&elems[1], symbols)
338
    {
339
1988
        return Some(elem);
340
3816
    }
341
11380
    elems.iter().find_map(|e| step_pair_element(e, symbols))
342
18612
}
343

            
344
3704
fn resolve_pair_element(expr: &Expr, symbols: &SymbolTable) -> Option<PairElement> {
345
    match expr {
346
1716
        Expr::WasmRuntime(ty) | Expr::WasmLocal(_, ty) => PairElement::from_wasm_type(*ty),
347
        // A numeric literal is a dimensionless Ratio (never a count), matching
348
        // `list::infer::literal_pair_element` and `infer_runtime_type`.
349
272
        Expr::Number(_) => Some(PairElement::Ratio),
350
        Expr::Bool(_) => Some(PairElement::Bool),
351
1716
        Expr::Symbol(name) => {
352
1716
            let sym = symbols.lookup(name)?;
353
1716
            let value = sym.value()?;
354
1716
            resolve_pair_element(value, symbols)
355
        }
356
        // A host-fn / form car hasn't reduced to a runtime placeholder yet;
357
        // eval on a CLONE so the accumulator element matches codegen. Admitted
358
        // only when the whole car subtree is pure-native (no user-callable
359
        // anywhere), so inference can't recurse into a (possibly
360
        // non-terminating) user lambda — native handlers eval their args, so a
361
        // top-level-head check alone wouldn't be enough.
362
        Expr::List(_) if is_pure_native_expr(expr, symbols) => {
363
            match eval_value(&mut symbols.clone(), expr).ok()? {
364
                Expr::List(_) => None,
365
                other => resolve_pair_element(&other, symbols),
366
            }
367
        }
368
        _ => None,
369
    }
370
3704
}
371

            
372
/// Whether `expr` is safe to `eval_value` purely for type inference: it
373
/// contains NO user-callable anywhere. Native eval handlers recurse into their
374
/// args, and higher-order natives (MAP/FILTER/FOLD) invoke a function-valued
375
/// argument, so the check spans the whole tree — every call head must be a
376
/// native/operator (no user shadow) and every operand pure-native, and a bare
377
/// symbol is admitted only when it doesn't name a user function/lambda value.
378
/// Conservative: any non-inert shape is rejected. Mirrors
379
/// `iteration::common::is_pure_native_expr`.
380
fn is_pure_native_expr(expr: &Expr, symbols: &SymbolTable) -> bool {
381
    match expr {
382
        Expr::Number(_)
383
        | Expr::Bool(_)
384
        | Expr::Nil
385
        | Expr::String(_)
386
        | Expr::Bytes(_)
387
        | Expr::Keyword(_)
388
        | Expr::Quote(_)
389
        | Expr::WasmRuntime(_)
390
        | Expr::WasmLocal(_, _) => true,
391
        Expr::Symbol(name) => !names_user_callable(name, symbols),
392
        Expr::List(elems) => match elems.first() {
393
            Some(Expr::Symbol(head)) => {
394
                head_is_native(head, symbols)
395
                    && elems[1..].iter().all(|e| is_pure_native_expr(e, symbols))
396
            }
397
            _ => false,
398
        },
399
        _ => false,
400
    }
401
}
402

            
403
fn head_is_native(head: &str, symbols: &SymbolTable) -> bool {
404
    match symbols.lookup(head) {
405
        Some(sym) => {
406
            sym.function().is_none()
407
                && matches!(sym.kind(), SymbolKind::Native | SymbolKind::Operator)
408
        }
409
        None => false,
410
    }
411
}
412

            
413
/// Whether `name` resolves to a user function / lambda value — passed to a
414
/// higher-order native it would be dereferenced and called during
415
/// `eval_value`, so it is NOT inert. Mirrors `iteration::common`'s sibling.
416
fn names_user_callable(name: &str, symbols: &SymbolTable) -> bool {
417
    match symbols.lookup(name) {
418
        Some(sym) => {
419
            sym.function().is_some()
420
                || matches!(sym.value(), Some(Expr::Lambda(_, _)))
421
                || matches!(sym.kind(), SymbolKind::Function | SymbolKind::Macro)
422
        }
423
        None => false,
424
    }
425
}