bugc: mark a function's return on its exit jump (#349)

Put the return context on the callee's JUMP back to the caller, with
the return value at stack slot 0, and leave the continuation JUMPDEST
with only the call site's code range. buildCallStack now pops a frame
on the return's postcondition step instead of one step later.

Author: gnidan

packages/bugc/CHANGELOG.md

@@ -62,6 +62,13 @@ support. Changes to the specification itself are tracked in the root
 - An array or bytes index and a slice bound must now be an unsigned
   integer, as in Solidity; cast a signed one, as in `a[i as uint256]`.
   Before, a signed index compiled ([#344]).
+- A function's `return` context is now on its exit: the `JUMP` back to
+  the caller, as the format's documentation describes. Its `data`
+  pointer (stack slot 0) is there when the function returns a value. The
+  caller's continuation `JUMPDEST` keeps only the call site's `code`
+  range. Before, the `return` was on that `JUMPDEST`, one instruction
+  later, so a debugger showed the function's frame for two steps after
+  control was back in the caller ([#349]).
 
 ### Fixed
 
@@ -269,3 +276,4 @@ First publication.
 [#344]: https://github.com/ethdebug/format/pull/344
 [#345]: https://github.com/ethdebug/format/pull/345
 [#347]: https://github.com/ethdebug/format/pull/347
+[#349]: https://github.com/ethdebug/format/pull/349

packages/bugc/src/evmgen/call-contexts.test.ts

@@ -106,18 +106,18 @@ code {
     },
   );
 
-  it("should emit return context on continuation JUMPDEST", async () => {
+  it("should emit return context on the callee's exit JUMP", async () => {
     const program = await compileProgram(source);
 
-    const returnJumpdests = findInstructionsWithContext(
+    const returnJumps = findInstructionsWithContext(
       program,
-      "JUMPDEST",
+      "JUMP",
       Context.isReturn,
     );
 
-    expect(returnJumpdests.length).toBeGreaterThanOrEqual(1);
+    expect(returnJumps.length).toBeGreaterThanOrEqual(1);
 
-    const { return: ret } = returnJumpdests[0].context;
+    const { return: ret } = returnJumps[0].context;
 
     expect(ret.identifier).toBe("add");
 
@@ -134,6 +134,14 @@ code {
     });
   });
 
+  it("should emit no return context on the continuation JUMPDEST", async () => {
+    const program = await compileProgram(source);
+
+    expect(
+      findInstructionsWithContext(program, "JUMPDEST", Context.isReturn),
+    ).toHaveLength(0);
+  });
+
   it(
     "should emit invoke context on callee entry " +
       "JUMPDEST with args and code target",
@@ -185,28 +193,26 @@ code {
     const program = await compileProgram(source);
 
     // The caller JUMP should come before the
-    // continuation JUMPDEST
+    // callee's exit JUMP
     const invokeJump = findInstructionsWithContext(
       program,
       "JUMP",
       Context.isInvoke,
     )[0];
 
-    const returnJumpdest = findInstructionsWithContext(
+    const returnJump = findInstructionsWithContext(
       program,
-      "JUMPDEST",
+      "JUMP",
       Context.isReturn,
     )[0];
 
     expect(invokeJump).toBeDefined();
-    expect(returnJumpdest).toBeDefined();
+    expect(returnJump).toBeDefined();
 
     // Invoke JUMP offset should be less than
-    // return JUMPDEST offset (caller comes first
+    // return JUMP offset (caller comes first
     // in bytecode)
-    expect(Number(invokeJump.offset)).toBeLessThan(
-      Number(returnJumpdest.offset),
-    );
+    expect(Number(invokeJump.offset)).toBeLessThan(Number(returnJump.offset));
   });
 
   describe("void function calls", () => {
@@ -238,15 +244,15 @@ code {
       async () => {
         const program = await compileProgram(voidSource);
 
-        const returnJumpdests = findInstructionsWithContext(
+        const returnJumps = findInstructionsWithContext(
           program,
-          "JUMPDEST",
+          "JUMP",
           Context.isReturn,
         );
 
-        expect(returnJumpdests.length).toBeGreaterThanOrEqual(1);
+        expect(returnJumps.length).toBeGreaterThanOrEqual(1);
 
-        const { return: ret } = returnJumpdests[0].context;
+        const { return: ret } = returnJumps[0].context;
         expect(ret.identifier).toBe("setVal");
         // Since setVal returns a value, data should
         // be present
@@ -310,15 +316,19 @@ code {
       expect(invokeIds).toContain("addThree");
       expect(invokeIds).toContain("add");
 
-      // Should have return contexts for all
-      // continuation points
-      const returnJumpdests = findInstructionsWithContext(
+      // Should have return contexts on both
+      // functions' exit JUMPs
+      const returnJumps = findInstructionsWithContext(
         program,
-        "JUMPDEST",
+        "JUMP",
         Context.isReturn,
       );
 
-      expect(returnJumpdests.length).toBeGreaterThanOrEqual(3);
+      const returnIds = returnJumps.map(
+        (instr) => instr.context.return.identifier,
+      );
+      expect(returnIds).toContain("addThree");
+      expect(returnIds).toContain("add");
     });
   });
 

packages/bugc/src/evmgen/call-timing.test.ts

@@ -0,0 +1,156 @@
+/**
+ * Call frame timing, checked against real traces. Instruction
+ * contexts are postconditions: the context of the instruction run
+ * at step i-1 holds at step i. Replaying the invoke and return
+ * contexts that way, a called function's frame must be on from the
+ * step after the caller's JUMP (the callee's entry JUMPDEST) through
+ * the step of the callee's exit JUMP, and off at the step after it
+ * (the caller's continuation JUMPDEST). At that step, the exit's
+ * `return` data pointer must read the returned value.
+ */
+import { describe, it, expect } from "vitest";
+import { createMachineState } from "@ethdebug/evm";
+import { dereference } from "@ethdebug/pointers";
+import { Program } from "@ethdebug/format";
+import type * as Format from "@ethdebug/format";
+import { traceLocals, type Level } from "#test/evm/locals";
+
+const { Context } = Program;
+
+const programs = [
+  {
+    name: "nested calls and a function with two exits",
+    source: `name Exits;
+define {
+  function add(a: uint256, b: uint256) -> uint256 { return a + b; };
+  function pick(a: uint256) -> uint256 {
+    if (a > 5) { return a; }
+    return 99;
+  };
+  function addThree(a: uint256) -> uint256 { return add(add(a, 1), 2); };
+}
+storage { [0] r: uint256; [1] s: uint256; [2] t: uint256; [3] u: uint256; }
+create { r = 0; }
+code {
+  r = add(10, 20);
+  s = pick(7);
+  t = pick(2);
+  u = addThree(5);
+}`,
+    // The values each function returns, in order, at O0 (higher
+    // levels inline some calls away)
+    returns: [30n, 7n, 99n, 6n, 8n, 8n],
+  },
+  {
+    name: "recursion",
+    source: `name Recursion;
+define {
+  function fact(n: uint256) -> uint256 {
+    if (n < 2) { return 1; }
+    return n * fact(n - 1);
+  };
+}
+storage { [0] r: uint256; }
+create { r = 0; }
+code { r = fact(4); }`,
+    returns: [1n, 2n, 6n, 24n],
+  },
+];
+
+/** A context's invoke and return leaves, outside inlined bodies */
+function events(context: Format.Program.Context | undefined): {
+  invoke?: string;
+  return?: string;
+} {
+  if (!context) return {};
+  if (Context.isTransform(context) && context.transform.includes("inline")) {
+    return {};
+  }
+  return {
+    ...(Context.isInvoke(context) ? { invoke: context.invoke.identifier } : {}),
+    ...(Context.isReturn(context) ? { return: context.return.identifier } : {}),
+  };
+}
+
+describe("call frame timing", () => {
+  for (const { name, source, returns } of programs) {
+    for (const level of [0, 1, 2, 3] as Level[]) {
+      it(`${name} at O${level}`, async () => {
+        const { program, executor, steps, instructionAt } = await traceLocals(
+          source,
+          level,
+        );
+
+        // Each function's body: from its entry JUMPDEST to the next
+        // function's entry (bugc lays functions out one after another,
+        // after the main code)
+        const entries = program.instructions
+          .filter(
+            (instruction) =>
+              instruction.operation?.mnemonic === "JUMPDEST" &&
+              events(instruction.context).invoke !== undefined,
+          )
+          .map((instruction) => ({
+            name: events(instruction.context).invoke!,
+            offset: Number(instruction.offset),
+          }))
+          .sort((a, b) => a.offset - b.offset);
+        const functionAt = (pc: number) =>
+          entries.filter(({ offset }) => offset <= pc).at(-1)?.name;
+
+        // Replay the frames by the postcondition rule. The caller's
+        // JUMP and the callee's JUMPDEST both carry the invoke, on
+        // consecutive steps: count it once.
+        const frames: { name: string; step: number }[] = [];
+        const read: bigint[] = [];
+        let calls = 0;
+        for (let i = 0; i < steps.length; i++) {
+          if (i > 0) {
+            const previous = instructionAt(steps[i - 1]);
+            const { invoke, return: returned } = events(previous?.context);
+            if (returned !== undefined && invoke !== undefined) {
+              // A tail call's back-edge: the frame is reused
+              frames[frames.length - 1] = { name: invoke, step: i };
+            } else if (returned !== undefined) {
+              expect(frames.pop()?.name).toBe(returned);
+              const ret = (previous!.context as Format.Program.Context.Return)
+                .return;
+              if (ret.data) {
+                const state = createMachineState(executor, {
+                  traceStep: steps[i],
+                });
+                const cursor = await dereference(ret.data.pointer, { state });
+                const view = await cursor.view(state);
+                read.push((await view.read(view.regions[0])).asUint());
+              }
+            } else if (invoke !== undefined) {
+              const top = frames[frames.length - 1];
+              if (top?.name === invoke && top.step === i - 1) {
+                top.step = i;
+              } else {
+                frames.push({ name: invoke, step: i });
+                calls++;
+              }
+            }
+          }
+
+          // The innermost frame is the function whose code this
+          // step runs: on at the callee's entry JUMPDEST, through
+          // its exit JUMP, and off at the caller's continuation
+          expect(
+            frames[frames.length - 1]?.name,
+            `step ${i} (pc ${steps[i].pc})`,
+          ).toBe(functionAt(steps[i].pc));
+        }
+
+        expect(frames).toEqual([]);
+        expect(read).toHaveLength(calls);
+        if (level === 0) {
+          expect(read).toEqual(returns);
+        } else {
+          for (const value of read) expect(returns).toContain(value);
+        }
+      });
+    }
+  }
+});

packages/bugc/src/evmgen/generation/block.ts

@@ -65,21 +65,16 @@ export function generate<S extends Stack>(
 
       // Set JUMPDEST for non-first blocks
       if (!isFirstBlock) {
-        // Check if this is a call continuation
-        let isContinuation = false;
-        let calledFunction = "";
+        // A call's continuation resumes at the call expression, so its
+        // JUMPDEST carries the call site's source range. The callee's
+        // exit JUMP already carries the `return`.
         let callSiteCode: Format.Program.Context.Code["code"] | undefined;
         if (func && predecessor) {
           const predBlock = func.blocks.get(predecessor);
           if (
             predBlock?.terminator.kind === "call" &&
             predBlock.terminator.continuation === block.id
           ) {
-            isContinuation = true;
-            calledFunction = predBlock.terminator.function;
-            // The continuation resumes at the call expression, so
-            // carry the call site's source range onto the return
-            // context (disjoint keys — flat composition).
             const ctx = predBlock.terminator.operationDebug?.context as
               | Record<string, unknown>
               | undefined;
@@ -89,39 +84,13 @@ export function generate<S extends Stack>(
           }
         }
 
-        // Add JUMPDEST with continuation annotation if applicable
-        if (isContinuation) {
-          // Return context describes state after JUMPDEST
-          // executes: TOS is the return value (if any).
-          // data pointer is required by the schema; for
-          // void returns, slot 0 is still valid (empty).
-          const calledFunc = functions?.get(calledFunction);
-          const declaration =
-            calledFunc?.loc && calledFunc?.sourceId
-              ? {
-                  source: { id: calledFunc.sourceId },
-                  range: calledFunc.loc,
-                }
-              : undefined;
-          const returnCtx: Format.Program.Context.Return = {
-            return: {
-              identifier: calledFunction,
-              ...(declaration ? { declaration } : {}),
-              data: {
-                pointer: {
-                  location: "stack" as const,
-                  slot: 0,
-                },
-              },
-            },
-          };
+        if (callSiteCode) {
           const entry = block.entryDebug;
           const continuationDebug = {
             ...entry,
             context: {
               ...entry?.context,
-              ...returnCtx,
-              ...(callSiteCode ? { code: callSiteCode } : {}),
+              code: callSiteCode,
             } as Format.Program.Context,
           };
           result = result.then(JUMPDEST({ debug: continuationDebug }));
@@ -260,7 +229,7 @@ export function generate<S extends Stack>(
       const terminate =
         term.kind === "call"
           ? generateCallTerminator(term, functions)
-          : generateTerminator(term, isLastBlock, isUserFunction);
+          : generateTerminator(term, isLastBlock, isUserFunction, func);
       // A terminator carries an inlined function's invoke or return
       // only when its block has no instruction to carry it (an
       // inlined function that starts with a call, or returns a call's