diff --git a/packages/bugc/src/evmgen/optimizer-contexts.test.ts b/packages/bugc/src/evmgen/optimizer-contexts.test.ts index cee43825a..51805d7ff 100644 --- a/packages/bugc/src/evmgen/optimizer-contexts.test.ts +++ b/packages/bugc/src/evmgen/optimizer-contexts.test.ts @@ -148,11 +148,18 @@ code { r = add(10, 20); }`; const program = await compileAt(source, level); const counts = countCallSites(program); - // One caller JUMP, one callee JUMPDEST, one - // continuation JUMPDEST — all naming "add". - expect(counts.invokeJump).toEqual({ add: 1 }); - expect(counts.invokeJumpdest).toEqual({ add: 1 }); - expect(counts.returnJumpdest).toEqual({ add: 1 }); + if (level >= 2) { + // `add` is a leaf single-return helper: inlining (L2+) + // splices its body into the caller, so there's no real + // caller JUMP for `add`. + expect(counts.invokeJump).toEqual({}); + } else { + // One caller JUMP, one callee JUMPDEST, one + // continuation JUMPDEST — all naming "add". + expect(counts.invokeJump).toEqual({ add: 1 }); + expect(counts.invokeJumpdest).toEqual({ add: 1 }); + expect(counts.returnJumpdest).toEqual({ add: 1 }); + } // Behavior is still correct. const result = await executeProgram(source, { @@ -185,9 +192,15 @@ code { r = add(2 + 3, 4 * 5); }`; const program = await compileAt(source, level); const counts = countCallSites(program); - expect(counts.invokeJump).toEqual({ add: 1 }); - expect(counts.invokeJumpdest).toEqual({ add: 1 }); - expect(counts.returnJumpdest).toEqual({ add: 1 }); + if (level >= 2) { + // `add` inlined at L2+ — its body is spliced in, no real + // caller JUMP. Constant-foldable args don't change that. + expect(counts.invokeJump).toEqual({}); + } else { + expect(counts.invokeJump).toEqual({ add: 1 }); + expect(counts.invokeJumpdest).toEqual({ add: 1 }); + expect(counts.returnJumpdest).toEqual({ add: 1 }); + } const result = await executeProgram(source, { calldata: "", @@ -224,9 +237,15 @@ code { const program = await compileAt(source, level); const counts = countCallSites(program); - expect(counts.invokeJump).toEqual({ dbl: 2 }); - expect(counts.invokeJumpdest).toEqual({ dbl: 1 }); - expect(counts.returnJumpdest).toEqual({ dbl: 2 }); + if (level >= 2) { + // Both `dbl` sites are leaf single-return calls, inlined + // at L2+ into the caller — no real caller JUMPs remain. + expect(counts.invokeJump).toEqual({}); + } else { + expect(counts.invokeJump).toEqual({ dbl: 2 }); + expect(counts.invokeJumpdest).toEqual({ dbl: 1 }); + expect(counts.returnJumpdest).toEqual({ dbl: 2 }); + } const result = await executeProgram(source, { calldata: "", @@ -349,18 +368,26 @@ code { r = addThree(1, 2, 3); }`; const program = await compileAt(source, level); const counts = countCallSites(program); - expect(counts.invokeJump).toEqual({ - addThree: 1, - add: 2, - }); - expect(counts.invokeJumpdest).toEqual({ - addThree: 1, - add: 1, - }); - expect(counts.returnJumpdest).toEqual({ - addThree: 1, - add: 2, - }); + if (level >= 2) { + // `add` (leaf) inlines into `addThree` at both sites; + // that makes `addThree` itself a leaf, so on a later + // fixpoint iteration it inlines into `main` too. End + // state: no real caller JUMPs remain. + expect(counts.invokeJump).toEqual({}); + } else { + expect(counts.invokeJump).toEqual({ + addThree: 1, + add: 2, + }); + expect(counts.invokeJumpdest).toEqual({ + addThree: 1, + add: 1, + }); + expect(counts.returnJumpdest).toEqual({ + addThree: 1, + add: 2, + }); + } const result = await executeProgram(source, { calldata: "", diff --git a/packages/bugc/src/optimizer/simple-optimizer.ts b/packages/bugc/src/optimizer/simple-optimizer.ts index b078c9638..1b52a0cdb 100644 --- a/packages/bugc/src/optimizer/simple-optimizer.ts +++ b/packages/bugc/src/optimizer/simple-optimizer.ts @@ -14,6 +14,7 @@ import { ReturnMergingStep, ReadWriteMergingStep, TailCallOptimizationStep, + InliningStep, } from "./steps/index.js"; /** @@ -58,9 +59,12 @@ function createOptimizationPipeline(level: number): OptimizationStep[] { ); } - // Level 2: Add CSE, tail call optimization, and jump optimization + // Level 2: Add inlining, CSE, tail call optimization, and + // jump optimization. Inlining runs first (after L1 fold) so + // TCO/CSE still apply to inlined code. if (level >= 2) { steps.push( + new InliningStep(), new CommonSubexpressionEliminationStep(), new TailCallOptimizationStep(), new JumpOptimizationStep(), diff --git a/packages/bugc/src/optimizer/steps/index.ts b/packages/bugc/src/optimizer/steps/index.ts index 75a02833f..640bac63c 100644 --- a/packages/bugc/src/optimizer/steps/index.ts +++ b/packages/bugc/src/optimizer/steps/index.ts @@ -7,3 +7,4 @@ export { BlockMergingStep } from "./block-merging.js"; export { ReturnMergingStep } from "./return-merging.js"; export { ReadWriteMergingStep } from "./read-write-merging.js"; export { TailCallOptimizationStep } from "./tail-call-optimization.js"; +export { InliningStep } from "./inlining.js"; diff --git a/packages/bugc/src/optimizer/steps/inlining.test.ts b/packages/bugc/src/optimizer/steps/inlining.test.ts new file mode 100644 index 000000000..668074784 --- /dev/null +++ b/packages/bugc/src/optimizer/steps/inlining.test.ts @@ -0,0 +1,157 @@ +/** + * Behavioral + structural tests for the function-inlining pass + * (level 2). + * + * Inlining must (a) preserve runtime behavior exactly, and + * (b) actually replace eligible calls with the callee's body — + * a fully-inlined callee is deleted and no `call` to it remains. + */ +import { describe, it, expect } from "vitest"; + +import { compile } from "#compiler"; +import * as Ir from "#ir"; +import { executeProgram } from "#test/evm/behavioral"; + +async function optimizedIr( + source: string, + level: 0 | 1 | 2 | 3, +): Promise { + const result = await compile({ + to: "ir", + source, + optimizer: { level }, + }); + if (!result.success) { + const errors = result.messages.error ?? []; + throw new Error( + "compile failed:\n" + + errors + .map((e: { message?: string }) => e.message ?? String(e)) + .join("\n"), + ); + } + return result.value.ir; +} + +/** All functions in the module, including
and . */ +function allFunctions(module: Ir.Module): Ir.Function[] { + return [ + module.main, + ...(module.create ? [module.create] : []), + ...module.functions.values(), + ]; +} + +/** Count `call` terminators targeting `callee` across the module. */ +function countCalls(module: Ir.Module, callee: string): number { + let n = 0; + for (const fn of allFunctions(module)) { + for (const block of fn.blocks.values()) { + const t = block.terminator; + if (t.kind === "call" && t.function === callee) n += 1; + } + } + return n; +} + +describe("function inlining (level 2)", () => { + describe("leaf helper, single return", () => { + const source = `name Demo; +define { + function add(a: uint256, b: uint256) -> uint256 { return a + b; }; +} +storage { [0] r: uint256; } +create {} +code { r = add(3, 4); }`; + + it("produces the same result at every level", async () => { + for (const level of [0, 1, 2, 3] as const) { + const res = await executeProgram(source, { + calldata: "", + optimizationLevel: level, + }); + expect(res.callSuccess).toBe(true); + expect(await res.getStorage(0n)).toBe(7n); + } + }); + + it("keeps the real call at levels 0-1", async () => { + for (const level of [0, 1] as const) { + const ir = await optimizedIr(source, level); + expect(ir.functions.has("add")).toBe(true); + expect(countCalls(ir, "add")).toBe(1); + } + }); + + it("inlines the call and deletes the callee at level 2", async () => { + const ir = await optimizedIr(source, 2); + expect(countCalls(ir, "add")).toBe(0); + expect(ir.functions.has("add")).toBe(false); + }); + }); + + describe("multiple call sites", () => { + const source = `name Multi; +define { + function dbl(x: uint256) -> uint256 { return x + x; }; +} +storage { [0] r: uint256; } +create { r = 0; } +code { + let a = dbl(5); + let b = dbl(10); + r = a + b; +}`; + + it("inlines every site and stays correct", async () => { + for (const level of [0, 1, 2, 3] as const) { + const res = await executeProgram(source, { + calldata: "", + optimizationLevel: level, + }); + expect(res.callSuccess).toBe(true); + expect(await res.getStorage(0n)).toBe(30n); + } + // Both sites inlined; callee deleted. + const ir = await optimizedIr(source, 2); + expect(countCalls(ir, "dbl")).toBe(0); + expect(ir.functions.has("dbl")).toBe(false); + }); + }); + + describe("does not inline into a tail-recursive function (protects TCO)", () => { + // `succ` is a leaf, but inlining it into `count`'s recursive + // call arguments would rewrite `count(succ(n))` into + // `count(n + 1)`, which the tail-call optimizer mishandles. + // The pass must leave recursive/TCO'd callers untouched. + const source = `name TailCall; +define { + function succ(n: uint256) -> uint256 { return n + 1; }; + function count(n: uint256, target: uint256) -> uint256 { + if (n < target) { return count(succ(n), target); } + else { return n; } + }; +} +storage { [0] r: uint256; } +create { r = 0; } +code { r = count(0, 5); }`; + + it("stays correct at every level", async () => { + for (const level of [0, 1, 2, 3] as const) { + const res = await executeProgram(source, { + calldata: "", + optimizationLevel: level, + }); + expect(res.callSuccess).toBe(true); + expect(await res.getStorage(0n)).toBe(5n); + } + }); + + it("does not inline succ into the recursive count", async () => { + // succ stays a real call so TCO can fire on count. + const ir = await optimizedIr(source, 2); + expect(ir.functions.has("succ")).toBe(true); + expect(countCalls(ir, "succ")).toBeGreaterThan(0); + }); + }); +}); diff --git a/packages/bugc/src/optimizer/steps/inlining.ts b/packages/bugc/src/optimizer/steps/inlining.ts new file mode 100644 index 000000000..e2bd4b712 --- /dev/null +++ b/packages/bugc/src/optimizer/steps/inlining.ts @@ -0,0 +1,480 @@ +/** + * Function inlining (level 2). + * + * Replaces calls to eligible internal functions with a copy of + * the callee's body spliced into the caller, so no runtime + * JUMP/frame is used. Inlined instructions keep their original + * source mapping (the callee's `operationDebug`), so inlined code + * still maps back to the callee. + * + * v1 eligibility: internal (user-defined), non-recursive callee + * that is either a leaf (calls nothing) or below a small size + * threshold. Applied at all call sites; a callee whose every + * site is inlined is deleted. + */ +import * as Ir from "#ir"; + +import { + BaseOptimizationStep, + type OptimizationContext, +} from "../optimizer.js"; + +/** Max IR-node count for a non-leaf callee to still inline. Tunable. */ +const INLINE_MAX_IR_NODES = 16; + +export class InliningStep extends BaseOptimizationStep { + name = "inlining"; + private siteCounter = 0; + + run(module: Ir.Module, _context: OptimizationContext): Ir.Module { + const optimized = this.cloneModule(module); + if (!optimized.functions || optimized.functions.size === 0) { + return optimized; + } + + const callGraph = buildCallGraph(optimized); + const eligible = new Set(); + for (const [name, fn] of optimized.functions) { + if (isEligible(name, fn, callGraph)) eligible.add(name); + } + if (eligible.size === 0) return optimized; + + // Track, per callee, how many sites remain un-inlined so we + // can delete fully-inlined callees afterward. + const remainingSites = new Map(); + for (const name of eligible) remainingSites.set(name, 0); + + // Named callers. Self-recursive callers are skipped: they are + // TailCall-optimized later, and inlining a helper into a + // self-recursive call's arguments (e.g. `count(succ(n))` -> + // `count(n + 1)`) rewrites the tail call into a computed-arg + // form that TCO mishandles, silently breaking the recursion. + // Correctness over coverage — inlining into recursive bodies is + // deferred. + const named: [string, Ir.Function][] = [ + ["
", optimized.main], + ...(optimized.create + ? ([["", optimized.create]] as [string, Ir.Function][]) + : []), + ...[...optimized.functions].map( + ([n, f]) => [n, f] as [string, Ir.Function], + ), + ]; + + for (const [callerName, caller] of named) { + // Self-recursive (pre-TCO) or already TCO'd (post-TCO, its + // self-call is now a jump-with-tailCall so the call graph no + // longer shows the recursion). Either way, don't inline into + // it. + if (reachableCallees(callerName, callGraph).has(callerName)) continue; + if (hasTailCallBackedge(caller)) continue; + this.inlineIntoFunction(caller, optimized, eligible, remainingSites); + } + + // Delete callees that no longer have any call site anywhere. + for (const name of eligible) { + if (!isCalledAnywhere(name, optimized)) { + optimized.functions.delete(name); + } + } + + return optimized; + } + + private inlineIntoFunction( + caller: Ir.Function, + module: Ir.Module, + eligible: Set, + _remainingSites: Map, + ): void { + // Snapshot block ids up front; we mutate the map as we splice. + let changed = true; + // Guard against pathological loops. + let guard = 0; + while (changed && guard++ < 1000) { + changed = false; + for (const [blockId, block] of caller.blocks) { + const term = block.terminator; + if (term.kind !== "call") continue; + if (!eligible.has(term.function)) continue; + const callee = module.functions?.get(term.function); + if (!callee) continue; + // Don't inline a function into itself. + if (callee === caller) continue; + + this.spliceCall(caller, blockId, block, term, callee); + changed = true; + break; // block map mutated — restart the scan + } + } + } + + private spliceCall( + caller: Ir.Function, + callBlockId: string, + callBlock: Ir.Block, + call: Extract, + callee: Ir.Function, + ): void { + const site = this.siteCounter++; + const prefix = `inl${site}$`; + + // --- build rename maps --- + const blockRename = new Map(); + for (const id of callee.blocks.keys()) { + blockRename.set(id, prefix + id); + } + + // param temp id -> bound argument Value + const paramSubst = new Map(); + callee.parameters.forEach((p, i) => { + const arg = call.arguments[i]; + if (arg) paramSubst.set(p.tempId, arg); + }); + + // every non-param temp defined in the callee -> fresh id + const idRename = new Map(); + for (const b of callee.blocks.values()) { + for (const phi of b.phis ?? []) rename(phi.dest); + for (const inst of b.instructions) { + if ("dest" in inst && typeof inst.dest === "string") { + rename(inst.dest); + } + } + } + function rename(id: string): void { + if (paramSubst.has(id)) return; + if (!idRename.has(id)) idRename.set(id, prefix + id); + } + + const remapValue = (v: Ir.Value): Ir.Value => { + if (v.kind !== "temp") return v; + const sub = paramSubst.get(v.id); + if (sub) return sub; + const nid = idRename.get(v.id); + return nid ? { ...v, id: nid } : v; + }; + + const entryBlockId = blockRename.get(callee.entry)!; + + // --- clone + remap callee blocks --- + for (const [origId, origBlock] of callee.blocks) { + const newId = blockRename.get(origId)!; + + // Cloned instructions keep their preserved source + // `operationDebug` (remapInstruction deep-clones it), so + // inlined code still maps back to the callee's source. + const instructions: Ir.Instruction[] = origBlock.instructions.map( + (inst) => remapInstruction(inst, remapValue, idRename), + ); + + const phis: Ir.Block.Phi[] = (origBlock.phis ?? []).map((phi) => + remapPhi(phi, remapValue, idRename, blockRename), + ); + + let terminator: Ir.Block.Terminator; + const t = origBlock.terminator; + if (t.kind === "return") { + // return -> jump to the caller's continuation, preserving + // the return's own source mapping. + terminator = { + kind: "jump", + target: call.continuation, + operationDebug: t.operationDebug, + }; + } else { + terminator = remapTerminator(t, remapValue, blockRename); + } + + caller.blocks.set(newId, { + id: newId, + phis, + instructions, + terminator, + predecessors: new Set(), + debug: origBlock.debug, + }); + } + + // --- wire the single return value into the caller --- + // v1 eligibility guarantees exactly one return. Substitute the + // call's dest temp with the (remapped) returned value across the + // whole caller — no phi, so it's robust to L3 block-merging. + if (call.dest) { + const returns = collectReturns(callee, blockRename, remapValue); + if (returns.length === 1) { + substituteTemp(caller, call.dest, returns[0].value); + } + } + + // --- rewire the calling block: call -> jump into inlined entry --- + callBlock.terminator = { + kind: "jump", + target: entryBlockId, + operationDebug: call.operationDebug, + }; + + void callBlockId; + recomputePredecessors(caller); + } +} + +// ---- helpers ---- + +function collectReturns( + callee: Ir.Function, + blockRename: Map, + remapValue: (v: Ir.Value) => Ir.Value, +): { block: string; value: Ir.Value }[] { + const out: { block: string; value: Ir.Value }[] = []; + for (const [origId, b] of callee.blocks) { + if (b.terminator.kind === "return" && b.terminator.value) { + out.push({ + block: blockRename.get(origId)!, + value: remapValue(b.terminator.value), + }); + } + } + return out; +} + +function buildCallGraph(module: Ir.Module): Map> { + const graph = new Map>(); + const fns: [string, Ir.Function][] = [ + ["
", module.main], + ...(module.create + ? ([["", module.create]] as [string, Ir.Function][]) + : []), + ...[...(module.functions ?? new Map())].map( + ([n, f]) => [n, f] as [string, Ir.Function], + ), + ]; + for (const [name, fn] of fns) { + const callees = new Set(); + for (const b of fn.blocks.values()) { + if (b.terminator.kind === "call") callees.add(b.terminator.function); + } + graph.set(name, callees); + } + return graph; +} + +function reachableCallees( + start: string, + graph: Map>, +): Set { + const seen = new Set(); + const stack = [...(graph.get(start) ?? [])]; + while (stack.length) { + const n = stack.pop()!; + if (seen.has(n)) continue; + seen.add(n); + for (const c of graph.get(n) ?? []) stack.push(c); + } + return seen; +} + +function functionSize(fn: Ir.Function): number { + let n = 0; + for (const b of fn.blocks.values()) { + n += b.instructions.length + 1; // + terminator + n += (b.phis ?? []).length; + } + return n; +} + +function hasTailCallBackedge(fn: Ir.Function): boolean { + for (const b of fn.blocks.values()) { + if (b.terminator.kind === "jump" && b.terminator.tailCall) return true; + } + return false; +} + +function returnCount(fn: Ir.Function): number { + let n = 0; + for (const b of fn.blocks.values()) { + if (b.terminator.kind === "return") n += 1; + } + return n; +} + +function isEligible( + name: string, + fn: Ir.Function, + graph: Map>, +): boolean { + const callees = graph.get(name) ?? new Set(); + // Non-recursive: name not reachable from itself. + if (reachableCallees(name, graph).has(name)) return false; + // v1: single return point only. Multi-return needs a phi at the + // continuation, which block-merging (L3) can turn into an + // invalid self-referential phi; deferred until that's handled. + if (returnCount(fn) !== 1) return false; + // Never inline a TCO-transformed function: after TailCall + // optimization a self-recursive function's back-edge becomes a + // `jump` with `tailCall`, which makes it look like a leaf + // single-return function on the next fixpoint iteration. Inlining + // it would clobber the tailcall showcase. A `tailCall` back-edge + // marks it as recursion, not a real leaf. + for (const b of fn.blocks.values()) { + if (b.terminator.kind === "jump" && b.terminator.tailCall) return false; + } + // v1: leaf callees only. Inlining a non-leaf callee whose own + // (eligible) calls also inline exposes a dest-substitution + // ordering bug in the nested chain; deferred. The size-threshold + // branch is kept for when that lands. + const isLeaf = callees.size === 0; + const smallEnough = functionSize(fn) <= INLINE_MAX_IR_NODES; + void smallEnough; + return isLeaf; +} + +function isCalledAnywhere(name: string, module: Ir.Module): boolean { + const fns = [ + module.main, + ...(module.create ? [module.create] : []), + ...(module.functions?.values() ?? []), + ]; + for (const fn of fns) { + for (const b of fn.blocks.values()) { + if (b.terminator.kind === "call" && b.terminator.function === name) { + return true; + } + } + } + return false; +} + +/** Deep-clone an instruction, remapping temp values and its dest. */ +function remapInstruction( + inst: Ir.Instruction, + remapValue: (v: Ir.Value) => Ir.Value, + idRename: Map, +): Ir.Instruction { + const cloned = structuredCloneValues(inst) as Ir.Instruction & { + dest?: string; + }; + remapValuesInPlace(cloned, remapValue); + if (typeof cloned.dest === "string") { + cloned.dest = idRename.get(cloned.dest) ?? cloned.dest; + } + return cloned; +} + +function remapPhi( + phi: Ir.Block.Phi, + remapValue: (v: Ir.Value) => Ir.Value, + idRename: Map, + blockRename: Map, +): Ir.Block.Phi { + const sources = new Map(); + for (const [pred, val] of phi.sources) { + sources.set(blockRename.get(pred) ?? pred, remapValue(val)); + } + return { + ...phi, + dest: idRename.get(phi.dest) ?? phi.dest, + sources, + }; +} + +function remapTerminator( + t: Ir.Block.Terminator, + remapValue: (v: Ir.Value) => Ir.Value, + blockRename: Map, +): Ir.Block.Terminator { + switch (t.kind) { + case "jump": + return { ...t, target: blockRename.get(t.target) ?? t.target }; + case "branch": + return { + ...t, + condition: remapValue(t.condition), + trueTarget: blockRename.get(t.trueTarget) ?? t.trueTarget, + falseTarget: blockRename.get(t.falseTarget) ?? t.falseTarget, + }; + case "return": + return { ...t, value: t.value ? remapValue(t.value) : undefined }; + case "call": + return { + ...t, + arguments: t.arguments.map(remapValue), + dest: t.dest, + continuation: blockRename.get(t.continuation) ?? t.continuation, + }; + } +} + +/** Deep clone that preserves nested plain objects and bigints. */ +function structuredCloneValues(obj: T): T { + return structuredClone(obj); +} + +/** Recursively rewrite any temp Value in place. */ +function remapValuesInPlace( + node: unknown, + remapValue: (v: Ir.Value) => Ir.Value, +): void { + if (!node || typeof node !== "object") return; + const obj = node as Record; + for (const key of Object.keys(obj)) { + const child = obj[key]; + if ( + child && + typeof child === "object" && + (child as { kind?: string }).kind === "temp" && + typeof (child as { id?: unknown }).id === "string" + ) { + obj[key] = remapValue(child as Ir.Value); + } else if (Array.isArray(child)) { + child.forEach((el, i) => { + if ( + el && + typeof el === "object" && + (el as { kind?: string }).kind === "temp" + ) { + child[i] = remapValue(el as Ir.Value); + } else { + remapValuesInPlace(el, remapValue); + } + }); + } else if (child && typeof child === "object") { + remapValuesInPlace(child, remapValue); + } + } +} + +/** Replace every use of temp `id` with `value` across the function. */ +function substituteTemp(fn: Ir.Function, id: string, value: Ir.Value): void { + const sub = (v: Ir.Value): Ir.Value => + v.kind === "temp" && v.id === id ? value : v; + for (const b of fn.blocks.values()) { + for (const inst of b.instructions) { + remapValuesInPlace(inst, sub); + } + for (const phi of b.phis ?? []) { + for (const [pred, val] of phi.sources) { + phi.sources.set(pred, sub(val)); + } + } + b.terminator = remapTerminator(b.terminator, sub, new Map()); + } +} + +function recomputePredecessors(fn: Ir.Function): void { + for (const b of fn.blocks.values()) b.predecessors = new Set(); + for (const [id, b] of fn.blocks) { + const t = b.terminator; + const targets: string[] = + t.kind === "jump" + ? [t.target] + : t.kind === "branch" + ? [t.trueTarget, t.falseTarget] + : t.kind === "call" + ? [t.continuation] + : []; + for (const tgt of targets) { + fn.blocks.get(tgt)?.predecessors.add(id); + } + } +}