Implements the Go to Definition keyboard shortcut (#2741)

Co-authored-by: Jonatan Kłosko <jonatanklosko@gmail.com>
Co-authored-by: José Valim <jose.valim@dashbit.co>

assets/js/hooks/cell.js

@@ -367,7 +367,7 @@ const Cell = {
 
     scrollIntoView(element, {
       scrollMode: "if-needed",
-      behavior: "smooth",
+      behavior: "instant",
       block: "center",
     });
   },

assets/js/hooks/cell_editor/live_editor.js

@@ -46,7 +46,7 @@ import { settingsStore } from "../../lib/settings";
 import Delta from "../../lib/delta";
 import Markdown from "../../lib/markdown";
 import { readOnlyHint } from "./live_editor/codemirror/read_only_hint";
-import { wait } from "../../lib/utils";
+import { isMacOS, wait } from "../../lib/utils";
 import Emitter from "../../lib/emitter";
 import CollabClient from "./live_editor/collab_client";
 import { languages } from "./live_editor/codemirror/languages";
@@ -363,14 +363,27 @@ export default class LiveEditor {
         settings.editor_mode === "emacs" ? [emacs()] : [],
         language ? language.support : [],
         EditorView.domEventHandlers({
+          click: this.handleEditorClick.bind(this),
           keydown: this.handleEditorKeydown.bind(this),
           blur: this.handleEditorBlur.bind(this),
           focus: this.handleEditorFocus.bind(this),
         }),
+        EditorView.clickAddsSelectionRange.of((event) => event.altKey),
       ],
     });
   }
 
+  /** @private */
+  handleEditorClick(event) {
+    const cmd = isMacOS() ? event.metaKey : event.ctrlKey;
+
+    if (cmd) {
+      this.jumpToDefinition(this.view);
+    }
+
+    return false;
+  }
+
   /** @private */
   handleEditorKeydown(event) {
     // We dispatch escape event, but only if it is not consumed by any
@@ -541,6 +554,29 @@ export default class LiveEditor {
       .catch(() => null);
   }
 
+  /** @private */
+  jumpToDefinition(view) {
+    const pos = view.state.selection.main.head;
+    const line = view.state.doc.lineAt(pos);
+    const lineLength = line.to - line.from;
+    const text = line.text;
+
+    const column = pos - line.from;
+    if (column < 1 || column > lineLength) return null;
+
+    return this.connection
+      .intellisenseRequest("definition", { line: text, column })
+      .then((response) => {
+        globalPubsub.broadcast("jump_to_editor", {
+          line: response.line,
+          file: response.file,
+        });
+
+        return true;
+      })
+      .catch(() => false);
+  }
+
   /** @private */
   signatureSource({ state, pos }) {
     const textUntilCursor = this.getSignatureHint(state, pos);

assets/js/hooks/session.js

@@ -157,6 +157,12 @@ const Session = {
       (event) => this.toggleCollapseAllSections(),
     );
 
+    this.subscriptions = [
+      globalPubsub.subscribe("jump_to_editor", ({ line, file }) =>
+        this.jumpToLine(file, line),
+      ),
+    ];
+
     this.initializeDragAndDrop();
 
     window.addEventListener(
@@ -270,6 +276,7 @@ const Session = {
       leaveChannel();
     }
 
+    this.subscriptions.forEach((subscription) => subscription.destroy());
     this.store.destroy();
   },
 
@@ -546,12 +553,9 @@ const Session = {
       const search = event.target.hash.replace("#go-to-definition", "");
       const params = new URLSearchParams(search);
       const line = parseInt(params.get("line"), 10);
-      const [_filename, cellId] = params.get("file").split("#cell:");
+      const file = params.get("file");
 
-      this.setFocusedEl(cellId);
-      this.setInsertMode(true);
-
-      globalPubsub.broadcast(`cells:${cellId}`, { type: "jump_to_line", line });
+      this.jumpToLine(file, line);
 
       event.preventDefault();
     }
@@ -1442,6 +1446,15 @@ const Session = {
   getElement(name) {
     return this.el.querySelector(`[data-el-${name}]`);
   },
+
+  jumpToLine(file, line) {
+    const [_filename, cellId] = file.split("#cell:");
+
+    this.setFocusedEl(cellId, { scroll: false });
+    this.setInsertMode(true);
+
+    globalPubsub.broadcast(`cells:${cellId}`, { type: "jump_to_line", line });
+  },
 };
 
 export default Session;

lib/livebook/intellisense.ex

@@ -53,6 +53,10 @@ defmodule Livebook.Intellisense do
     get_details(line, column, context, node)
   end
 
+  def handle_request({:definition, line, column}, context, node) do
+    get_definitions(line, column, context, node)
+  end
+
   def handle_request({:signature, hint}, context, node) do
     get_signature_items(hint, context, node)
   end
@@ -452,7 +456,7 @@ defmodule Livebook.Intellisense do
        ) do
     join_with_divider([
       code(inspect(module)),
-      format_definition_link(module, context),
+      format_definition_link(module, context, {:module, module}),
       format_docs_link(module),
       format_documentation(documentation, :all)
     ])
@@ -539,23 +543,9 @@ defmodule Livebook.Intellisense do
     """
   end
 
-  defp format_definition_link(module, context, function_or_type \\ nil) do
-    if context.ebin_path do
-      path = Path.join(context.ebin_path, "#{module}.beam")
-
-      identifier =
-        if function_or_type,
-          do: function_or_type,
-          else: {:module, module}
-
-      with true <- File.exists?(path),
-           {:ok, line} <- Docs.locate_definition(path, identifier) do
-        file = module.module_info(:compile)[:source]
-        query_string = URI.encode_query(%{file: to_string(file), line: line})
-        "[Go to definition](#go-to-definition?#{query_string})"
-      else
-        _otherwise -> nil
-      end
+  defp format_definition_link(module, context, identifier) do
+    if query = get_definition_location(module, context, identifier) do
+      "[Go to definition](#go-to-definition?#{URI.encode_query(query)})"
     end
   end
 
@@ -710,6 +700,56 @@ defmodule Livebook.Intellisense do
     raise "unknown documentation format #{inspect(format)}"
   end
 
+  @doc """
+  Returns the identifier definition located in `column` in `line`.
+  """
+  @spec get_definitions(String.t(), pos_integer(), context(), node()) ::
+          Runtime.definition_response() | nil
+  def get_definitions(line, column, context, node) do
+    case IdentifierMatcher.locate_identifier(line, column, context, node) do
+      %{matches: []} ->
+        nil
+
+      %{matches: matches, range: range} ->
+        matches
+        |> Enum.sort_by(& &1[:arity], :asc)
+        |> Enum.flat_map(&List.wrap(get_definition_location(&1, context)))
+        |> case do
+          [%{file: file, line: line} | _] -> %{range: range, file: file, line: line}
+          _ -> nil
+        end
+    end
+  end
+
+  defp get_definition_location(%{kind: :module, module: module}, context) do
+    get_definition_location(module, context, {:module, module})
+  end
+
+  defp get_definition_location(
+         %{kind: :function, module: module, name: name, arity: arity},
+         context
+       ) do
+    get_definition_location(module, context, {:function, name, arity})
+  end
+
+  defp get_definition_location(%{kind: :type, module: module, name: name, arity: arity}, context) do
+    get_definition_location(module, context, {:type, name, arity})
+  end
+
+  defp get_definition_location(module, context, identifier) do
+    if context.ebin_path do
+      path = Path.join(context.ebin_path, "#{module}.beam")
+
+      with true <- File.exists?(path),
+           {:ok, line} <- Docs.locate_definition(path, identifier) do
+        file = module.module_info(:compile)[:source]
+        %{file: to_string(file), line: line}
+      else
+        _otherwise -> nil
+      end
+    end
+  end
+
   # Erlang HTML AST
   # See https://erlang.org/doc/apps/erl_docgen/doc_storage.html#erlang-documentation-format
 

lib/livebook/intellisense/identifier_matcher.ex

@@ -627,19 +627,35 @@ defmodule Livebook.Intellisense.IdentifierMatcher do
   end
 
   defp get_matching_modules(hint, ctx) do
-    ctx.node
+    ctx
     |> get_modules()
     |> Enum.filter(&ctx.matcher.(Atom.to_string(&1), hint))
     |> Enum.uniq()
   end
 
-  defp get_modules(node) do
-    modules = cached_all_loaded(node)
-
+  defp get_modules(%{node: node} = ctx) do
+    # On interactive mode, we load modules from the application
+    # and then the ones from runtime. For a remote node, ideally
+    # we would get the applications one, but there is no cheap
+    # way to do such, so we get :code.all_loaded and cache it
+    # instead (which includes all modules anyway on embedded mode).
     if node == node() and :code.get_mode() == :interactive do
-      modules ++ get_modules_from_applications()
+      runtime_modules(ctx.intellisense_context.ebin_path) ++ get_modules_from_applications()
     else
-      modules
+      cached_all_loaded(node)
+    end
+  end
+
+  defp runtime_modules(path) do
+    with true <- is_binary(path),
+         {:ok, beams} <- File.ls(path) do
+      for beam <- beams, String.ends_with?(beam, ".beam") do
+        beam
+        |> binary_slice(0..-6//1)
+        |> String.to_atom()
+      end
+    else
+      _ -> []
     end
   end
 

lib/livebook/runtime.ex

@@ -503,6 +503,7 @@ defprotocol Livebook.Runtime do
   @type intellisense_request ::
           completion_request()
           | details_request()
+          | definition_request()
           | signature_request()
           | format_request()
 
@@ -516,6 +517,7 @@ defprotocol Livebook.Runtime do
           nil
           | completion_response()
           | details_response()
+          | definition_response()
           | signature_response()
           | format_response()
 
@@ -553,6 +555,21 @@ defprotocol Livebook.Runtime do
           contents: list(String.t())
         }
 
+  @typedoc """
+  Looks up more the definition about an identifier found in `column` in
+  `line`.
+  """
+  @type definition_request :: {:definition, line :: String.t(), column :: pos_integer()}
+
+  @type definition_response :: %{
+          range: %{
+            from: non_neg_integer(),
+            to: non_neg_integer()
+          },
+          line: pos_integer(),
+          file: String.t()
+        }
+
   @typedoc """
   Looks up a list of function signatures matching the given hint.
 

lib/livebook_web/live/session_live.ex

@@ -616,6 +616,10 @@ defmodule LivebookWeb.SessionLive do
           column = Text.JS.js_column_to_elixir(column, line)
           {:details, line, column}
 
+        %{"type" => "definition", "line" => line, "column" => column} ->
+          column = Text.JS.js_column_to_elixir(column, line)
+          {:definition, line, column}
+
         %{"type" => "signature", "hint" => hint} ->
           {:signature, hint}
 

test/livebook/intellisense_test.exs

@@ -251,16 +251,6 @@ defmodule Livebook.IntellisenseTest do
              ] = Intellisense.get_completion_items("RuntimeE", context, node())
     end
 
-    test "caches all loaded modules" do
-      context = eval(do: nil)
-      Intellisense.get_completion_items("Hub", context, node())
-
-      key = {Intellisense.IdentifierMatcher, node()}
-      assert [_ | _] = :persistent_term.get(key, :error)
-      Intellisense.IdentifierMatcher.clear_all_loaded(node())
-      assert :error = :persistent_term.get(key, :error)
-    end
-
     test "Elixir struct completion lists nested options" do
       context = eval(do: nil)
 
@@ -999,11 +989,16 @@ defmodule Livebook.IntellisenseTest do
              ] = Intellisense.get_completion_items("^my_va", context, node())
     end
 
-    defmodule SublevelTest.LevelA.LevelB do
-    end
+    @tag :tmp_dir
+    test "Elixir completion sublevel", %{tmp_dir: tmp_dir} do
+      context =
+        eval tmp_dir do
+        end
 
-    test "Elixir completion sublevel" do
-      context = eval(do: nil)
+      compile_and_save_bytecode(tmp_dir, ~S'''
+      defmodule Livebook.IntellisenseTest.SublevelTest.LevelA.LevelB do
+      end
+      ''')
 
       assert [%{label: "LevelA"}] =
                Intellisense.get_completion_items(
@@ -1100,12 +1095,17 @@ defmodule Livebook.IntellisenseTest do
       :code.delete(Sample)
     end
 
-    defmodule MyStruct do
-      defstruct [:my_val]
-    end
+    @tag :tmp_dir
+    test "completion for struct names", %{tmp_dir: tmp_dir} do
+      context =
+        eval tmp_dir do
+        end
 
-    test "completion for struct names" do
-      context = eval(do: nil)
+      compile_and_save_bytecode(tmp_dir, ~S'''
+      defmodule Livebook.IntellisenseTest.MyStruct do
+        defstruct [:my_val]
+      end
+      ''')
 
       assert [
                %{label: "MyStruct"}
@@ -1117,9 +1117,16 @@ defmodule Livebook.IntellisenseTest do
                )
     end
 
-    test "completion for struct keys" do
+    @tag :tmp_dir
+    test "completion for struct keys", %{tmp_dir: tmp_dir} do
+      compile_and_save_bytecode(tmp_dir, ~S'''
+      defmodule Livebook.IntellisenseTest.MyStruct do
+        defstruct [:my_val]
+      end
+      ''')
+
       context =
-        eval do
+        eval tmp_dir do
           struct = %Livebook.IntellisenseTest.MyStruct{}
         end
 
@@ -1128,8 +1135,17 @@ defmodule Livebook.IntellisenseTest do
              ] = Intellisense.get_completion_items("struct.my", context, node())
     end
 
-    test "completion for struct keys inside struct" do
-      context = eval(do: nil)
+    @tag :tmp_dir
+    test "completion for struct keys inside struct", %{tmp_dir: tmp_dir} do
+      context =
+        eval tmp_dir do
+        end
+
+      compile_and_save_bytecode(tmp_dir, ~S'''
+      defmodule Livebook.IntellisenseTest.MyStruct do
+        defstruct [:my_val]
+      end
+      ''')
 
       assert [
                %{
@@ -1154,8 +1170,18 @@ defmodule Livebook.IntellisenseTest do
                )
     end
 
-    test "completion for struct keys inside struct removes filled keys" do
-      context = eval(do: nil)
+    @tag :tmp_dir
+    test "completion for struct keys inside struct removes filled keys",
+         %{tmp_dir: tmp_dir} do
+      context =
+        eval tmp_dir do
+        end
+
+      compile_and_save_bytecode(tmp_dir, ~S'''
+      defmodule Livebook.IntellisenseTest.MyStruct do
+        defstruct [:my_val]
+      end
+      ''')
 
       assert [] =
                Intellisense.get_completion_items(
@@ -1173,8 +1199,17 @@ defmodule Livebook.IntellisenseTest do
       refute Enum.find(completions, &match?(%{label: "__exception__"}, &1))
     end
 
-    test "completion for struct keys in update syntax" do
-      context = eval(do: nil)
+    @tag :tmp_dir
+    test "completion for struct keys in update syntax", %{tmp_dir: tmp_dir} do
+      context =
+        eval tmp_dir do
+        end
+
+      compile_and_save_bytecode(tmp_dir, ~S'''
+      defmodule Livebook.IntellisenseTest.MyStruct do
+        defstruct [:my_val]
+      end
+      ''')
 
       assert [
                %{
@@ -1581,77 +1616,75 @@ defmodule Livebook.IntellisenseTest do
 
       assert content =~ ~r"https://www.erlang.org/doc/man/string.html#uppercase-1"
     end
+  end
 
-    @tag :tmp_dir
-    test "returns the go to definition link", %{tmp_dir: tmp_dir} do
-      Code.put_compiler_option(:debug_info, true)
+  @tag :tmp_dir
+  test "get_definitions/4 returns the go to definition query string", %{tmp_dir: tmp_dir} do
+    Code.put_compiler_option(:debug_info, true)
 
-      context =
-        eval tmp_dir do
-          alias Livebook.IntellisenseTest.GoToDefinition
-        end
+    context =
+      eval tmp_dir do
+        alias Livebook.IntellisenseTest.GoToDefinition
+      end
 
-      code = ~S'''
-      defmodule Livebook.IntellisenseTest.GoToDefinition do
-        @type t :: term()
-        @type foo :: foo(:bar)
-        @type foo(var) :: {var, t()}
+    code = ~S'''
+    defmodule Livebook.IntellisenseTest.GoToDefinition do
+      @type t :: term()
+      @type foo :: foo(:bar)
+      @type foo(var) :: {var, t()}
 
-        defmacro with_logging(do: block) do
-          quote do
-            require Logger
-            Logger.debug("Running code")
-            result = unquote(block)
-            Logger.debug("Result: #{inspect(result)}")
-            result
-          end
-        end
-
-        @spec hello(var :: term()) :: foo(term())
-        def hello(message) do
-          {:bar, message}
+      defmacro with_logging(do: block) do
+        quote do
+          require Logger
+          Logger.debug("Running code")
+          result = unquote(block)
+          Logger.debug("Result: #{inspect(result)}")
+          result
         end
       end
-      '''
 
-      file = "#{__ENV__.file}#cell:#{Livebook.Utils.random_short_id()}"
-      path = Path.join(tmp_dir, "Elixir.Livebook.IntellisenseTest.GoToDefinition.beam")
-
-      [{_module, bytecode}] = Code.compile_string(code, file)
-      File.write!(path, bytecode)
-      Code.prepend_path(tmp_dir)
-
-      assert %{contents: [content]} =
-               Intellisense.get_details("GoToDefinition", 14, context, node())
-
-      assert content =~ "#go-to-definition?#{URI.encode_query(%{file: file, line: 1})}"
-
-      assert %{contents: [content]} =
-               Intellisense.get_details("GoToDefinition.t", 16, context, node())
-
-      assert content =~ "#go-to-definition?#{URI.encode_query(%{file: file, line: 2})}"
-
-      assert %{contents: [arity_0_content, arity_1_content]} =
-               Intellisense.get_details("GoToDefinition.foo", 18, context, node())
-
-      assert arity_0_content =~ "#go-to-definition?#{URI.encode_query(%{file: file, line: 3})}"
-      assert arity_1_content =~ "#go-to-definition?#{URI.encode_query(%{file: file, line: 4})}"
-
-      assert %{contents: [content]} =
-               Intellisense.get_details("GoToDefinition.with_logging", 20, context, node())
-
-      assert content =~ "#go-to-definition?#{URI.encode_query(%{file: file, line: 6})}"
-
-      assert %{contents: [content]} =
-               Intellisense.get_details("GoToDefinition.hello", 18, context, node())
-
-      assert content =~ "#go-to-definition?#{URI.encode_query(%{file: file, line: 17})}"
-    after
-      Code.put_compiler_option(:debug_info, false)
-      Code.delete_path(tmp_dir)
-      :code.purge(:"Elixir.Livebook.IntellisenseTest.GoToDefinition")
-      :code.delete(:"Elixir.Livebook.IntellisenseTest.GoToDefinition")
+      @spec hello(var :: term()) :: foo(term())
+      def hello(message) do
+        {:bar, message}
+      end
     end
+    '''
+
+    file = "#{__ENV__.file}#cell:#{Livebook.Utils.random_short_id()}"
+    compile_and_save_bytecode(tmp_dir, code, file)
+
+    assert Intellisense.get_definitions("GoToDefinition", 14, context, node()) == %{
+             line: 1,
+             file: file,
+             range: %{to: 15, from: 1}
+           }
+
+    assert Intellisense.get_definitions("GoToDefinition.t", 16, context, node()) == %{
+             line: 2,
+             file: file,
+             range: %{to: 17, from: 1}
+           }
+
+    # For now, we aren't fetching the expected arity but we will address it later.
+    assert Intellisense.get_definitions("GoToDefinition.foo", 18, context, node()) == %{
+             line: 3,
+             file: file,
+             range: %{to: 19, from: 1}
+           }
+
+    assert Intellisense.get_definitions("GoToDefinition.with_logging", 20, context, node()) == %{
+             line: 6,
+             file: file,
+             range: %{to: 28, from: 1}
+           }
+
+    assert Intellisense.get_definitions("GoToDefinition.hello", 18, context, node()) == %{
+             line: 17,
+             file: file,
+             range: %{to: 21, from: 1}
+           }
+  after
+    Code.put_compiler_option(:debug_info, false)
   end
 
   describe "get_signature_items/3" do
@@ -2032,4 +2065,18 @@ defmodule Livebook.IntellisenseTest do
       assert content =~ "No documentation available"
     end
   end
+
+  defp compile_and_save_bytecode(dir, code, file \\ "nofile") do
+    [{module, bytecode}] = Code.compile_string(code, file)
+    path = Path.join(dir, "#{module}.beam")
+
+    File.write!(path, bytecode)
+    Code.prepend_path(dir)
+
+    on_exit(fn ->
+      Code.delete_path(dir)
+      :code.purge(module)
+      :code.delete(module)
+    end)
+  end
 end

test/test_helper.exs

@@ -6,6 +6,9 @@ Livebook.Runtime.ErlDist.NodeManager.start(
   unload_modules_on_termination: false
 )
 
+# We load erts so we can access its modules for completion.
+Application.load(:erts)
+
 # Use the embedded runtime in tests by default, so they are cheaper
 # to run. Other runtimes can be tested by setting them explicitly
 Application.put_env(:livebook, :default_runtime, Livebook.Runtime.Embedded.new())