Context-aware "use" completions (#336)

* Conext-aware "use" completions

When using `use`, we would complete any module, even those without
`__using__` macros defined.

This PR filters any suggestions inside a `use` directive to ensure
that any suggested modules (or a submodule) contain a __using__ macro
of arity 1.

One thing that cropped up is that not all modules are loaded when
queried, so I had to use `:code.all_available` to generate the module
list. This wasn't super duper fast, so I had to put a cache in front
of it, which is the large part of this PR.

Another oddity is that I couldn't just send anonymous functions as
predicates, so I had to come up with a way to implement filtering
logic across distribution, so I did a matchspec-like thingy, which was
a bit ugly, so I made some syntax to make it look like function captures.

apps/remote_control/lib/lexical/remote_control/api.ex

@@ -60,4 +60,14 @@ defmodule Lexical.RemoteControl.Api do
       position
     ])
   end
+
+  def modules_with_prefix(%Project{} = project, prefix)
+      when is_binary(prefix) or is_atom(prefix) do
+    RemoteControl.call(project, RemoteControl.Modules, :with_prefix, [prefix])
+  end
+
+  def modules_with_prefix(%Project{} = project, prefix, predicate)
+      when is_binary(prefix) or is_atom(prefix) do
+    RemoteControl.call(project, RemoteControl.Modules, :with_prefix, [prefix, predicate])
+  end
 end

apps/remote_control/lib/lexical/remote_control/modules.ex

@@ -0,0 +1,158 @@
+defmodule Lexical.RemoteControl.Modules do
+  @moduledoc """
+  Utilities for dealing with modules on the remote control node
+  """
+  defmodule Predicate.Syntax do
+    @moduledoc """
+    Syntax helpers for the predicate syntax
+    """
+    defmacro __using__(_) do
+      quote do
+        import unquote(__MODULE__), only: [predicate: 1]
+      end
+    end
+
+    defmacro predicate(call) do
+      predicate_mfa =
+        case call do
+          {:&, _, [{{:., _, [{:__aliases__, _, module}, fn_name]}, _, args}]} ->
+            # This represents the syntax of &Kernel.foo(&1, :a)
+            {Module.concat(module), fn_name, capture_to_placeholder(args)}
+
+          {:&, _, [{fn_name, _, args}]} ->
+            # This represents foo(:a, :b)
+            {Kernel, fn_name, capture_to_placeholder(args)}
+
+          _ ->
+            message = """
+            Invalid predicate.
+
+              Predicates should look like function captures, i.e.
+              predicate(&Module.function(&1, :other)).
+
+              Instead, I got predicate(#{Macro.to_string(call)})
+            """
+
+            raise CompileError, description: message, file: __CALLER__.file, line: __CALLER__.line
+        end
+
+      Macro.escape(predicate_mfa)
+    end
+
+    defp capture_to_placeholder(args) do
+      Enum.map(args, fn
+        {:&, _, [1]} -> :"$1"
+        arg -> arg
+      end)
+    end
+  end
+
+  @cache_timeout Application.compile_env(:remote_control, :modules_cache_expiry, {10, :second})
+
+  @doc """
+  Returns all modules matching a prefix
+
+  `with_prefix` returns all modules on the node on which it runs that start with the given prefix.
+  It's worth noting that it will return _all modules_ regardless if they have been loaded or not.
+
+  You can optionally pass a predicate function to further select which modules are returned, but
+  it's important to understand that the predicate can only be a function reference to a function that
+  exists on the `remote_control` node. I.e. you CANNOT pass anonymous functions to this module.
+
+  To ease things, there is a syntax helper in the `Predicate.Syntax` module that allows you to specify
+  predicates via a syntax that looks like function captures.
+  """
+  def with_prefix(prefix_module, predicate_mfa \\ {Function, :identity, [:"$1"]})
+
+  def with_prefix(prefix_module, mfa) when is_atom(prefix_module) do
+    prefix_module
+    |> to_string()
+    |> with_prefix(mfa)
+  end
+
+  def with_prefix("Elixir." <> _ = prefix, mfa) do
+    results =
+      for {module_string, already_loaded?} <- all_modules(),
+          String.starts_with?(module_string, prefix),
+          module = Module.concat([module_string]),
+          ensure_loaded?(module, already_loaded?),
+          apply_predicate(module, mfa) do
+        {module_string, module}
+      end
+
+    {module_strings, modules_with_prefix} = Enum.unzip(results)
+
+    mark_loaded(module_strings)
+
+    modules_with_prefix
+  end
+
+  def with_prefix(prefix, mfa) do
+    with_prefix("Elixir." <> prefix, mfa)
+  end
+
+  defp apply_predicate(module_arg, {invoked_module, function, args}) do
+    args =
+      Enum.map(args, fn
+        :"$1" ->
+          module_arg
+
+        other ->
+          other
+      end)
+
+    apply(invoked_module, function, args)
+  end
+
+  defp ensure_loaded?(_, true), do: true
+  defp ensure_loaded?(module, _), do: Code.ensure_loaded?(module)
+
+  defp mark_loaded(modules) when is_list(modules) do
+    newly_loaded = Map.new(modules, &{&1, true})
+    {expires, all_loaded} = :persistent_term.get(__MODULE__)
+    updated = Map.merge(all_loaded, newly_loaded)
+
+    :persistent_term.put(__MODULE__, {expires, updated})
+  end
+
+  defp all_modules do
+    case term() do
+      {:ok, modules} ->
+        modules
+
+      :error ->
+        {_expires, modules} = cache = rebuild_cache()
+        :persistent_term.put(__MODULE__, cache)
+        modules
+    end
+  end
+
+  defp term do
+    {expires_at, modules} = :persistent_term.get(__MODULE__, {nil, []})
+
+    if expired?(expires_at) do
+      :error
+    else
+      {:ok, modules}
+    end
+  end
+
+  defp expired?(nil), do: true
+
+  defp expired?(expires) do
+    DateTime.compare(DateTime.utc_now(), expires) == :gt
+  end
+
+  defp rebuild_cache do
+    {amount, unit} = @cache_timeout
+
+    expires = DateTime.add(DateTime.utc_now(), amount, unit)
+
+    module_map =
+      Map.new(:code.all_available(), fn {module_charlist, _path, already_loaded?} ->
+        {to_string(module_charlist), already_loaded?}
+      end)
+
+    {expires, module_map}
+  end
+end

apps/remote_control/test/lexical/remote_control/modules_test.exs

@@ -0,0 +1,66 @@
+defmodule Lexical.RemoteControl.ModulesTest do
+  alias Lexical.RemoteControl.Modules
+  use Modules.Predicate.Syntax
+
+  use ExUnit.Case
+  use Lexical.Test.EventualAssertions
+
+  describe "simple prefixes" do
+    test "specifying a prefix with a string" do
+      found = Modules.with_prefix("En")
+
+      for module <- found,
+          string = module |> Module.split() |> Enum.join(".") do
+        assert String.starts_with?(string, "En")
+      end
+    end
+
+    test "specifying a prefix with a module" do
+      found = Modules.with_prefix(Enum)
+
+      for module <- found,
+          string = module |> Module.split() |> Enum.join(".") do
+        assert String.starts_with?(string, "Enum")
+      end
+    end
+
+    test "new modules are added after expiry" do
+      assert [] = Modules.with_prefix(DoesntExistYet)
+      Module.create(DoesntExistYet, quote(do: nil), file: "foo.ex")
+      assert_eventually [DoesntExistYet] = Modules.with_prefix(DoesntExistYet)
+    end
+
+    test "finds unloaded modules" do
+      modules = "GenEvent" |> Modules.with_prefix() |> Enum.map(&to_string/1)
+      assert "Elixir.GenEvent" in modules
+      assert "Elixir.GenEvent.Stream" in modules
+
+      # ensure it loads the module
+      assert "GenEvent" |> List.wrap() |> Module.concat() |> Code.ensure_loaded?()
+    end
+
+    test "not finding anything" do
+      assert [] = Modules.with_prefix("LexicalIsTheBest")
+    end
+  end
+
+  describe "using predicate descriptors" do
+    test "it should place the argument where you specify" do
+      assert [module] =
+               Modules.with_prefix("GenEvent", {Kernel, :macro_exported?, [:"$1", :__using__, 1]})
+
+      assert to_string(module) == "Elixir.GenEvent"
+    end
+
+    test "it should work with the predicate syntax helpers" do
+      assert [GenServer] =
+               Modules.with_prefix("GenServer", predicate(&macro_exported?(&1, :__using__, 1)))
+
+      assert [GenServer] =
+               Modules.with_prefix(
+                 "GenServer",
+                 predicate(&Kernel.macro_exported?(&1, :__using__, 1))
+               )
+    end
+  end
+end

apps/server/lib/lexical/server/code_intelligence/completion.ex

@@ -9,10 +9,12 @@ defmodule Lexical.Server.CodeIntelligence.Completion do
   alias Lexical.Protocol.Types.InsertTextFormat
   alias Lexical.RemoteControl
   alias Lexical.RemoteControl.Completion.Candidate
+  alias Lexical.RemoteControl.Modules.Predicate
   alias Lexical.Server.CodeIntelligence.Completion.Env
   alias Lexical.Server.Project.Intelligence
   alias Mix.Tasks.Namespace
 
+  use Predicate.Syntax
   require InsertTextFormat
   require Logger
 
@@ -186,6 +188,22 @@ defmodule Lexical.Server.CodeIntelligence.Completion do
           Candidate.Struct
         ]
 
+      Env.in_context?(env, :use) ->
+        case result do
+          %{full_name: full_name} ->
+            with_prefix =
+              RemoteControl.Api.modules_with_prefix(
+                env.project,
+                full_name,
+                predicate(&macro_exported?(&1, :__using__, 1))
+              )
+
+            not Enum.empty?(with_prefix)
+
+          _ ->
+            false
+        end
+
       true ->
         true
     end

apps/server/lib/lexical/server/code_intelligence/completion/translations/macro.ex

@@ -7,7 +7,7 @@ defmodule Lexical.Server.CodeIntelligence.Completion.Translations.Macro do
 
   use Translatable.Impl, for: Candidate.Macro
 
-  @snippet_macros ~w(def defp defmacro defmacrop defimpl defmodule defprotocol defguard defguardp defexception test)
+  @snippet_macros ~w(def defp defmacro defmacrop defimpl defmodule defprotocol defguard defguardp defexception test use)
 
   def translate(%Candidate.Macro{name: name}, _builder, _env)
       when name in ["__before_compile__", "__using__", "__after_compile__"] do
@@ -248,6 +248,18 @@ defmodule Lexical.Server.CodeIntelligence.Completion.Translations.Macro do
     |> builder.boost()
   end
 
+  def translate(%Candidate.Macro{name: "use", arity: 1}, builder, env) do
+    label = "use (invoke another module's __using__ macro)"
+    snippet = "use $0"
+
+    env
+    |> builder.snippet(snippet,
+      kind: :class,
+      label: label
+    )
+    |> builder.boost()
+  end
+
   def translate(%Candidate.Macro{name: "require" <> _, arity: 2} = macro, builder, env) do
     label = "#{macro.name} (require a module's macros)"
 

apps/server/test/lexical/server/code_intelligence/completion/translations/macro_test.exs

@@ -277,6 +277,17 @@ defmodule Lexical.Server.CodeIntelligence.Completion.Translations.MacroTest do
       assert completion.insert_text == "alias $0"
     end
 
+    test "use returns a snippet", %{project: project} do
+      assert {:ok, completion} =
+               project
+               |> complete("us|")
+               |> fetch_completion("use ")
+
+      assert completion.label == "use (invoke another module's __using__ macro)"
+      assert completion.insert_text_format == :snippet
+      assert completion.insert_text == "use $0"
+    end
+
     test "import returns a snippet", %{project: project} do
       assert {:ok, completion} =
                project

apps/server/test/lexical/server/code_intelligence/completion/translations/module_or_behaviour_test.exs

@@ -226,5 +226,19 @@ defmodule Lexical.Server.CodeIntelligence.Completion.Translations.ModuleOrBehavi
       assert order.kind == :struct
       assert apply_completion(order) =~ "%Order{$1}"
     end
+
+    test "completions in use directives only return modules who have a descendent that has __using__",
+         %{
+           project: project
+         } do
+      source = ~q[
+        use En|
+      ]
+
+      completions = complete(project, source)
+      assert length(completions) == 2
+      assert {:ok, _} = fetch_completion(completions, insert_text: "ExUnit")
+      assert {:ok, _} = fetch_completion(completions, insert_text: "ExUnitProperties")
+    end
   end
 end

config/test.exs

@@ -3,7 +3,11 @@ alias Lexical.Server.JsonRpc
 alias Lexical.Test.Transport.NoOp
 
 config :logger, level: :none
-config :remote_control, edit_window_millis: 10
+
+config :remote_control,
+  edit_window_millis: 10,
+  modules_cache_expiry: {50, :millisecond}
+
 config :server, transport: NoOp
 config :stream_data, initial_size: 50