Repository: marimo-team/marimo
Pull request: Go to definition functionality added to hotkey and context menu
Fixes https://github.com/marimo-team/marimo/issues/1126

For this one, I set the hotkey to Mod+Shift+, but let me know if there's a better one.

I'm unsure if this is structured correctly, specifically, the utility functions specified in `cell-editor.tsx`. Let me know if there is a better way to structure it. Basically, I wasn't sure how to ensure we could use the same functionality with the right-click menu and the hotkeys. The solution works, but there is a bit of code repetition.

Created this with the help of [Glide](https://glide.agenticlabs.com)! The change in this document only covers adding the hotkey, I manually added the right-click after as it was a small change:
https://glide.agenticlabs.com/sharing?userId=65bae76b7fbfb1c20ac1f3b8&taskId=ghM1mqX

Problem statement:
Go to cell defining something
### Description

Quickly go to the cell defining a variable or a function.

### Suggested solution

- An item in the right click menu.
- A keyboard shortcut.

### Alternative

_No response_

### Additional context

_No response_

Hints from issue discussion:
Could this be a simple matter to code now that #1365 is fixed?
I didn't factor https://github.com/marimo-team/marimo/pull/1365 in when making these changes.

Here's the plan I used to make the change:
https://glide.agenticlabs.com/sharing?userId=65bae76b7fbfb1c20ac1f3b8&taskId=ghM1mqX

Do we think this needs to be updated based on the new functionality of `goToDefinition`? It looked like from my testing that `goToDefinition` just focused on the first instance of the variable within the cell, not across cells?
I thought thanks to marimo's design we already know the cell defining a
variable, so the remainder is going to the (single) definition in that cell.
Thanks for introducing me to glide, but I think I don't have the time and
understanding to contribute right now!

On Sun, May 19, 2024 at 3:15 PM Robert McKnight ***@***.***>
wrote:

> I didn't factor #1365 <https://github.com/marimo-team/marimo/pull/1365>
> in when making these changes.
>
> Here's the plan I used to make the change:
>
> https://glide.agenticlabs.com/sharing?userId=65bae76b7fbfb1c20ac1f3b8&taskId=ghM1mqX
>
> Do we think this needs to be updated based on the new functionality of
> goToDefinition? It looked like from my testing that goToDefinition just
> focused on the first instance of the variable within the cell, not across
> cells?
>
> —
> Reply to this email directly, view it on GitHub
> <https://github.com/marimo-team/marimo/issues/1126#issuecomment-2119214990>,
> or unsubscribe
> <https://github.com/notifications/unsubscribe-auth/AYWFZPD2VSJJH6W7IPOVXMTZDCJU7AVCNFSM6AAAAABGEPXL76VHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZDCMJZGIYTIOJZGA>
> .
> You are receiving this because you commented.Message ID:
> ***@***.***>
>


-- 
Alon Levy
Software Developer
Greenvibe
054-2395317

Commit message: Remove the call through

Patch to review:
diff --git a/frontend/src/components/editor/cell/cell-context-menu.tsx b/frontend/src/components/editor/cell/cell-context-menu.tsx
index f2a97d3b151..1efd618b6b4 100644
--- a/frontend/src/components/editor/cell/cell-context-menu.tsx
+++ b/frontend/src/components/editor/cell/cell-context-menu.tsx
@@ -18,7 +18,10 @@ import {
   CopyIcon,
   ImageIcon,
   ScissorsIcon,
+  SearchIcon,
 } from "lucide-react";
+import { useVariables } from "@/core/variables/state";
+import { goToDefinition } from "@/core/codemirror/go-to-definition";
 
 interface Props extends CellActionButtonProps {
   children: React.ReactNode;
@@ -26,6 +29,7 @@ interface Props extends CellActionButtonProps {
 
 export const CellActionsContextMenu = ({ children, ...props }: Props) => {
   const actions = useCellActionButtons({ cell: props });
+  const variables = useVariables();
   const [imageRightClicked, setImageRightClicked] =
     React.useState<HTMLImageElement>();
 
@@ -81,6 +85,17 @@ export const CellActionsContextMenu = ({ children, ...props }: Props) => {
         }
       },
     },
+    {
+      label: "Go to Definition",
+      icon: <SearchIcon size={13} strokeWidth={1.5} />,
+      handle: () => {
+        const { getEditorView } = props;
+        const editorView = getEditorView();
+        if (editorView) {
+          goToDefinition(editorView, variables);
+        }
+      },
+    },
   ];
 
   return (
diff --git a/frontend/src/core/cells/cells.ts b/frontend/src/core/cells/cells.ts
index 8f983637eb8..160559be733 100644
--- a/frontend/src/core/cells/cells.ts
+++ b/frontend/src/core/cells/cells.ts
@@ -156,7 +156,12 @@ const {
 } = createReducerAndAtoms(initialNotebookState, {
   createNewCell: (
     state,
-    action: { cellId: CellId | "__end__"; before: boolean; code?: string, newCellId?: CellId},
+    action: {
+      cellId: CellId | "__end__";
+      before: boolean;
+      code?: string;
+      newCellId?: CellId;
+    },
   ) => {
     const { cellId, before, code } = action;
     let { newCellId } = action;
@@ -249,6 +254,7 @@ const {
       cell: state.cellHandles[focusCellId],
       config: state.cellData[focusCellId].config,
       codeFocus: before ? "bottom" : "top",
+      variableName: undefined,
     });
     return state;
   },
@@ -556,6 +562,7 @@ const {
         cell: state.cellHandles[nextCellId],
         config: state.cellData[nextCellId].config,
         codeFocus: before ? "bottom" : "top",
+        variableName: undefined,
       });
       return state;
     }
@@ -581,6 +588,7 @@ const {
           cell: state.cellHandles[nextCellId],
           config: state.cellData[nextCellId].config,
           codeFocus: undefined,
+          variableName: undefined,
         });
       }
 
diff --git a/frontend/src/core/cells/scrollCellIntoView.ts b/frontend/src/core/cells/scrollCellIntoView.ts
index 573668633b2..4ebc3f2ceee 100644
--- a/frontend/src/core/cells/scrollCellIntoView.ts
+++ b/frontend/src/core/cells/scrollCellIntoView.ts
@@ -4,17 +4,20 @@ 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";
 
 export function focusAndScrollCellIntoView({
   cellId,
   cell,
   config,
   codeFocus,
+  variableName,
 }: {
   cellId: CellId;
   cell: RefObject<CellHandle>;
   config: CellConfig;
   codeFocus: "top" | "bottom" | undefined;
+  variableName: string | undefined;
 }) {
   if (!cell) {
     return;
@@ -53,6 +56,8 @@ export function focusAndScrollCellIntoView({
           head: lastLine.from,
         },
       });
+    } else if (variableName) {
+      goToDefinition(editor, variableName);
     }
   }
 
diff --git a/frontend/src/core/codemirror/cells/extensions.ts b/frontend/src/core/codemirror/cells/extensions.ts
index 3cfed77918e..f2c99ace950 100644
--- a/frontend/src/core/codemirror/cells/extensions.ts
+++ b/frontend/src/core/codemirror/cells/extensions.ts
@@ -9,6 +9,8 @@ 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 {useVariables} from "@/core/variables/state";
 
 export interface MovementCallbacks
   extends Pick<
@@ -50,6 +52,7 @@ export function cellMovementBundle(
     toggleHideCode,
     aiCellCompletion,
   } = callbacks;
+  const variables = useVariables()
 
   const hotkeys: KeyBinding[] = [
     {
@@ -240,6 +243,15 @@ export function cellMovementBundle(
         return true;
       },
     },
+    {
+      key: HOTKEYS.getHotkey("cell.goToDefinition").key,
+      preventDefault: true,
+      stopPropagation: true,
+      run: (ev) => {
+        goToDefinition(ev, variables);
+        return true;
+      },
+    },
     {
       key: HOTKEYS.getHotkey("cell.splitCell").key,
       preventDefault: true,
@@ -251,7 +263,6 @@ export function cellMovementBundle(
           ev.contentDOM.blur();
           moveToNextCell({ cellId, before: false }); // focus new cell
         });
-
         return true;
       },
     },
diff --git a/frontend/src/core/codemirror/go-to-definition.ts b/frontend/src/core/codemirror/go-to-definition.ts
new file mode 100644
index 00000000000..16d44afb9bb
--- /dev/null
+++ b/frontend/src/core/codemirror/go-to-definition.ts
@@ -0,0 +1,48 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+import { EditorState } from "@codemirror/state";
+import { EditorView } from "@codemirror/view";
+import { getPositionAtWordBounds } from "../codemirror/completion/hints";
+import { Variables, VariableName } from "../variables/types";
+import { focusAndScrollCellIntoView } from "../cells/scrollCellIntoView";
+import { store } from "../state/jotai";
+import { notebookAtom } from "../cells/cells";
+
+const getWordUnderCursor = (state: EditorState) => {
+  const { from, to } = state.selection.main;
+  let variableName: string;
+
+  if (from === to) {
+    const { startToken, endToken } = getPositionAtWordBounds(state.doc, from);
+    variableName = state.doc.sliceString(startToken, endToken);
+  } else {
+    variableName = state.doc.sliceString(from, to);
+  }
+
+  return variableName;
+};
+
+const getCellIdOfDefinition = (variables: Variables, variableName: string) => {
+  const variable = variables[variableName as VariableName];
+  if (!variable || variable.declaredBy.length === 0) {
+    return null;
+  }
+  const focusCellId = variable.declaredBy[0];
+  return focusCellId;
+};
+
+export function goToDefinition(view: EditorView, variables: Variables) {
+  const state = view.state;
+  const variableName = getWordUnderCursor(state);
+  const focusCellId = getCellIdOfDefinition(variables, variableName);
+
+  if (focusCellId) {
+    const notebookState = store.get(notebookAtom);
+    focusAndScrollCellIntoView({
+      cellId: focusCellId,
+      cell: notebookState.cellHandles[focusCellId],
+      config: notebookState.cellData[focusCellId].config,
+      codeFocus: undefined,
+      variableName,
+    });
+  }
+}
diff --git a/frontend/src/core/codemirror/language/utils.ts b/frontend/src/core/codemirror/language/utils.ts
index 30d84473007..14a36365b42 100644
--- a/frontend/src/core/codemirror/language/utils.ts
+++ b/frontend/src/core/codemirror/language/utils.ts
@@ -8,7 +8,7 @@ import { EditorState } from "@codemirror/state";
  * Handles when the editor has a different language adapter
  */
 export function getEditorCodeAsPython(
-  editor: { state: EditorState},
+  editor: { state: EditorState },
   fromPos?: number,
   toPos?: number,
 ): string {
diff --git a/frontend/src/core/hotkeys/__tests__/shortcuts.test.ts b/frontend/src/core/hotkeys/__tests__/shortcuts.test.ts
index bc46620fe7b..b8d33ff009a 100644
--- a/frontend/src/core/hotkeys/__tests__/shortcuts.test.ts
+++ b/frontend/src/core/hotkeys/__tests__/shortcuts.test.ts
@@ -33,6 +33,12 @@ describe("parseShortcut", () => {
     expect(shortcut(event)).toBe(true);
   });
 
+  it("should recognize cell.goToDefinition hotkey", () => {
+    const shortcut = parseShortcut("F12");
+    const event = new KeyboardEvent("keydown", { key: "F12" });
+    expect(shortcut(event)).toBe(true);
+  });
+
   it("should not recognize incorrect single key shortcuts", () => {
     const shortcut = parseShortcut("a");
     const event = new KeyboardEvent("keydown", { key: "b" });
diff --git a/frontend/src/core/hotkeys/hotkeys.ts b/frontend/src/core/hotkeys/hotkeys.ts
index b4449ce4891..dc5e14a6ab0 100644
--- a/frontend/src/core/hotkeys/hotkeys.ts
+++ b/frontend/src/core/hotkeys/hotkeys.ts
@@ -301,6 +301,11 @@ const DEFAULT_HOT_KEY = {
     group: "Navigation",
     key: "Mod-Shift-s",
   },
+  "cell.goToDefinition": {
+    name: "Go to Definition",
+    group: "Navigation",
+    key: "F12",
+  },
 } satisfies Record<string, Hotkey>;
 
 export type HotkeyAction = keyof typeof DEFAULT_HOT_KEY;


Code review comments:
diff_hunk: @@ -52,6 +52,7 @@ export function cellMovementBundle(
     toggleHideCode,
     aiCellCompletion,
   } = callbacks;
+  const variables = useVariables()
path: frontend/src/core/codemirror/cells/extensions.ts
text: @user1:
you can use `store.get(variables)` and can likely push this into `goToDefinition`

@author:
Will add, thanks for the tip