feat: go-to-definition with cmd+click

frontend/src/components/dependency-graph/panels.tsx

@@ -1,6 +1,6 @@
 /* Copyright 2024 Marimo. All rights reserved. */
 import React, { memo } from "react";
-import { Edge, Panel } from "reactflow";
+import { type Edge, Panel } from "reactflow";
 import { Button } from "../ui/button";
 import {
   Rows3Icon,
@@ -13,18 +13,18 @@ import {
   SettingsIcon,
   MoreVerticalIcon,
 } from "lucide-react";
-import { GraphLayoutView, GraphSelection, GraphSettings } from "./types";
+import type { GraphLayoutView, GraphSelection, GraphSettings } from "./types";
 import { CellLink } from "../editor/links/cell-link";
 import { CellLinkList } from "../editor/links/cell-link-list";
 import { VariableName } from "../variables/common";
-import { Variable, Variables } from "@/core/variables/types";
+import type { Variable, Variables } from "@/core/variables/types";
 import { Popover, PopoverContent, PopoverTrigger } from "../ui/popover";
 import { Checkbox } from "../ui/checkbox";
 import { Label } from "../ui/label";
 import { ConnectionCellActionsDropdown } from "../editor/cell/cell-actions";
 import { getCellEditorView } from "@/core/cells/cells";
-import { CellId } from "@/core/cells/ids";
-import { goToDefinition } from "@/core/codemirror/find-replace/search-highlight";
+import type { CellId } from "@/core/cells/ids";
+import { goToVariableDefinition } from "@/core/codemirror/go-to-definition/commands";
 
 interface Props {
   view: GraphLayoutView;
@@ -120,7 +120,7 @@ export const GraphSelectionPanel: React.FC<{
   const highlightInCell = (cellId: CellId, variableName: string) => {
     const editorView = getCellEditorView(cellId);
     if (editorView) {
-      goToDefinition(editorView, variableName);
+      goToVariableDefinition(editorView, variableName);
     }
   };
 

frontend/src/components/editor/cell/cell-context-menu.tsx

@@ -1,7 +1,7 @@
 /* Copyright 2024 Marimo. All rights reserved. */
 import React, { Fragment } from "react";
 import {
-  CellActionButtonProps,
+  type CellActionButtonProps,
   useCellActionButtons,
 } from "../actions/useCellActionButton";
 import {
@@ -12,15 +12,15 @@ import {
   ContextMenuSeparator,
 } from "@/components/ui/context-menu";
 import { renderMinimalShortcut } from "@/components/shortcuts/renderShortcut";
-import { ActionButton } from "../actions/types";
+import type { ActionButton } from "../actions/types";
 import {
   ClipboardPasteIcon,
   CopyIcon,
   ImageIcon,
   ScissorsIcon,
   SearchIcon,
 } from "lucide-react";
-import { goToDefinition } from "@/core/codemirror/go-to-definition";
+import { goToDefinitionAtCursorPosition } from "@/core/codemirror/go-to-definition/utils";
 
 interface Props extends CellActionButtonProps {
   children: React.ReactNode;
@@ -90,7 +90,7 @@ export const CellActionsContextMenu = ({ children, ...props }: Props) => {
         const { getEditorView } = props;
         const editorView = getEditorView();
         if (editorView) {
-          goToDefinition(editorView);
+          goToDefinitionAtCursorPosition(editorView);
         }
       },
     },

frontend/src/components/variables/variables-table.tsx

@@ -8,25 +8,25 @@ import {
   TableCell,
   Table,
 } from "../ui/table";
-import { Variable, Variables } from "@/core/variables/types";
-import { CellId } from "@/core/cells/ids";
+import type { Variable, Variables } from "@/core/variables/types";
+import type { CellId } from "@/core/cells/ids";
 import { CellLink } from "@/components/editor/links/cell-link";
 import { cn } from "@/utils/cn";
 import { SquareEqualIcon, WorkflowIcon } from "lucide-react";
 import {
   useReactTable,
   getCoreRowModel,
   flexRender,
-  ColumnDef,
-  SortingState,
+  type ColumnDef,
+  type SortingState,
   getSortedRowModel,
-  ColumnSort,
+  type ColumnSort,
   getFilteredRowModel,
 } from "@tanstack/react-table";
 import { DataTableColumnHeader } from "../data-table/column-header";
 import { sortBy } from "lodash-es";
 import { getCellEditorView } from "@/core/cells/cells";
-import { goToDefinition } from "@/core/codemirror/find-replace/search-highlight";
+import { goToVariableDefinition } from "@/core/codemirror/go-to-definition/commands";
 import { SearchInput } from "../ui/input";
 import { CellLinkList } from "../editor/links/cell-link-list";
 import { VariableName } from "./common";
@@ -122,7 +122,7 @@ const COLUMNS = [
       const highlightInCell = (cellId: CellId) => {
         const editorView = getCellEditorView(cellId);
         if (editorView) {
-          goToDefinition(editorView, name);
+          goToVariableDefinition(editorView, name);
         }
       };
 

frontend/src/core/cells/cells.ts

@@ -34,7 +34,7 @@ import { clamp } from "@/utils/math";
 import type { LayoutState } from "../layout/layout";
 import { notebookIsRunning } from "./utils";
 import {
-  getEditorCodeAsPython,
+  splitEditor,
   updateEditorCodeFromPython,
 } from "../codemirror/language/utils";
 
@@ -627,38 +627,18 @@ const {
       cellLogs: [],
     };
   },
-  splitCell: (state, action: { cellId: CellId; cursorPos: number }) => {
-    const { cellId, cursorPos } = action;
+  splitCell: (state, action: { cellId: CellId }) => {
+    const { cellId } = action;
     const index = state.cellIds.indexOf(cellId);
     const cell = state.cellData[cellId];
     const cellHandle = state.cellHandles[cellId].current;
 
     if (cellHandle?.editorView == null) {
-      // TODO: Because of this we can't do  the reducer tests like for the other functions
       return state;
     }
 
-    // Figure out if we're at the start or end of a line to adjust the cursor positions
-    const isCursorAtLineStart =
-      cell.code.length > 0 && cell.code[cursorPos - 1] === "\n";
-    const isCursorAtLineEnd =
-      cell.code.length > 0 && cell.code[cursorPos] === "\n";
-
-    const beforeAdjustedCursorPos = isCursorAtLineStart
-      ? cursorPos - 1
-      : cursorPos;
-    const afterAdjustedCursorPos = isCursorAtLineEnd
-      ? cursorPos + 1
-      : cursorPos;
-
-    const beforeCursorCode = getEditorCodeAsPython(
-      cellHandle.editorView,
-      0,
-      beforeAdjustedCursorPos,
-    );
-    const afterCursorCode = getEditorCodeAsPython(
+    const { beforeCursorCode, afterCursorCode } = splitEditor(
       cellHandle.editorView,
-      afterAdjustedCursorPos,
     );
 
     updateEditorCodeFromPython(cellHandle.editorView, beforeCursorCode);
@@ -673,7 +653,9 @@ const {
         [cellId]: {
           ...cell,
           code: beforeCursorCode,
-          edited: beforeCursorCode.trim() !== cell.lastCodeRun,
+          edited:
+            Boolean(beforeCursorCode) &&
+            beforeCursorCode.trim() !== cell.lastCodeRun?.trim(),
         },
         [newCellId]: createCell({
           id: newCellId,

frontend/src/core/cells/scrollCellIntoView.ts

@@ -1,10 +1,10 @@
 /* Copyright 2024 Marimo. All rights reserved. */
-import { RefObject } from "react";
+import type { RefObject } from "react";
 import { Logger } from "../../utils/Logger";
-import { CellId, HTMLCellId } from "./ids";
-import { CellHandle } from "@/components/editor/Cell";
-import { CellConfig } from "./types";
-import { goToDefinition } from "../codemirror/find-replace/search-highlight";
+import { type CellId, HTMLCellId } from "./ids";
+import type { CellHandle } from "@/components/editor/Cell";
+import type { CellConfig } from "./types";
+import { goToVariableDefinition } from "../codemirror/go-to-definition/commands";
 
 export function focusAndScrollCellIntoView({
   cellId,
@@ -57,7 +57,7 @@ export function focusAndScrollCellIntoView({
         },
       });
     } else if (variableName) {
-      goToDefinition(editor, variableName);
+      goToVariableDefinition(editor, variableName);
     }
   }
 

frontend/src/core/codemirror/cells/extensions.ts

@@ -9,7 +9,7 @@ import { getEditorCodeAsPython } from "../language/utils";
 import { formattingChangeEffect } from "../format";
 import { closeCompletion, completionStatus } from "@codemirror/autocomplete";
 import { isAtEndOfEditor, isAtStartOfEditor } from "../utils";
-import { goToDefinition } from "../go-to-definition";
+import { goToDefinitionAtCursorPosition } from "../go-to-definition/utils";
 
 export interface MovementCallbacks
   extends Pick<
@@ -246,7 +246,7 @@ export function cellMovementBundle(
       preventDefault: true,
       stopPropagation: true,
       run: (ev) => {
-        goToDefinition(ev);
+        goToDefinitionAtCursorPosition(ev);
         return true;
       },
     },
@@ -255,8 +255,7 @@ export function cellMovementBundle(
       preventDefault: true,
       stopPropagation: true,
       run: (ev) => {
-        const cursorPos = ev.state.selection.main.head;
-        splitCell({ cellId, cursorPos });
+        splitCell({ cellId });
         requestAnimationFrame(() => {
           ev.contentDOM.blur();
           moveToNextCell({ cellId, before: false }); // focus new cell

frontend/src/core/codemirror/cm.ts

@@ -34,20 +34,20 @@ import {
   EditorView,
 } from "@codemirror/view";
 
-import { EditorState, Extension, Prec } from "@codemirror/state";
+import { EditorState, type Extension, Prec } from "@codemirror/state";
 import { oneDark } from "@codemirror/theme-one-dark";
 
-import { CompletionConfig, KeymapConfig } from "../config/config-schema";
-import { Theme } from "../../theme/useTheme";
+import type { CompletionConfig, KeymapConfig } from "../config/config-schema";
+import type { Theme } from "../../theme/useTheme";
 
 import { findReplaceBundle } from "./find-replace/extension";
 import {
-  CodeCallbacks,
-  MovementCallbacks,
+  type CodeCallbacks,
+  type MovementCallbacks,
   cellCodeEditingBundle,
   cellMovementBundle,
 } from "./cells/extensions";
-import { CellId } from "../cells/ids";
+import type { CellId } from "../cells/ids";
 import { keymapBundle } from "./keymaps/keymaps";
 import { scrollActiveLineIntoView } from "./extensions";
 import { copilotBundle } from "./copilot/extension";
@@ -58,6 +58,7 @@ import {
   clickablePlaceholderExtension,
   smartPlaceholderExtension,
 } from "./placeholder/extensions";
+import { goToDefinitionBundle } from "./go-to-definition/extension";
 
 export interface CodeMirrorSetupOpts {
   cellId: CellId;
@@ -91,6 +92,8 @@ export const setupCodeMirror = ({
     cellCodeEditingBundle(cellId, cellCodeCallbacks),
     // Comes last so that it can be overridden
     basicBundle(completionConfig, theme),
+    // Underline cmd+clickable placeholder
+    goToDefinitionBundle(),
     showPlaceholder
       ? Prec.highest(smartPlaceholderExtension("import marimo as mo"))
       : enableAI

frontend/src/core/codemirror/find-replace/search-highlight.ts

@@ -4,15 +4,14 @@ import { RangeSetBuilder, StateEffect, StateField } from "@codemirror/state";
 import {
   Decoration,
   ViewPlugin,
-  DecorationSet,
+  type DecorationSet,
   EditorView,
-  ViewUpdate,
+  type ViewUpdate,
 } from "@codemirror/view";
-import { QueryType, asQueryCreator } from "./query";
+import { type QueryType, asQueryCreator } from "./query";
 import { store } from "@/core/state/jotai";
 import { findReplaceAtom } from "./state";
 import { getAllEditorViews } from "@/core/cells/cells";
-import { syntaxTree } from "@codemirror/language";
 
 const setSearchQuery = StateEffect.define<SearchQuery>();
 
@@ -147,57 +146,3 @@ export const highlightTheme = EditorView.baseTheme({
     backgroundColor: "#6199ff88 !important",
   },
 });
-
-/**
- * This function will select the first occurrence of the given variable name.
- */
-export function goToDefinition(
-  view: EditorView,
-  variableName: string,
-): boolean {
-  const state = view.state;
-  const tree = syntaxTree(state);
-
-  let found = false;
-  let from = 0;
-
-  tree.iterate({
-    enter: (node) => {
-      if (found) {
-        return false;
-      } // Stop traversal if found
-
-      // Check if the node is an identifier and matches the variable name
-      if (
-        node.name === "VariableName" &&
-        state.doc.sliceString(node.from, node.to) === variableName
-      ) {
-        from = node.from;
-        found = true;
-        return false; // Stop traversal
-      }
-
-      // Skip comments and strings
-      if (node.name === "Comment" || node.name === "String") {
-        return false;
-      }
-    },
-  });
-
-  if (found) {
-    view.focus();
-    view.dispatch({
-      selection: {
-        anchor: from,
-        head: from,
-      },
-      // Unfortunately, EditorView.scrollIntoView does
-      // not support smooth scrolling.
-      effects: EditorView.scrollIntoView(from, {
-        y: "center",
-      }),
-    });
-    return true;
-  }
-  return false;
-}