logicaffeine_compile/vm/wasm/
func.rs

1//! Per-function WebAssembly lowering: a region of VM bytecode (`&[Op]`, a register machine)
2//! → a WebAssembly function (a stack machine with i64 locals).
3//!
4//! ## Control flow
5//!
6//! VM bytecode uses absolute jumps; WebAssembly has only *structured* control flow. The
7//! lowering splits the region into basic blocks and emits the standard **dispatch loop**:
8//! a `loop` wrapping `block`s nested one-per-basic-block, with a `br_table` on a "next
9//! block" local at the bottom. Each block runs its arithmetic, then sets the next-block
10//! index and re-dispatches (`br` to the loop) — or `return`s. This translates *any*
11//! reducible (and irreducible) control flow correctly, not just recognized patterns.
12//!
13//! ## Eligibility (Phase 6 deopt-at-concurrency, mirrored)
14//!
15//! Only integer regions are emitted: const/move/arith/compare/jump/return. Any op outside
16//! that set — every concurrency / channel / networking op included — makes
17//! [`compile_region_to_wasm`] return `None`, so the region stays on the bytecode tier (a
18//! yield-free emitted module by construction, exactly like the native backend's deny-list).
19//!
20//! This is the byte emitter both the browser WASM-JIT tier (`super::region_jit`) and the
21//! direct AOT backend consume; it depends only on [`super::encode`], not on the `wasm-jit`
22//! feature.
23
24use super::cfg::{assemble_dispatch_loop, Blocks};
25use super::encode::*;
26use crate::vm::instruction::{CompiledProgram, Constant, Op};
27
28/// Whether a region op is WASM-JIT-eligible (the integer fragment). Everything else —
29/// notably every concurrency / channel / networking op — is rejected so the region stays
30/// on the bytecode tier.
31fn is_supported(op: &Op) -> bool {
32    matches!(
33        op,
34        Op::LoadConst { .. }
35            | Op::Move { .. }
36            | Op::Add { .. }
37            | Op::Sub { .. }
38            | Op::Mul { .. }
39            | Op::Div { .. }
40            | Op::Mod { .. }
41            | Op::BitXor { .. }
42            | Op::BitAnd { .. }
43            | Op::BitOr { .. }
44            | Op::Shl { .. }
45            | Op::Shr { .. }
46            | Op::Lt { .. }
47            | Op::Gt { .. }
48            | Op::LtEq { .. }
49            | Op::GtEq { .. }
50            | Op::Eq { .. }
51            | Op::NotEq { .. }
52            | Op::Jump { .. }
53            | Op::JumpIfFalse { .. }
54            | Op::JumpIfTrue { .. }
55            | Op::Return { .. }
56            | Op::ReturnNothing
57    )
58}
59
60/// Lower an integer region to a WebAssembly module exporting one function `f` of
61/// `num_params` `i64` parameters returning `i64`. VM registers map 1:1 to WASM locals
62/// (`0..num_params` params, `num_params..num_regs` declared i64 locals); one extra i32
63/// local at index `num_regs` is the dispatch "next block" index.
64///
65/// Returns `None` — the region stays on the bytecode tier — for any op outside the integer
66/// fragment (control flow, calls, and *all* concurrency / channel / networking ops), a
67/// non-`Int` constant, an out-of-range jump target, or a region without a reachable
68/// `Return`.
69pub fn compile_region_to_wasm(
70    ops: &[Op],
71    constants: &[Constant],
72    num_params: u32,
73    num_regs: u32,
74) -> Option<Vec<u8>> {
75    let n = ops.len();
76    if n == 0 {
77        return None;
78    }
79
80    // ---- 1. Eligibility: the integer fragment, and a reachable `Return` ----
81    let mut has_return = false;
82    for op in ops {
83        if !is_supported(op) {
84            return None;
85        }
86        if matches!(op, Op::Return { .. }) {
87            has_return = true;
88        }
89    }
90    if !has_return {
91        return None;
92    }
93
94    let blocks = Blocks::new(ops)?;
95    let num_blocks = blocks.num_blocks();
96    let pc_local = num_regs; // the dispatch "next block" index lives just past the registers
97
98    // ---- 2. Emit each block's body + terminator ----
99    let mut blocks_code: Vec<Vec<u8>> = Vec::with_capacity(num_blocks);
100    for k in 0..num_blocks {
101        let start = blocks.start(k);
102        let end = blocks.end(k);
103        let mut code = Vec::new();
104        let mut terminated = false;
105        for pc in start..end {
106            match ops[pc] {
107                Op::LoadConst { dst, idx } => {
108                    let v = match constants.get(idx as usize)? {
109                        Constant::Int(x) => *x,
110                        _ => return None,
111                    };
112                    code.push(0x42); // i64.const
113                    leb_i64(&mut code, v);
114                    local_set(&mut code, dst as u32);
115                }
116                Op::Move { dst, src } => {
117                    local_get(&mut code, src as u32);
118                    local_set(&mut code, dst as u32);
119                }
120                // Add/Sub/Mul are CHECKED: they trap on signed overflow so the task falls
121                // back to the VM (which promotes to BigInt) — never a silent wrap. This keeps
122                // every tier in sync on the EXACT-integer contract.
123                Op::Add { dst, lhs, rhs } => emit_checked_addsub(&mut code, false, dst, lhs, rhs),
124                Op::Sub { dst, lhs, rhs } => emit_checked_addsub(&mut code, true, dst, lhs, rhs),
125                Op::Mul { dst, lhs, rhs } => emit_checked_mul(&mut code, dst, lhs, rhs),
126                // Div/Mod already trap on their edge cases (div-by-zero, i64::MIN / -1) in WASM
127                // → host error → VM fallback, so no extra guard is needed.
128                Op::Div { dst, lhs, rhs } => arith(&mut code, 0x7F, dst, lhs, rhs), // i64.div_s
129                Op::Mod { dst, lhs, rhs } => arith(&mut code, 0x81, dst, lhs, rhs), // i64.rem_s
130                // Bitwise / shift: pure wrapping, no overflow concern. The VM masks the shift
131                // count mod 64 (`wrapping_shl`/`wrapping_shr`); WASM's i64.shl/shr_s do the
132                // same, and shr is ARITHMETIC (signed) to match `wrapping_shr` on i64.
133                Op::BitXor { dst, lhs, rhs } => arith(&mut code, 0x85, dst, lhs, rhs), // i64.xor
134                Op::BitAnd { dst, lhs, rhs } => arith(&mut code, 0x83, dst, lhs, rhs), // i64.and
135                Op::BitOr { dst, lhs, rhs } => arith(&mut code, 0x84, dst, lhs, rhs),  // i64.or
136                Op::Shl { dst, lhs, rhs } => arith(&mut code, 0x86, dst, lhs, rhs),   // i64.shl
137                Op::Shr { dst, lhs, rhs } => arith(&mut code, 0x87, dst, lhs, rhs),   // i64.shr_s
138                Op::Lt { dst, lhs, rhs } => compare(&mut code, 0x53, dst, lhs, rhs), // i64.lt_s
139                Op::Gt { dst, lhs, rhs } => compare(&mut code, 0x55, dst, lhs, rhs), // i64.gt_s
140                Op::LtEq { dst, lhs, rhs } => compare(&mut code, 0x57, dst, lhs, rhs), // i64.le_s
141                Op::GtEq { dst, lhs, rhs } => compare(&mut code, 0x59, dst, lhs, rhs), // i64.ge_s
142                Op::Eq { dst, lhs, rhs } => compare(&mut code, 0x51, dst, lhs, rhs), // i64.eq
143                Op::NotEq { dst, lhs, rhs } => compare(&mut code, 0x52, dst, lhs, rhs), // i64.ne
144                Op::Jump { target } => {
145                    code.push(0x41); // i32.const next-block
146                    leb_u32(&mut code, blocks.block_of(target) as u32);
147                    local_set(&mut code, pc_local);
148                    code.push(0x0C); // br $loop
149                    leb_u32(&mut code, blocks.br_loop(k));
150                    terminated = true;
151                    break;
152                }
153                Op::JumpIfFalse { cond, target } => {
154                    emit_cond_jump(&mut code, cond, true, blocks.block_of(target), blocks.block_of(pc + 1), pc_local, blocks.br_loop(k));
155                    terminated = true;
156                    break;
157                }
158                Op::JumpIfTrue { cond, target } => {
159                    emit_cond_jump(&mut code, cond, false, blocks.block_of(target), blocks.block_of(pc + 1), pc_local, blocks.br_loop(k));
160                    terminated = true;
161                    break;
162                }
163                Op::Return { src } => {
164                    local_get(&mut code, src as u32);
165                    code.push(0x0F); // return
166                    terminated = true;
167                    break;
168                }
169                Op::ReturnNothing => {
170                    // The compiler appends a `ReturnNothing` after a value `Return` as a
171                    // fallthrough safety terminator. For an i64-returning function it is
172                    // unreachable; encode it as a trap so any (erroneous) reachable path
173                    // diverges from the VM and is caught by the differential.
174                    code.push(0x00); // unreachable
175                    terminated = true;
176                    break;
177                }
178                _ => return None,
179            }
180        }
181        if !terminated {
182            // Fell off the block end into the next block: set next, re-dispatch.
183            let next = blocks.block_of(end);
184            code.push(0x41);
185            leb_u32(&mut code, next as u32);
186            local_set(&mut code, pc_local);
187            code.push(0x0C);
188            leb_u32(&mut code, blocks.br_loop(k));
189        }
190        blocks_code.push(code);
191    }
192
193    // ---- 3. Assemble the function body: dispatch loop ----
194    let mut body = assemble_dispatch_loop(pc_local, &blocks_code);
195    body.push(0x00); // unreachable (only reached via $exit / default — never on a Return path)
196    body.push(0x0B); // end function
197
198    // ---- 4. Module sections ----
199    let mut module = vec![0x00, 0x61, 0x73, 0x6D, 0x01, 0x00, 0x00, 0x00];
200
201    let mut ty = Vec::new();
202    leb_u32(&mut ty, 1);
203    ty.push(0x60);
204    leb_u32(&mut ty, num_params);
205    for _ in 0..num_params {
206        ty.push(I64);
207    }
208    leb_u32(&mut ty, 1);
209    ty.push(I64);
210    section(&mut module, 1, &ty);
211
212    let mut func = Vec::new();
213    leb_u32(&mut func, 1);
214    leb_u32(&mut func, 0);
215    section(&mut module, 3, &func);
216
217    let mut export = Vec::new();
218    leb_u32(&mut export, 1);
219    leb_u32(&mut export, 1);
220    export.push(b'f');
221    export.push(0x00);
222    leb_u32(&mut export, 0);
223    section(&mut module, 7, &export);
224
225    // Code section: locals = (num_regs - num_params) i64 + one i32 (the dispatch index).
226    let num_i64_locals = num_regs.saturating_sub(num_params);
227    let mut entry = Vec::new();
228    leb_u32(&mut entry, if num_i64_locals > 0 { 2 } else { 1 }); // local groups
229    if num_i64_locals > 0 {
230        leb_u32(&mut entry, num_i64_locals);
231        entry.push(I64);
232    }
233    leb_u32(&mut entry, 1);
234    entry.push(I32);
235    entry.extend_from_slice(&body);
236    let mut code_sec = Vec::new();
237    leb_u32(&mut code_sec, 1);
238    leb_u32(&mut code_sec, entry.len() as u32);
239    code_sec.extend_from_slice(&entry);
240    section(&mut module, 10, &code_sec);
241
242    Some(module)
243}
244
245/// Lower function `fi` of a compiled program to a WebAssembly module — the integration
246/// entry the VM's tier-up uses. A function's body lives in the shared `program.code` from
247/// its `entry_pc` until the next function's entry (or the program end), with **absolute**
248/// jump targets; this extracts that slice and rebases every jump into the region-local
249/// 0-based space [`compile_region_to_wasm`] expects. Returns `None` (stay on bytecode) if
250/// the body leaves the integer fragment or a jump escapes the function.
251pub fn compile_function_to_wasm(program: &CompiledProgram, fi: usize) -> Option<Vec<u8>> {
252    let f = program.functions.get(fi)?;
253    // The emitted module returns an i64 that `WasmTier::on_call`'s caller boxes as an `Int`.
254    // That is only sound when the function actually returns an Int — a `Bool`-returning
255    // function (a comparison result) would be mis-boxed as `Int(1)` instead of `Bool(true)`
256    // (`Show` would print `1`, not `true`), and `Float` rides different bits. So tier ONLY
257    // declared-`Int` returns; everything else (Bool / Float / inferred-`None`) deopts to the
258    // VM, which types the result correctly. Keeps every tier in sync on return types.
259    if f.ret_kind != Some(crate::vm::native_tier::SlotKind::Int) {
260        return None;
261    }
262    let entry = f.entry_pc;
263    // The function spans [entry, next_entry) — functions are concatenated after Main.
264    let end = program
265        .functions
266        .iter()
267        .map(|g| g.entry_pc)
268        .filter(|&e| e > entry)
269        .min()
270        .unwrap_or(program.code.len());
271    if entry >= end || end > program.code.len() {
272        return None;
273    }
274    let rebase = |t: usize| -> Option<usize> {
275        if t >= entry && t < end {
276            Some(t - entry)
277        } else {
278            None // a jump out of the function body — not a self-contained region
279        }
280    };
281    let mut region: Vec<Op> = Vec::with_capacity(end - entry);
282    for &op in &program.code[entry..end] {
283        region.push(match op {
284            Op::Jump { target } => Op::Jump { target: rebase(target)? },
285            Op::JumpIfFalse { cond, target } => Op::JumpIfFalse { cond, target: rebase(target)? },
286            Op::JumpIfTrue { cond, target } => Op::JumpIfTrue { cond, target: rebase(target)? },
287            other => other,
288        });
289    }
290    compile_region_to_wasm(
291        &region,
292        &program.constants,
293        f.param_count as u32,
294        f.register_count as u32,
295    )
296}
297
298#[cfg(all(test, feature = "wasm-jit", not(target_arch = "wasm32")))]
299mod tests {
300    use super::*;
301
302    /// Run the emitted module's `f` through the pure-Rust wasmi interpreter (instantiation
303    /// also validates the bytes), returning `f(arg)` — the end-to-end codegen proof.
304    fn run(module: &[u8], arg: i64) -> i64 {
305        let engine = wasmi::Engine::default();
306        let m = wasmi::Module::new(&engine, module).expect("emitted bytes are valid wasm");
307        let mut store = wasmi::Store::new(&engine, ());
308        let instance = wasmi::Linker::<()>::new(&engine)
309            .instantiate(&mut store, &m)
310            .unwrap()
311            .start(&mut store)
312            .unwrap();
313        let f = instance.get_typed_func::<i64, i64>(&store, "f").unwrap();
314        f.call(&mut store, arg).unwrap()
315    }
316
317    #[test]
318    fn wasm_jit_emits_and_runs_straight_line() {
319        // f(x) = x*x + x
320        let ops = vec![
321            Op::Mul { dst: 1, lhs: 0, rhs: 0 },
322            Op::Add { dst: 2, lhs: 1, rhs: 0 },
323            Op::Return { src: 2 },
324        ];
325        let module = compile_region_to_wasm(&ops, &[], 1, 3).expect("emits");
326        assert_eq!(run(&module, 5), 30);
327        assert_eq!(run(&module, 7), 56);
328    }
329
330    #[test]
331    fn wasm_jit_emits_and_runs_loop() {
332        // f(n) = n + (n-1) + ... + 1   (triangular number) via a counted loop:
333        //   acc = 0; one = 1; zero = 0;
334        //   loop:  cond = n > zero;  if !cond goto exit;  acc += n;  n -= one;  goto loop;
335        //   exit:  return acc
336        // regs: 0=n, 1=acc, 2=one, 3=zero, 4=cond. consts: [0, 1].
337        let ops = vec![
338            Op::LoadConst { dst: 1, idx: 0 }, // acc = 0
339            Op::LoadConst { dst: 2, idx: 1 }, // one = 1
340            Op::LoadConst { dst: 3, idx: 0 }, // zero = 0
341            Op::Gt { dst: 4, lhs: 0, rhs: 3 }, // [pc 3] cond = n > 0
342            Op::JumpIfFalse { cond: 4, target: 8 }, // if !cond goto 8
343            Op::Add { dst: 1, lhs: 1, rhs: 0 }, // acc += n
344            Op::Sub { dst: 0, lhs: 0, rhs: 2 }, // n -= 1
345            Op::Jump { target: 3 },             // goto 3
346            Op::Return { src: 1 },              // [pc 8] return acc
347        ];
348        let consts = vec![Constant::Int(0), Constant::Int(1)];
349        let module = compile_region_to_wasm(&ops, &consts, 1, 5).expect("loop region emits");
350        assert_eq!(run(&module, 5), 15); // 5+4+3+2+1
351        assert_eq!(run(&module, 10), 55);
352        assert_eq!(run(&module, 1), 1);
353        assert_eq!(run(&module, 0), 0);
354        assert_eq!(run(&module, 100), 5050);
355    }
356
357    #[test]
358    fn wasm_jit_emits_and_runs_branch() {
359        // f(x) = if x < 10 { x * 2 } else { x + 100 }
360        // regs: 0=x, 1=ten, 2=cond, 3=two, 4=result, 5=hundred. consts: [10, 2, 100].
361        let ops = vec![
362            Op::LoadConst { dst: 1, idx: 0 },       // ten = 10
363            Op::Lt { dst: 2, lhs: 0, rhs: 1 },      // cond = x < 10
364            Op::JumpIfFalse { cond: 2, target: 6 }, // if !cond goto else(6)
365            Op::LoadConst { dst: 3, idx: 1 },       // two = 2
366            Op::Mul { dst: 4, lhs: 0, rhs: 3 },     // result = x * 2
367            Op::Return { src: 4 },
368            Op::LoadConst { dst: 5, idx: 2 },       // [pc 6] hundred = 100
369            Op::Add { dst: 4, lhs: 0, rhs: 5 },     // result = x + 100
370            Op::Return { src: 4 },
371        ];
372        let consts = vec![Constant::Int(10), Constant::Int(2), Constant::Int(100)];
373        let module = compile_region_to_wasm(&ops, &consts, 1, 6).expect("branch region emits");
374        assert_eq!(run(&module, 3), 6);
375        assert_eq!(run(&module, 9), 18);
376        assert_eq!(run(&module, 10), 110);
377        assert_eq!(run(&module, 50), 150);
378    }
379
380    #[test]
381    fn wasm_jit_rejects_concurrency_op() {
382        let ops = vec![Op::ChanRecv { dst: 1, chan: 0 }, Op::Return { src: 1 }];
383        assert!(
384            compile_region_to_wasm(&ops, &[], 1, 2).is_none(),
385            "concurrency ops are WASM-JIT-ineligible"
386        );
387    }
388
389    #[test]
390    fn wasm_jit_requires_a_return() {
391        let ops = vec![Op::Add { dst: 1, lhs: 0, rhs: 0 }];
392        assert!(compile_region_to_wasm(&ops, &[], 1, 2).is_none());
393    }
394
395    /// `Not` negates TRUTHINESS — for a raw i64 register the WASM-JIT cannot
396    /// tell a Bool 0/1 from a general Int without kind context, so it MUST
397    /// reject it (region deopts to the VM) rather than guess. This locks the
398    /// soundness boundary: the integer fragment stops where semantics need kinds.
399    #[test]
400    fn wasm_jit_rejects_type_ambiguous_ops() {
401        for op in [
402            Op::Not { dst: 1, src: 0 },
403        ] {
404            let ops = vec![op, Op::Return { src: 1 }];
405            assert!(
406                compile_region_to_wasm(&ops, &[], 1, 2).is_none(),
407                "type-overloaded op must be WASM-JIT-ineligible (deopt, not miscompile): {:?}",
408                ops[0]
409            );
410        }
411    }
412}
logicaffeine_compile/vm/wasm/func.rs

logicaffeine_compile/vm/wasm/
func.rs
