diff --git a/packages/bugc/src/evmgen/generation/control-flow/terminator.ts b/packages/bugc/src/evmgen/generation/control-flow/terminator.ts index 9334d11a1d..6a8d5331d1 100644 --- a/packages/bugc/src/evmgen/generation/control-flow/terminator.ts +++ b/packages/bugc/src/evmgen/generation/control-flow/terminator.ts @@ -472,8 +472,22 @@ function buildTailCallJumpOptions(tailCall: Ir.Block.TailCall): { } : undefined; + // The call site's own source range, composed flat alongside the + // invoke/return (disjoint keys). This maps the back-edge JUMP back + // to the recursive call expression, matching the `{invoke, code}` a + // real caller JUMP carries and the `{return, code}` on the + // continuation JUMPDEST for a non-TCO call. + const callSiteCode = + tailCall.callSiteLoc && tailCall.callSiteSourceId + ? { + source: { id: tailCall.callSiteSourceId }, + range: tailCall.callSiteLoc, + } + : undefined; + const combined: Format.Program.Context.Return & - Format.Program.Context.Invoke = { + Format.Program.Context.Invoke & + Partial = { return: { identifier: tailCall.function, ...(declaration ? { declaration } : {}), @@ -490,6 +504,7 @@ function buildTailCallJumpOptions(tailCall: Ir.Block.TailCall): { }, }, }, + ...(callSiteCode ? { code: callSiteCode } : {}), }; // Route through the shared helper so all transform emission diff --git a/packages/bugc/src/evmgen/tco-callsite.test.ts b/packages/bugc/src/evmgen/tco-callsite.test.ts new file mode 100644 index 0000000000..f00bcdc7c6 --- /dev/null +++ b/packages/bugc/src/evmgen/tco-callsite.test.ts @@ -0,0 +1,103 @@ +/** + * Tail-call optimization preserves the call site's source range. + * + * TCO (level 2+) rewrites a tail-recursive call into a back-edge JUMP + * that carries a flat `{return, invoke, transform: ["tailcall"]}` + * context. Without the call-site range, that JUMP maps only to the + * callee declaration, not the recursive call expression the user + * wrote. This carries the call site's `code` range flat alongside the + * invoke/return (disjoint keys) — matching the `{invoke, code}` a real + * caller JUMP gets and the `{return, code}` on a non-TCO continuation + * JUMPDEST. + */ +import { describe, it, expect } from "vitest"; + +import { compile } from "#compiler"; +import { executeProgram } from "#test/evm/behavioral"; +import { Program } from "@ethdebug/format"; +import type * as Format from "@ethdebug/format"; + +const { Context } = Program; + +const source = `name TailSum; +define { + function sum(n: uint256, acc: uint256) -> uint256 { + if (n == 0) { return acc; } + else { return sum(n - 1, acc + n); } + }; +} +storage { [0] r: uint256; } +create { r = 0; } +code { r = sum(5, 0); }`; + +async function runtimeProgram(level: 2 | 3): Promise { + const result = await compile({ + to: "bytecode", + source, + optimizer: { level }, + }); + if (!result.success) throw new Error("compile failed"); + return result.value.bytecode.runtimeProgram; +} + +/** Flatten a context into leaves, unwrapping gather/pick. */ +function leaves(ctx: Format.Program.Context): Format.Program.Context[] { + if (Context.isGather(ctx)) return ctx.gather.flatMap(leaves); + if ("pick" in ctx && Array.isArray((ctx as { pick: unknown[] }).pick)) { + return (ctx as { pick: Format.Program.Context[] }).pick.flatMap(leaves); + } + return [ctx]; +} + +function hasCodeRange(ctx: Format.Program.Context): boolean { + return [ctx, ...leaves(ctx)].some((c) => { + const code = (c as { code?: { range?: { offset?: unknown } } }).code; + return !!code && typeof code.range?.offset === "number"; + }); +} + +/** Find the TCO back-edge JUMP: a JUMP tagged transform ["tailcall"]. */ +function tailCallJump(program: Format.Program) { + return program.instructions.find( + (i) => + i.operation?.mnemonic === "JUMP" && + i.context !== undefined && + [i.context, ...leaves(i.context)].some( + (c) => Context.isTransform(c) && c.transform.includes("tailcall"), + ), + ); +} + +describe("call-site source range on the TCO back-edge JUMP", () => { + for (const level of [2, 3] as const) { + it(`the back-edge JUMP carries a call-site code range (level ${level})`, async () => { + const program = await runtimeProgram(level); + const jump = tailCallJump(program); + expect(jump, "tailcall JUMP").toBeDefined(); + expect(hasCodeRange(jump!.context!)).toBe(true); + }); + + it(`the back-edge JUMP still carries invoke and return identity (level ${level})`, async () => { + const program = await runtimeProgram(level); + const jump = tailCallJump(program); + const parts = [jump!.context!, ...leaves(jump!.context!)]; + expect(parts.some((c) => Context.isInvoke(c))).toBe(true); + expect(parts.some((c) => Context.isReturn(c))).toBe(true); + }); + + it(`preserves behavior vs unoptimized (level ${level})`, async () => { + const base = await executeProgram(source, { + calldata: "", + optimizationLevel: 0, + }); + const opt = await executeProgram(source, { + calldata: "", + optimizationLevel: level, + }); + expect(base.callSuccess).toBe(true); + expect(opt.callSuccess).toBe(true); + // TCO must not change the computed result. + expect(await opt.getStorage(0n)).toBe(await base.getStorage(0n)); + }); + } +}); diff --git a/packages/bugc/src/ir/spec/block.ts b/packages/bugc/src/ir/spec/block.ts index 182bc35f85..ce64cd62b7 100644 --- a/packages/bugc/src/ir/spec/block.ts +++ b/packages/bugc/src/ir/spec/block.ts @@ -47,6 +47,10 @@ export namespace Block { declarationLoc?: Ast.SourceLocation; /** Source ID for the declaration (inherited from module) */ declarationSourceId?: string; + /** Source location of the call site (the recursive call expression) */ + callSiteLoc?: Ast.SourceLocation; + /** Source ID for the call site (inherited from module) */ + callSiteSourceId?: string; } /** diff --git a/packages/bugc/src/optimizer/steps/tail-call-optimization.ts b/packages/bugc/src/optimizer/steps/tail-call-optimization.ts index 251cd78eab..e18e7d35f0 100644 --- a/packages/bugc/src/optimizer/steps/tail-call-optimization.ts +++ b/packages/bugc/src/optimizer/steps/tail-call-optimization.ts @@ -1,4 +1,5 @@ import * as Ir from "#ir"; +import type * as Format from "@ethdebug/format"; import { BaseOptimizationStep, type OptimizationContext, @@ -152,10 +153,25 @@ export class TailCallOptimizationStep extends BaseOptimizationStep { // debugger can still see the recursive call in // the trace, even though the implementation is // now a block-internal jump. + // + // The call expression's own source range lives on the + // original call terminator's operationDebug. Carry it on + // the tailCall so the back-edge JUMP can map back to the + // call site (mirroring the flat `{invoke, code}` a real + // caller JUMP gets), not just the callee declaration. + const callSiteCode = ( + callTerm.operationDebug?.context as + | { code?: Format.Program.Context.Code["code"] } + | undefined + )?.code; const tailCall: Ir.Block.TailCall = { function: funcName, ...(func.loc ? { declarationLoc: func.loc } : {}), ...(func.sourceId ? { declarationSourceId: func.sourceId } : {}), + ...(callSiteCode?.range ? { callSiteLoc: callSiteCode.range } : {}), + ...(typeof callSiteCode?.source?.id === "string" + ? { callSiteSourceId: callSiteCode.source.id } + : {}), }; block.terminator = {