@marimo-team/islands 0.23.7-dev50 → 0.23.7-dev52

package/src/core/codemirror/go-to-definition/commands.ts

@@ -1,7 +1,31 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
 import { syntaxTree } from "@codemirror/language";
+import type { EditorState } from "@codemirror/state";
 import { EditorView } from "@codemirror/view";
+import type { SyntaxNode, Tree, TreeCursor } from "@lezer/common";
+
+const SCOPE_CREATING_NODES = new Set([
+  "FunctionDefinition",
+  "LambdaExpression",
+  "ArrayComprehensionExpression",
+  "SetComprehensionExpression",
+  "DictionaryComprehensionExpression",
+  "ComprehensionExpression",
+  "ClassDefinition",
+]);
+
+const POSITION_SENSITIVE_SCOPES = new Set(["ClassDefinition"]);
+
+interface ScopeContext {
+  id: number;
+  type: string;
+}
+
+interface VariableDeclaration {
+  from: number;
+  scopeId: number;
+}
 
 function goToPosition(view: EditorView, from: number): void {
   view.focus();
@@ -22,50 +46,386 @@ function goToPosition(view: EditorView, from: number): void {
   });
 }
 
-/**
- * This function will select the first occurrence of the given variable name,
- * for a given editor view.
- * @param view The editor view which contains the variable name.
- * @param variableName The name of the variable to select, if found in the editor.
- */
-export function goToVariableDefinition(
-  view: EditorView,
+function findFirstMatchingVariable(
+  state: EditorState,
   variableName: string,
-): boolean {
-  const { state } = view;
+): number | null {
   const tree = syntaxTree(state);
 
-  let found = false;
-  let from = 0;
+  let from: number | null = null;
 
   tree.iterate({
     enter: (node) => {
-      if (found) {
+      if (from !== null) {
         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
+        return false;
       }
 
-      // Skip comments and strings
       if (node.name === "Comment" || node.name === "String") {
         return false;
       }
+
+      return undefined;
     },
   });
 
-  if (found) {
-    goToPosition(view, from);
-    return true;
+  return from;
+}
+
+function getScopeChain(tree: Tree, usagePosition: number): ScopeContext[] {
+  const scopeChain: ScopeContext[] = [];
+  let currentNode: SyntaxNode | null = tree.resolveInner(usagePosition, 0);
+
+  while (currentNode) {
+    if (SCOPE_CREATING_NODES.has(currentNode.name)) {
+      // Skip ClassDefinition if we've already seen a function/lambda.
+      const inFunctionLikeScope = scopeChain.some(
+        (scope) =>
+          scope.type === "FunctionDefinition" ||
+          scope.type === "LambdaExpression",
+      );
+      if (!(inFunctionLikeScope && currentNode.name === "ClassDefinition")) {
+        scopeChain.push({
+          id: currentNode.from,
+          type: currentNode.name,
+        });
+      }
+    }
+    currentNode = currentNode.parent;
+  }
+
+  scopeChain.push({ id: -1, type: "global" });
+  return scopeChain;
+}
+
+function addDeclaration(
+  declarations: VariableDeclaration[],
+  scopeId: number,
+  from: number,
+) {
+  declarations.push({ scopeId, from });
+}
+
+function traverseChildren(
+  cursor: TreeCursor,
+  callback: (node: SyntaxNode) => void,
+) {
+  if (cursor.firstChild()) {
+    do {
+      callback(cursor.node);
+    } while (cursor.nextSibling());
+  }
+}
+
+function collectMatchingTargets(
+  cursor: TreeCursor,
+  state: EditorState,
+  variableName: string,
+  scopeId: number,
+  declarations: VariableDeclaration[],
+) {
+  switch (cursor.name) {
+    case "VariableName":
+      if (state.doc.sliceString(cursor.from, cursor.to) === variableName) {
+        addDeclaration(declarations, scopeId, cursor.from);
+      }
+      break;
+
+    case "TupleExpression":
+    case "ArrayExpression": {
+      const childCursor = cursor.node.cursor();
+      childCursor.firstChild();
+      do {
+        collectMatchingTargets(
+          childCursor,
+          state,
+          variableName,
+          scopeId,
+          declarations,
+        );
+      } while (childCursor.nextSibling());
+      break;
+    }
+    default:
+      break;
   }
-  return false;
+}
+
+function collectFunctionParameters(
+  node: SyntaxNode | Tree,
+  state: EditorState,
+  variableName: string,
+  scopeId: number,
+  declarations: VariableDeclaration[],
+) {
+  const cursor = node.cursor();
+  cursor.firstChild();
+  do {
+    if (cursor.name !== "ParamList") {
+      continue;
+    }
+
+    const paramCursor = cursor.node.cursor();
+    paramCursor.firstChild();
+    do {
+      if (
+        paramCursor.name === "VariableName" &&
+        state.doc.sliceString(paramCursor.from, paramCursor.to) === variableName
+      ) {
+        addDeclaration(declarations, scopeId, paramCursor.from);
+      }
+    } while (paramCursor.nextSibling());
+  } while (cursor.nextSibling());
+}
+
+function collectForTargets(
+  node: SyntaxNode | Tree,
+  state: EditorState,
+  variableName: string,
+  scopeId: number,
+  declarations: VariableDeclaration[],
+) {
+  const cursor = node.cursor();
+  cursor.firstChild();
+  let foundFor = false;
+  do {
+    if (cursor.name === "for") {
+      foundFor = true;
+    } else if (foundFor && cursor.name === "in") {
+      break;
+    } else if (foundFor) {
+      collectMatchingTargets(
+        cursor,
+        state,
+        variableName,
+        scopeId,
+        declarations,
+      );
+    }
+  } while (cursor.nextSibling());
+}
+
+function collectMatchingDeclarations(
+  node: SyntaxNode | Tree,
+  state: EditorState,
+  variableName: string,
+  scopeStack: number[],
+  declarations: VariableDeclaration[],
+) {
+  const cursor = node.cursor();
+  const nodeName = cursor.name;
+  const nodeStart = cursor.from;
+
+  const isNewScope = SCOPE_CREATING_NODES.has(nodeName);
+  const currentScopeStack = isNewScope
+    ? [...scopeStack, nodeStart]
+    : scopeStack;
+  const currentScope = currentScopeStack[currentScopeStack.length - 1] ?? -1;
+
+  switch (nodeName) {
+    case "FunctionDefinition":
+    case "ClassDefinition": {
+      const subCursor = node.cursor();
+      subCursor.firstChild();
+      do {
+        if (
+          subCursor.name === "VariableName" &&
+          state.doc.sliceString(subCursor.from, subCursor.to) === variableName
+        ) {
+          const parentScope = scopeStack[scopeStack.length - 1] ?? -1;
+          addDeclaration(declarations, parentScope, subCursor.from);
+          break;
+        }
+      } while (subCursor.nextSibling());
+
+      if (nodeName === "FunctionDefinition") {
+        collectFunctionParameters(
+          node,
+          state,
+          variableName,
+          nodeStart,
+          declarations,
+        );
+      }
+      break;
+    }
+    case "LambdaExpression":
+      collectFunctionParameters(
+        node,
+        state,
+        variableName,
+        nodeStart,
+        declarations,
+      );
+      break;
+
+    case "ArrayComprehensionExpression":
+    case "DictionaryComprehensionExpression":
+    case "SetComprehensionExpression":
+    case "ComprehensionExpression":
+    case "ForStatement":
+      collectForTargets(node, state, variableName, currentScope, declarations);
+      break;
+
+    case "AssignStatement": {
+      const assignOpPositions: number[] = [];
+      const subCursor = node.cursor();
+      subCursor.firstChild();
+      do {
+        if (subCursor.name === "AssignOp") {
+          assignOpPositions.push(subCursor.from);
+        }
+      } while (subCursor.nextSibling());
+
+      const lastAssignOpPosition =
+        assignOpPositions[assignOpPositions.length - 1];
+      if (lastAssignOpPosition === undefined) {
+        break;
+      }
+
+      const targetCursor = node.cursor();
+      targetCursor.firstChild();
+      do {
+        if (targetCursor.from < lastAssignOpPosition) {
+          collectMatchingTargets(
+            targetCursor,
+            state,
+            variableName,
+            currentScope,
+            declarations,
+          );
+        }
+      } while (targetCursor.nextSibling());
+      break;
+    }
+    case "ImportStatement": {
+      const subCursor = node.cursor();
+      subCursor.firstChild();
+      do {
+        if (
+          subCursor.name === "VariableName" &&
+          state.doc.sliceString(subCursor.from, subCursor.to) === variableName
+        ) {
+          addDeclaration(declarations, currentScope, subCursor.from);
+        }
+      } while (subCursor.nextSibling());
+      break;
+    }
+    case "ImportFromStatement": {
+      const subCursor = node.cursor();
+      subCursor.firstChild();
+      let foundImport = false;
+      do {
+        if (subCursor.name === "import") {
+          foundImport = true;
+        } else if (
+          foundImport &&
+          subCursor.name === "VariableName" &&
+          state.doc.sliceString(subCursor.from, subCursor.to) === variableName
+        ) {
+          addDeclaration(declarations, currentScope, subCursor.from);
+        }
+      } while (subCursor.nextSibling());
+      break;
+    }
+    case "TryStatement":
+    case "WithStatement": {
+      const subCursor = node.cursor();
+      subCursor.firstChild();
+      let foundAs = false;
+      do {
+        if (subCursor.name === "as") {
+          foundAs = true;
+        } else if (
+          foundAs &&
+          subCursor.name === "VariableName" &&
+          state.doc.sliceString(subCursor.from, subCursor.to) === variableName
+        ) {
+          addDeclaration(declarations, currentScope, subCursor.from);
+          foundAs = false;
+        }
+      } while (subCursor.nextSibling());
+      break;
+    }
+    default:
+      break;
+  }
+
+  traverseChildren(cursor, (childNode) => {
+    collectMatchingDeclarations(
+      childNode,
+      state,
+      variableName,
+      currentScopeStack,
+      declarations,
+    );
+  });
+}
+
+function findScopedDefinitionPosition(
+  state: EditorState,
+  variableName: string,
+  usagePosition: number,
+): number | null {
+  const tree = syntaxTree(state);
+  const declarations: VariableDeclaration[] = [];
+
+  collectMatchingDeclarations(tree, state, variableName, [], declarations);
+
+  const clampedUsagePosition = Math.max(
+    0,
+    Math.min(usagePosition, state.doc.length),
+  );
+
+  for (const scope of getScopeChain(tree, clampedUsagePosition)) {
+    const match = declarations
+      .filter((declaration) => declaration.scopeId === scope.id)
+      .filter((declaration) => {
+        return POSITION_SENSITIVE_SCOPES.has(scope.type)
+          ? declaration.from <= clampedUsagePosition
+          : true;
+      })
+      .toSorted((left, right) => left.from - right.from)[0];
+
+    if (match) {
+      return match.from;
+    }
+  }
+
+  return null;
+}
+
+/**
+ * This function will select the first occurrence of the given variable name,
+ * for a given editor view.
+ * @param view The editor view which contains the variable name.
+ * @param variableName The name of the variable to select, if found in the editor.
+ * @param usagePosition The position of the variable usage, if available.
+ */
+export function goToVariableDefinition(
+  view: EditorView,
+  variableName: string,
+  usagePosition?: number,
+): boolean {
+  const { state } = view;
+  const from =
+    (usagePosition !== undefined
+      ? findScopedDefinitionPosition(state, variableName, usagePosition)
+      : null) ?? findFirstMatchingVariable(state, variableName);
+
+  if (from === null) {
+    return false;
+  }
+
+  goToPosition(view, from);
+  return true;
 }
 
 /**

package/src/core/codemirror/go-to-definition/utils.ts

@@ -18,10 +18,16 @@ function getWordUnderCursor(state: EditorState) {
   const { from, to } = state.selection.main;
   if (from === to) {
     const { startToken, endToken } = getPositionAtWordBounds(state.doc, from);
-    return state.doc.sliceString(startToken, endToken);
+    return {
+      position: startToken,
+      word: state.doc.sliceString(startToken, endToken),
+    };
   }
 
-  return state.doc.sliceString(from, to);
+  return {
+    position: from,
+    word: state.doc.sliceString(from, to),
+  };
 }
 
 /**
@@ -51,15 +57,15 @@ function isPrivateVariable(variableName: string) {
  */
 export function goToDefinitionAtCursorPosition(view: EditorView): boolean {
   const { state } = view;
-  const variableName = getWordUnderCursor(state);
-  if (!variableName) {
+  const { position, word } = getWordUnderCursor(state);
+  if (!word) {
     return false;
   }
   // Close popovers/tooltips
   closeCompletion(view);
   view.dispatch({ effects: closeHoverTooltips });
 
-  return goToDefinition(view, variableName);
+  return goToDefinition(view, word, position);
 }
 
 /**
@@ -69,7 +75,19 @@ export function goToDefinitionAtCursorPosition(view: EditorView): boolean {
 export function goToDefinition(
   view: EditorView,
   variableName: string,
+  usagePosition?: number,
 ): boolean {
+  if (usagePosition !== undefined) {
+    const foundLocally = goToVariableDefinition(
+      view,
+      variableName,
+      usagePosition,
+    );
+    if (foundLocally) {
+      return true;
+    }
+  }
+
   // The variable may exist in another cell
   const editorWithVariable = getEditorForVariable(view, variableName);
   if (!editorWithVariable) {