Skip to main content

pepl_codegen/
stmt.rs

1//! Statement code generation.
2//!
3//! Statements do not produce a value on the stack (they consume their results
4//! or have side effects like `set`).
5
6use pepl_types::ast::*;
7use wasm_encoder::{BlockType, Function, Instruction, ValType};
8
9use crate::compiler::FuncContext;
10use crate::error::CodegenResult;
11use crate::expr::emit_expr;
12use crate::gas;
13use crate::runtime::*;
14use crate::types::*;
15
16/// Emit a slice of statements.
17pub fn emit_stmts(stmts: &[Stmt], ctx: &mut FuncContext, f: &mut Function) -> CodegenResult<()> {
18    for stmt in stmts {
19        emit_stmt(stmt, ctx, f)?;
20    }
21    Ok(())
22}
23
24/// Emit a single statement.
25pub fn emit_stmt(stmt: &Stmt, ctx: &mut FuncContext, f: &mut Function) -> CodegenResult<()> {
26    match stmt {
27        Stmt::Set(set) => emit_set(set, ctx, f),
28        Stmt::Let(let_bind) => emit_let(let_bind, ctx, f),
29        Stmt::If(if_expr) => emit_if_stmt(if_expr, ctx, f),
30        Stmt::For(for_expr) => emit_for_stmt(for_expr, ctx, f),
31        Stmt::Match(match_expr) => emit_match_stmt(match_expr, ctx, f),
32        Stmt::Return(_) => emit_return(f),
33        Stmt::Assert(assert_stmt) => emit_assert(assert_stmt, ctx, f),
34        Stmt::Expr(expr_stmt) => emit_expr_stmt(&expr_stmt.expr, ctx, f),
35    }
36}
37
38// ══════════════════════════════════════════════════════════════════════════════
39// Set statement
40// ══════════════════════════════════════════════════════════════════════════════
41
42fn emit_set(set: &SetStmt, ctx: &mut FuncContext, f: &mut Function) -> CodegenResult<()> {
43    // `set field = value` updates the state record.
44    // For nested paths `set a.b.c = x`, we must do immutable record update:
45    //   state.a = { ...state.a, b: { ...state.a.b, c: x } }
46    // For now we handle single-level and emit inline for nested.
47
48    // Evaluate the new value
49    let val_local = ctx.alloc_local(ValType::I32);
50    emit_expr(&set.value, ctx, f)?;
51    f.instruction(&Instruction::LocalSet(val_local));
52
53    if set.target.len() == 1 {
54        // Simple: update one field in the state record
55        let field_name = &set.target[0].name;
56        emit_state_field_set(field_name, val_local, ctx, f)?;
57    } else {
58        // Nested set: a.b.c = x
59        // We need to build the chain from inside out.
60        // For now, handle 2-level: set a.b = x
61        // Read the outer record, create new record with updated field, store back.
62        emit_nested_set(&set.target, val_local, ctx, f)?;
63    }
64    Ok(())
65}
66
67/// Set a single state field.
68pub(crate) fn emit_state_field_set(
69    field_name: &str,
70    val_local: u32,
71    ctx: &mut FuncContext,
72    f: &mut Function,
73) -> CodegenResult<()> {
74    // We rebuild the state record with the updated field.
75    // This is expensive but correct for the immutable-update model.
76    // Strategy: iterate state fields, for each field:
77    //   if name matches → use val_local
78    //   else → copy from old state
79
80    let state_fields = ctx.state_field_names.clone();
81    let field_count = state_fields.len();
82    let entries_local = ctx.alloc_local(ValType::I32);
83
84    // Allocate entries array
85    f.instruction(&Instruction::I32Const((field_count * 12) as i32));
86    f.instruction(&Instruction::Call(rt_func_idx(RT_ALLOC)));
87    f.instruction(&Instruction::LocalSet(entries_local));
88
89    for (i, sf) in state_fields.iter().enumerate() {
90        let (key_ptr, key_len) = ctx.intern_string(sf);
91        let base = (i * 12) as u64;
92
93        // key_offset
94        f.instruction(&Instruction::LocalGet(entries_local));
95        f.instruction(&Instruction::I32Const(key_ptr as i32));
96        f.instruction(&Instruction::I32Store(memarg(base, 2)));
97        // key_len
98        f.instruction(&Instruction::LocalGet(entries_local));
99        f.instruction(&Instruction::I32Const(key_len as i32));
100        f.instruction(&Instruction::I32Store(memarg(base + 4, 2)));
101
102        // value: either the new value or the old one
103        if sf == field_name {
104            f.instruction(&Instruction::LocalGet(entries_local));
105            f.instruction(&Instruction::LocalGet(val_local));
106            f.instruction(&Instruction::I32Store(memarg(base + 8, 2)));
107        } else {
108            // Read from old state
109            let old_val = ctx.alloc_local(ValType::I32);
110            f.instruction(&Instruction::GlobalGet(GLOBAL_STATE_PTR));
111            f.instruction(&Instruction::I32Const(key_ptr as i32));
112            f.instruction(&Instruction::I32Const(key_len as i32));
113            f.instruction(&Instruction::Call(rt_func_idx(RT_VAL_RECORD_GET)));
114            f.instruction(&Instruction::LocalSet(old_val));
115            f.instruction(&Instruction::LocalGet(entries_local));
116            f.instruction(&Instruction::LocalGet(old_val));
117            f.instruction(&Instruction::I32Store(memarg(base + 8, 2)));
118        }
119    }
120
121    // Build new state record
122    f.instruction(&Instruction::LocalGet(entries_local));
123    f.instruction(&Instruction::I32Const(field_count as i32));
124    f.instruction(&Instruction::Call(rt_func_idx(RT_VAL_RECORD)));
125    f.instruction(&Instruction::GlobalSet(GLOBAL_STATE_PTR));
126
127    Ok(())
128}
129
130/// Handle nested set: `set a.b.c = x` for arbitrary depth.
131///
132/// Algorithm (e.g. `set a.b.c = x`, target = [a, b, c]):
133///   Phase 1 — Walk down:  old_a = state.a,  old_b = old_a.b
134///   Phase 2 — Rebuild up: new_b = { ...old_b, c: x },  new_a = { ...old_a, b: new_b }
135///   Phase 3 — Set root:   state.a = new_a
136fn emit_nested_set(
137    target: &[Ident],
138    val_local: u32,
139    ctx: &mut FuncContext,
140    f: &mut Function,
141) -> CodegenResult<()> {
142    let depth = target.len(); // e.g. [a, b, c] → 3
143
144    // Phase 1: Walk down — read each intermediate record.
145    // For depth=3 (a.b.c): intermediates[0] = state.a, intermediates[1] = old_a.b
146    // We read depth-1 intermediates (don't read the last field — we're replacing it).
147    let mut intermediates: Vec<u32> = Vec::with_capacity(depth - 1);
148
149    for i in 0..depth - 1 {
150        let (key_ptr, key_len) = ctx.intern_string(&target[i].name);
151        let record_local = ctx.alloc_local(ValType::I32);
152
153        if i == 0 {
154            // First level reads from global state
155            f.instruction(&Instruction::GlobalGet(GLOBAL_STATE_PTR));
156        } else {
157            // Subsequent levels read from the previous intermediate
158            f.instruction(&Instruction::LocalGet(intermediates[i - 1]));
159        }
160        f.instruction(&Instruction::I32Const(key_ptr as i32));
161        f.instruction(&Instruction::I32Const(key_len as i32));
162        f.instruction(&Instruction::Call(rt_func_idx(RT_VAL_RECORD_GET)));
163        f.instruction(&Instruction::LocalSet(record_local));
164
165        intermediates.push(record_local);
166    }
167
168    // Phase 2: Rebuild from inside out.
169    // Start with val_local, then wrap it in successive record rebuilds
170    // working from the deepest level back up to the root.
171    let mut current_val = val_local;
172
173    for i in (0..depth - 1).rev() {
174        let old_record = intermediates[i];
175        let field_to_replace = &target[i + 1].name;
176
177        current_val =
178            emit_record_field_replace(old_record, field_to_replace, current_val, ctx, f)?;
179    }
180
181    // Phase 3: Set the root state field to the fully-rebuilt record.
182    emit_state_field_set(&target[0].name, current_val, ctx, f)?;
183
184    Ok(())
185}
186
187/// Clone a record with one field replaced, returning the local holding the new record.
188///
189/// Copies all entries from `old_record_local`, replacing the entry whose key matches
190/// `field_name` with `new_val_local`. Returns the local index of the new record pointer.
191fn emit_record_field_replace(
192    old_record_local: u32,
193    field_name: &str,
194    new_val_local: u32,
195    ctx: &mut FuncContext,
196    f: &mut Function,
197) -> CodegenResult<u32> {
198    // Read old record's entries pointer and field count
199    let old_entries_ptr = ctx.alloc_local(ValType::I32);
200    let old_count = ctx.alloc_local(ValType::I32);
201    f.instruction(&Instruction::LocalGet(old_record_local));
202    f.instruction(&Instruction::I32Load(memarg(4, 2))); // entries_ptr (w1)
203    f.instruction(&Instruction::LocalSet(old_entries_ptr));
204    f.instruction(&Instruction::LocalGet(old_record_local));
205    f.instruction(&Instruction::I32Load(memarg(8, 2))); // field_count (w2)
206    f.instruction(&Instruction::LocalSet(old_count));
207
208    // Allocate new entries array: count * 12 bytes
209    let new_entries = ctx.alloc_local(ValType::I32);
210    f.instruction(&Instruction::LocalGet(old_count));
211    f.instruction(&Instruction::I32Const(12));
212    f.instruction(&Instruction::I32Mul);
213    f.instruction(&Instruction::Call(rt_func_idx(RT_ALLOC)));
214    f.instruction(&Instruction::LocalSet(new_entries));
215
216    // Intern the target field key
217    let (field_key_ptr, field_key_len) = ctx.intern_string(field_name);
218
219    // Loop: copy each entry, replacing the target field's value
220    let i_local = ctx.alloc_local(ValType::I32);
221    let src_local = ctx.alloc_local(ValType::I32);
222    let dst_local = ctx.alloc_local(ValType::I32);
223    f.instruction(&Instruction::I32Const(0));
224    f.instruction(&Instruction::LocalSet(i_local));
225
226    f.instruction(&Instruction::Block(BlockType::Empty)); // break target
227    f.instruction(&Instruction::Loop(BlockType::Empty));
228
229    // if i >= count → break
230    f.instruction(&Instruction::LocalGet(i_local));
231    f.instruction(&Instruction::LocalGet(old_count));
232    f.instruction(&Instruction::I32GeU);
233    f.instruction(&Instruction::BrIf(1));
234
235    // src = old_entries + i * 12
236    f.instruction(&Instruction::LocalGet(old_entries_ptr));
237    f.instruction(&Instruction::LocalGet(i_local));
238    f.instruction(&Instruction::I32Const(12));
239    f.instruction(&Instruction::I32Mul);
240    f.instruction(&Instruction::I32Add);
241    f.instruction(&Instruction::LocalSet(src_local));
242
243    // dst = new_entries + i * 12
244    f.instruction(&Instruction::LocalGet(new_entries));
245    f.instruction(&Instruction::LocalGet(i_local));
246    f.instruction(&Instruction::I32Const(12));
247    f.instruction(&Instruction::I32Mul);
248    f.instruction(&Instruction::I32Add);
249    f.instruction(&Instruction::LocalSet(dst_local));
250
251    // Copy key_offset and key_len
252    f.instruction(&Instruction::LocalGet(dst_local));
253    f.instruction(&Instruction::LocalGet(src_local));
254    f.instruction(&Instruction::I32Load(memarg(0, 2))); // src.key_offset
255    f.instruction(&Instruction::I32Store(memarg(0, 2))); // dst.key_offset
256    f.instruction(&Instruction::LocalGet(dst_local));
257    f.instruction(&Instruction::LocalGet(src_local));
258    f.instruction(&Instruction::I32Load(memarg(4, 2))); // src.key_len
259    f.instruction(&Instruction::I32Store(memarg(4, 2))); // dst.key_len
260
261    // Check if this entry's key matches field_name
262    // Compare lengths first, then memcmp
263    f.instruction(&Instruction::LocalGet(src_local));
264    f.instruction(&Instruction::I32Load(memarg(4, 2))); // entry.key_len
265    f.instruction(&Instruction::I32Const(field_key_len as i32));
266    f.instruction(&Instruction::I32Eq);
267    f.instruction(&Instruction::If(BlockType::Empty));
268    // Lengths match — memcmp key data
269    f.instruction(&Instruction::LocalGet(src_local));
270    f.instruction(&Instruction::I32Load(memarg(0, 2))); // entry.key_offset
271    f.instruction(&Instruction::I32Const(field_key_ptr as i32)); // target key_ptr
272    f.instruction(&Instruction::I32Const(field_key_len as i32)); // key_len
273    f.instruction(&Instruction::Call(rt_func_idx(RT_MEMCMP)));
274    f.instruction(&Instruction::If(BlockType::Empty));
275    // Match → write new value
276    f.instruction(&Instruction::LocalGet(dst_local));
277    f.instruction(&Instruction::LocalGet(new_val_local));
278    f.instruction(&Instruction::I32Store(memarg(8, 2)));
279    f.instruction(&Instruction::Else);
280    // Lengths matched but content didn't → copy old value
281    f.instruction(&Instruction::LocalGet(dst_local));
282    f.instruction(&Instruction::LocalGet(src_local));
283    f.instruction(&Instruction::I32Load(memarg(8, 2)));
284    f.instruction(&Instruction::I32Store(memarg(8, 2)));
285    f.instruction(&Instruction::End); // end memcmp check
286    f.instruction(&Instruction::Else);
287    // Lengths don't match → copy old value
288    f.instruction(&Instruction::LocalGet(dst_local));
289    f.instruction(&Instruction::LocalGet(src_local));
290    f.instruction(&Instruction::I32Load(memarg(8, 2)));
291    f.instruction(&Instruction::I32Store(memarg(8, 2)));
292    f.instruction(&Instruction::End); // end length check
293
294    // i += 1
295    f.instruction(&Instruction::LocalGet(i_local));
296    f.instruction(&Instruction::I32Const(1));
297    f.instruction(&Instruction::I32Add);
298    f.instruction(&Instruction::LocalSet(i_local));
299    f.instruction(&Instruction::Br(0)); // continue loop
300
301    f.instruction(&Instruction::End); // end loop
302    f.instruction(&Instruction::End); // end block
303
304    // Build new record from entries
305    let new_record = ctx.alloc_local(ValType::I32);
306    f.instruction(&Instruction::LocalGet(new_entries));
307    f.instruction(&Instruction::LocalGet(old_count));
308    f.instruction(&Instruction::Call(rt_func_idx(RT_VAL_RECORD)));
309    f.instruction(&Instruction::LocalSet(new_record));
310
311    Ok(new_record)
312}
313
314// ══════════════════════════════════════════════════════════════════════════════
315// Let binding
316// ══════════════════════════════════════════════════════════════════════════════
317
318fn emit_let(let_bind: &LetBinding, ctx: &mut FuncContext, f: &mut Function) -> CodegenResult<()> {
319    emit_expr(&let_bind.value, ctx, f)?;
320
321    match &let_bind.name {
322        Some(ident) => {
323            let local = ctx.alloc_local(ValType::I32);
324            f.instruction(&Instruction::LocalSet(local));
325            ctx.push_local(&ident.name, local);
326        }
327        None => {
328            // `let _ = expr` — discard
329            f.instruction(&Instruction::Drop);
330        }
331    }
332    Ok(())
333}
334
335// ══════════════════════════════════════════════════════════════════════════════
336// If / For / Match as statements (values are discarded)
337// ══════════════════════════════════════════════════════════════════════════════
338
339fn emit_if_stmt(if_expr: &IfExpr, ctx: &mut FuncContext, f: &mut Function) -> CodegenResult<()> {
340    // Evaluate condition
341    emit_expr(&if_expr.condition, ctx, f)?;
342    f.instruction(&Instruction::I32Load(memarg(4, 2)));
343
344    f.instruction(&Instruction::If(BlockType::Empty));
345    emit_stmts(&if_expr.then_block.stmts, ctx, f)?;
346
347    match &if_expr.else_branch {
348        Some(ElseBranch::Block(block)) => {
349            f.instruction(&Instruction::Else);
350            emit_stmts(&block.stmts, ctx, f)?;
351        }
352        Some(ElseBranch::ElseIf(elif)) => {
353            f.instruction(&Instruction::Else);
354            emit_if_stmt(elif, ctx, f)?;
355        }
356        None => {}
357    }
358
359    f.instruction(&Instruction::End);
360    Ok(())
361}
362
363fn emit_for_stmt(for_expr: &ForExpr, ctx: &mut FuncContext, f: &mut Function) -> CodegenResult<()> {
364    // Same as emit_for_expr but discard the result
365    let list_local = ctx.alloc_local(ValType::I32);
366    let arr_local = ctx.alloc_local(ValType::I32);
367    let count_local = ctx.alloc_local(ValType::I32);
368    let i_local = ctx.alloc_local(ValType::I32);
369
370    emit_expr(&for_expr.iterable, ctx, f)?;
371    f.instruction(&Instruction::LocalSet(list_local));
372
373    f.instruction(&Instruction::LocalGet(list_local));
374    f.instruction(&Instruction::I32Load(memarg(4, 2)));
375    f.instruction(&Instruction::LocalSet(arr_local));
376
377    f.instruction(&Instruction::LocalGet(list_local));
378    f.instruction(&Instruction::I32Load(memarg(8, 2)));
379    f.instruction(&Instruction::LocalSet(count_local));
380
381    f.instruction(&Instruction::I32Const(0));
382    f.instruction(&Instruction::LocalSet(i_local));
383
384    let item_local = ctx.alloc_local(ValType::I32);
385    let index_local = if for_expr.index.is_some() {
386        Some(ctx.alloc_local(ValType::I32))
387    } else {
388        None
389    };
390
391    ctx.push_local(&for_expr.item.name, item_local);
392    if let (Some(idx_ident), Some(idx_local)) = (&for_expr.index, index_local) {
393        ctx.push_local(&idx_ident.name, idx_local);
394    }
395
396    f.instruction(&Instruction::Block(BlockType::Empty));
397    f.instruction(&Instruction::Loop(BlockType::Empty));
398
399    gas::emit_gas_tick(f, ctx.data.gas_exhausted_ptr, ctx.data.gas_exhausted_len);
400
401    f.instruction(&Instruction::LocalGet(i_local));
402    f.instruction(&Instruction::LocalGet(count_local));
403    f.instruction(&Instruction::I32GeU);
404    f.instruction(&Instruction::BrIf(1));
405
406    f.instruction(&Instruction::LocalGet(arr_local));
407    f.instruction(&Instruction::LocalGet(i_local));
408    f.instruction(&Instruction::I32Const(4));
409    f.instruction(&Instruction::I32Mul);
410    f.instruction(&Instruction::I32Add);
411    f.instruction(&Instruction::I32Load(memarg(0, 2)));
412    f.instruction(&Instruction::LocalSet(item_local));
413
414    if let Some(idx_local) = index_local {
415        f.instruction(&Instruction::I32Const(VALUE_SIZE as i32));
416        f.instruction(&Instruction::Call(rt_func_idx(RT_ALLOC)));
417        f.instruction(&Instruction::LocalTee(idx_local));
418        f.instruction(&Instruction::I32Const(TAG_NUMBER));
419        f.instruction(&Instruction::I32Store(memarg(0, 2)));
420        f.instruction(&Instruction::LocalGet(idx_local));
421        f.instruction(&Instruction::LocalGet(i_local));
422        f.instruction(&Instruction::F64ConvertI32U);
423        f.instruction(&Instruction::F64Store(memarg(4, 3)));
424    }
425
426    emit_stmts(&for_expr.body.stmts, ctx, f)?;
427
428    f.instruction(&Instruction::LocalGet(i_local));
429    f.instruction(&Instruction::I32Const(1));
430    f.instruction(&Instruction::I32Add);
431    f.instruction(&Instruction::LocalSet(i_local));
432    f.instruction(&Instruction::Br(0));
433
434    f.instruction(&Instruction::End);
435    f.instruction(&Instruction::End);
436
437    ctx.pop_local(&for_expr.item.name);
438    if let Some(idx_ident) = &for_expr.index {
439        ctx.pop_local(&idx_ident.name);
440    }
441
442    Ok(())
443}
444
445fn emit_match_stmt(
446    match_expr: &MatchExpr,
447    ctx: &mut FuncContext,
448    f: &mut Function,
449) -> CodegenResult<()> {
450    // Emit as expr then drop the result
451    crate::expr::emit_expr(
452        &Expr::new(
453            ExprKind::Match(Box::new(match_expr.clone())),
454            match_expr.span,
455        ),
456        ctx,
457        f,
458    )?;
459    f.instruction(&Instruction::Drop);
460    Ok(())
461}
462
463fn emit_return(f: &mut Function) -> CodegenResult<()> {
464    // In dispatch_action (returns i32), we need a value on the stack.
465    // Push a nil value as the return value for early return.
466    f.instruction(&Instruction::Call(rt_func_idx(RT_VAL_NIL)));
467    f.instruction(&Instruction::Return);
468    Ok(())
469}
470
471fn emit_assert(
472    assert_stmt: &AssertStmt,
473    ctx: &mut FuncContext,
474    f: &mut Function,
475) -> CodegenResult<()> {
476    // Evaluate condition
477    emit_expr(&assert_stmt.condition, ctx, f)?;
478    // Extract bool
479    f.instruction(&Instruction::I32Load(memarg(4, 2)));
480    f.instruction(&Instruction::I32Eqz);
481    f.instruction(&Instruction::If(BlockType::Empty));
482
483    // Assertion failed → trap with message
484    if let Some(msg) = &assert_stmt.message {
485        let (ptr, len) = ctx.intern_string(msg);
486        f.instruction(&Instruction::I32Const(ptr as i32));
487        f.instruction(&Instruction::I32Const(len as i32));
488    } else {
489        f.instruction(&Instruction::I32Const(ctx.data.assert_failed_ptr as i32));
490        f.instruction(&Instruction::I32Const(ctx.data.assert_failed_len as i32));
491    }
492    f.instruction(&Instruction::Call(IMPORT_TRAP));
493    f.instruction(&Instruction::Unreachable);
494    f.instruction(&Instruction::End);
495    Ok(())
496}
497
498fn emit_expr_stmt(expr: &Expr, ctx: &mut FuncContext, f: &mut Function) -> CodegenResult<()> {
499    emit_expr(expr, ctx, f)?;
500    f.instruction(&Instruction::Drop);
501    Ok(())
502}
503
504/// Create a `MemArg`.
505fn memarg(offset: u64, align: u32) -> wasm_encoder::MemArg {
506    wasm_encoder::MemArg {
507        offset,
508        align,
509        memory_index: 0,
510    }
511}