Provide go to definition of modules and functions defined in the notebook (#2730)

Co-authored-by: Jonatan Kłosko <jonatanklosko@gmail.com>

assets/js/hooks/cell.js

@@ -159,6 +159,8 @@ const Cell = {
   handleCellEvent(event) {
     if (event.type === "dispatch_queue_evaluation") {
       this.handleDispatchQueueEvaluation(event.dispatch);
+    } else if (event.type === "jump_to_line") {
+      this.handleJumpToLine(event.line);
     }
   },
 
@@ -175,6 +177,12 @@ const Cell = {
     }
   },
 
+  handleJumpToLine(line) {
+    if (this.isFocused) {
+      this.currentEditor().moveCursorToLine(line);
+    }
+  },
+
   handleCellEditorCreated(tag, liveEditor) {
     this.liveEditors[tag] = liveEditor;
 

assets/js/hooks/cell_editor/live_editor.js

@@ -11,7 +11,7 @@ import {
   lineNumbers,
   highlightActiveLineGutter,
 } from "@codemirror/view";
-import { EditorState } from "@codemirror/state";
+import { EditorState, EditorSelection } from "@codemirror/state";
 import {
   indentOnInput,
   bracketMatching,
@@ -56,6 +56,7 @@ import {
 } from "./live_editor/codemirror/commands";
 import { ancestorNode, closestNode } from "./live_editor/codemirror/tree_utils";
 import { selectingClass } from "./live_editor/codemirror/selecting_class";
+import { globalPubsub } from "../../lib/pubsub";
 
 /**
  * Mounts cell source editor with real-time collaboration mechanism.
@@ -179,6 +180,17 @@ export default class LiveEditor {
     this.view.focus();
   }
 
+  /**
+   * Updates editor selection such that cursor points to the given line.
+   */
+  moveCursorToLine(lineNumber) {
+    const line = this.view.state.doc.line(lineNumber);
+
+    this.view.dispatch({
+      selection: EditorSelection.single(line.from),
+    });
+  }
+
   /**
    * Removes focus from the editor.
    */
@@ -515,6 +527,7 @@ export default class LiveEditor {
               item.classList.add("cm-hoverDocsContent");
               item.classList.add("cm-markdown");
               dom.appendChild(item);
+
               new Markdown(item, content, {
                 defaultCodeLanguage: this.language,
                 useDarkTheme: this.usesDarkTheme(),

assets/js/hooks/session.js

@@ -539,6 +539,23 @@ const Session = {
       this.setInsertMode(false);
     }
 
+    if (
+      event.target.matches("a") &&
+      event.target.hash.startsWith("#go-to-definition")
+    ) {
+      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:");
+
+      this.setFocusedEl(cellId);
+      this.setInsertMode(true);
+
+      globalPubsub.broadcast(`cells:${cellId}`, { type: "jump_to_line", line });
+
+      event.preventDefault();
+    }
+
     const evalButton = event.target.closest(
       `[data-el-queue-cell-evaluation-button]`,
     );

lib/livebook/intellisense.ex

@@ -20,6 +20,7 @@ defmodule Livebook.Intellisense do
   """
   @type context :: %{
           env: Macro.Env.t(),
+          ebin_path: String.t() | nil,
           map_binding: (Code.binding() -> any())
         }
 
@@ -413,7 +414,11 @@ defmodule Livebook.Intellisense do
         nil
 
       matches ->
-        contents = Enum.map(matches, &format_details_item/1)
+        contents =
+          matches
+          |> Enum.sort_by(& &1[:arity], :asc)
+          |> Enum.map(&format_details_item(&1, context))
+
         %{range: range, contents: contents}
     end
   end
@@ -422,13 +427,13 @@ defmodule Livebook.Intellisense do
   defp include_in_details?(%{kind: :bitstring_modifier}), do: false
   defp include_in_details?(_), do: true
 
-  defp format_details_item(%{kind: :variable, name: name}), do: code(name)
+  defp format_details_item(%{kind: :variable, name: name}, _context), do: code(name)
 
-  defp format_details_item(%{kind: :map_field, name: name}), do: code(name)
+  defp format_details_item(%{kind: :map_field, name: name}, _context), do: code(name)
 
-  defp format_details_item(%{kind: :in_map_field, name: name}), do: code(name)
+  defp format_details_item(%{kind: :in_map_field, name: name}, _context), do: code(name)
 
-  defp format_details_item(%{kind: :in_struct_field, name: name, default: default}) do
+  defp format_details_item(%{kind: :in_struct_field, name: name, default: default}, _context) do
     join_with_divider([
       code(name),
       """
@@ -441,27 +446,35 @@ defmodule Livebook.Intellisense do
     ])
   end
 
-  defp format_details_item(%{kind: :module, module: module, documentation: documentation}) do
+  defp format_details_item(
+         %{kind: :module, module: module, documentation: documentation},
+         context
+       ) do
     join_with_divider([
       code(inspect(module)),
+      format_definition_link(module, context),
       format_docs_link(module),
       format_documentation(documentation, :all)
     ])
   end
 
-  defp format_details_item(%{
+  defp format_details_item(
-         kind: :function,
+         %{
-         module: module,
+           kind: :function,
-         name: name,
+           module: module,
-         arity: arity,
+           name: name,
-         documentation: documentation,
+           arity: arity,
-         signatures: signatures,
+           documentation: documentation,
-         specs: specs,
+           signatures: signatures,
-         meta: meta
+           specs: specs,
-       }) do
+           meta: meta
+         },
+         context
+       ) do
     join_with_divider([
       format_signatures(signatures, module) |> code(),
       join_with_middle_dot([
+        format_definition_link(module, context, {:function, name, arity}),
         format_docs_link(module, {:function, name, arity}),
         format_meta(:since, meta)
       ]),
@@ -471,23 +484,30 @@ defmodule Livebook.Intellisense do
     ])
   end
 
-  defp format_details_item(%{
+  defp format_details_item(
-         kind: :type,
+         %{
-         module: module,
+           kind: :type,
-         name: name,
+           module: module,
-         arity: arity,
+           name: name,
-         documentation: documentation,
+           arity: arity,
-         type_spec: type_spec
+           documentation: documentation,
-       }) do
+           type_spec: type_spec
+         },
+         context
+       ) do
     join_with_divider([
       format_type_signature(type_spec, module) |> code(),
+      format_definition_link(module, context, {:type, name, arity}),
       format_docs_link(module, {:type, name, arity}),
       format_type_spec(type_spec, @extended_line_length) |> code(),
       format_documentation(documentation, :all)
     ])
   end
 
-  defp format_details_item(%{kind: :module_attribute, name: name, documentation: documentation}) do
+  defp format_details_item(
+         %{kind: :module_attribute, name: name, documentation: documentation},
+         _context
+       ) do
     join_with_divider([
       code("@#{name}"),
       format_documentation(documentation, :all)
@@ -519,14 +539,29 @@ 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
+    end
+  end
+
   defp format_docs_link(module, function_or_type \\ nil) do
     app = Application.get_application(module)
-
+    module_name = module_name(module)
-    module_name =
-      case Atom.to_string(module) do
-        "Elixir." <> name -> name
-        name -> name
-      end
 
     is_otp? =
       case :code.which(module) do
@@ -868,4 +903,11 @@ defmodule Livebook.Intellisense do
        ) do
     group_type_list_items(items, [{:li, [], prev_content ++ [{:p, [], content}]} | acc])
   end
+
+  defp module_name(module) do
+    case Atom.to_string(module) do
+      "Elixir." <> name -> name
+      name -> name
+    end
+  end
 end

lib/livebook/intellisense/docs.ex

@@ -40,6 +40,9 @@ defmodule Livebook.Intellisense.Docs do
   @type type_spec() :: {type_kind(), term()}
   @type type_kind() :: :type | :opaque
 
+  @type definition ::
+          {:module, module()} | {:function | :type, name :: atom(), arity :: pos_integer()}
+
   @doc """
   Fetches documentation for the given module if available.
   """
@@ -174,4 +177,57 @@ defmodule Livebook.Intellisense.Docs do
   # so we explicitly list it.
   defp ensure_loaded?(Elixir), do: false
   defp ensure_loaded?(module), do: Code.ensure_loaded?(module)
+
+  @doc """
+  Extracts the location about an identifier found.
+
+  The function returns the line where the identifier is located.
+  """
+  @spec locate_definition(String.t(), definition()) :: {:ok, pos_integer()} | :error
+  def locate_definition(path, identifier)
+
+  def locate_definition(path, {:module, module}) do
+    with {:ok, {:raw_abstract_v1, annotations}} <- beam_lib_chunks(path, :abstract_code) do
+      {:attribute, anno, :module, ^module} =
+        Enum.find(annotations, &match?({:attribute, _, :module, _}, &1))
+
+      {:ok, :erl_anno.line(anno)}
+    end
+  end
+
+  def locate_definition(path, {:function, name, arity}) do
+    with {:ok, {:debug_info_v1, _, {:elixir_v1, meta, _}}} <- beam_lib_chunks(path, :debug_info),
+         {_pair, _kind, kw, _body} <- keyfind(meta.definitions, {name, arity}) do
+      Keyword.fetch(kw, :line)
+    end
+  end
+
+  def locate_definition(path, {:type, name, arity}) do
+    with {:ok, {:raw_abstract_v1, annotations}} <- beam_lib_chunks(path, :abstract_code) do
+      fetch_type_line(annotations, name, arity)
+    end
+  end
+
+  defp fetch_type_line(annotations, name, arity) do
+    for {:attribute, anno, :type, {^name, _, vars}} <- annotations, length(vars) == arity do
+      :erl_anno.line(anno)
+    end
+    |> case do
+      [] -> :error
+      lines -> {:ok, Enum.min(lines)}
+    end
+  end
+
+  defp beam_lib_chunks(path, key) do
+    path = String.to_charlist(path)
+
+    case :beam_lib.chunks(path, [key]) do
+      {:ok, {_, [{^key, value}]}} -> {:ok, value}
+      _ -> :error
+    end
+  end
+
+  defp keyfind(list, key) do
+    List.keyfind(list, key, 0) || :error
+  end
 end

lib/livebook/intellisense/signature_matcher.ex

@@ -13,7 +13,7 @@ defmodule Livebook.Intellisense.SignatureMatcher do
   Evaluation binding and environment is used to expand aliases,
   imports, access variable values, etc.
   """
-  @spec get_matching_signatures(String.t(), Livebook.Intellisense.intellisense_context(), node()) ::
+  @spec get_matching_signatures(String.t(), Livebook.Intellisense.context(), node()) ::
           {:ok, list(signature_info()), active_argument :: non_neg_integer()} | :error
   def get_matching_signatures(hint, intellisense_context, node) do
     %{env: env} = intellisense_context

lib/livebook/runtime/evaluator.ex

@@ -183,17 +183,17 @@ defmodule Livebook.Runtime.Evaluator do
   @doc """
   Returns an empty intellisense context.
   """
-  @spec intellisense_context() :: Livebook.Intellisense.intellisense_context()
+  @spec intellisense_context() :: Livebook.Intellisense.context()
   def intellisense_context() do
     env = Code.env_for_eval([])
     map_binding = fn fun -> fun.([]) end
-    %{env: env, map_binding: map_binding}
+    %{env: env, ebin_path: ebin_path(), map_binding: map_binding}
   end
 
   @doc """
   Builds intellisense context from the given evaluation.
   """
-  @spec intellisense_context(t(), list(ref())) :: Livebook.Intellisense.intellisense_context()
+  @spec intellisense_context(t(), list(ref())) :: Livebook.Intellisense.context()
   def intellisense_context(evaluator, parent_refs) do
     {:dictionary, dictionary} = Process.info(evaluator.pid, :dictionary)
 
@@ -210,7 +210,11 @@ defmodule Livebook.Runtime.Evaluator do
 
     map_binding = fn fun -> map_binding(evaluator, parent_refs, fun) end
 
-    %{env: env, map_binding: map_binding}
+    %{
+      env: env,
+      ebin_path: find_in_dictionary(dictionary, @ebin_path_key),
+      map_binding: map_binding
+    }
   end
 
   defp find_in_dictionary(dictionary, key) do

lib/livebook/session.ex

@@ -129,7 +129,7 @@ defmodule Livebook.Session do
           app_pid: pid() | nil,
           auto_shutdown_ms: non_neg_integer() | nil,
           auto_shutdown_timer_ref: reference() | nil,
-          started_by: Livebook.User.t() | nil
+          started_by: Livebook.Users.User.t() | nil
         }
 
   @type memory_usage ::
@@ -201,7 +201,7 @@ defmodule Livebook.Session do
   @doc """
   Fetches session information from the session server.
   """
-  @spec get_by_pid(pid()) :: Session.t()
+  @spec get_by_pid(pid()) :: t()
   def get_by_pid(pid) do
     GenServer.call(pid, :describe_self, @timeout)
   end

test/livebook/intellisense_test.exs

@@ -5,7 +5,7 @@ defmodule Livebook.IntellisenseTest do
 
   # Returns intellisense context resulting from evaluating
   # the given block of code in a fresh context.
-  defmacrop eval(do: block) do
+  defmacrop eval(ebin_path \\ System.tmp_dir!(), do: block) do
     quote do
       block = unquote(Macro.escape(block))
       binding = []
@@ -14,6 +14,7 @@ defmodule Livebook.IntellisenseTest do
 
       %{
         env: env,
+        ebin_path: unquote(ebin_path),
         map_binding: fn fun -> fun.(binding) end
       }
     end
@@ -1580,6 +1581,72 @@ defmodule Livebook.IntellisenseTest do
 
       assert content =~ ~r"https://www.erlang.org/doc/man/string.html#uppercase-1"
     end
+
+    test "returns the go to definition link" do
+      Code.put_compiler_option(:debug_info, true)
+
+      [runtime_path | _] = :code.get_path()
+      runtime_path = to_string(runtime_path)
+
+      context = eval(runtime_path, do: nil)
+
+      code = ~S'''
+      defmodule 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}
+        end
+      end
+      '''
+
+      file = "#{__ENV__.file}#cell:#{Livebook.Utils.random_short_id()}"
+      path = Path.join(runtime_path, "Elixir.GoToDefinition.beam")
+
+      [{_module, bytecode}] = Code.compile_string(code, file)
+      File.write!(path, bytecode)
+
+      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)
+    end
   end
 
   describe "get_signature_items/3" do
@@ -1891,18 +1958,17 @@ defmodule Livebook.IntellisenseTest do
             '''
 
           Livebook.Runtime.evaluate_code(runtime, :elixir, code, {:c1, :e1}, [])
-
           receive do: ({:runtime_evaluation_response, :e1, _, _} -> :ok)
           send(parent, :continue)
 
-          Process.sleep(:infinity)
+          receive do: (:done -> :ok)
         end
       })
 
     receive do: (:continue -> :ok)
 
     on_exit(fn ->
-      Process.exit(runtime_owner_pid, :kill)
+      send(runtime_owner_pid, :done)
     end)
 
     [node: runtime.node]