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
9092
pub(super) fn parse_do_vars(context: &str, expr: &Expr) -> Result<Vec<DoVar>> {
50
9092
    let list = match expr {
51
9092
        Expr::List(elems) => elems,
52
        _ => {
53
            return Err(Error::Compile(format!(
54
                "{context}: expected variable list, got {expr:?}"
55
            )));
56
        }
57
    };
58
9092
    list.iter()
59
11268
        .map(|spec| match spec {
60
            Expr::Symbol(name) => Ok(DoVar {
61
                name: name.clone(),
62
                init: None,
63
                step: None,
64
            }),
65
11268
            Expr::List(elems) if !elems.is_empty() && elems.len() <= 3 => {
66
11268
                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
11268
                Ok(DoVar {
73
11268
                    name: name.to_string(),
74
11268
                    init: elems.get(1).cloned(),
75
11268
                    step: elems.get(2).cloned(),
76
11268
                })
77
            }
78
            _ => Err(Error::Compile(format!(
79
                "{context}: malformed variable spec: {spec:?}"
80
            ))),
81
11268
        })
82
9092
        .collect()
83
9092
}
84

            
85
9092
pub(super) fn parse_end_clause<'a>(
86
9092
    context: &str,
87
9092
    expr: &'a Expr,
88
9092
) -> Result<(&'a Expr, &'a [Expr])> {
89
9092
    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
9092
    if elems.is_empty() {
95
        return Err(Error::Compile(format!(
96
            "{context}: end clause must have a test form"
97
        )));
98
9092
    }
99
9092
    Ok((&elems[0], &elems[1..]))
100
9092
}
101

            
102
13135
pub(super) fn infer_wasm_type(
103
13135
    expr: &Expr,
104
13135
    step_hint: Option<&Expr>,
105
13135
    symbols: &SymbolTable,
106
13135
) -> 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
13135
    if let Expr::Nil = expr
111
1500
        && let Some(step) = step_hint
112
1499
        && let Some(elem) = step_pair_element(step, symbols)
113
    {
114
953
        return WasmType::PairRef(elem);
115
12182
    }
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
12182
    if let Expr::Nil = expr
124
547
        && let Some(step) = step_hint
125
546
        && is_pure_native_expr(step, symbols)
126
342
        && let Ok(resolved) = eval_value(&mut symbols.clone(), step)
127
136
        && let Some(ty) = crate::compiler::expr::classify_stack_type(&resolved)
128
    {
129
136
        return ty;
130
12046
    }
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
12046
    if let Expr::Bytes(_) = expr {
134
        return WasmType::StringRef;
135
12046
    }
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
12046
    crate::compiler::expr::classify_stack_type(expr).unwrap_or(WasmType::I32)
141
13135
}
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
4180
pub(super) fn infer_result_pair_element(
150
4180
    result_form: &Expr,
151
4180
    body: &[Expr],
152
4180
    stepped: &[(String, Option<Expr>)],
153
4180
    symbols: &SymbolTable,
154
4180
) -> Option<PairElement> {
155
4180
    let Expr::Symbol(target) = result_form else {
156
        return None;
157
    };
158
4180
    let from_body = collect_setf_targets(body)
159
4180
        .into_iter()
160
4180
        .find_map(|(name, rhs)| {
161
3160
            (name == *target)
162
3160
                .then_some(rhs)
163
3160
                .flatten()
164
3160
                .as_ref()
165
3160
                .and_then(|r| step_pair_element(r, symbols))
166
3160
        });
167
4180
    if from_body.is_some() {
168
2337
        return from_body;
169
1843
    }
170
2591
    stepped.iter().find_map(|(name, step)| {
171
2591
        if name != target {
172
1571
            return None;
173
1020
        }
174
1020
        step.as_ref().and_then(|s| step_pair_element(s, symbols))
175
2591
    })
176
4180
}
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
14682
fn step_pair_element(expr: &Expr, symbols: &SymbolTable) -> Option<PairElement> {
184
14682
    let Expr::List(elems) = expr else {
185
7371
        return None;
186
    };
187
7311
    if let Some(Expr::Symbol(name)) = elems.first()
188
7175
        && name == "CONS"
189
4448
        && elems.len() == 3
190
4448
        && let Some(elem) = resolve_pair_element_from_expr(&elems[1], symbols)
191
    {
192
3902
        return Some(elem);
193
3409
    }
194
9003
    elems.iter().find_map(|e| step_pair_element(e, symbols))
195
14682
}
196

            
197
8348
fn resolve_pair_element_from_expr(expr: &Expr, symbols: &SymbolTable) -> Option<PairElement> {
198
1360
    match expr {
199
3900
        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
3086
        Expr::Symbol(name) => {
206
3086
            let sym = symbols.lookup(name)?;
207
2880
            let value = sym.value()?;
208
2880
            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
1360
        Expr::List(_) if is_pure_native_expr(expr, symbols) => {
218
1020
            let resolved = eval_value(&mut symbols.clone(), expr).ok()?;
219
1020
            match resolved {
220
                Expr::List(_) => None,
221
1020
                other => resolve_pair_element_from_expr(&other, symbols),
222
            }
223
        }
224
340
        _ => None,
225
    }
226
8348
}
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
5242
fn is_pure_native_expr(expr: &Expr, symbols: &SymbolTable) -> bool {
242
5242
    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
1836
        | Expr::WasmLocal(_, _) => true,
252
412
        Expr::Symbol(name) => !names_user_callable(name, symbols),
253
2994
        Expr::List(elems) => match elems.first() {
254
2994
            Some(Expr::Symbol(head)) => {
255
2994
                head_is_native(head, symbols)
256
3336
                    && elems[1..].iter().all(|e| is_pure_native_expr(e, symbols))
257
            }
258
            _ => false,
259
        },
260
        _ => false,
261
    }
262
5242
}
263

            
264
2994
fn head_is_native(head: &str, symbols: &SymbolTable) -> bool {
265
2994
    match symbols.lookup(head) {
266
2994
        Some(sym) => {
267
2994
            sym.function().is_none()
268
2654
                && matches!(sym.kind(), SymbolKind::Native | SymbolKind::Operator)
269
        }
270
        None => false,
271
    }
272
2994
}
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
412
fn names_user_callable(name: &str, symbols: &SymbolTable) -> bool {
278
412
    match symbols.lookup(name) {
279
206
        Some(sym) => {
280
206
            sym.function().is_some()
281
206
                || matches!(sym.value(), Some(Expr::Lambda(_, _)))
282
206
                || matches!(sym.kind(), SymbolKind::Function | SymbolKind::Macro)
283
        }
284
206
        None => false,
285
    }
286
412
}
287

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

            
298
233202
fn collect_setf_targets_inner(expr: &Expr, targets: &mut Vec<(String, Option<Expr>)>) {
299
233202
    if let Expr::List(elems) = expr {
300
78862
        if let Some(Expr::Symbol(name)) = elems.first()
301
78862
            && name == "SETF"
302
        {
303
11028
            for pair in elems[1..].chunks(2) {
304
11028
                if let Some(Expr::Symbol(var)) = pair.first()
305
11028
                    && !targets.iter().any(|(n, _)| n == var)
306
10892
                {
307
10892
                    targets.push((var.clone(), pair.get(1).cloned()));
308
10892
                }
309
            }
310
67834
        }
311
217128
        for elem in elems {
312
217128
            collect_setf_targets_inner(elem, targets);
313
217128
        }
314
154340
    }
315
233202
}
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
3836
pub(in crate::compiler::special) fn mark_runtime_setf_targets(
326
3836
    symbols: &mut SymbolTable,
327
3836
    body: &[Expr],
328
3836
    exclude: &[&str],
329
3836
) {
330
3836
    for (name, rhs) in collect_setf_targets(body) {
331
3220
        if exclude.contains(&name.as_str()) {
332
            continue;
333
3220
        }
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
3220
        let Some(current) = symbols.lookup(&name).and_then(|s| s.value().cloned()) else {
339
            continue;
340
        };
341
3220
        if matches!(current, Expr::WasmRuntime(_) | Expr::WasmLocal(_, _)) {
342
            continue;
343
3220
        }
344
3220
        let ty = infer_wasm_type(&current, rhs.as_ref(), symbols);
345
3220
        symbols.define(Symbol::new(&name, SymbolKind::Variable).with_value(Expr::WasmRuntime(ty)));
346
    }
347
3836
}
348

            
349
4512
pub(in crate::compiler::special) fn promote_to_wasm_local(
350
4512
    ctx: &mut CompileContext,
351
4512
    emit: &mut FunctionEmitter,
352
4512
    symbols: &mut SymbolTable,
353
4512
    name: &str,
354
4512
    step_hint: Option<&Expr>,
355
4512
) -> Result<()> {
356
4512
    let sym = symbols
357
4512
        .lookup(name)
358
4512
        .ok_or_else(|| Error::UndefinedSymbol(name.to_string()))?;
359
    if matches!(
360
4512
        sym.value(),
361
        Some(Expr::WasmLocal(_, _) | Expr::WasmRuntime(_))
362
    ) {
363
4512
        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
4512
}
377

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