From 2da7cd3cfb9bc824b6b3ba864a25e0c233ab8143 Mon Sep 17 00:00:00 2001 From: Cam Pedersen Date: Mon, 10 Aug 2026 07:59:52 -0400 Subject: [PATCH 1/2] EIR: make tail calls real, from lowering through every backend MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `End::Tail` existed in the IR and all three backends carried code for it, but nothing ever constructed it — the lowering emits every call as `Op::Call` into a register plus a jump to a merge block that returns it. Because the variant was unreachable, the backend paths handling it had rotted unnoticed. Lowering: a new `eir::tailcall` pass recognizes tail position as a shape in the finished IR rather than threading a flag through `lower_expr`. Return threading turns `Jmp(j, args)` into `Ret(args[i])` when `j` does nothing but return its own `i`th parameter; tail marking then turns a trailing `Call(r, f, args)` + `Ret(r)` into `End::Tail(f, args)`. Run to fixpoint, this covers if/match/do/when and multi-clause arity dispatch at once. `Op::Invoke` is deliberately left alone. `End::TailInvoke` is not simply "Invoke in tail position": with a continuation callee it takes the tail-resume path instead of establishing a prompt, and `handle` lowers its body to a thunk called through `Op::Invoke` precisely to create the frame that delimits captured continuations. A static rewrite cannot distinguish those from an ordinary closure call. The call must also be its block's last op, so a following `PopHandler` is never skipped. VM: `End::Tail` left the *caller's* captures installed, so a tail-called function's `Op::Upval` read the wrong frame's upvalues — latent while nothing emitted `End::Tail`, live as soon as the pass fires. Captures are now an explicit argument to `enter_tail_frame`: empty for `Tail` (mirroring `Op::Call`), the closure's own list for `TailInvoke`. wasm: `End::Tail` emitted `call` + `return` rather than `return_call`, so tail recursion grew the wasm stack; the legacy codegen had done this correctly all along. Single-block functions were worse — they fell through to a catch-all returning Unit, so `[fn f [n] [g n]]` never called `g`. native: tail calls were plain calls, and `End::TailInvoke` returned Unit without performing the call at all. Loon functions now use Cranelift's `tail` calling convention so `End::Tail` lowers to `return_call`, with a C-ABI trampoline for the entry point; `TailInvoke` is a hard error until the backend models closures. VM dispatch, independent of the above: the register-file size was recomputed by scanning every block and op of the callee on every call and every tail call, making a tail-recursive loop cost O(code size) per iteration. It is now computed once per function. The interpreter loop also cloned each op before executing it — a heap allocation per instruction for every call, vector, map, ADT and builtin — which the borrow checker no longer requires now that the block is read out of the loop's own `Rc` clone. 10M-deep mutual recursion on the VM: 2.67 GB peak RSS and 1.55s before, 9.8 MB and 0.35s after. All 19 samples produce byte-identical output against the baseline binary (except `replay-demo.oo`, which is documented to branch on the wall clock). Co-Authored-By: Claude Opus 5 --- crates/loon-lang/src/eir/lower.rs | 60 +++++- crates/loon-lang/src/eir/mod.rs | 1 + crates/loon-lang/src/eir/native.rs | 216 ++++++++++++++++++++-- crates/loon-lang/src/eir/tailcall.rs | 266 +++++++++++++++++++++++++++ crates/loon-lang/src/eir/vm.rs | 260 ++++++++++++++++++-------- crates/loon-lang/src/eir/wasm.rs | 183 ++++++++++++++++-- 6 files changed, 873 insertions(+), 113 deletions(-) create mode 100644 crates/loon-lang/src/eir/tailcall.rs diff --git a/crates/loon-lang/src/eir/lower.rs b/crates/loon-lang/src/eir/lower.rs index 68a69bf..19f8ba4 100644 --- a/crates/loon-lang/src/eir/lower.rs +++ b/crates/loon-lang/src/eir/lower.rs @@ -21,7 +21,11 @@ use std::path::PathBuf; pub fn lower(checker: &Checker) -> Module { let mut ctx = Lower::new(checker); ctx.lower_program(); - ctx.finish() + let mut module = ctx.finish(); + // Calls are emitted uniformly as `Op::Call` + a jump to a merge block; + // which of them are in tail position is recognized on the finished IR. + super::tailcall::mark_tail_calls(&mut module); + module } // ─── Lowering context ────────────────────────────────────────────────────── @@ -2815,6 +2819,60 @@ mod tests { assert!(module.funcs.len() >= 2); // __main + add } + /// `[recur ...]` is the one tail construct the lowering does emit: it + /// becomes `End::Recur`, a jump back to block 0, so it costs no frame. + #[test] + fn recur_lowers_to_end_recur() { + let module = lower_src( + r#" + [fn countdown [n] [if [= n 0] :done [recur [- n 1]]]] + [countdown 3] + "#, + ); + let countdown = module + .funcs + .iter() + .find(|f| f.name.as_deref() == Some("countdown")) + .expect("countdown should be lowered"); + assert!( + countdown + .blocks + .iter() + .any(|b| matches!(b.end, End::Recur(_))), + "recur should lower to End::Recur" + ); + } + + /// A call in tail position becomes `End::Tail`. The lowering itself still + /// emits every call the same way — `Op::Call` into a register plus a jump + /// to a merge block — and `eir::tailcall` recognizes the tail-position + /// shape afterwards. This test guards the end-to-end result. + #[test] + fn mutual_recursion_lowers_to_a_tail_call() { + let module = lower_src( + r#" + [fn even? [n] [if [= n 0] true [odd? [- n 1]]]] + [fn odd? [n] [if [= n 0] false [even? [- n 1]]]] + [even? 10] + "#, + ); + for name in ["even?", "odd?"] { + let f = module + .funcs + .iter() + .find(|f| f.name.as_deref() == Some(name)) + .unwrap_or_else(|| panic!("{name} should be lowered")); + assert_eq!( + f.blocks + .iter() + .filter(|b| matches!(b.end, End::Tail(..))) + .count(), + 1, + "{name} should end in a tail call" + ); + } + } + #[test] fn lower_if() { let module = lower_src("[if true 1 2]"); diff --git a/crates/loon-lang/src/eir/mod.rs b/crates/loon-lang/src/eir/mod.rs index 2329142..62955c6 100644 --- a/crates/loon-lang/src/eir/mod.rs +++ b/crates/loon-lang/src/eir/mod.rs @@ -9,6 +9,7 @@ pub mod lower; pub mod native; pub mod net; pub mod replay; +pub mod tailcall; pub mod trace; pub mod value64; pub mod vm; diff --git a/crates/loon-lang/src/eir/native.rs b/crates/loon-lang/src/eir/native.rs index 3ea5b36..19b28bf 100644 --- a/crates/loon-lang/src/eir/native.rs +++ b/crates/loon-lang/src/eir/native.rs @@ -10,7 +10,7 @@ //! - Arithmetic, comparison, and logic binary ops //! - Unary ops (neg, not) //! - Mov, branches, jumps, returns -//! - Function calls (direct) +//! - Function calls (direct), including tail calls (`return_call`) //! - Builtin println (via extern) //! //! Not yet implemented (fall back to VM): @@ -18,12 +18,18 @@ //! - Collection construction (Vec, Map, Set, Tuple, ADT) //! - Field access, tag extraction //! - Effect operations (perform, push/pop handler) -//! - Tail calls (compiled as regular calls + return) //! - String operations +//! +//! Loon functions are compiled with Cranelift's `tail` calling convention so +//! that `End::Tail` can lower to a real `return_call` (constant stack for +//! mutual tail recursion). That convention is not the platform C ABI, so the +//! entry point is reached through a small C-ABI trampoline — see +//! `ENTRY_TRAMPOLINE`. use cranelift_codegen::ir::condcodes::IntCC; use cranelift_codegen::ir::types::I64; use cranelift_codegen::ir::{AbiParam, Function, InstBuilder, Signature, UserFuncName}; +use cranelift_codegen::isa::CallConv; use cranelift_codegen::settings::{self, Configurable}; use cranelift_codegen::Context; use cranelift_frontend::{FunctionBuilder, FunctionBuilderContext, Variable}; @@ -50,6 +56,9 @@ const VAL_TRUE: u64 = BASE | TAG_IMM | 1; const VAL_FALSE: u64 = BASE | TAG_IMM | 2; const VAL_NONE: u64 = BASE | TAG_IMM | 3; +/// Symbol name of the C-ABI shim that calls the module's entry function. +const ENTRY_TRAMPOLINE: &str = "loon_entry_trampoline"; + // ─── Runtime helper functions ─────────────────────────────────────────────── /// Runtime: println a NaN-boxed value. Called from compiled code. @@ -192,12 +201,16 @@ impl NativeModule { let mut jit = JITModule::new(builder); - // Create the default calling convention signature: all args and return are i64. - let call_conv = jit.isa().default_call_conv(); + // The platform C ABI, used for the runtime helpers (plain Rust + // `extern "C"` functions) and for the entry trampoline. + let c_call_conv = jit.isa().default_call_conv(); + // Loon functions use the `tail` convention instead: `return_call` + // requires caller and callee to share a tail-call-capable convention. + let call_conv = CallConv::Tail; // Declare runtime helper functions. let rt_println_sig = { - let mut sig = Signature::new(call_conv); + let mut sig = Signature::new(c_call_conv); sig.params.push(AbiParam::new(I64)); sig.returns.push(AbiParam::new(I64)); sig @@ -302,15 +315,53 @@ impl NativeModule { })?; } + // Entry trampoline: `execute()` calls the module through a plain C + // function pointer, but the entry itself uses the `tail` convention, + // which is not the C ABI. Bridge the two with a C-ABI shim that calls + // the entry and returns its result. + let entry_cl_id = func_map.funcs[&eir_module.entry.0]; + let trampoline_id = { + let mut sig = Signature::new(c_call_conv); + sig.returns.push(AbiParam::new(I64)); + let id = jit + .declare_function(ENTRY_TRAMPOLINE, Linkage::Local, &sig) + .map_err(|e| Error { + message: format!("declare {ENTRY_TRAMPOLINE}: {e}"), + phase: "native:declare", + })?; + + let mut cl_func = Function::with_name_signature( + UserFuncName::user(0, eir_module.funcs.len() as u32), + sig, + ); + { + let mut builder = FunctionBuilder::new(&mut cl_func, &mut fb_ctx); + let block = builder.create_block(); + builder.switch_to_block(block); + builder.seal_block(block); + let entry_ref = jit.declare_func_in_func(entry_cl_id, builder.func); + let call = builder.ins().call(entry_ref, &[]); + let result = builder.inst_results(call)[0]; + builder.ins().return_(&[result]); + builder.finalize(); + } + let mut ctx = Context::for_function(cl_func); + jit.define_function(id, &mut ctx).map_err(|e| Error { + message: format!("define {ENTRY_TRAMPOLINE}: {e}"), + phase: "native:codegen", + })?; + id + }; + // Finalize all definitions. jit.finalize_definitions().map_err(|e| Error { message: format!("finalize: {e}"), phase: "native:finalize", })?; - // Get entry function pointer. - let entry_cl_id = func_map.funcs[&eir_module.entry.0]; - let entry_fn = jit.get_finalized_function(entry_cl_id); + // Get entry function pointer (the C-ABI trampoline, not the entry + // itself — see above). + let entry_fn = jit.get_finalized_function(trampoline_id); Ok(NativeModule { _jit: jit, @@ -922,8 +973,9 @@ fn compile_terminator( } End::Tail(func_id, args) => { - // Compile tail calls as regular calls + return (no TCO in Cranelift - // for our calling convention yet). + // A real tail call: `return_call` replaces the current frame, so + // mutual tail recursion runs in constant stack. This is why loon + // functions are compiled with `CallConv::Tail`. let cl_func_id = func_map.funcs.get(&func_id.0).ok_or_else(|| Error { message: format!("unknown tail call target {}", func_id.0), phase: "native:compile", @@ -933,15 +985,19 @@ fn compile_terminator( .iter() .map(|r| builder.use_var(vars[r.0 as usize])) .collect(); - let call = builder.ins().call(func_ref, &arg_vals); - let result = builder.inst_results(call)[0]; - builder.ins().return_(&[result]); + builder.ins().return_call(func_ref, &arg_vals); } End::TailInvoke(_callee, _args) => { - // Indirect tail calls need closure support. - let unit = builder.ins().iconst(I64, VAL_UNIT as i64); - builder.ins().return_(&[unit]); + // Indirect tail calls need closures, which this backend does not + // represent yet (`Op::Close`/`Op::Invoke` are still stubs). Fail + // loudly rather than returning Unit — a silent wrong answer is far + // worse than a missing feature. + return Err(Error { + message: "tail call to a closure is not supported by the native backend yet" + .to_string(), + phase: "native:compile", + }); } End::Recur(args) => { @@ -1087,6 +1143,134 @@ mod tests { assert_eq!(result.as_int(), 120); } + /// Mutual tail recursion from source must run in constant stack. Compiled + /// as a plain call + return this recurses a million frames deep and + /// overflows, so reaching the assert is the evidence that `End::Tail` is + /// both emitted by the lowering and lowered to `return_call` here. + #[test] + fn tail_calls_run_in_constant_stack_from_source() { + let src = r#" + [fn even? [n] [if [= n 0] true [odd? [- n 1]]]] + [fn odd? [n] [if [= n 0] false [even? [- n 1]]]] + [even? 1000000] + "#; + assert_eq!(eval_native(src).unwrap(), Val::TRUE); + } + + /// The same property stated directly against `End::Tail`, independent of + /// what the lowering happens to produce. + #[test] + fn tail_calls_run_in_constant_stack() { + use crate::eir::{Block, BlockId, Func, FuncId, Ty}; + use crate::syntax::Span; + + // `even(n) = n == 0 ? true : odd(n - 1)`, and vice versa. + let parity = |id: u32, other: u32, base: bool| Func { + id: FuncId(id), + name: Some(format!("parity{id}")), + params: vec![Ty::Int], + ret: Ty::Bool, + evidence: vec![], + captures: vec![], + blocks: vec![ + Block { + id: BlockId(0), + params: vec![Reg(0)], + ops: vec![ + Op::Lit(Reg(1), Lit::Int(0), Span::ZERO), + Op::Bin(Reg(2), BinOp::Eq, Reg(0), Reg(1), Span::ZERO), + ], + end: End::Br(Reg(2), BlockId(1), BlockId(2)), + }, + Block { + id: BlockId(1), + params: vec![], + ops: vec![Op::Lit(Reg(3), Lit::Bool(base), Span::ZERO)], + end: End::Ret(Reg(3)), + }, + Block { + id: BlockId(2), + params: vec![], + ops: vec![ + Op::Lit(Reg(4), Lit::Int(1), Span::ZERO), + Op::Bin(Reg(5), BinOp::Sub, Reg(0), Reg(4), Span::ZERO), + ], + end: End::Tail(FuncId(other), vec![Reg(5)]), + }, + ], + span: Span::ZERO, + is_closure: false, + }; + + let module = crate::eir::Module { + funcs: vec![ + parity(0, 1, true), + parity(1, 0, false), + Func { + id: FuncId(2), + name: Some("__main".to_string()), + params: vec![], + ret: Ty::Bool, + evidence: vec![], + captures: vec![], + blocks: vec![Block { + id: BlockId(0), + params: vec![], + ops: vec![ + Op::Lit(Reg(0), Lit::Int(1_000_000), Span::ZERO), + Op::Call(Reg(1), FuncId(0), vec![Reg(0)], Span::ZERO), + ], + end: End::Ret(Reg(1)), + }], + span: Span::ZERO, + is_closure: false, + }, + ], + strings: vec![], + ctors: vec![], + entry: FuncId(2), + }; + + let mut backend = NativeBackend; + let native = backend.compile(&module).expect("compilation failed"); + assert_eq!(native.execute(), Val::TRUE); + } + + /// A tail call to a closure has no lowering yet. It must fail loudly — + /// it used to return Unit without performing the call at all. + #[test] + fn tail_invoke_is_an_error_not_a_silent_unit() { + use crate::eir::{Block, BlockId, Func, FuncId, Ty}; + use crate::syntax::Span; + + let module = crate::eir::Module { + funcs: vec![Func { + id: FuncId(0), + name: Some("__main".to_string()), + params: vec![], + ret: Ty::Any, + evidence: vec![], + captures: vec![], + blocks: vec![Block { + id: BlockId(0), + params: vec![], + ops: vec![Op::Lit(Reg(0), Lit::Int(1), Span::ZERO)], + end: End::TailInvoke(Reg(0), vec![Reg(0)]), + }], + span: Span::ZERO, + is_closure: false, + }], + strings: vec![], + ctors: vec![], + entry: FuncId(0), + }; + + match NativeBackend.compile(&module) { + Ok(_) => panic!("tail invoke should not compile silently"), + Err(e) => assert!(e.message.contains("closure"), "unexpected error: {e}"), + } + } + #[test] fn compile_division() { let result = eval_native("[/ 10 3]").unwrap(); diff --git a/crates/loon-lang/src/eir/tailcall.rs b/crates/loon-lang/src/eir/tailcall.rs new file mode 100644 index 0000000..ee0a26f --- /dev/null +++ b/crates/loon-lang/src/eir/tailcall.rs @@ -0,0 +1,266 @@ +//! Tail-call marking: rewrite calls in tail position to `End::Tail`. +//! +//! The lowering builds every call the same way — `Op::Call` into a register, +//! then a jump to whatever block merges the surrounding `if`/`match`/clause +//! arms, which returns that register. Tail position is therefore not a +//! property the lowering tracks; it is a *shape* in the finished IR: +//! +//! ```text +//! b2: ...; Call(r6, odd?, [r5]) b2: ... +//! Jmp(b3, [r6]) => Tail(odd?, [r5]) +//! b3(p7): Ret(p7) b3(p7): Ret(p7) (now unreachable) +//! ``` +//! +//! Recognizing the shape after the fact covers every construct at once — `if`, +//! `match`, `do`, `when`, multi-clause arity dispatch — without threading a +//! tail-position flag through all of `lower_expr`. +//! +//! Two rewrites, applied to fixpoint: +//! +//! 1. **Return threading** — a `Jmp` to a block that does nothing but return +//! one of its own parameters becomes a `Ret` of the corresponding argument. +//! 2. **Tail marking** — a block whose last op is a `Call` producing exactly +//! the register the block then returns becomes an `End::Tail`, dropping the +//! op. +//! +//! ## Why only `Op::Call` +//! +//! `Op::Invoke` (closure / function-pointer calls) is deliberately left alone. +//! `End::TailInvoke` is not simply "Invoke in tail position" in the VM: when +//! its callee is a continuation it takes the *tail-resume* path, which reuses +//! the frame below instead of establishing a fresh prompt. `handle` also +//! lowers its body to a thunk called through `Op::Invoke` precisely to create +//! a prompt frame — eliding that frame would move the boundary that delimits +//! captured continuations. A static rewrite cannot tell those cases apart from +//! an ordinary closure call, so it must not try. + +use super::{End, Func, Module, Op}; + +/// Rewrite tail-position calls across every function in the module. +pub fn mark_tail_calls(module: &mut Module) { + for func in &mut module.funcs { + // Threading can expose new tail calls and marking can expose new + // threadable jumps, so alternate until neither fires. + while thread_returns(func) | mark_calls(func) {} + } +} + +/// Rewrite `Jmp(j, args)` to `Ret(args[i])` when `j` is a pure return block: +/// no ops, and its terminator returns its own `i`th parameter. +/// +/// This is what collapses the merge block that `if`/`match` arms jump to. The +/// merge block itself is left in place — other predecessors may still use it, +/// and an unreachable block is harmless. +fn thread_returns(func: &mut Func) -> bool { + let mut changed = false; + + for idx in 0..func.blocks.len() { + let End::Jmp(target, ref args) = func.blocks[idx].end else { + continue; + }; + let target = target.0 as usize; + if target == idx || target >= func.blocks.len() { + continue; + } + + let dest = &func.blocks[target]; + // Only a block that does nothing but return a parameter is safe to + // inline into its predecessor. Returning a non-parameter register would + // move the read to a block where that register may not be defined. + if !dest.ops.is_empty() || dest.params.len() != args.len() { + continue; + } + let End::Ret(returned) = dest.end else { + continue; + }; + let Some(param_idx) = dest.params.iter().position(|p| *p == returned) else { + continue; + }; + + let arg = args[param_idx]; + func.blocks[idx].end = End::Ret(arg); + changed = true; + } + + changed +} + +/// Rewrite a trailing `Call(r, f, args)` + `Ret(r)` into `End::Tail(f, args)`. +/// +/// The call must be the block's *last* op: anything emitted after it (a +/// `PopHandler`, say) still has to run before the function returns, and a tail +/// call would skip it. +fn mark_calls(func: &mut Func) -> bool { + let mut changed = false; + + for block in &mut func.blocks { + let End::Ret(returned) = block.end else { + continue; + }; + let Some(Op::Call(dst, callee, args, _)) = block.ops.last() else { + continue; + }; + if *dst != returned { + continue; + } + + block.end = End::Tail(*callee, args.clone()); + block.ops.pop(); + changed = true; + } + + changed +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::check::Checker; + use crate::eir::{BlockId, FuncId, Reg}; + use crate::parser::parse; + + fn lower_src(src: &str) -> Module { + let exprs = parse(src).expect("parse failed"); + let mut checker = Checker::new(); + let _ = checker.check_program(&exprs); + crate::eir::lower::lower(&checker) + } + + fn func<'a>(module: &'a Module, name: &str) -> &'a Func { + module + .funcs + .iter() + .find(|f| f.name.as_deref() == Some(name)) + .unwrap_or_else(|| panic!("no function named {name}")) + } + + fn tail_targets(f: &Func) -> Vec { + f.blocks + .iter() + .filter_map(|b| match &b.end { + End::Tail(fid, _) => Some(*fid), + _ => None, + }) + .collect() + } + + #[test] + fn mutual_recursion_becomes_tail_calls() { + let module = lower_src( + r#" + [fn even? [n] [if [= n 0] true [odd? [- n 1]]]] + [fn odd? [n] [if [= n 0] false [even? [- n 1]]]] + [even? 10] + "#, + ); + let even = func(&module, "even?"); + let odd = func(&module, "odd?"); + assert_eq!(tail_targets(even).len(), 1, "even? should tail-call odd?"); + assert_eq!(tail_targets(odd).len(), 1, "odd? should tail-call even?"); + assert_eq!(tail_targets(even)[0], odd.id); + assert_eq!(tail_targets(odd)[0], even.id); + } + + #[test] + fn tail_call_through_match_is_marked() { + let module = lower_src( + r#" + [fn go [n] [match n 0 :done _ [go [- n 1]]]] + [go 3] + "#, + ); + let go = func(&module, "go"); + assert_eq!(tail_targets(go), vec![go.id]); + } + + #[test] + fn tail_call_through_do_is_marked() { + let module = lower_src( + r#" + [fn go [n] [if [= n 0] :done [do [let _ 0] [go [- n 1]]]]] + [go 3] + "#, + ); + assert_eq!(tail_targets(func(&module, "go")).len(), 1); + } + + /// A call whose result is used is not in tail position and must stay a + /// `Call` — turning it into a tail call would skip the multiply. + #[test] + fn non_tail_call_is_left_alone() { + let module = lower_src( + r#" + [fn fact [n] [if [<= n 1] 1 [* n [fact [- n 1]]]]] + [fact 5] + "#, + ); + let fact = func(&module, "fact"); + assert!( + tail_targets(fact).is_empty(), + "a call feeding an arithmetic op is not a tail call" + ); + assert!( + fact.blocks + .iter() + .any(|b| b.ops.iter().any(|o| matches!(o, Op::Call(..)))), + "the recursive call should survive as an Op::Call" + ); + } + + /// `Op::Invoke` is never rewritten — see the module docs. + #[test] + fn closure_calls_are_not_marked() { + let module = lower_src( + r#" + [fn apply-it [f x] [f x]] + [apply-it [fn [y] [+ y 1]] 1] + "#, + ); + for f in &module.funcs { + for b in &f.blocks { + assert!( + !matches!(b.end, End::TailInvoke(..)), + "invoke must not be rewritten to a tail invoke" + ); + } + } + } + + /// Ops emitted after the call still have to run, so the call is not in + /// tail position even though its result is returned. + #[test] + fn call_followed_by_another_op_is_not_marked() { + use crate::eir::{Block, Lit, Ty}; + use crate::syntax::Span; + + let mut module = Module { + funcs: vec![Func { + id: FuncId(0), + name: Some("f".to_string()), + params: vec![], + ret: Ty::Any, + evidence: vec![], + captures: vec![], + blocks: vec![Block { + id: BlockId(0), + params: vec![], + ops: vec![ + Op::Call(Reg(0), FuncId(0), vec![], Span::ZERO), + Op::PopHandler(Span::ZERO), + Op::Lit(Reg(1), Lit::Unit, Span::ZERO), + ], + end: End::Ret(Reg(0)), + }], + span: Span::ZERO, + is_closure: false, + }], + strings: vec![], + ctors: vec![], + entry: FuncId(0), + }; + + mark_tail_calls(&mut module); + assert!(matches!(module.funcs[0].blocks[0].end, End::Ret(_))); + assert_eq!(module.funcs[0].blocks[0].ops.len(), 3); + } +} diff --git a/crates/loon-lang/src/eir/vm.rs b/crates/loon-lang/src/eir/vm.rs index ee32bdd..a3942f3 100644 --- a/crates/loon-lang/src/eir/vm.rs +++ b/crates/loon-lang/src/eir/vm.rs @@ -279,6 +279,27 @@ pub struct Vm { /// immediate singleton `Val::NONE` at every construction site, so `None` /// never allocates and `is_truthy` stays a pure bit test. none_tag: Option, + /// Register-file size per function, indexed by `FuncId`. Computed once at + /// VM construction (and extended for funcs appended later) because the + /// scan is O(function size) — doing it per call made every tail call cost + /// proportional to the *callee's code size* rather than O(1), which defeats + /// the point of having proper tail calls in the first place. + reg_counts: Vec, +} + +/// Number of registers a function's frame needs: one past the highest register +/// mentioned as an op destination or a block parameter. +fn reg_count_of(func: &Func) -> usize { + func.blocks + .iter() + .flat_map(|b| { + b.ops + .iter() + .map(|op| op.dst().0 + 1) + .chain(b.params.iter().map(|r| r.0 + 1)) + }) + .max() + .unwrap_or(0) as usize } /// Heap allocation statistics collected during VM execution. @@ -313,6 +334,7 @@ impl Vm { .rev() .find(|c| c.name == "None") .map(|c| c.tag); + let reg_counts = module.funcs.iter().map(reg_count_of).collect(); Self { module: Rc::new(module), heap: Vec::new(), @@ -334,9 +356,44 @@ impl Vm { runtime_syms: Vec::new(), rand_state: None, none_tag, + reg_counts, } } + /// Register-file size for `func`, from the table built at construction. + /// Falls back to computing (and caching) it for functions appended to the + /// module after the VM was created, so the table can never go stale. + fn reg_count(&mut self, func: FuncId) -> usize { + let idx = func.0 as usize; + if idx >= self.reg_counts.len() { + let module = Rc::clone(&self.module); + while self.reg_counts.len() <= idx { + let f = &module.funcs[self.reg_counts.len()]; + self.reg_counts.push(reg_count_of(f)); + } + } + self.reg_counts[idx] + } + + /// Resize the current register file for a frame that is being *replaced* + /// in place (a tail call), load `vals` into the argument registers, and + /// install the callee's captures. + /// + /// `captures` must be what the equivalent non-tail call would have passed: + /// empty for `End::Tail` (mirroring `Op::Call`), the closure's own capture + /// list for `End::TailInvoke` (mirroring `Op::Invoke`). Leaving the + /// *caller's* captures in place would let the callee's `Op::Upval` read + /// the wrong frame's upvalues. + fn enter_tail_frame(&mut self, func: FuncId, vals: &[Val], captures: Vec) { + let needed = self.reg_count(func).max(vals.len()) + 16; + self.regs.resize(needed, Val::UNIT); + self.regs[..vals.len()].copy_from_slice(vals); + self.captures = captures; + self.func = func; + self.block = BlockId(0); + self.ip = 0; + } + /// Construct an ADT value, normalizing the nullary `None` constructor to /// the immediate singleton `Val::NONE`. Every ADT construction site MUST /// go through this (never `alloc(Obj::Adt(..))` directly), otherwise a @@ -577,19 +634,7 @@ impl Vm { } // Set up new frame - let func = &self.module.funcs[func_id.0 as usize]; - let reg_count = func - .blocks - .iter() - .flat_map(|b| { - b.ops - .iter() - .map(|op| op.dst().0 + 1) - .chain(b.params.iter().map(|r| r.0 + 1)) - }) - .max() - .unwrap_or(0) as usize; - let reg_count = reg_count.max(args.len()); + let reg_count = self.reg_count(func_id).max(args.len()); self.regs = vec![Val::UNIT; reg_count + 16]; // padding for safety for (i, &val) in args.iter().enumerate() { @@ -725,25 +770,29 @@ impl Vm { let module = Rc::clone(&self.module); loop { - // Fetch current instruction or terminator by index (no stale refs) + // Borrow the current block out of our own `Rc` clone rather + // than out of `self`, so ops and terminators can be read by + // reference while `self` is mutably borrowed. `module` keeps the + // code alive for the whole loop, so these borrows stay valid across + // `exec_op` — cloning each op just to satisfy the borrow checker + // cost a heap allocation per instruction with a `Vec` operand + // (every call, vector, map, ADT and builtin). let func_idx = self.func.0 as usize; let block_idx = self.block.0 as usize; - let ops_len = module.funcs[func_idx].blocks[block_idx].ops.len(); + let block = &module.funcs[func_idx].blocks[block_idx]; - if self.ip < ops_len { - // Clone the op to avoid borrowing module across exec_op - let op = module.funcs[func_idx].blocks[block_idx].ops[self.ip].clone(); + if self.ip < block.ops.len() { + let op = &block.ops[self.ip]; self.current_span = op.span(); self.ip += 1; - self.exec_op(&op)?; + self.exec_op(op)?; continue; } // Execute terminator - let end = module.funcs[func_idx].blocks[block_idx].end.clone(); - match end { + match &block.end { End::Ret(reg) => { - let val = self.r(reg); + let val = self.r(*reg); if self.frames.is_empty() { return Ok(val); } @@ -757,26 +806,24 @@ impl Vm { } } - End::Jmp(target, ref args) => { + End::Jmp(target, args) => { let vals = self.read_regs(args); - let params: Vec = module.funcs[func_idx].blocks[target.0 as usize] - .params - .clone(); + let params = &module.funcs[func_idx].blocks[target.0 as usize].params; for (param, val) in params.iter().zip(vals.iter()) { self.w(*param, *val); } - self.block = target; + self.block = *target; self.ip = 0; } End::Br(cond, then_b, else_b) => { - let v = self.r(cond); - self.block = if v.is_truthy() { then_b } else { else_b }; + let v = self.r(*cond); + self.block = if v.is_truthy() { *then_b } else { *else_b }; self.ip = 0; } - End::Switch(scrutinee, ref cases, default) => { - let v = self.r(scrutinee); + End::Switch(scrutinee, cases, default) => { + let v = self.r(*scrutinee); let tag = if v.is_none() { self.none_tag.unwrap_or(0) } else if let Some(Obj::Adt(t, _)) = self.get_obj(v) { @@ -790,37 +837,18 @@ impl Vm { .iter() .find(|(t, _)| *t == tag) .map(|(_, b)| *b) - .unwrap_or(default); + .unwrap_or(*default); self.block = target; self.ip = 0; } - End::Tail(func_id, ref args) => { + End::Tail(func_id, args) => { let vals = self.read_regs(args); - let f = &module.funcs[func_id.0 as usize]; - let reg_count = f - .blocks - .iter() - .flat_map(|b| { - b.ops - .iter() - .map(|op| op.dst().0 + 1) - .chain(b.params.iter().map(|r| r.0 + 1)) - }) - .max() - .unwrap_or(0) as usize; - let needed = reg_count.max(vals.len()) + 16; - self.regs.resize(needed, Val::UNIT); - for (i, &val) in vals.iter().enumerate() { - self.regs[i] = val; - } - self.func = func_id; - self.block = BlockId(0); - self.ip = 0; + self.enter_tail_frame(*func_id, &vals, Vec::new()); } - End::TailInvoke(callee, ref args) => { - let func_val = self.r(callee); + End::TailInvoke(callee, args) => { + let func_val = self.r(*callee); let vals = self.read_regs(args); if matches!(self.get_obj(func_val), Some(Obj::Continuation { .. })) { // Tail resume (`[resume v]` as the handler's whole body): @@ -830,27 +858,7 @@ impl Vm { let v = vals.first().copied().unwrap_or(Val::UNIT); self.resume_continuation(func_val, v, None)?; } else if let Some(Obj::Closure(fid, caps)) = self.get_obj(func_val).cloned() { - let f = &module.funcs[fid.0 as usize]; - let reg_count = f - .blocks - .iter() - .flat_map(|b| { - b.ops - .iter() - .map(|op| op.dst().0 + 1) - .chain(b.params.iter().map(|r| r.0 + 1)) - }) - .max() - .unwrap_or(0) as usize; - let needed = reg_count.max(vals.len()) + 16; - self.regs.resize(needed, Val::UNIT); - for (i, &val) in vals.iter().enumerate() { - self.regs[i] = val; - } - self.captures = caps; - self.func = fid; - self.block = BlockId(0); - self.ip = 0; + self.enter_tail_frame(fid, &vals, caps); } else { return Err( VmError::new(VmErrorKind::NotCallable).with_span(self.current_span) @@ -858,9 +866,9 @@ impl Vm { } } - End::Recur(ref args) => { + End::Recur(args) => { let vals = self.read_regs(args); - let params: Vec = module.funcs[func_idx].blocks[0].params.clone(); + let params = &module.funcs[func_idx].blocks[0].params; for (param, val) in params.iter().zip(vals.iter()) { self.w(*param, *val); } @@ -3676,6 +3684,102 @@ mod tests { eval_eir(src).expect("vm error").output } + /// `End::Tail` must be exactly `Op::Call` + `End::Ret` with the frame + /// reused — including captures. `Op::Call` enters the callee with an empty + /// capture list; a tail call that left the *caller's* captures installed + /// would let the callee's `Op::Upval` read the caller's upvalues instead. + /// + /// Built by hand: the lowering only marks tail calls to named top-level + /// functions, which never carry captures, so source cannot express this. + #[test] + fn vm_tail_call_does_not_inherit_caller_captures() { + use crate::eir::{Block, Ty}; + use crate::syntax::Span; + + // callee() = upval 0 — it has no captures of its own, so this must + // read as Unit however it was entered. + let callee = |id: u32| Func { + id: FuncId(id), + name: Some("callee".to_string()), + params: vec![], + ret: Ty::Any, + evidence: vec![], + captures: vec![], + blocks: vec![Block { + id: BlockId(0), + params: vec![], + ops: vec![Op::Upval(Reg(0), 0, Span::ZERO)], + end: End::Ret(Reg(0)), + }], + span: Span::ZERO, + is_closure: false, + }; + + // A closure capturing 7 that reaches `callee` either way. + let caller = |tail: bool| Func { + id: FuncId(1), + name: Some("caller".to_string()), + params: vec![], + ret: Ty::Any, + evidence: vec![], + captures: vec![], + blocks: vec![Block { + id: BlockId(0), + params: vec![], + ops: if tail { + vec![] + } else { + vec![Op::Call(Reg(0), FuncId(0), vec![], Span::ZERO)] + }, + end: if tail { + End::Tail(FuncId(0), vec![]) + } else { + End::Ret(Reg(0)) + }, + }], + span: Span::ZERO, + is_closure: true, + }; + + let build = |tail: bool| Module { + funcs: vec![ + callee(0), + caller(tail), + Func { + id: FuncId(2), + name: Some("__main".to_string()), + params: vec![], + ret: Ty::Any, + evidence: vec![], + captures: vec![], + blocks: vec![Block { + id: BlockId(0), + params: vec![], + ops: vec![ + Op::Lit(Reg(0), Lit::Int(7), Span::ZERO), + Op::Close(Reg(1), FuncId(1), vec![Reg(0)], Span::ZERO), + Op::Invoke(Reg(2), Reg(1), vec![], Span::ZERO), + ], + end: End::Ret(Reg(2)), + }], + span: Span::ZERO, + is_closure: false, + }, + ], + strings: vec![], + ctors: vec![], + entry: FuncId(2), + }; + + let via_call = Vm::new(build(false)).run().expect("vm error").value; + let via_tail = Vm::new(build(true)).run().expect("vm error").value; + assert_eq!(via_call, Val::UNIT); + assert_eq!( + via_tail, via_call, + "a tail call must not leak the caller's captures into the callee" + ); + } + #[test] fn vm_durable_resume() { // Durable execution as a handler: a RESUME tower replays a journal of diff --git a/crates/loon-lang/src/eir/wasm.rs b/crates/loon-lang/src/eir/wasm.rs index a3f6278..4b4623c 100644 --- a/crates/loon-lang/src/eir/wasm.rs +++ b/crates/loon-lang/src/eir/wasm.rs @@ -115,6 +115,12 @@ enum WasmInstr { GlobalSet(u32), Call(u32), CallIndirect(u32), + /// `return_call` — a proper tail call. Replaces the current frame instead + /// of stacking a new one, so tail recursion runs in constant stack. + ReturnCall(u32), + /// `return_call_indirect` — the closure/function-pointer form (arity keys + /// the shared `(i64…) -> i64` type, same as `CallIndirect`). + ReturnCallIndirect(u32), Return, Drop, Unreachable, @@ -762,6 +768,11 @@ impl<'a> CompileCtx<'a> { out.push(WasmInstr::LocalGet(reg.0)); out.push(WasmInstr::Return); } + // A single-block function can still end in a tail call — e.g. + // `[fn f [n] [g n]]`. Falling through to the catch-all below would + // return Unit and never call `g` at all. + End::Tail(fid, args) => emit_tail(*fid, args, out), + End::TailInvoke(callee, args) => emit_tail_invoke(*callee, args, out), End::Trap => { out.push(WasmInstr::Unreachable); } @@ -838,24 +849,9 @@ impl<'a> CompileCtx<'a> { out.push(WasmInstr::Br(dispatch_depth)); } - End::Tail(fid, args) => { - for arg in args { - out.push(WasmInstr::LocalGet(arg.0)); - } - out.push(WasmInstr::Call(HOST_IMPORT_COUNT + fid.0)); - out.push(WasmInstr::Return); - } + End::Tail(fid, args) => emit_tail(*fid, args, out), - End::TailInvoke(callee, args) => { - for arg in args { - out.push(WasmInstr::LocalGet(arg.0)); - } - out.push(WasmInstr::LocalGet(callee.0)); - emit_unbox_int(out); - out.push(WasmInstr::I32WrapI64); - out.push(WasmInstr::CallIndirect(args.len() as u32)); - out.push(WasmInstr::Return); - } + End::TailInvoke(callee, args) => emit_tail_invoke(*callee, args, out), End::Recur(args) => { let entry_block = &func.blocks[0]; @@ -907,7 +903,8 @@ impl<'a> CompileCtx<'a> { let mut indirect_arities: Vec = Vec::new(); for func in &self.functions { for instr in &func.body { - if let WasmInstr::CallIndirect(arity) = instr { + if let WasmInstr::CallIndirect(arity) | WasmInstr::ReturnCallIndirect(arity) = instr + { if !indirect_arities.contains(arity) { indirect_arities.push(*arity); } @@ -1057,6 +1054,28 @@ fn align4(n: u32) -> u32 { (n + 3) & !3 } +/// Emit `End::Tail` as a wasm `return_call`. Every EIR function has the same +/// `(i64…) -> i64` shape, so the caller/callee result types always match and +/// the tail call validates unconditionally. +fn emit_tail(fid: FuncId, args: &[Reg], out: &mut Vec) { + for arg in args { + out.push(WasmInstr::LocalGet(arg.0)); + } + out.push(WasmInstr::ReturnCall(HOST_IMPORT_COUNT + fid.0)); +} + +/// Emit `End::TailInvoke` as a `return_call_indirect` through the function +/// table, unboxing the callee value into a table index. +fn emit_tail_invoke(callee: Reg, args: &[Reg], out: &mut Vec) { + for arg in args { + out.push(WasmInstr::LocalGet(arg.0)); + } + out.push(WasmInstr::LocalGet(callee.0)); + emit_unbox_int(out); + out.push(WasmInstr::I32WrapI64); + out.push(WasmInstr::ReturnCallIndirect(args.len() as u32)); +} + /// Unbox NaN-boxed int: extract 48-bit payload, sign-extend to i64. /// Consumes one i64, produces one i64. fn emit_unbox_int(out: &mut Vec) { @@ -1210,6 +1229,16 @@ fn emit_wasm_instr( table_index: 0, }); } + WasmInstr::ReturnCall(i) => { + f.instruction(&Instruction::ReturnCall(*i)); + } + WasmInstr::ReturnCallIndirect(arity) => { + let type_idx = indirect_type_map.get(arity).copied().unwrap_or(0); + f.instruction(&Instruction::ReturnCallIndirect { + type_index: type_idx, + table_index: 0, + }); + } WasmInstr::Return => { f.instruction(&Instruction::Return); } @@ -1610,6 +1639,124 @@ mod tests { } } + /// Mutual tail recursion, `even`/`odd` shaped: + /// - func 0 is multi-block and tail-calls func 1 from its else branch + /// - func 1 is single-block and does nothing but tail-call func 0 + /// + /// Both terminators must become real `return_call`s; the single-block form + /// used to fall through to a catch-all that returned Unit without calling. + fn make_mutual_tail_module() -> super::super::Module { + super::super::Module { + funcs: vec![ + Func { + id: FuncId(0), + name: Some("even".to_string()), + params: vec![Ty::Int], + ret: Ty::Bool, + evidence: vec![], + captures: vec![], + blocks: vec![ + Block { + id: BlockId(0), + params: vec![Reg(0)], + ops: vec![ + Op::Lit(Reg(1), Lit::Int(0), Span::ZERO), + Op::Bin(Reg(2), BinOp::Eq, Reg(0), Reg(1), Span::ZERO), + ], + end: End::Br(Reg(2), BlockId(1), BlockId(2)), + }, + Block { + id: BlockId(1), + params: vec![], + ops: vec![Op::Lit(Reg(3), Lit::Bool(true), Span::ZERO)], + end: End::Ret(Reg(3)), + }, + Block { + id: BlockId(2), + params: vec![], + ops: vec![ + Op::Lit(Reg(4), Lit::Int(1), Span::ZERO), + Op::Bin(Reg(5), BinOp::Sub, Reg(0), Reg(4), Span::ZERO), + ], + end: End::Tail(FuncId(1), vec![Reg(5)]), + }, + ], + span: Span::ZERO, + is_closure: false, + }, + Func { + id: FuncId(1), + name: Some("odd".to_string()), + params: vec![Ty::Int], + ret: Ty::Bool, + evidence: vec![], + captures: vec![], + blocks: vec![Block { + id: BlockId(0), + params: vec![Reg(0)], + ops: vec![], + end: End::Tail(FuncId(0), vec![Reg(0)]), + }], + span: Span::ZERO, + is_closure: false, + }, + ], + strings: vec![], + ctors: vec![], + entry: FuncId(0), + } + } + + /// Tail calls lower to `return_call`, not `call` + `return`. Without this + /// the wasm stack grows one frame per tail call and deep tail recursion + /// overflows — the exact thing `End::Tail` exists to prevent. + #[test] + fn tail_calls_lower_to_return_call() { + let module = make_mutual_tail_module(); + let mut ctx = CompileCtx::new(&module); + ctx.compile_module().expect("compilation failed"); + + for (idx, f) in ctx.functions.iter().enumerate() { + assert!( + f.body.iter().any(|i| matches!(i, WasmInstr::ReturnCall(_))), + "func {idx} should tail-call via return_call, got {:?}", + f.body + ); + assert!( + !f.body.iter().any(|i| matches!(i, WasmInstr::Call(_))), + "func {idx} should not emit a plain call for a tail call" + ); + } + } + + /// `End::TailInvoke` (closure in tail position) uses the indirect form. + #[test] + fn tail_invoke_lowers_to_return_call_indirect() { + let mut module = make_mutual_tail_module(); + module.funcs[1].blocks[0].end = End::TailInvoke(Reg(0), vec![Reg(0)]); + let mut ctx = CompileCtx::new(&module); + ctx.compile_module().expect("compilation failed"); + + assert!( + ctx.functions[1] + .body + .iter() + .any(|i| matches!(i, WasmInstr::ReturnCallIndirect(_))), + "tail invoke should use return_call_indirect" + ); + } + + /// The emitted binary must still validate with tail calls enabled. + #[test] + fn tail_call_module_validates() { + let module = make_mutual_tail_module(); + let mut backend = WasmBackend; + let bytes = backend.compile(&module).expect("compilation failed"); + wasmparser::Validator::new() + .validate_all(&bytes) + .expect("tail-calling module should validate"); + } + #[test] fn compile_produces_valid_wasm() { let module = make_int_module(42); From b65f63726f7666b7debd50752568d92e22687195 Mon Sep 17 00:00:00 2001 From: Cam Pedersen Date: Mon, 10 Aug 2026 08:07:35 -0400 Subject: [PATCH 2/2] native: enable frame pointers for the tail calling convention MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Cranelift's x64 backend asserts when emitting a tail call without frame pointers ("the current implementation relies on them being present"). Every Loon function now uses `CallConv::Tail`, so this hit ordinary calls too — `compile_function_call` and `compile_recursive_function` failed alongside the tail-call tests. aarch64 maintains a frame pointer unconditionally, which is why it only showed up on CI's x86_64 runner. Verified both ways on x86_64 via `--target x86_64-apple-darwin` under Rosetta: all 454 lib tests pass with the flag, and the four native tests reproduce the exact CI assertion without it. Also addresses three review nits: - `enter_tail_frame` cleared only the argument registers, so a frame smaller than its caller's kept the caller's values in the rest of the file — `resize` fills only the slots it adds. `Op::Call` installs a freshly zeroed file, so a tail call has to hand over the same; this was the register-file twin of the capture leak already fixed here. Covered by `vm_tail_call_does_not_leak_caller_registers`, confirmed to fail without the clear. - Spell out which of the three return-threading guards carries which part of the safety argument. - Give the entry trampoline its own `UserFuncName` namespace instead of an index one past the EIR functions'. Co-Authored-By: Claude Opus 5 --- crates/loon-lang/src/eir/native.rs | 20 ++++-- crates/loon-lang/src/eir/tailcall.rs | 11 +++- crates/loon-lang/src/eir/vm.rs | 93 ++++++++++++++++++++++++++-- 3 files changed, 112 insertions(+), 12 deletions(-) diff --git a/crates/loon-lang/src/eir/native.rs b/crates/loon-lang/src/eir/native.rs index 19b28bf..263a73b 100644 --- a/crates/loon-lang/src/eir/native.rs +++ b/crates/loon-lang/src/eir/native.rs @@ -181,6 +181,18 @@ impl NativeModule { phase: "native:setup", })?; } + // Required by `CallConv::Tail`, which every Loon function uses so that + // `End::Tail` can lower to `return_call`. Cranelift's x64 tail-call + // emitter asserts on a missing frame pointer ("frame pointers aren't + // fundamentally required for tail calls, but the current + // implementation relies on them being present"); aarch64 maintains one + // unconditionally, so this only bites on x86_64. + flag_builder + .set("preserve_frame_pointers", "true") + .map_err(|e| Error { + message: format!("cranelift flag error: {e}"), + phase: "native:setup", + })?; let isa_builder = cranelift_native::builder().map_err(|msg| Error { message: format!("unsupported host: {msg}"), phase: "native:setup", @@ -330,10 +342,10 @@ impl NativeModule { phase: "native:declare", })?; - let mut cl_func = Function::with_name_signature( - UserFuncName::user(0, eir_module.funcs.len() as u32), - sig, - ); + // Namespace 1: EIR functions occupy namespace 0, indexed by + // FuncId. The trampoline is not an EIR function, so it gets its + // own namespace rather than an index just past the end of theirs. + let mut cl_func = Function::with_name_signature(UserFuncName::user(1, 0), sig); { let mut builder = FunctionBuilder::new(&mut cl_func, &mut fb_ctx); let block = builder.create_block(); diff --git a/crates/loon-lang/src/eir/tailcall.rs b/crates/loon-lang/src/eir/tailcall.rs index ee0a26f..ecbf198 100644 --- a/crates/loon-lang/src/eir/tailcall.rs +++ b/crates/loon-lang/src/eir/tailcall.rs @@ -64,9 +64,14 @@ fn thread_returns(func: &mut Func) -> bool { } let dest = &func.blocks[target]; - // Only a block that does nothing but return a parameter is safe to - // inline into its predecessor. Returning a non-parameter register would - // move the read to a block where that register may not be defined. + // Only a block that does nothing but return one of its own parameters + // is safe to inline into its predecessor. Three conditions, each load + // bearing: + // - no ops, or inlining would skip work the callee still owes; + // - arity match, so `args[i]` really is what `params[i]` binds to; + // - the returned register is a *parameter* (the `position` lookup + // below), because returning a register defined elsewhere would move + // the read into a block where it may not be defined. if !dest.ops.is_empty() || dest.params.len() != args.len() { continue; } diff --git a/crates/loon-lang/src/eir/vm.rs b/crates/loon-lang/src/eir/vm.rs index a3942f3..f89d416 100644 --- a/crates/loon-lang/src/eir/vm.rs +++ b/crates/loon-lang/src/eir/vm.rs @@ -379,13 +379,22 @@ impl Vm { /// in place (a tail call), load `vals` into the argument registers, and /// install the callee's captures. /// - /// `captures` must be what the equivalent non-tail call would have passed: - /// empty for `End::Tail` (mirroring `Op::Call`), the closure's own capture - /// list for `End::TailInvoke` (mirroring `Op::Invoke`). Leaving the - /// *caller's* captures in place would let the callee's `Op::Upval` read - /// the wrong frame's upvalues. + /// A tail call must be indistinguishable from the equivalent non-tail call + /// followed by a return, so the callee has to see exactly the frame that + /// `call_func_with_captures` would have built for it: + /// + /// - **A clean register file.** The non-tail path installs a fresh + /// `vec![Val::UNIT; _]`; reusing the buffer with a bare `resize` would + /// leave the caller's values in every register past the arguments, since + /// `resize` only fills slots it *adds*. Clearing first reuses the + /// allocation while still handing over a blank frame. + /// - **The callee's own captures**: empty for `End::Tail` (mirroring + /// `Op::Call`), the closure's capture list for `End::TailInvoke` + /// (mirroring `Op::Invoke`). Leaving the *caller's* captures in place + /// would let the callee's `Op::Upval` read the wrong frame's upvalues. fn enter_tail_frame(&mut self, func: FuncId, vals: &[Val], captures: Vec) { let needed = self.reg_count(func).max(vals.len()) + 16; + self.regs.clear(); self.regs.resize(needed, Val::UNIT); self.regs[..vals.len()].copy_from_slice(vals); self.captures = captures; @@ -3780,6 +3789,80 @@ mod tests { ); } + /// The register-file half of the same rule: `Op::Call` hands the callee a + /// freshly zeroed frame, so a tail call reusing the caller's buffer must + /// clear it. Otherwise a register the callee reads before writing sees the + /// caller's leftovers instead of Unit. + #[test] + fn vm_tail_call_does_not_leak_caller_registers() { + use crate::eir::{Block, Ty}; + use crate::syntax::Span; + + // callee() returns a register it never writes. + let callee = || Func { + id: FuncId(0), + name: Some("callee".to_string()), + params: vec![], + ret: Ty::Any, + evidence: vec![], + captures: vec![], + blocks: vec![Block { + id: BlockId(0), + params: vec![], + ops: vec![], + end: End::Ret(Reg(3)), + }], + span: Span::ZERO, + is_closure: false, + }; + + // __main stashes 7 in that same register, then reaches callee either + // way. The caller's frame is the larger one, so a bare `resize` would + // not overwrite Reg(3). + let build = |tail: bool| Module { + funcs: vec![ + callee(), + Func { + id: FuncId(1), + name: Some("__main".to_string()), + params: vec![], + ret: Ty::Any, + evidence: vec![], + captures: vec![], + blocks: vec![Block { + id: BlockId(0), + params: vec![], + ops: { + let mut ops = vec![Op::Lit(Reg(3), Lit::Int(7), Span::ZERO)]; + if !tail { + ops.push(Op::Call(Reg(9), FuncId(0), vec![], Span::ZERO)); + } + ops + }, + end: if tail { + End::Tail(FuncId(0), vec![]) + } else { + End::Ret(Reg(9)) + }, + }], + span: Span::ZERO, + is_closure: false, + }, + ], + strings: vec![], + ctors: vec![], + entry: FuncId(1), + }; + + let via_call = Vm::new(build(false)).run().expect("vm error").value; + let via_tail = Vm::new(build(true)).run().expect("vm error").value; + assert_eq!(via_call, Val::UNIT); + assert_eq!( + via_tail, via_call, + "a tail call must not leak the caller's registers into the callee" + ); + } + #[test] fn vm_durable_resume() { // Durable execution as a handler: a RESUME tower replays a journal of