web: show storage as of the current step in the trace viewer

packages/web/src/theme/ProgramExample/TraceDrawer.tsx

@@ -25,6 +25,7 @@ import {
 } from "@ethdebug/bugc-react";
 import { Executor, createTraceCollector, type TraceStep } from "@ethdebug/evm";
 import { dereference, Data, type Machine } from "@ethdebug/pointers";
+import { storageByStep } from "./storageByStep";
 import {
   buildCallStack,
   effectiveContextForStep,
@@ -82,7 +83,14 @@ function TraceDrawerContent(): JSX.Element {
   const [currentStep, setCurrentStep] = useState(0);
   const [isTracing, setIsTracing] = useState(false);
   const [traceError, setTraceError] = useState<string | null>(null);
-  const [storage, setStorage] = useState<Record<string, string>>({});
+  const [initialStorage, setInitialStorage] = useState<Record<string, string>>(
+    {},
+  );
+  const storages = useMemo(
+    () => storageByStep(trace, initialStorage),
+    [trace, initialStorage],
+  );
+  const storage = storages[currentStep] ?? initialStorage;
   const [showInstructionObject, setShowInstructionObject] = useState(false);
   const [objectHeight, setObjectHeight] = useState(OBJECT_DEFAULT_HEIGHT);
   const [isResizingObject, setIsResizingObject] = useState(false);
@@ -342,7 +350,7 @@ function TraceDrawerContent(): JSX.Element {
       const step = trace[frame.stepIndex];
       if (!step) return;
 
-      const state = traceStepToState(step, storage);
+      const state = traceStepToState(step, storages[frame.stepIndex] ?? {});
       const args: ResolvedArg[] = ptrs.map((_, i) => ({
         name: names?.[i] ?? `_${i}`,
       }));
@@ -376,7 +384,7 @@ function TraceDrawerContent(): JSX.Element {
     return () => {
       cancelled = true;
     };
-  }, [callStack, trace, storage]);
+  }, [callStack, trace, storages]);
 
   // Resolve the current instruction's variable values by
   // dereferencing each variable's pointer against the step
@@ -437,7 +445,7 @@ function TraceDrawerContent(): JSX.Element {
       setTrace([]);
       setCurrentStep(0);
       setTraceError(null);
-      setStorage({});
+      setInitialStorage({});
 
       let bytecode: BytecodeOutput | undefined;
 
@@ -489,13 +497,7 @@ function TraceDrawerContent(): JSX.Element {
           await executor.deploy(createHex);
         }
 
-        const [handler, getTrace] = createTraceCollector();
-        await executor.execute({}, handler);
-
-        const collectedTrace = getTrace();
-        setTrace(collectedTrace.steps);
-        setCurrentStep(0);
-
+        // storage before execution (e.g. set by the constructor)
         const storageEntries: Record<string, string> = {};
         for (let i = 0n; i < 16n; i++) {
           const value = await executor.getStorage(i);
@@ -504,7 +506,14 @@ function TraceDrawerContent(): JSX.Element {
             storageEntries[slot] = `0x${value.toString(16).padStart(64, "0")}`;
           }
         }
-        setStorage(storageEntries);
+
+        const [handler, getTrace] = createTraceCollector();
+        await executor.execute({}, handler);
+
+        const collectedTrace = getTrace();
+        setTrace(collectedTrace.steps);
+        setCurrentStep(0);
+        setInitialStorage(storageEntries);
       } catch (e) {
         setTraceError(e instanceof Error ? e.message : String(e));
       } finally {

packages/web/src/theme/ProgramExample/storageByStep.test.ts

@@ -0,0 +1,49 @@
+import { describe, it, expect } from "vitest";
+
+import { storageByStep } from "./storageByStep";
+
+const step = (opcode: string, stack: bigint[] = []) => ({
+  pc: 0,
+  opcode,
+  stack,
+});
+
+describe("storageByStep", () => {
+  const slot0 = `0x${"00".padStart(2, "0")}`;
+  const word = (n: number) => `0x${n.toString(16).padStart(64, "0")}`;
+
+  // stack arrays list the top of the stack last
+  const trace = [
+    step("PUSH1"),
+    step("SSTORE", [1n, 0n]),
+    step("PUSH1"),
+    step("SSTORE", [2n, 0n]),
+    step("STOP"),
+  ];
+
+  it("shows storage as of each step, not the final storage", () => {
+    const states = storageByStep(trace, {});
+
+    expect(states[0]).toEqual({});
+    expect(states[1]).toEqual({});
+    expect(states[2]).toEqual({ [slot0]: word(1) });
+    expect(states[3]).toEqual({ [slot0]: word(1) });
+    expect(states[4]).toEqual({ [slot0]: word(2) });
+  });
+
+  it("starts from the given initial storage", () => {
+    const states = storageByStep(trace, { "0x05": word(9) });
+
+    expect(states[0]).toEqual({ "0x05": word(9) });
+    expect(states[4]).toEqual({ "0x05": word(9), [slot0]: word(2) });
+  });
+
+  it("removes slots that are written with zero", () => {
+    const states = storageByStep([step("SSTORE", [0n, 0n]), step("STOP")], {
+      [slot0]: word(7),
+    });
+
+    expect(states[0]).toEqual({ [slot0]: word(7) });
+    expect(states[1]).toEqual({});
+  });
+});

packages/web/src/theme/ProgramExample/storageByStep.ts

@@ -0,0 +1,39 @@
+import type { TraceStep } from "@ethdebug/evm";
+
+/**
+ * Compute the storage that holds at each step of a trace: the
+ * initial storage plus the SSTOREs executed before that step.
+ * Like the stack and memory in a step, this is the state before
+ * the step's own instruction runs.
+ *
+ * Trace steps do not carry storage, so each SSTORE is read from
+ * its operands on the stack (key on top, then value).
+ */
+export function storageByStep(
+  trace: TraceStep[],
+  initial: Record<string, string>,
+): Record<string, string>[] {
+  const states: Record<string, string>[] = [];
+  let current = initial;
+
+  for (const step of trace) {
+    states.push(current);
+
+    if (step.opcode !== "SSTORE" || step.stack.length < 2) {
+      continue;
+    }
+
+    const key = step.stack[step.stack.length - 1];
+    const value = step.stack[step.stack.length - 2];
+    const slot = `0x${key.toString(16).padStart(2, "0")}`;
+
+    current = { ...current };
+    if (value === 0n) {
+      delete current[slot];
+    } else {
+      current[slot] = `0x${value.toString(16).padStart(64, "0")}`;
+    }
+  }
+
+  return states;
+}