ui: add source-level stepping to trace viewer controls

packages/programs-react/src/components/TraceContext.tsx

@@ -23,6 +23,54 @@ import {
 } from "#utils/mockTrace";
 import { traceStepToMachineState } from "#utils/traceState";
 
+/**
+ * Compute a key representing an instruction's source range,
+ * used to detect when stepping has moved to a new source
+ * location. Returns empty string for instructions without
+ * source ranges.
+ */
+function sourceRangeKey(instruction: Program.Instruction | undefined): string {
+  if (!instruction?.context) return "";
+
+  const ctx = instruction.context as Record<string, unknown>;
+  const ranges = collectCodeRanges(ctx);
+  if (ranges.length === 0) return "";
+
+  return ranges.map((r) => `${r.offset}:${r.length}`).join(",");
+}
+
+function collectCodeRanges(
+  ctx: Record<string, unknown>,
+): Array<{ offset: number; length: number }> {
+  if ("code" in ctx && typeof ctx.code === "object") {
+    const code = ctx.code as Record<string, unknown>;
+    if (code.range && typeof code.range === "object") {
+      const r = code.range as Record<string, number>;
+      if (typeof r.offset === "number" && typeof r.length === "number") {
+        return [{ offset: r.offset, length: r.length }];
+      }
+    }
+  }
+
+  if ("gather" in ctx && Array.isArray(ctx.gather)) {
+    return ctx.gather.flatMap((item: unknown) =>
+      item && typeof item === "object"
+        ? collectCodeRanges(item as Record<string, unknown>)
+        : [],
+    );
+  }
+
+  if ("pick" in ctx && Array.isArray(ctx.pick)) {
+    return ctx.pick.flatMap((item: unknown) =>
+      item && typeof item === "object"
+        ? collectCodeRanges(item as Record<string, unknown>)
+        : [],
+    );
+  }
+
+  return [];
+}
+
 /**
  * A variable with its resolved value.
  */
@@ -98,10 +146,14 @@ export interface TraceState {
   /** Whether we're at the last step */
   isAtEnd: boolean;
 
-  /** Move to the next step */
+  /** Move to the next trace step */
   stepForward(): void;
-  /** Move to the previous step */
+  /** Move to the previous trace step */
   stepBackward(): void;
+  /** Step to the next different source range */
+  stepToNextSource(): void;
+  /** Step to the previous different source range */
+  stepToPrevSource(): void;
   /** Jump to a specific step */
   jumpToStep(index: number): void;
   /** Reset to the first step */
@@ -377,6 +429,41 @@ export function TraceProvider({
     setCurrentStepIndex((prev) => Math.max(prev - 1, 0));
   }, []);
 
+  const stepToNextSource = useCallback(() => {
+    setCurrentStepIndex((prev) => {
+      const currentKey = sourceRangeKey(pcToInstruction.get(trace[prev]?.pc));
+      for (let i = prev + 1; i < trace.length; i++) {
+        const instr = pcToInstruction.get(trace[i].pc);
+        const key = sourceRangeKey(instr);
+        if (key !== currentKey && key !== "") {
+          return i;
+        }
+      }
+      return trace.length - 1;
+    });
+  }, [trace, pcToInstruction]);
+
+  const stepToPrevSource = useCallback(() => {
+    setCurrentStepIndex((prev) => {
+      const currentKey = sourceRangeKey(pcToInstruction.get(trace[prev]?.pc));
+      // First skip past all steps with the same range
+      let i = prev - 1;
+      while (i > 0) {
+        const key = sourceRangeKey(pcToInstruction.get(trace[i].pc));
+        if (key !== currentKey && key !== "") break;
+        i--;
+      }
+      // Now find the start of that source range
+      const targetKey = sourceRangeKey(pcToInstruction.get(trace[i]?.pc));
+      while (i > 0) {
+        const prevKey = sourceRangeKey(pcToInstruction.get(trace[i - 1]?.pc));
+        if (prevKey !== targetKey) break;
+        i--;
+      }
+      return Math.max(0, i);
+    });
+  }, [trace, pcToInstruction]);
+
   const jumpToStep = useCallback(
     (index: number) => {
       setCurrentStepIndex(Math.max(0, Math.min(index, trace.length - 1)));
@@ -406,6 +493,8 @@ export function TraceProvider({
     isAtEnd: currentStepIndex >= trace.length - 1,
     stepForward,
     stepBackward,
+    stepToNextSource,
+    stepToPrevSource,
     jumpToStep,
     reset,
     jumpToEnd,

packages/programs-react/src/components/TraceControls.tsx

@@ -33,6 +33,8 @@ export function TraceControls({
     isAtEnd,
     stepBackward,
     stepForward,
+    stepToPrevSource,
+    stepToNextSource,
     reset,
     jumpToEnd,
   } = useTraceContext();
@@ -47,25 +49,43 @@ export function TraceControls({
           title="Reset to start"
           type="button"
         >
-          ⏮
+          ⏮
+        </button>
+        <button
+          className="trace-control-btn trace-control-prev-source"
+          onClick={stepToPrevSource}
+          disabled={isAtStart}
+          title="Previous source location"
+          type="button"
+        >
+          &#x25C0;
         </button>
         <button
           className="trace-control-btn trace-control-prev"
           onClick={stepBackward}
           disabled={isAtStart}
-          title="Previous step"
+          title="Previous trace step"
           type="button"
         >
-          ←
+          &#x25C1;
         </button>
         <button
           className="trace-control-btn trace-control-next"
           onClick={stepForward}
           disabled={isAtEnd}
-          title="Next step"
+          title="Next trace step"
+          type="button"
+        >
+          &#x25B7;
+        </button>
+        <button
+          className="trace-control-btn trace-control-next-source"
+          onClick={stepToNextSource}
+          disabled={isAtEnd}
+          title="Next source location"
           type="button"
         >
-          →
+          &#x25B6;
         </button>
         <button
           className="trace-control-btn trace-control-end"
@@ -74,7 +94,7 @@ export function TraceControls({
           title="Jump to end"
           type="button"
         >
-          ⏭
+          &#x23ED;
         </button>
       </div>
 

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

@@ -258,6 +258,45 @@ function TraceDrawerContent(): JSX.Element {
     setCurrentStep((prev) => Math.max(prev - 1, 0));
   };
 
+  const rangeKey = (stepIdx: number): string => {
+    const step = trace[stepIdx];
+    if (!step) return "";
+    const instr = pcToInstruction.get(step.pc);
+    if (!instr?.debug?.context) return "";
+    const ranges = extractSourceRange(instr.debug.context);
+    if (ranges.length === 0) return "";
+    return ranges.map((r) => `${r.offset}:${r.length}`).join(",");
+  };
+
+  const stepToNextSource = () => {
+    setCurrentStep((prev) => {
+      const currentKey = rangeKey(prev);
+      for (let i = prev + 1; i < trace.length; i++) {
+        const key = rangeKey(i);
+        if (key !== currentKey && key !== "") return i;
+      }
+      return trace.length - 1;
+    });
+  };
+
+  const stepToPrevSource = () => {
+    setCurrentStep((prev) => {
+      const currentKey = rangeKey(prev);
+      let i = prev - 1;
+      while (i > 0) {
+        const key = rangeKey(i);
+        if (key !== currentKey && key !== "") break;
+        i--;
+      }
+      const targetKey = rangeKey(i);
+      while (i > 0) {
+        if (rangeKey(i - 1) !== targetKey) break;
+        i--;
+      }
+      return Math.max(0, i);
+    });
+  };
+
   const jumpToStart = () => setCurrentStep(0);
   const jumpToEnd = () => setCurrentStep(trace.length - 1);
 
@@ -332,21 +371,37 @@ function TraceDrawerContent(): JSX.Element {
                   &#x23EE;
                 </button>
                 <button
-                  onClick={stepBackward}
+                  onClick={stepToPrevSource}
                   disabled={currentStep === 0}
                   className="trace-nav-btn"
-                  title="Step backward"
+                  title="Previous source location"
                 >
                   &#x25C0;
                 </button>
+                <button
+                  onClick={stepBackward}
+                  disabled={currentStep === 0}
+                  className="trace-nav-btn trace-nav-btn-step"
+                  title="Previous trace step"
+                >
+                  &#x25C1;
+                </button>
                 <span className="trace-step-info">
                   {currentStep + 1} / {trace.length}
                 </span>
                 <button
                   onClick={stepForward}
                   disabled={currentStep >= trace.length - 1}
+                  className="trace-nav-btn trace-nav-btn-step"
+                  title="Next trace step"
+                >
+                  &#x25B7;
+                </button>
+                <button
+                  onClick={stepToNextSource}
+                  disabled={currentStep >= trace.length - 1}
                   className="trace-nav-btn"
-                  title="Step forward"
+                  title="Next source location"
                 >
                   &#x25B6;
                 </button>