1
//! Shared types and helpers across the DO/DO*/DOLIST family.
2
//!
3
//! Three concerns live here:
4
//! - parsing the form's surface syntax (`parse_do_vars`,
5
//!   `parse_end_clause`, `DoVar`, `DoLoop`);
6
//! - static-time termination + type inference
7
//!   (`static_loop_terminates`, `infer_wasm_type`,
8
//!   `infer_result_pair_element`, `step_pair_element`);
9
//! - the codegen-side setf-target machinery shared by all three
10
//!   constructs (`collect_setf_targets`, `promote_to_wasm_local`).
11
//!
12
//! Everything is `pub(super)` so the sibling `do_loop` and `dolist`
13
//! modules can use it; nothing escapes to the broader codebase.
14

            
15
use crate::ast::{Expr, PairElement, WasmType};
16
use crate::compiler::context::CompileContext;
17
use crate::compiler::emit::FunctionEmitter;
18
use crate::compiler::expr::{compile_for_stack, eval_value};
19
use crate::error::{Error, Result};
20
use crate::runtime::{Symbol, SymbolKind, SymbolTable};
21

            
22
/// Per-binding spec parsed from the variable list of a DO/DO* form.
23
/// `(name init step)` — both `init` and `step` are optional in source
24
/// syntax; the runtime treats absent init as `nil` and absent step as
25
/// "leave the binding alone for this iteration."
26
pub(super) struct DoVar {
27
    pub(super) name: String,
28
    pub(super) init: Option<Expr>,
29
    pub(super) step: Option<Expr>,
30
}
31

            
32
/// Owns the raw shape of a do-loop after the surface form has been
33
/// validated. Reused by both the runtime codegen entry points
34
/// (`compile_do_runtime{,_for_effect,_for_stack}`) and the iterators
35
/// inside each variant.
36
pub(super) struct DoLoop<'a> {
37
    pub(super) vars: &'a [DoVar],
38
    pub(super) end_test: &'a Expr,
39
    pub(super) result_forms: &'a [Expr],
40
    pub(super) body: &'a [Expr],
41
    /// True for DO* (sequential steps), false for DO (parallel).
42
    pub(super) sequential: bool,
43
}
44

            
45
pub(super) const MAX_STATIC_LOOP_ITERS: usize = 64;
46

            
47
11209
pub(super) fn parse_do_vars(context: &str, expr: &Expr) -> Result<Vec<DoVar>> {
48
11209
    let list = expr.as_list().ok_or_else(|| {
49
        Error::Compile(format!("{context}: expected variable list, got {expr:?}"))
50
    })?;
51
11209
    list.iter()
52
13782
        .map(|spec| match spec {
53
            Expr::Symbol(name) => Ok(DoVar {
54
                name: name.clone(),
55
                init: None,
56
                step: None,
57
            }),
58
13782
            Expr::List(elems) if !elems.is_empty() && elems.len() <= 3 => {
59
13782
                let name = elems[0].as_symbol().ok_or_else(|| {
60
                    Error::Compile(format!(
61
                        "{context}: variable name must be a symbol, got {:?}",
62
                        elems[0]
63
                    ))
64
                })?;
65
13782
                Ok(DoVar {
66
13782
                    name: name.to_string(),
67
13782
                    init: elems.get(1).cloned(),
68
13782
                    step: elems.get(2).cloned(),
69
13782
                })
70
            }
71
            _ => Err(Error::Compile(format!(
72
                "{context}: malformed variable spec: {spec:?}"
73
            ))),
74
13782
        })
75
11209
        .collect()
76
11209
}
77

            
78
11209
pub(super) fn parse_end_clause<'a>(
79
11209
    context: &str,
80
11209
    expr: &'a Expr,
81
11209
) -> Result<(&'a Expr, &'a [Expr])> {
82
11209
    let elems = expr.as_list().ok_or_else(|| {
83
        Error::Compile(format!(
84
            "{context}: end clause must be a list, got {expr:?}"
85
        ))
86
    })?;
87
11209
    if elems.is_empty() {
88
        return Err(Error::Compile(format!(
89
            "{context}: end clause must have a test form"
90
        )));
91
11209
    }
92
11209
    Ok((&elems[0], &elems[1..]))
93
11209
}
94

            
95
16030
pub(super) fn infer_wasm_type(
96
16030
    expr: &Expr,
97
16030
    step_hint: Option<&Expr>,
98
16030
    symbols: &SymbolTable,
99
16030
) -> WasmType {
100
    // A `nil` DO-var init whose step builds a `(cons V acc)` accumulator is
101
    // sized from the cons car's element — checked before the shared classifier
102
    // (which would type a bare nil as `Bool`).
103
16030
    if let Expr::Nil = expr
104
1828
        && let Some(step) = step_hint
105
1827
        && let Some(elem) = step_pair_element(step, symbols)
106
    {
107
1163
        return WasmType::PairRef(elem);
108
14867
    }
109
    // A `nil`-init var whose step assigns a runtime value of some OTHER type —
110
    // e.g. `(setf out (get-input-entities))` (PairRef) or `(setf out (some-
111
    // ratio))` — must be sized from the step's actual produced type, not the
112
    // `Bool` a bare nil classifies to. Otherwise the local is sized i32-shaped
113
    // while the step pushes a ref, and the emitted wasm fails validation. Only
114
    // a pure-native step is probed (eval on a CLONE — no live-table mutation,
115
    // no recursion into user code), matching `resolve_pair_element_from_expr`.
116
14867
    if let Expr::Nil = expr
117
665
        && let Some(step) = step_hint
118
664
        && is_pure_native_expr(step, symbols)
119
415
        && let Ok(resolved) = eval_value(&mut symbols.clone(), step)
120
166
        && let Some(ty) = crate::compiler::expr::classify_stack_type(&resolved)
121
    {
122
166
        return ty;
123
14701
    }
124
    // Bytes share StringRef's representation but aren't a `compile_for_stack`
125
    // literal, so they're classified here rather than in `classify_stack_type`.
126
14701
    if let Expr::Bytes(_) = expr {
127
        return WasmType::StringRef;
128
14701
    }
129
    // Everything else mirrors `compile_for_stack` via the shared classifier;
130
    // List / Quote / Symbol / Lambda fall back to `I32` (the historical
131
    // integer-counter default — the do-var's own `compile_for_stack` declares
132
    // its actual type and a mismatch surfaces as a hard emit error).
133
14701
    crate::compiler::expr::classify_stack_type(expr).unwrap_or(WasmType::I32)
134
16030
}
135

            
136
/// Looks for an accumulator `(cons V <result-name>)` that determines
137
/// the runtime `PairElement` of the do's result. Two patterns: body
138
/// `(setf acc (cons V acc))` and step `(acc nil (cons V acc))` (or
139
/// nested via if/cond). Without inspecting the step path, a do whose
140
/// accumulator lives in its var-step gets the placeholder I32, and
141
/// consumer dolists downcast to i31 — runtime cast trap.
142
5067
pub(super) fn infer_result_pair_element(
143
5067
    result_form: &Expr,
144
5067
    body: &[Expr],
145
5067
    stepped: &[(String, Option<Expr>)],
146
5067
    symbols: &SymbolTable,
147
5067
) -> Option<PairElement> {
148
5067
    let Expr::Symbol(target) = result_form else {
149
        return None;
150
    };
151
5067
    let from_body = collect_setf_targets(body)
152
5067
        .into_iter()
153
5067
        .find_map(|(name, rhs)| {
154
3822
            (name == *target)
155
3822
                .then_some(rhs)
156
3822
                .flatten()
157
3822
                .as_ref()
158
3822
                .and_then(|r| step_pair_element(r, symbols))
159
3822
        });
160
5067
    if from_body.is_some() {
161
2825
        return from_body;
162
2242
    }
163
3155
    stepped.iter().find_map(|(name, step)| {
164
3155
        if name != target {
165
1910
            return None;
166
1245
        }
167
1245
        step.as_ref().and_then(|s| step_pair_element(s, symbols))
168
3155
    })
169
5067
}
170

            
171
/// If `step` is (or contains) a `(CONS car cdr)` whose car has a
172
/// derivable wasm type, returns the corresponding `PairElement`. Used by
173
/// nil-init type inference so a let-bound accumulator like
174
/// `(setf result (cons i result))` picks up `i`'s element type rather
175
/// than defaulting to a placeholder.
176
17853
fn step_pair_element(expr: &Expr, symbols: &SymbolTable) -> Option<PairElement> {
177
17853
    let Expr::List(elems) = expr else {
178
8967
        return None;
179
    };
180
8886
    if let Some(Expr::Symbol(name)) = elems.first()
181
8720
        && name == "CONS"
182
5399
        && elems.len() == 3
183
5399
        && let Some(elem) = resolve_pair_element_from_expr(&elems[1], symbols)
184
    {
185
4735
        return Some(elem);
186
4151
    }
187
10959
    elems.iter().find_map(|e| step_pair_element(e, symbols))
188
17853
}
189

            
190
10132
fn resolve_pair_element_from_expr(expr: &Expr, symbols: &SymbolTable) -> Option<PairElement> {
191
1660
    match expr {
192
4733
        Expr::WasmRuntime(ty) | Expr::WasmLocal(_, ty) => PairElement::from_wasm_type(*ty),
193
1
        Expr::Number(_) => crate::compiler::native::literal_pair_element(expr),
194
1
        Expr::T => Some(PairElement::Bool),
195
3737
        Expr::Symbol(name) => {
196
3737
            let sym = symbols.lookup(name)?;
197
3488
            let value = sym.value()?;
198
3488
            resolve_pair_element_from_expr(value, symbols)
199
        }
200
        // A host-fn / form car (e.g. `(transaction-tag-count i)`) hasn't been
201
        // reduced to a runtime placeholder yet. Eval it on a CLONE (pure — no
202
        // live-table mutation) so the accumulator's element type matches what
203
        // codegen will push, instead of falling through to the scalar default.
204
        // Gated on a native/operator head: those reduce to a runtime
205
        // placeholder without executing user code, so this can't recurse into
206
        // a (possibly non-terminating) user lambda during type inference.
207
1660
        Expr::List(_) if is_pure_native_expr(expr, symbols) => {
208
1245
            let resolved = eval_value(&mut symbols.clone(), expr).ok()?;
209
1245
            match resolved {
210
                Expr::List(_) => None,
211
1245
                other => resolve_pair_element_from_expr(&other, symbols),
212
            }
213
        }
214
415
        _ => None,
215
    }
216
10132
}
217

            
218
/// Whether `expr` is safe to `eval_value` purely for type inference — i.e.
219
/// it contains NO user-callable anywhere in its tree. Native eval handlers
220
/// (e.g. `+`) recurse into their arguments, and higher-order natives
221
/// (MAP/FILTER/FOLD) invoke a function-valued argument, so checking only the
222
/// outer head is insufficient: `(+ 1 (user-recursive i))` or
223
/// `(map (function user-recursive) xs)` could still execute user code and
224
/// non-terminate. This walks the WHOLE expression and admits it only when
225
/// every call head AND every operand is itself pure-native. A bare symbol is
226
/// admitted only when it does not name a user function/lambda value (a
227
/// fn-valued symbol could be dereferenced+called by a HOF native); literals
228
/// and runtime placeholders are always inert. Conservative by construction —
229
/// any shape not explicitly inert (Lambda, Quote, FUNCTION-forms, …) is
230
/// rejected. Mirrored in `binding::infer`.
231
6391
fn is_pure_native_expr(expr: &Expr, symbols: &SymbolTable) -> bool {
232
6391
    match expr {
233
        Expr::Number(_)
234
        | Expr::T
235
        | Expr::Nil
236
        | Expr::String(_)
237
        | Expr::Bytes(_)
238
        | Expr::Keyword(_)
239
        | Expr::Quote(_)
240
        | Expr::WasmRuntime(_)
241
2241
        | Expr::WasmLocal(_, _) => true,
242
498
        Expr::Symbol(name) => !names_user_callable(name, symbols),
243
3652
        Expr::List(elems) => match elems.first() {
244
3652
            Some(Expr::Symbol(head)) => {
245
3652
                head_is_native(head, symbols)
246
4067
                    && elems[1..].iter().all(|e| is_pure_native_expr(e, symbols))
247
            }
248
            _ => false,
249
        },
250
        _ => false,
251
    }
252
6391
}
253

            
254
3652
fn head_is_native(head: &str, symbols: &SymbolTable) -> bool {
255
3652
    match symbols.lookup(head) {
256
3652
        Some(sym) => {
257
3652
            sym.function().is_none()
258
3237
                && matches!(sym.kind(), SymbolKind::Native | SymbolKind::Operator)
259
        }
260
        None => false,
261
    }
262
3652
}
263

            
264
/// Whether `name` resolves to a user function / lambda value — such a symbol,
265
/// passed to a higher-order native, would be dereferenced and called during
266
/// `eval_value`, so it is NOT inert for purity purposes.
267
498
fn names_user_callable(name: &str, symbols: &SymbolTable) -> bool {
268
498
    match symbols.lookup(name) {
269
249
        Some(sym) => {
270
249
            sym.function().is_some()
271
249
                || matches!(sym.value(), Some(Expr::Lambda(_, _)))
272
249
                || matches!(sym.kind(), SymbolKind::Function | SymbolKind::Macro)
273
        }
274
249
        None => false,
275
    }
276
498
}
277

            
278
19758
pub(in crate::compiler::special) fn collect_setf_targets(
279
19758
    exprs: &[Expr],
280
19758
) -> Vec<(String, Option<Expr>)> {
281
19758
    let mut targets = Vec::new();
282
19758
    for expr in exprs {
283
19592
        collect_setf_targets_inner(expr, &mut targets);
284
19592
    }
285
19758
    targets
286
19758
}
287

            
288
283224
fn collect_setf_targets_inner(expr: &Expr, targets: &mut Vec<(String, Option<Expr>)>) {
289
283224
    if let Expr::List(elems) = expr {
290
95790
        if let Some(Expr::Symbol(name)) = elems.first()
291
95790
            && name == "SETF"
292
        {
293
13450
            for pair in elems[1..].chunks(2) {
294
13450
                if let Some(Expr::Symbol(var)) = pair.first()
295
13450
                    && !targets.iter().any(|(n, _)| n == var)
296
13284
                {
297
13284
                    targets.push((var.clone(), pair.get(1).cloned()));
298
13284
                }
299
            }
300
82340
        }
301
263632
        for elem in elems {
302
263632
            collect_setf_targets_inner(elem, targets);
303
263632
        }
304
187434
    }
305
283224
}
306

            
307
/// Eval-path mirror of [`promote_to_wasm_local`]: when a *runtime* loop body
308
/// `setf`s an OUTER variable, the loop's runtime iteration mutates it, but the
309
/// eval (const-fold) surface never runs that body, so the variable keeps its
310
/// stale const init. Mark each such target as a `WasmRuntime` placeholder of
311
/// its codegen-inferred type so the enclosing eval (an IF-test classification,
312
/// a defun's tail value) sees it as runtime — exactly what codegen's
313
/// accumulator promotion does. Idempotent: re-marking an already-runtime
314
/// target is a no-op, keeping the two-surface re-eval design sound.
315
4648
pub(in crate::compiler::special) fn mark_runtime_setf_targets(
316
4648
    symbols: &mut SymbolTable,
317
4648
    body: &[Expr],
318
4648
    exclude: &[&str],
319
4648
) {
320
4648
    for (name, rhs) in collect_setf_targets(body) {
321
3901
        if exclude.contains(&name.as_str()) {
322
            continue;
323
3901
        }
324
        // Only re-type a var that is ALREADY bound in this scope — never
325
        // synthesize a missing one. A `setf` to an undefined / nested-bound
326
        // name is codegen's concern (`promote_to_wasm_local` errors on it); the
327
        // eval mirror must not mask that by defining a spurious outer symbol.
328
3901
        let Some(current) = symbols.lookup(&name).and_then(|s| s.value().cloned()) else {
329
            continue;
330
        };
331
3901
        if matches!(current, Expr::WasmRuntime(_) | Expr::WasmLocal(_, _)) {
332
            continue;
333
3901
        }
334
3901
        let ty = infer_wasm_type(&current, rhs.as_ref(), symbols);
335
3901
        symbols.define(Symbol::new(&name, SymbolKind::Variable).with_value(Expr::WasmRuntime(ty)));
336
    }
337
4648
}
338

            
339
5561
pub(in crate::compiler::special) fn promote_to_wasm_local(
340
5561
    ctx: &mut CompileContext,
341
5561
    emit: &mut FunctionEmitter,
342
5561
    symbols: &mut SymbolTable,
343
5561
    name: &str,
344
5561
    step_hint: Option<&Expr>,
345
5561
) -> Result<()> {
346
5561
    let sym = symbols
347
5561
        .lookup(name)
348
5561
        .ok_or_else(|| Error::UndefinedSymbol(name.to_string()))?;
349
    if matches!(
350
5561
        sym.value(),
351
        Some(Expr::WasmLocal(_, _) | Expr::WasmRuntime(_))
352
    ) {
353
5561
        return Ok(());
354
    }
355
    let val = sym.value().cloned().unwrap_or(Expr::Nil);
356
    let ty = infer_wasm_type(&val, step_hint, symbols);
357
    let idx = ctx.alloc_local(ty)?;
358
    if matches!(ty, WasmType::PairRef(_)) && matches!(val, Expr::Nil) {
359
        emit.ref_null(ctx.ids.ty_pair);
360
    } else {
361
        compile_for_stack(ctx, emit, symbols, &val)?;
362
    }
363
    emit.local_set(idx);
364
    symbols.define(Symbol::new(name, SymbolKind::Variable).with_value(Expr::WasmLocal(idx, ty)));
365
    Ok(())
366
5561
}
367

            
368
2905
pub(super) fn static_loop_terminates(
369
2905
    local: &SymbolTable,
370
2905
    end_test: &Expr,
371
2905
    stepped: &[(String, Option<Expr>)],
372
2905
    sequential: bool,
373
2905
) -> bool {
374
2905
    let mut sim = local.clone();
375
2905
    for _ in 0..MAX_STATIC_LOOP_ITERS {
376
35856
        match eval_value(&mut sim, end_test) {
377
35856
            Ok(test) if !test.is_nil() => return true,
378
33366
            Ok(_) => {}
379
            Err(_) => return false,
380
        }
381
33366
        if sequential {
382
22991
            for (name, step) in stepped {
383
22991
                let Some(s) = step else { continue };
384
22991
                match eval_value(&mut sim, s) {
385
22991
                    Ok(val) => {
386
22991
                        if let Some(sym) = sim.lookup_mut(name) {
387
22991
                            sym.set_value(val);
388
22991
                        }
389
                    }
390
                    Err(_) => return false,
391
                }
392
            }
393
        } else {
394
21497
            let new_vals: Vec<_> = stepped
395
21497
                .iter()
396
40172
                .filter_map(|(name, step)| {
397
40172
                    step.as_ref()
398
40172
                        .and_then(|s| eval_value(&mut sim, s).ok().map(|v| (name.clone(), v)))
399
40172
                })
400
21497
                .collect();
401
34611
            for (name, val) in new_vals {
402
34611
                if let Some(sym) = sim.lookup_mut(&name) {
403
34611
                    sym.set_value(val);
404
34611
                }
405
            }
406
        }
407
    }
408
415
    false
409
2905
}