1
//! `MAP` eval + compile.
2
//!
3
//! Three paths:
4
//! - **Constant-fold**: when the function and every list arg resolves
5
//!   at compile time, `map_fn` walks element-by-element via the eval
6
//!   `call` pipeline and returns the resulting list literal.
7
//! - **Runtime closure, literal list**: emit a per-element FUNCALL,
8
//!   stash each result in a fresh local, then walk the locals in
9
//!   reverse and prepend each car onto the accumulator via `pair_new`.
10
//! - **Runtime closure or runtime `PairRef` list**: walk the input
11
//!   chain at runtime, calling the function on each car, prepend onto
12
//!   a reversed accumulator, and run `emit_reverse_loop` once at the
13
//!   end so the output preserves input order.
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::{call, compile_expr, compile_for_stack, eval_value, format_expr};
19
use crate::error::{Error, Result};
20
use crate::runtime::SymbolTable;
21

            
22
use super::reverse::emit_reverse_loop;
23

            
24
994
pub(super) fn map_fn(symbols: &mut SymbolTable, args: &[Expr]) -> Result<Expr> {
25
994
    if args.len() < 2 {
26
71
        return Err(Error::Arity {
27
71
            name: "MAP".to_string(),
28
71
            expected: 2,
29
71
            actual: args.len(),
30
71
        });
31
923
    }
32
923
    let fn_resolved = eval_value(symbols, &args[0])?;
33
923
    if let Some(WasmType::Closure(sig)) = fn_resolved.wasm_type() {
34
        // MAP's OUTPUT element is the closure's RESULT type (it transforms each
35
        // element) — NOT the input element. The closure-emit site records the
36
        // result so this eval prediction agrees with codegen (`compile_map_*`
37
        // uses the per-row closure result). Fall back to the input element only
38
        // when the result isn't recorded or isn't a pair-cell type (e.g. a
39
        // nested-pair result, which codegen rejects anyway).
40
142
        let elem = symbols
41
142
            .closure_result(sig)
42
142
            .and_then(PairElement::from_wasm_type)
43
142
            .or(runtime_pair_input_element(symbols, &args[1..])?)
44
142
            .unwrap_or(PairElement::Ratio);
45
142
        return Ok(Expr::WasmRuntime(WasmType::PairRef(elem)));
46
781
    }
47
781
    if list_args_have_runtime_pair(symbols, &args[1..])? {
48
        let elem = runtime_pair_input_element(symbols, &args[1..])?
49
            .ok_or_else(|| Error::Compile("MAP: runtime list element type unknown".to_string()))?;
50
        return Ok(Expr::WasmRuntime(WasmType::PairRef(elem)));
51
781
    }
52
    // Static fn + constant list(s): try to fully reduce. If the body produces
53
    // a runtime value per element (a side-effecting native, or a runtime
54
    // result), MAP ISN'T constant-foldable — surface a runtime placeholder
55
    // typed by the body's per-element result so the compile path lowers each
56
    // element call at runtime instead of baking placeholders into a list.
57
781
    let folded = constant_fold_map(symbols, &args[0], &args[1..])?;
58
710
    if let Some(elem) = mapped_runtime_element(&folded) {
59
        return Ok(Expr::WasmRuntime(WasmType::PairRef(elem)));
60
710
    }
61
710
    Ok(folded)
62
994
}
63

            
64
/// If a folded MAP result is a list whose elements are runtime placeholders
65
/// (the body produced runtime values), returns the unified `PairElement` of
66
/// those results — the signal that MAP must lower at runtime. `None` when the
67
/// fold produced a genuine constant list.
68
1207
fn mapped_runtime_element(folded: &Expr) -> Option<PairElement> {
69
1207
    let elems = match folded {
70
1207
        Expr::Quote(inner) => match inner.as_ref() {
71
1207
            Expr::List(elems) => elems,
72
            _ => return None,
73
        },
74
        Expr::List(elems) => elems,
75
        _ => return None,
76
    };
77
1207
    if !elems.iter().any(Expr::is_wasm_runtime) {
78
923
        return None;
79
284
    }
80
284
    let mut elem: Option<PairElement> = None;
81
639
    for e in elems {
82
639
        let pe = e
83
639
            .wasm_type()
84
639
            .and_then(PairElement::from_wasm_type)
85
639
            .or_else(|| super::infer::literal_pair_element(e))
86
639
            .unwrap_or(PairElement::AnyRef);
87
639
        elem = Some(match elem {
88
355
            Some(prev) => prev.widen(pe),
89
284
            None => pe,
90
        });
91
    }
92
284
    Some(elem.unwrap_or(PairElement::AnyRef))
93
1207
}
94

            
95
1562
fn constant_fold_map(
96
1562
    symbols: &mut SymbolTable,
97
1562
    function_arg: &Expr,
98
1562
    list_args: &[Expr],
99
1562
) -> Result<Expr> {
100
1562
    let lists: Vec<Vec<Expr>> = list_args
101
1562
        .iter()
102
1917
        .map(|arg| {
103
1917
            let resolved = eval_value(symbols, arg)?;
104
1917
            extract_list_elements(&resolved)
105
1917
        })
106
1562
        .collect::<Result<_>>()?;
107

            
108
1491
    let min_len = lists.iter().map(std::vec::Vec::len).min().unwrap_or(0);
109

            
110
1491
    let mut results = Vec::with_capacity(min_len);
111
3195
    for i in 0..min_len {
112
3195
        let mut call_args = vec![function_arg.clone()];
113
3976
        for list in &lists {
114
3976
            call_args.push(as_literal_arg(&list[i]));
115
3976
        }
116
3195
        let result = call(symbols, &call_args)?;
117
2911
        let resolved_result = eval_value(symbols, &result)?;
118
2911
        results.push(resolved_result);
119
    }
120
1207
    Ok(Expr::Quote(Box::new(Expr::List(results))))
121
1562
}
122

            
123
5467
fn as_literal_arg(expr: &Expr) -> Expr {
124
5467
    match expr {
125
497
        Expr::Symbol(_) | Expr::List(_) | Expr::Cons(_, _) => Expr::Quote(Box::new(expr.clone())),
126
4970
        _ => expr.clone(),
127
    }
128
5467
}
129

            
130
1988
fn list_args_have_runtime_pair(symbols: &mut SymbolTable, list_args: &[Expr]) -> Result<bool> {
131
2343
    for arg in list_args {
132
2343
        let resolved = eval_value(symbols, arg)?;
133
2343
        if matches!(resolved.wasm_type(), Some(WasmType::PairRef(_))) {
134
213
            return Ok(true);
135
2130
        }
136
    }
137
1775
    Ok(false)
138
1988
}
139

            
140
142
fn runtime_pair_input_element(
141
142
    symbols: &mut SymbolTable,
142
142
    list_args: &[Expr],
143
142
) -> Result<Option<PairElement>> {
144
142
    for arg in list_args {
145
142
        let resolved = eval_value(symbols, arg)?;
146
142
        if let Some(WasmType::PairRef(elem)) = resolved.wasm_type() {
147
            return Ok(Some(elem));
148
142
        }
149
    }
150
142
    Ok(None)
151
142
}
152

            
153
4899
pub(super) fn extract_list_elements(expr: &Expr) -> Result<Vec<Expr>> {
154
4899
    match expr {
155
        Expr::List(elems) => Ok(elems.clone()),
156
284
        Expr::Nil => Ok(vec![]),
157
4544
        Expr::Quote(inner) => match inner.as_ref() {
158
4544
            Expr::List(elems) => Ok(elems.clone()),
159
            Expr::Nil => Ok(vec![]),
160
            other => Err(Error::Compile(format!(
161
                "MAP expects list arguments, got quoted {}",
162
                format_expr(other)
163
            ))),
164
        },
165
71
        other => Err(Error::Compile(format!(
166
71
            "MAP expects list arguments, got {}",
167
71
            format_expr(other)
168
71
        ))),
169
    }
170
4899
}
171

            
172
1207
pub(super) fn compile_map(
173
1207
    ctx: &mut CompileContext,
174
1207
    emit: &mut FunctionEmitter,
175
1207
    symbols: &mut SymbolTable,
176
1207
    args: &[Expr],
177
1207
) -> Result<()> {
178
1207
    if args.len() < 2 {
179
        return Err(Error::Arity {
180
            name: "MAP".to_string(),
181
            expected: 2,
182
            actual: args.len(),
183
        });
184
1207
    }
185
1207
    let fn_resolved = eval_value(symbols, &args[0])?;
186
1207
    let runtime_closure = matches!(fn_resolved.wasm_type(), Some(WasmType::Closure(_)));
187
1207
    let runtime_list = list_args_have_runtime_pair(symbols, &args[1..])?;
188

            
189
1207
    if runtime_closure || runtime_list {
190
426
        let ty = compile_map_to_stack(ctx, emit, symbols, args)?;
191
426
        return crate::compiler::expr::serialize_stack_to_output(ctx, emit, ty);
192
781
    }
193
781
    let result = constant_fold_map(symbols, &args[0], &args[1..])?;
194
497
    if mapped_runtime_element(&result).is_some() {
195
        // Body isn't constant-foldable — lower the static fn + constant list
196
        // element-by-element via the runtime path.
197
284
        let ty = compile_map_to_stack(ctx, emit, symbols, args)?;
198
284
        return crate::compiler::expr::serialize_stack_to_output(ctx, emit, ty);
199
213
    }
200
213
    compile_expr(ctx, emit, symbols, &result)
201
1207
}
202

            
203
781
pub(super) fn compile_map_to_stack(
204
781
    ctx: &mut CompileContext,
205
781
    emit: &mut FunctionEmitter,
206
781
    symbols: &mut SymbolTable,
207
781
    args: &[Expr],
208
781
) -> Result<WasmType> {
209
781
    if args.len() < 2 {
210
        return Err(Error::Compile(
211
            "MAP: stack-position lowering requires a function and at least one list".to_string(),
212
        ));
213
781
    }
214
    // Single runtime `PairRef` list: walk the chain at runtime.
215
781
    if args.len() == 2
216
710
        && let Some(WasmType::PairRef(elem)) = eval_value(symbols, &args[1])?.wasm_type()
217
    {
218
213
        return compile_map_runtime_list(ctx, emit, symbols, &args[0], &args[1], elem);
219
568
    }
220
    // One or more constant lists with a callable whose body isn't constant-
221
    // foldable (runtime/side-effecting). Zip the constant lists and apply the
222
    // function per row via FUNCALL — serves a closure OR a bare lambda, and any
223
    // arity of lists. (A multi-list MAP over RUNTIME pair chains is not
224
    // supported here — only the single-list runtime chain is.)
225
568
    let lists: Vec<Vec<Expr>> = args[1..]
226
568
        .iter()
227
639
        .map(|arg| {
228
639
            let resolved = eval_value(symbols, arg)?;
229
639
            if matches!(resolved.wasm_type(), Some(WasmType::PairRef(_))) {
230
                return Err(Error::Compile(
231
                    "MAP: a multi-list mapping requires constant lists; a runtime list \
232
                     is only supported as the sole list argument"
233
                        .to_string(),
234
                ));
235
639
            }
236
639
            extract_list_elements(&resolved)
237
639
        })
238
568
        .collect::<Result<_>>()?;
239
568
    compile_map_literal(ctx, emit, symbols, &args[0], &lists)
240
781
}
241

            
242
/// Lower `(map <fn> '(a0 a1 ...) '(b0 b1 ...) ...)` row-by-row over one or more
243
/// constant lists (zipped to the shortest). Each per-row call rides FUNCALL (→
244
/// `call_ref` for a closure, inline for a bare lambda); results land in fresh
245
/// per-row locals. We then walk the locals in reverse and prepend each car onto
246
/// the accumulator via `pair_new` so the output keeps input order without an
247
/// extra reverse pass.
248
568
fn compile_map_literal(
249
568
    ctx: &mut CompileContext,
250
568
    emit: &mut FunctionEmitter,
251
568
    symbols: &mut SymbolTable,
252
568
    fn_arg: &Expr,
253
568
    lists: &[Vec<Expr>],
254
568
) -> Result<WasmType> {
255
568
    let rows = lists.iter().map(Vec::len).min().unwrap_or(0);
256
568
    if rows == 0 {
257
        emit.ref_null(ctx.ids.ty_pair);
258
        return Ok(WasmType::PairRef(PairElement::Ratio));
259
568
    }
260
568
    let mut element_locals: Vec<(u32, PairElement)> = Vec::with_capacity(rows);
261
568
    let mut shared_elem: Option<PairElement> = None;
262
1349
    for row in 0..rows {
263
1349
        let mut call = vec![Expr::Symbol("FUNCALL".to_string()), fn_arg.clone()];
264
1491
        call.extend(lists.iter().map(|list| as_literal_arg(&list[row])));
265
1349
        let funcall = Expr::List(call);
266
1349
        let ty = compile_for_stack(ctx, emit, symbols, &funcall)?;
267
1349
        let elem = PairElement::from_wasm_type(ty).ok_or_else(|| {
268
            Error::Compile(format!(
269
                "MAP: closure result type {ty} can't ride a typed pair; \
270
                 flatten via let-bind first"
271
            ))
272
        })?;
273
        // Heterogeneous per-element results widen the chain to `AnyRef` (the
274
        // ADR-0025 escape hatch) instead of erroring — each car is boxed to
275
        // anyref at prepend. Matches the eval-side `mapped_runtime_element`.
276
1349
        shared_elem = Some(match shared_elem {
277
781
            Some(prev) => prev.widen(elem),
278
568
            None => elem,
279
        });
280
1349
        let local = ctx.alloc_local(elem.as_wasm_type())?;
281
1349
        emit.local_set(local);
282
1349
        element_locals.push((local, elem));
283
    }
284
568
    let Some(chain_elem) = shared_elem else {
285
        return Err(Error::Compile(
286
            "MAP: empty literal-list closure mapping reached prepend phase".to_string(),
287
        ));
288
    };
289
568
    let pair_idx = ctx.ids.ty_pair;
290
568
    let acc_local = ctx.alloc_local(WasmType::PairRef(chain_elem))?;
291
568
    emit.ref_null(pair_idx);
292
568
    emit.local_set(acc_local);
293
1349
    for (local, elem_ty) in element_locals.iter().rev() {
294
1349
        emit.local_get(*local);
295
1349
        // Box each car for the CHAIN's element slot: an AnyRef chain needs
296
1349
        // every car widened to anyref (i31 for i32/bool); a homogeneous chain
297
1349
        // boxes per the element's own type.
298
1349
        box_for_pair_car(emit, chain_car_box(chain_elem, *elem_ty));
299
1349
        emit.local_get(acc_local);
300
1349
        emit.call(ctx.ids.pair_new);
301
1349
        emit.local_set(acc_local);
302
1349
    }
303
568
    emit.local_get(acc_local);
304
568
    Ok(WasmType::PairRef(chain_elem))
305
568
}
306

            
307
/// The `PairElement` to box a car AS when prepending into a chain whose slot is
308
/// `chain_elem`. For an `AnyRef` chain, an i32/bool car must still be i31-boxed
309
/// (so `box_for_pair_car` sees its own value type); ref-typed cars are anyref
310
/// subtypes already. For a homogeneous chain the car boxes per the chain type.
311
1349
fn chain_car_box(chain_elem: PairElement, car_elem: PairElement) -> PairElement {
312
1349
    match chain_elem {
313
142
        PairElement::AnyRef => car_elem,
314
1207
        other => other,
315
    }
316
1349
}
317

            
318
/// Walks a runtime `PairRef(elem)` input. Builds the result reversed
319
/// (one prepend per car) by calling the function on each car, then
320
/// reverses once at the end so the output preserves input order.
321
213
fn compile_map_runtime_list(
322
213
    ctx: &mut CompileContext,
323
213
    emit: &mut FunctionEmitter,
324
213
    symbols: &mut SymbolTable,
325
213
    fn_arg: &Expr,
326
213
    list_expr: &Expr,
327
213
    elem: PairElement,
328
213
) -> Result<WasmType> {
329
213
    let result_elem = elem;
330
213
    let pair_idx = ctx.ids.ty_pair;
331
213
    let pair_local = ctx.alloc_local(WasmType::PairRef(result_elem))?;
332
213
    let acc_local = ctx.alloc_local(WasmType::PairRef(result_elem))?;
333
213
    let car_local = ctx.alloc_local(elem.as_wasm_type())?;
334

            
335
213
    compile_for_stack(ctx, emit, symbols, list_expr)?;
336
213
    emit.local_set(pair_local);
337

            
338
213
    emit.ref_null(pair_idx);
339
213
    emit.local_set(acc_local);
340

            
341
213
    emit.block_start();
342
213
    emit.loop_start();
343

            
344
213
    emit.local_get(pair_local);
345
213
    emit.ref_is_null();
346
213
    emit.br_if(1);
347

            
348
213
    emit.local_get(pair_local);
349
213
    emit.struct_get(pair_idx, 0);
350
213
    crate::compiler::native::list::emit_pair_car_downcast(ctx, emit, elem);
351
213
    emit.local_set(car_local);
352

            
353
213
    let mapped_ty = compile_map_call_with_local(ctx, emit, symbols, fn_arg, car_local, elem)?;
354
213
    let actual_elem = PairElement::from_wasm_type(mapped_ty).ok_or_else(|| {
355
        Error::Compile(format!(
356
            "MAP: closure result type {mapped_ty} can't ride a typed pair; \
357
             flatten via let-bind first"
358
        ))
359
    })?;
360
213
    if actual_elem != result_elem {
361
        return Err(Error::Compile(format!(
362
            "MAP: closure result element {actual_elem} doesn't match input element {result_elem}; \
363
             heterogeneous mapping isn't supported yet"
364
        )));
365
213
    }
366
213
    box_for_pair_car(emit, actual_elem);
367
213
    emit.local_get(acc_local);
368
213
    emit.call(ctx.ids.pair_new);
369
213
    emit.local_set(acc_local);
370

            
371
213
    emit.local_get(pair_local);
372
213
    emit.struct_get(pair_idx, 1);
373
213
    emit.local_set(pair_local);
374

            
375
213
    emit.br(0);
376
213
    emit.block_end();
377
213
    emit.block_end();
378

            
379
    // The accumulator is in reverse order — walk it once more through
380
    // the standard reverse loop so the output preserves input order.
381
213
    let acc_expr = Expr::WasmLocal(acc_local, WasmType::PairRef(result_elem));
382
213
    emit_reverse_loop(ctx, emit, symbols, &acc_expr, result_elem)?;
383
213
    Ok(WasmType::PairRef(result_elem))
384
213
}
385

            
386
1562
fn box_for_pair_car(emit: &mut FunctionEmitter, elem: PairElement) {
387
    // I32 and Bool share the i31-boxed car representation.
388
1562
    if matches!(elem, PairElement::I32 | PairElement::Bool) {
389
426
        emit.ref_i31();
390
1136
    }
391
1562
}
392

            
393
/// Emit a per-iteration call of `fn_arg` against the value held in
394
/// `car_local`. Mirrors the FUNCALL stack-position emit but inlined so
395
/// we don't have to allocate a temporary `Expr` for the local.
396
213
fn compile_map_call_with_local(
397
213
    ctx: &mut CompileContext,
398
213
    emit: &mut FunctionEmitter,
399
213
    symbols: &mut SymbolTable,
400
213
    fn_arg: &Expr,
401
213
    car_local: u32,
402
213
    elem: PairElement,
403
213
) -> Result<WasmType> {
404
213
    let car_expr = Expr::WasmLocal(car_local, elem.as_wasm_type());
405
213
    let funcall = Expr::List(vec![
406
213
        Expr::Symbol("FUNCALL".to_string()),
407
213
        fn_arg.clone(),
408
213
        car_expr,
409
213
    ]);
410
213
    compile_for_stack(ctx, emit, symbols, &funcall)
411
213
}