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
4402
fn tail(elems: &[Expr], from: usize) -> &[Expr] {
21
4402
    elems.get(from..).unwrap_or(&[])
22
4402
}
23

            
24
7534
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
7534
    if let Expr::Bytes(_) = expr {
29
        return WasmType::StringRef;
30
7534
    }
31
7534
    match crate::compiler::expr::classify_stack_type(expr) {
32
7392
        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
142
        None => WasmType::PairRef(PairElement::I32),
39
    }
40
7534
}
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
10153
pub(super) fn assigned_var_runtime_type(
51
10153
    body: &[Expr],
52
10153
    name: &str,
53
10153
    init_val: &Expr,
54
10153
    symbols: &SymbolTable,
55
10153
) -> Option<WasmType> {
56
10863
    if !body.iter().any(|form| form_assigns_var(form, name)) {
57
2982
        return None;
58
7171
    }
59
7171
    let mut env = symbols.clone();
60
7171
    if let Some(elem) = walk_for_setf_pair_element(body, name, &mut env) {
61
2059
        return Some(WasmType::PairRef(elem));
62
5112
    }
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
5112
    collect_setf_rhs(body, name)
77
5112
        .into_iter()
78
5325
        .find_map(|rhs| {
79
5325
            eval_value(&mut env.clone(), rhs)
80
5325
                .ok()
81
5325
                .and_then(|v| v.wasm_type())
82
5325
        })
83
5112
        .or_else(|| Some(infer_runtime_type(init_val)))
84
10153
}
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
5112
fn collect_setf_rhs<'a>(body: &'a [Expr], name: &str) -> Vec<&'a Expr> {
91
5112
    let mut out = Vec::new();
92
5112
    collect_setf_rhs_into(body, name, &mut out);
93
5112
    out
94
5112
}
95

            
96
41038
fn collect_setf_rhs_into<'a>(body: &'a [Expr], name: &str, out: &mut Vec<&'a Expr>) {
97
110547
    for form in body {
98
110547
        let Expr::List(elems) = form else { continue };
99
35926
        if let Some(Expr::Symbol(head)) = elems.first()
100
34222
            && (head == "SETF" || head == "SET!")
101
        {
102
6603
            for pair in elems[1..].chunks(2) {
103
6603
                if let (Some(Expr::Symbol(place)), Some(rhs)) = (pair.first(), pair.get(1))
104
6603
                    && place == name
105
5467
                {
106
5467
                    out.push(rhs);
107
5467
                }
108
            }
109
29394
        }
110
        // Don't descend into a binder that rebinds `name` — its body's setfs
111
        // target a DIFFERENT (shadowing) variable.
112
35926
        if !binder_rebinds(elems, name) {
113
35926
            collect_setf_rhs_into(elems, name, out);
114
35926
        }
115
    }
116
41038
}
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
130782
fn form_assigns_var(form: &Expr, name: &str) -> bool {
122
130782
    let Expr::List(elems) = form else {
123
77461
        return false;
124
    };
125
53321
    if let Some(Expr::Symbol(head)) = elems.first()
126
48564
        && (head == "SETF" || head == "SET!")
127
    {
128
7668
        for pair in elems[1..].chunks(2) {
129
7668
            if let Some(Expr::Symbol(place)) = pair.first()
130
7668
                && place == name
131
            {
132
7171
                return true;
133
497
            }
134
        }
135
45653
    }
136
46150
    if binder_rebinds(elems, name) {
137
213
        return false;
138
45937
    }
139
119919
    elems.iter().any(|e| form_assigns_var(e, name))
140
130782
}
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
124960
fn binder_rebinds(elems: &[Expr], name: &str) -> bool {
147
124960
    let Some(Expr::Symbol(head)) = elems.first() else {
148
6674
        return false;
149
    };
150
118286
    let in_specs = |spec: Option<&Expr>| {
151
6461
        matches!(spec, Some(Expr::List(items)) if items.iter().any(|b| match b {
152
497
            Expr::Symbol(s) => s == name,
153
5964
            Expr::List(parts) => matches!(parts.first(), Some(Expr::Symbol(s)) if s == name),
154
            _ => false,
155
6461
        }))
156
6177
    };
157
118286
    match head.to_ascii_uppercase().as_str() {
158
118286
        "LAMBDA" | "LET" | "LET*" | "DO" | "DO*" => in_specs(elems.get(1)),
159
112109
        "DOLIST" => matches!(
160
5751
            elems.get(1),
161
5751
            Some(Expr::List(spec)) if matches!(spec.first(), Some(Expr::Symbol(s)) if s == name)
162
        ),
163
106358
        _ => false,
164
    }
165
124960
}
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
639
pub(super) fn infer_runtime_type_with_body(
177
639
    val: &Expr,
178
639
    result_expr: &Expr,
179
639
    body: &[Expr],
180
639
    symbols: &SymbolTable,
181
639
) -> WasmType {
182
639
    if let Expr::Symbol(name) = result_expr {
183
142
        let mut env = symbols.clone();
184
142
        if let Some(elem) = walk_for_setf_pair_element(body, name, &mut env) {
185
            return WasmType::PairRef(elem);
186
142
        }
187
497
    }
188
639
    infer_runtime_type(val)
189
639
}
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
7313
fn walk_for_setf_pair_element(
196
7313
    body: &[Expr],
197
7313
    target: &str,
198
7313
    env: &mut SymbolTable,
199
7313
) -> Option<PairElement> {
200
11857
    for form in body {
201
11857
        if let Some(elem) = walk_form(form, target, env) {
202
2059
            return Some(elem);
203
9798
        }
204
    }
205
5254
    None
206
7313
}
207

            
208
113529
fn walk_form(form: &Expr, target: &str, env: &mut SymbolTable) -> Option<PairElement> {
209
113529
    let Expr::List(elems) = form else {
210
70645
        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
42884
    if binder_rebinds(elems, target) {
215
        return None;
216
42884
    }
217
42884
    let head = elems.first().and_then(|e| match e {
218
42671
        Expr::Symbol(s) => Some(s.as_str()),
219
213
        _ => None,
220
42884
    });
221
42884
    match head {
222
42671
        Some("SETF") => {
223
8662
            for pair in elems[1..].chunks(2) {
224
8662
                if let Some(Expr::Symbol(var)) = pair.first()
225
8662
                    && var == target
226
7526
                    && let Some(rhs) = pair.get(1)
227
7526
                    && let Some(elem) = step_pair_element(rhs, env)
228
                {
229
2059
                    return Some(elem);
230
6603
                }
231
            }
232
13206
            for e in &elems[1..] {
233
13206
                if let Some(found) = walk_form(e, target, env) {
234
                    return Some(found);
235
13206
                }
236
            }
237
6532
            None
238
        }
239
34080
        Some("DOLIST") => {
240
1988
            if let Some(Expr::List(var_spec)) = elems.get(1)
241
1988
                && let Some(Expr::Symbol(var_name)) = var_spec.first()
242
1988
                && let Some(list_expr) = var_spec.get(1)
243
            {
244
1988
                let mut probe = env.clone();
245
1988
                if let Ok(list_val) = eval_value(&mut probe, list_expr)
246
1988
                    && let Some(WasmType::PairRef(elem)) = list_val.wasm_type()
247
1420
                {
248
1420
                    env.define(
249
1420
                        Symbol::new(var_name, SymbolKind::Variable)
250
1420
                            .with_value(Expr::WasmRuntime(elem.as_wasm_type())),
251
1420
                    );
252
1420
                }
253
            }
254
2130
            for e in tail(elems, 2) {
255
2130
                if let Some(found) = walk_form(e, target, env) {
256
213
                    return Some(found);
257
1917
                }
258
            }
259
1775
            None
260
        }
261
32092
        Some("DO") | Some("DO*") => {
262
2201
            if let Some(Expr::List(vars)) = elems.get(1) {
263
2201
                for v in vars {
264
2201
                    if let Expr::List(vparts) = v
265
2201
                        && let Some(Expr::Symbol(vname)) = vparts.first()
266
2201
                        && let Some(init) = vparts.get(1)
267
                    {
268
2201
                        let mut probe = env.clone();
269
2201
                        if let Ok(init_val) = eval_value(&mut probe, init) {
270
2201
                            // A DO var is a RUNTIME value of the init's stack
271
2201
                            // type, not the const init itself: an integer init
272
2201
                            // (`(i 0 …)`) makes `i` a runtime Index (I32), so
273
2201
                            // `(cons i …)` is an I32 cell. Binding the const `0`
274
2201
                            // would resolve its pair element as Ratio (the
275
2201
                            // numeric-literal cell slot) and mis-size the
276
2201
                            // accumulator → a CAR ref.cast trap.
277
2201
                            env.define(
278
2201
                                Symbol::new(vname, SymbolKind::Variable)
279
2201
                                    .with_value(Expr::WasmRuntime(infer_runtime_type(&init_val))),
280
2201
                            );
281
2201
                        }
282
                    }
283
                }
284
            }
285
2201
            for e in tail(elems, 3) {
286
2201
                if let Some(found) = walk_form(e, target, env) {
287
1562
                    return Some(found);
288
639
                }
289
            }
290
639
            None
291
        }
292
29891
        Some("LET") | Some("LET*") => {
293
213
            if let Some(Expr::List(bindings)) = elems.get(1) {
294
213
                for b in bindings {
295
213
                    if let Expr::List(parts) = b
296
213
                        && let Some(Expr::Symbol(bname)) = parts.first()
297
                    {
298
213
                        let init_val = match parts.get(1) {
299
213
                            Some(init) => {
300
213
                                let mut probe = env.clone();
301
213
                                eval_value(&mut probe, init).unwrap_or(Expr::Nil)
302
                            }
303
                            None => Expr::Nil,
304
                        };
305
213
                        env.define(Symbol::new(bname, SymbolKind::Variable).with_value(init_val));
306
                    }
307
                }
308
            }
309
213
            for e in tail(elems, 2) {
310
213
                if let Some(found) = walk_form(e, target, env) {
311
                    return Some(found);
312
213
                }
313
            }
314
213
            None
315
        }
316
        _ => {
317
83922
            for e in elems {
318
83922
                if let Some(found) = walk_form(e, target, env) {
319
1775
                    return Some(found);
320
82147
                }
321
            }
322
28116
            None
323
        }
324
    }
325
113529
}
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
19383
fn step_pair_element(expr: &Expr, symbols: &SymbolTable) -> Option<PairElement> {
331
19383
    let Expr::List(elems) = expr else {
332
13348
        return None;
333
    };
334
6035
    if let Some(Expr::Symbol(name)) = elems.first()
335
6035
        && name == "CONS"
336
2059
        && elems.len() == 3
337
2059
        && let Some(elem) = resolve_pair_element(&elems[1], symbols)
338
    {
339
2059
        return Some(elem);
340
3976
    }
341
11857
    elems.iter().find_map(|e| step_pair_element(e, symbols))
342
19383
}
343

            
344
3834
fn resolve_pair_element(expr: &Expr, symbols: &SymbolTable) -> Option<PairElement> {
345
    match expr {
346
1775
        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
284
        Expr::Number(_) => Some(PairElement::Ratio),
350
        Expr::Bool(_) => Some(PairElement::Bool),
351
1775
        Expr::Symbol(name) => {
352
1775
            let sym = symbols.lookup(name)?;
353
1775
            let value = sym.value()?;
354
1775
            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
3834
}
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
}