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
use super::super::control::is_truthy;
23

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

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

            
47
pub(super) const MAX_STATIC_LOOP_ITERS: usize = 64;
48

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

            
85
9447
pub(super) fn parse_end_clause<'a>(
86
9447
    context: &str,
87
9447
    expr: &'a Expr,
88
9447
) -> Result<(&'a Expr, &'a [Expr])> {
89
9447
    let elems = expr.as_list().ok_or_else(|| {
90
        Error::Compile(format!(
91
            "{context}: end clause must be a list, got {expr:?}"
92
        ))
93
    })?;
94
9447
    if elems.is_empty() {
95
        return Err(Error::Compile(format!(
96
            "{context}: end clause must have a test form"
97
        )));
98
9447
    }
99
9447
    Ok((&elems[0], &elems[1..]))
100
9447
}
101

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

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

            
178
/// If `step` is (or contains) a `(CONS car cdr)` whose car has a
179
/// derivable wasm type, returns the corresponding `PairElement`. Used by
180
/// nil-init type inference so a let-bound accumulator like
181
/// `(setf result (cons i result))` picks up `i`'s element type rather
182
/// than defaulting to a placeholder.
183
15273
fn step_pair_element(expr: &Expr, symbols: &SymbolTable) -> Option<PairElement> {
184
15273
    let Expr::List(elems) = expr else {
185
7671
        return None;
186
    };
187
7602
    if let Some(Expr::Symbol(name)) = elems.first()
188
7460
        && name == "CONS"
189
4619
        && elems.len() == 3
190
4619
        && let Some(elem) = resolve_pair_element_from_expr(&elems[1], symbols)
191
    {
192
4051
        return Some(elem);
193
3551
    }
194
9375
    elems.iter().find_map(|e| step_pair_element(e, symbols))
195
15273
}
196

            
197
8668
fn resolve_pair_element_from_expr(expr: &Expr, symbols: &SymbolTable) -> Option<PairElement> {
198
1420
    match expr {
199
4049
        Expr::WasmRuntime(ty) | Expr::WasmLocal(_, ty) => PairElement::from_wasm_type(*ty),
200
        // A numeric literal is a dimensionless Ratio (never a count), matching
201
        // `list::infer::literal_pair_element` and `infer_wasm_type` — so it
202
        // agrees with what CONS codegen emits for the same car.
203
1
        Expr::Number(_) => Some(PairElement::Ratio),
204
1
        Expr::Bool(_) => Some(PairElement::Bool),
205
3197
        Expr::Symbol(name) => {
206
3197
            let sym = symbols.lookup(name)?;
207
2984
            let value = sym.value()?;
208
2984
            resolve_pair_element_from_expr(value, symbols)
209
        }
210
        // A host-fn / form car (e.g. `(transaction-tag-count i)`) hasn't been
211
        // reduced to a runtime placeholder yet. Eval it on a CLONE (pure — no
212
        // live-table mutation) so the accumulator's element type matches what
213
        // codegen will push, instead of falling through to the scalar default.
214
        // Gated on a native/operator head: those reduce to a runtime
215
        // placeholder without executing user code, so this can't recurse into
216
        // a (possibly non-terminating) user lambda during type inference.
217
1420
        Expr::List(_) if is_pure_native_expr(expr, symbols) => {
218
1065
            let resolved = eval_value(&mut symbols.clone(), expr).ok()?;
219
1065
            match resolved {
220
                Expr::List(_) => None,
221
1065
                other => resolve_pair_element_from_expr(&other, symbols),
222
            }
223
        }
224
355
        _ => None,
225
    }
226
8668
}
227

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

            
264
3124
fn head_is_native(head: &str, symbols: &SymbolTable) -> bool {
265
3124
    match symbols.lookup(head) {
266
3124
        Some(sym) => {
267
3124
            sym.function().is_none()
268
2769
                && matches!(sym.kind(), SymbolKind::Native | SymbolKind::Operator)
269
        }
270
        None => false,
271
    }
272
3124
}
273

            
274
/// Whether `name` resolves to a user function / lambda value — such a symbol,
275
/// passed to a higher-order native, would be dereferenced and called during
276
/// `eval_value`, so it is NOT inert for purity purposes.
277
426
fn names_user_callable(name: &str, symbols: &SymbolTable) -> bool {
278
426
    match symbols.lookup(name) {
279
213
        Some(sym) => {
280
213
            sym.function().is_some()
281
213
                || matches!(sym.value(), Some(Expr::Lambda(_, _)))
282
213
                || matches!(sym.kind(), SymbolKind::Function | SymbolKind::Macro)
283
        }
284
213
        None => false,
285
    }
286
426
}
287

            
288
16760
pub(in crate::compiler::special) fn collect_setf_targets(
289
16760
    exprs: &[Expr],
290
16760
) -> Vec<(String, Option<Expr>)> {
291
16760
    let mut targets = Vec::new();
292
16760
    for expr in exprs {
293
16689
        collect_setf_targets_inner(expr, &mut targets);
294
16689
    }
295
16760
    targets
296
16760
}
297

            
298
241712
fn collect_setf_targets_inner(expr: &Expr, targets: &mut Vec<(String, Option<Expr>)>) {
299
241712
    if let Expr::List(elems) = expr {
300
81729
        if let Some(Expr::Symbol(name)) = elems.first()
301
81729
            && name == "SETF"
302
        {
303
11435
            for pair in elems[1..].chunks(2) {
304
11435
                if let Some(Expr::Symbol(var)) = pair.first()
305
11435
                    && !targets.iter().any(|(n, _)| n == var)
306
11293
                {
307
11293
                    targets.push((var.clone(), pair.get(1).cloned()));
308
11293
                }
309
            }
310
70294
        }
311
225023
        for elem in elems {
312
225023
            collect_setf_targets_inner(elem, targets);
313
225023
        }
314
159983
    }
315
241712
}
316

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

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

            
378
2485
pub(super) fn static_loop_terminates(
379
2485
    local: &SymbolTable,
380
2485
    end_test: &Expr,
381
2485
    stepped: &[(String, Option<Expr>)],
382
2485
    sequential: bool,
383
2485
) -> bool {
384
2485
    let mut sim = local.clone();
385
2485
    for _ in 0..MAX_STATIC_LOOP_ITERS {
386
30672
        match eval_value(&mut sim, end_test) {
387
30672
            Ok(test) if is_truthy(&test) => return true,
388
28542
            Ok(_) => {}
389
            Err(_) => return false,
390
        }
391
28542
        if sequential {
392
19667
            for (name, step) in stepped {
393
19667
                let Some(s) = step else { continue };
394
19667
                match eval_value(&mut sim, s) {
395
19667
                    Ok(val) => {
396
19667
                        if let Some(sym) = sim.lookup_mut(name) {
397
19667
                            sym.set_value(val);
398
19667
                        }
399
                    }
400
                    Err(_) => return false,
401
                }
402
            }
403
        } else {
404
18389
            let new_vals: Vec<_> = stepped
405
18389
                .iter()
406
34364
                .filter_map(|(name, step)| {
407
34364
                    step.as_ref()
408
34364
                        .and_then(|s| eval_value(&mut sim, s).ok().map(|v| (name.clone(), v)))
409
34364
                })
410
18389
                .collect();
411
29607
            for (name, val) in new_vals {
412
29607
                if let Some(sym) = sim.lookup_mut(&name) {
413
29607
                    sym.set_value(val);
414
29607
                }
415
            }
416
        }
417
    }
418
355
    false
419
2485
}