Add reactive notebook minimap for navigating and visualizing dataflow

frontend/src/components/data-table/hooks/use-panel-ownership.ts

@@ -7,7 +7,7 @@ import {
   isCellAwareAtom,
 } from "@/components/editor/chrome/panels/context-aware-panel/atoms";
 import type { PanelType } from "@/components/editor/chrome/panels/context-aware-panel/context-aware-panel";
-import { lastFocusedCellIdAtom } from "@/core/cells/focus";
+import { useCellFocusActions, useLastFocusedCellId } from "@/core/cells/focus";
 import type { CellId } from "@/core/cells/ids";
 import { Logger } from "@/utils/Logger";
 
@@ -21,9 +21,8 @@ export function usePanelOwnership(
   cellId?: CellId | null,
 ): PanelOwnershipResult {
   let isPanelCellAware = useAtomValue(isCellAwareAtom);
-  const [lastFocusedCellId, setLastFocusedCellId] = useAtom(
-    lastFocusedCellIdAtom,
-  );
+  const { focusCell } = useCellFocusActions();
+  const lastFocusedCellId = useLastFocusedCellId();
   const [panelType, setPanelType] = useAtom(contextAwarePanelType);
   const [panelOwner, setPanelOwner] = useAtom(contextAwarePanelOwner);
   const [isContextAwarePanelOpen, setContextAwarePanelOpen] = useAtom(
@@ -64,7 +63,7 @@ export function usePanelOwnership(
       setPanelOwner(panelId);
       // if cell-aware, we want to focus on this cell when toggled open
       if (isPanelCellAware && cellId) {
-        setLastFocusedCellId(cellId);
+        focusCell({ cellId });
       }
       setContextAwarePanelOpen(true);
       setPanelType(panelType);

frontend/src/components/editor/chrome/state.ts

@@ -10,6 +10,7 @@ export interface ChromeState {
   selectedPanel: PanelType | undefined;
   isSidebarOpen: boolean;
   isTerminalOpen: boolean;
+  isMinimapOpen: boolean;
 }
 
 const KEY = "marimo:sidebar";
@@ -21,6 +22,7 @@ const storage = new ZodLocalStorage<ChromeState>(
       .transform((v) => v as PanelType),
     isSidebarOpen: z.boolean(),
     isTerminalOpen: z.boolean(),
+    isMinimapOpen: z.boolean(),
   }),
   initialState,
 );
@@ -30,6 +32,7 @@ function initialState(): ChromeState {
     selectedPanel: "variables", // initial panel
     isSidebarOpen: false,
     isTerminalOpen: false,
+    isMinimapOpen: false,
   };
 }
 
@@ -71,6 +74,14 @@ const {
       ...state,
       isTerminalOpen: isOpen,
     }),
+    toggleMinimap: (state) => ({
+      ...state,
+      isMinimapOpen: !state.isMinimapOpen,
+    }),
+    setIsMinimapOpen: (state, isOpen: boolean) => ({
+      ...state,
+      isMinimapOpen: isOpen,
+    }),
   },
   [(_prevState, newState) => storage.set(KEY, newState)],
 );

frontend/src/components/editor/chrome/wrapper/__tests__/minimap-state.test.ts

@@ -0,0 +1,209 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+import { describe, expect, it } from "vitest";
+import type { CellId } from "@/core/cells/ids";
+import type { VariableName, Variables } from "@/core/variables/types";
+import { buildCellGraph } from "../minimap-state";
+
+const cellId = (id: string) => id as CellId;
+const varName = (name: string) => name as VariableName;
+
+describe("buildCellGraph", () => {
+  it("builds graph for simple linear dependency", () => {
+    // cell1 declares x, cell2 uses x
+    const cellIds = [cellId("cell1"), cellId("cell2")];
+    const variables: Variables = {
+      [varName("x")]: {
+        name: varName("x"),
+        declaredBy: [cellId("cell1")],
+        usedBy: [cellId("cell2")],
+      },
+    };
+
+    const graph = buildCellGraph(cellIds, variables);
+
+    expect(graph).toMatchInlineSnapshot(`
+      {
+        "cell1": {
+          "ancestors": Set {},
+          "children": Set {
+            "cell2",
+          },
+          "descendants": Set {
+            "cell2",
+          },
+          "parents": Set {},
+          "variables": [
+            "x",
+          ],
+        },
+        "cell2": {
+          "ancestors": Set {
+            "cell1",
+          },
+          "children": Set {},
+          "descendants": Set {},
+          "parents": Set {
+            "cell1",
+          },
+          "variables": [],
+        },
+      }
+    `);
+  });
+
+  it("builds graph for diamond pattern", () => {
+    // cell1 declares x
+    // cell2 and cell3 both use x and declare y and z
+    // cell4 uses y and z
+    const cellIds = [
+      cellId("cell1"),
+      cellId("cell2"),
+      cellId("cell3"),
+      cellId("cell4"),
+    ];
+    const variables: Variables = {
+      [varName("x")]: {
+        name: varName("x"),
+        declaredBy: [cellId("cell1")],
+        usedBy: [cellId("cell2"), cellId("cell3")],
+      },
+      [varName("y")]: {
+        name: varName("y"),
+        declaredBy: [cellId("cell2")],
+        usedBy: [cellId("cell4")],
+      },
+      [varName("z")]: {
+        name: varName("z"),
+        declaredBy: [cellId("cell3")],
+        usedBy: [cellId("cell4")],
+      },
+    };
+
+    const graph = buildCellGraph(cellIds, variables);
+
+    expect(graph).toMatchInlineSnapshot(`
+      {
+        "cell1": {
+          "ancestors": Set {},
+          "children": Set {
+            "cell2",
+            "cell3",
+          },
+          "descendants": Set {
+            "cell2",
+            "cell4",
+            "cell3",
+          },
+          "parents": Set {},
+          "variables": [
+            "x",
+          ],
+        },
+        "cell2": {
+          "ancestors": Set {
+            "cell1",
+          },
+          "children": Set {
+            "cell4",
+          },
+          "descendants": Set {
+            "cell4",
+          },
+          "parents": Set {
+            "cell1",
+          },
+          "variables": [
+            "y",
+          ],
+        },
+        "cell3": {
+          "ancestors": Set {
+            "cell1",
+          },
+          "children": Set {
+            "cell4",
+          },
+          "descendants": Set {
+            "cell4",
+          },
+          "parents": Set {
+            "cell1",
+          },
+          "variables": [
+            "z",
+          ],
+        },
+        "cell4": {
+          "ancestors": Set {
+            "cell2",
+            "cell1",
+            "cell3",
+          },
+          "children": Set {},
+          "descendants": Set {},
+          "parents": Set {
+            "cell2",
+            "cell3",
+          },
+          "variables": [],
+        },
+      }
+    `);
+  });
+
+  it("handles self-referencing and isolated cells", () => {
+    const cellIds = [cellId("cell1"), cellId("cell2"), cellId("cell3")];
+    const variables: Variables = {
+      [varName("x")]: {
+        name: varName("x"),
+        declaredBy: [cellId("cell1")],
+        usedBy: [cellId("cell1"), cellId("cell2")], // self-reference + downstream
+      },
+      [varName("y")]: {
+        name: varName("y"),
+        declaredBy: [cellId("cell3")],
+        usedBy: [], // isolated variable
+      },
+    };
+
+    const graph = buildCellGraph(cellIds, variables);
+
+    expect(graph).toMatchInlineSnapshot(`
+      {
+        "cell1": {
+          "ancestors": Set {},
+          "children": Set {
+            "cell2",
+          },
+          "descendants": Set {
+            "cell2",
+          },
+          "parents": Set {},
+          "variables": [
+            "x",
+          ],
+        },
+        "cell2": {
+          "ancestors": Set {
+            "cell1",
+          },
+          "children": Set {},
+          "descendants": Set {},
+          "parents": Set {
+            "cell1",
+          },
+          "variables": [],
+        },
+        "cell3": {
+          "ancestors": Set {},
+          "children": Set {},
+          "descendants": Set {},
+          "parents": Set {},
+          "variables": [
+            "y",
+          ],
+        },
+      }
+    `);
+  });
+});

frontend/src/components/editor/chrome/wrapper/app-chrome.tsx

@@ -32,6 +32,7 @@ import { SnippetsPanel } from "../panels/snippets-panel";
 import { TracingPanel } from "../panels/tracing-panel";
 import { VariablePanel } from "../panels/variable-panel";
 import { useChromeActions, useChromeState } from "../state";
+import { Minimap } from "./minimap";
 import { PanelsWrapper } from "./panels";
 import { createStorage } from "./storage";
 import { handleDragging } from "./utils";
@@ -244,6 +245,7 @@ export const AppChrome: React.FC<PropsWithChildren> = ({ children }) => {
         </Panel>
         <ContextAwarePanel />
       </PanelGroup>
+      <Minimap />
       <ErrorBoundary>
         <TooltipProvider>
           <Footer />

frontend/src/components/editor/chrome/wrapper/footer-items/minimap-status.tsx

@@ -0,0 +1,21 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+import { MapIcon } from "lucide-react";
+import { useChromeActions, useChromeState } from "../../state";
+import { FooterItem } from "../footer-item";
+
+export const MinimapStatusIcon: React.FC = () => {
+  const { isMinimapOpen } = useChromeState();
+  const { toggleMinimap } = useChromeActions();
+
+  return (
+    <FooterItem
+      tooltip="Toggle Minimap"
+      selected={isMinimapOpen}
+      onClick={() => toggleMinimap()}
+      data-testid="footer-minimap"
+    >
+      <MapIcon className="h-4 w-4" />
+    </FooterItem>
+  );
+};

frontend/src/components/editor/chrome/wrapper/footer.tsx

@@ -17,6 +17,7 @@ import { AIStatusIcon } from "./footer-items/ai-status";
 import { BackendConnection } from "./footer-items/backend-status";
 import { CopilotStatusIcon } from "./footer-items/copilot-status";
 import { MachineStats } from "./footer-items/machine-stats";
+import { MinimapStatusIcon } from "./footer-items/minimap-status";
 import { RTCStatus } from "./footer-items/rtc-status";
 import { RuntimeSettings } from "./footer-items/runtime-settings";
 
@@ -78,6 +79,7 @@ export const Footer: React.FC = () => {
 
       <div className="flex items-center flex-shrink-0 min-w-0">
         <MachineStats />
+        <MinimapStatusIcon />
         <AIStatusIcon />
         <CopilotStatusIcon />
         <RTCStatus />