Optimize protocol consolidation

* Allow directory listings to be pre-computed
* Do not require struct modules to be loaded

The current implementation would require us to
recompute protocols whenever a struct changed,
something that is not currently available.
Furthermore, the fact we had to load modules
naturally made consolidation more expensive.
Therefore we need a lazy type resolution,
and until we have such feature, we will only
track struct names.

Author: josevalim

lib/elixir/lib/module/types/apply.ex

@@ -1469,6 +1469,17 @@ defmodule Module.Types.Apply do
                   """
               end
 
+            # The protocol has, at the moment, simplified clauses, so we build the complete one
+            clauses =
+              case mod.__protocol__(:impls) do
+                {:consolidated, mods} ->
+                  domain = mods |> Enum.map(&Module.Types.Of.impl/1) |> Enum.reduce(&union/2)
+                  [{[domain], dynamic()}]
+
+                _ ->
+                  clauses
+              end
+
             {"""
              but expected a type that implements the #{inspect(mod)} protocol.
              #{fix}\

lib/elixir/lib/module/types/of.ex

@@ -142,16 +142,25 @@ defmodule Module.Types.Of do
     {Any, term()}
   ]
 
+  @doc """
+  Currently, for protocol implementations, we only store
+  the open struct definition. This is because we don't want
+  to reconsolidate whenever the struct changes, but at the
+  moment we can't store references either. Ideally struct
+  types on protocol dispatches would be lazily resolved.
+  """
+  def impl(for, mode \\ :closed)
+
   for {for, type} <- impls do
-    def impl(unquote(for)), do: unquote(Macro.escape(type))
+    def impl(unquote(for), _mode), do: unquote(Macro.escape(type))
   end
 
-  def impl(struct) do
+  def impl(struct, mode) do
     # Elixir did not strictly require the implementation to be available,
     # so we need to deal with such cases accordingly.
     # TODO: Assume implementation is available on Elixir v2.0.
     # A warning is emitted since v1.19+.
-    if info = Code.ensure_loaded?(struct) && struct.__info__(:struct) do
+    if info = mode == :closed && Code.ensure_loaded?(struct) && struct.__info__(:struct) do
       struct_type(struct, info)
     else
       open_map(__struct__: atom([struct]))

lib/elixir/lib/protocol.ex

@@ -453,7 +453,7 @@ defmodule Protocol do
       true
 
   """
-  @spec extract_protocols([charlist | String.t()]) :: [atom]
+  @spec extract_protocols([charlist | String.t() | {charlist, [charlist]}]) :: [atom]
   def extract_protocols(paths) do
     extract_matching_by_attribute(paths, [?E, ?l, ?i, ?x, ?i, ?r, ?.], fn module, attributes ->
       case attributes[:__protocol__] do
@@ -482,7 +482,7 @@ defmodule Protocol do
       true
 
   """
-  @spec extract_impls(module, [charlist | String.t()]) :: [atom]
+  @spec extract_impls(module, [charlist | String.t() | {charlist, [charlist]}]) :: [atom]
   def extract_impls(protocol, paths) when is_atom(protocol) do
     prefix = Atom.to_charlist(protocol) ++ [?.]
 
@@ -496,17 +496,25 @@ defmodule Protocol do
 
   defp extract_matching_by_attribute(paths, prefix, callback) do
     for path <- paths,
-        # Do not use protocols as they may be consolidating
-        path = if(is_list(path), do: path, else: String.to_charlist(path)),
-        file <- list_dir(path),
+        {path, files} = list_dir(path),
+        file <- files,
         mod = extract_from_file(path, file, prefix, callback),
         do: mod
   end
 
+  # Do not use protocols as they may be consolidating
+  defp list_dir({path, files}) when is_list(path) and is_list(files) do
+    {path, files}
+  end
+
+  defp list_dir(path) when is_binary(path) do
+    list_dir(String.to_charlist(path))
+  end
+
   defp list_dir(path) when is_list(path) do
     case :file.list_dir(path) do
-      {:ok, files} -> files
-      _ -> []
+      {:ok, files} -> {path, files}
+      _ -> {path, []}
     end
   end
 
@@ -633,7 +641,7 @@ defmodule Protocol do
         checker =
           if checker do
             update_in(checker.exports, fn exports ->
-              signatures = new_signatures(definitions, protocol_funs, protocol, types)
+              signatures = new_signatures(definitions, protocol_funs, protocol, types, structs)
 
               for {fun, info} <- exports do
                 if sig = Map.get(signatures, fun) do
@@ -652,14 +660,13 @@ defmodule Protocol do
     end
   end
 
-  defp new_signatures(definitions, protocol_funs, protocol, types) do
+  defp new_signatures(definitions, protocol_funs, protocol, types, structs) do
     alias Module.Types.Descr
+    types_minus_any = List.delete(types, Any)
 
     clauses =
-      types
-      |> List.delete(Any)
-      |> Enum.map(fn impl ->
-        {[Module.Types.Of.impl(impl)], Descr.atom([__concat__(protocol, impl)])}
+      Enum.map(types_minus_any, fn impl ->
+        {[Module.Types.Of.impl(impl, :open)], Descr.atom([__concat__(protocol, impl)])}
       end)
 
     {domain, impl_for, impl_for!} =
@@ -674,10 +681,16 @@ defmodule Protocol do
           end
 
         _ ->
+          structs_domain =
+            case structs do
+              [] -> Descr.none()
+              _ -> Descr.open_map(__struct__: Descr.atom(structs))
+            end
+
           domain =
-            clauses
-            |> Enum.map(fn {[domain], _} -> domain end)
-            |> Enum.reduce(&Descr.union/2)
+            Enum.reduce(types_minus_any -- structs, structs_domain, fn impl, acc ->
+              Descr.union(Module.Types.Of.impl(impl, :open), acc)
+            end)
 
           not_domain = Descr.negation(domain)
 

lib/elixir/test/elixir/protocol/consolidation_test.exs

@@ -183,21 +183,23 @@ defmodule Protocol.ConsolidationTest do
       assert domain == [term()]
 
       assert clauses == [
-               {[Of.impl(ImplStruct)], atom([Sample.Protocol.ConsolidationTest.ImplStruct])},
-               {[negation(Of.impl(ImplStruct))], atom([nil])}
+               {[Of.impl(ImplStruct, :open)],
+                atom([Sample.Protocol.ConsolidationTest.ImplStruct])},
+               {[negation(Of.impl(ImplStruct, :open))], atom([nil])}
              ]
 
       assert %{{:impl_for!, 1} => %{sig: {:strong, domain, clauses}}} = exports
-      assert domain == [Of.impl(ImplStruct)]
+      assert domain == [Of.impl(ImplStruct, :open)]
 
       assert clauses == [
-               {[Of.impl(ImplStruct)], atom([Sample.Protocol.ConsolidationTest.ImplStruct])}
+               {[Of.impl(ImplStruct, :open)],
+                atom([Sample.Protocol.ConsolidationTest.ImplStruct])}
              ]
 
       assert %{{:ok, 1} => %{sig: {:strong, nil, clauses}}} = exports
 
       assert clauses == [
-               {[Of.impl(ImplStruct)], dynamic()}
+               {[Of.impl(ImplStruct, :open)], dynamic()}
              ]
     end
 
@@ -212,9 +214,11 @@ defmodule Protocol.ConsolidationTest do
       assert domain == [term()]
 
       assert clauses == [
-               {[Of.impl(Map)], atom([WithAny.Map])},
-               {[Of.impl(ImplStruct)], atom([WithAny.Protocol.ConsolidationTest.ImplStruct])},
-               {[negation(union(Of.impl(ImplStruct), Of.impl(Map)))], atom([WithAny.Any])}
+               {[Of.impl(Map, :open)], atom([WithAny.Map])},
+               {[Of.impl(ImplStruct, :open)],
+                atom([WithAny.Protocol.ConsolidationTest.ImplStruct])},
+               {[negation(union(Of.impl(ImplStruct, :open), Of.impl(Map, :open)))],
+                atom([WithAny.Any])}
              ]
 
       assert %{{:ok, 2} => %{sig: {:strong, nil, clauses}}} = exports
@@ -275,18 +279,34 @@ defmodule Protocol.ConsolidationTest do
       assert Inspect in protos
     end
 
+    test "protocols with expanded path" do
+      path = to_charlist(Application.app_dir(:elixir, "ebin"))
+      {:ok, mods} = :file.list_dir(path)
+      protos = Protocol.extract_protocols([{path, mods}])
+      assert Enumerable in protos
+      assert Inspect in protos
+    end
+
     test "implementations with charlist path" do
-      protos =
+      impls =
         Protocol.extract_impls(Enumerable, [to_charlist(Application.app_dir(:elixir, "ebin"))])
 
-      assert List in protos
-      assert Function in protos
+      assert List in impls
+      assert Function in impls
     end
 
     test "implementations with binary path" do
-      protos = Protocol.extract_impls(Enumerable, [Application.app_dir(:elixir, "ebin")])
-      assert List in protos
-      assert Function in protos
+      impls = Protocol.extract_impls(Enumerable, [Application.app_dir(:elixir, "ebin")])
+      assert List in impls
+      assert Function in impls
+    end
+
+    test "implementations with expanded path" do
+      path = to_charlist(Application.app_dir(:elixir, "ebin"))
+      {:ok, mods} = :file.list_dir(path)
+      impls = Protocol.extract_impls(Enumerable, [{path, mods}])
+      assert List in impls
+      assert Function in impls
     end
   end
 end

lib/mix/lib/mix/compilers/protocol.ex

@@ -126,11 +126,16 @@ defmodule Mix.Compilers.Protocol do
   end
 
   defp consolidation_paths do
-    filter_otp(:code.get_path(), :code.lib_dir())
-  end
-
-  defp filter_otp(paths, otp) do
-    Enum.filter(paths, &(not :lists.prefix(otp, &1)))
+    otp = :code.lib_dir()
+
+    # Pre-list the paths in each directory for performance
+    for path <- :code.get_path(), not :lists.prefix(otp, path) do
+      {path,
+       case :file.list_dir(path) do
+         {:ok, files} -> files
+         _ -> []
+       end}
+    end
   end
 
   defp consolidate([], _paths, output, _opts) do