Recompose multiple orelse as Kernel.in for performance, closes #15108

Author: josevalim

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

@@ -367,8 +367,8 @@ defmodule Module.Types.Apply do
     right_literal? = Macro.quoted_literal?(right)
 
     case {left_literal?, right_literal?} do
-      {true, false} -> custom_compare(name, right, left, expected, expr, stack, context, of_fun)
-      {false, true} -> custom_compare(name, left, right, expected, expr, stack, context, of_fun)
+      {true, false} -> literal_compare(name, right, left, expected, expr, stack, context, of_fun)
+      {false, true} -> literal_compare(name, left, right, expected, expr, stack, context, of_fun)
       {literal?, _} -> compare(name, left, right, literal?, expr, stack, context, of_fun)
     end
   end
@@ -495,7 +495,7 @@ defmodule Module.Types.Apply do
             when (fun in [:length, :map_size] and is_integer(literal) and literal >= 0) or
                    (fun in [:tuple_size] and literal in 0..15)
 
-  defp custom_compare(
+  defp literal_compare(
          name,
          {{:., _, [:erlang, fun]}, _, [arg]} = left,
          literal,
@@ -547,14 +547,18 @@ defmodule Module.Types.Apply do
     end
   end
 
-  defp custom_compare(name, arg, literal, expected, expr, stack, context, of_fun) do
+  defp literal_compare(name, arg, literal, expected, expr, stack, context, of_fun) do
+    {type, context} = of_fun.(literal, term(), expr, stack, context)
+    literal_compare(name, arg, type, singleton?(type), expected, expr, stack, context, of_fun)
+  end
+
+  def literal_compare(name, arg, type, singleton?, expected, expr, stack, context, of_fun) do
     case booleaness(expected) do
       booleaness when booleaness in [:maybe_both, :none] ->
-        compare(name, arg, literal, false, expr, stack, context, of_fun)
+        {arg_type, context} = of_fun.(arg, term(), expr, stack, context)
+        return_compare(name, arg_type, type, boolean(), false, expr, stack, context)
 
       {boolean, _maybe_or_always} ->
-        {literal_type, context} = of_fun.(literal, term(), expr, stack, context)
-
         {polarity, return} =
           case boolean do
             true -> {name in [:==, :"=:="], @atom_true}
@@ -563,28 +567,28 @@ defmodule Module.Types.Apply do
 
         # This logic mirrors the code in `Pattern.of_pattern_tree`
         # If it is a singleton, we can always be precise
-        if singleton?(literal_type) do
-          expected = if polarity, do: literal_type, else: negation(literal_type)
+        if singleton? do
+          expected = if polarity, do: type, else: negation(type)
           {arg_type, context} = of_fun.(arg, expected, expr, stack, context)
           result = if subtype?(arg_type, expected), do: return, else: boolean()
 
           # Because reverse polarity means we will infer negated types
           # (which are naturally disjoint), we skip checks in such cases
           skip_check? = not polarity
-          return_compare(name, arg_type, literal_type, result, skip_check?, expr, stack, context)
+          return_compare(name, arg_type, type, result, skip_check?, expr, stack, context)
         else
           expected =
             cond do
               # We are checking for `not x == 1` or similar, we can't say anything about x
               polarity == false -> term()
               # We are checking for `x == 1`, make sure x is integer or float
-              name in [:==, :"/="] -> numberize(literal_type)
+              name in [:==, :"/="] -> numberize(type)
               # Otherwise we have the literal type as is
-              true -> literal_type
+              true -> type
             end
 
           {arg_type, context} = of_fun.(arg, expected, expr, stack, context)
-          return_compare(name, arg_type, literal_type, boolean(), false, expr, stack, context)
+          return_compare(name, arg_type, type, boolean(), false, expr, stack, context)
         end
     end
   end

lib/elixir/lib/module/types/pattern.ex

@@ -1046,6 +1046,26 @@ defmodule Module.Types.Pattern do
     of_remote(fun, args, call, expected, stack, context)
   end
 
+  # This is reconstructed as part of orelse
+  def of_guard({{:., _, [Kernel, :in]}, _meta, [left, right]} = call, expected, _, stack, context) do
+    {right, context} = Enum.map_reduce(right, context, &of_guard(&1, term(), call, stack, &2))
+    {singleton, non_singleton} = Enum.split_with(right, &singleton?/1)
+
+    fun = fn
+      [], _singleton?, context ->
+        {none(), context}
+
+      types, singleton?, context ->
+        type = Enum.reduce(types, &union/2)
+        fun = &of_guard/5
+        Apply.literal_compare(:"=:=", left, type, singleton?, expected, call, stack, context, fun)
+    end
+
+    {singleton_type, context} = fun.(singleton, true, context)
+    {non_singleton_type, context} = fun.(non_singleton, false, context)
+    {union(singleton_type, non_singleton_type), context}
+  end
+
   # var
   def of_guard({_, meta, _} = var, expected, expr, stack, context) when is_var(var) do
     version = Keyword.fetch!(meta, :version)
@@ -1136,27 +1156,10 @@ defmodule Module.Types.Pattern do
     # If we have multiple operations in a row,
     # we unpack them into a single pass, to avoid
     # building nested conditional environments.
-    [left | right] = unpack_op(call, fun, [])
-
-    # In case the right side is a sequence of comparisons with imprecise literals, collapse them.
-    right =
-      with {{:., _, [:erlang, comp_op]}, _, [{var_name, _, var_ctx}, literal]}
-           when comp_op in @comp_op and is_atom(var_name) and is_atom(var_ctx) and
-                  (is_number(literal) or is_binary(literal)) <- left,
-           true <-
-             Enum.all?(right, fn
-               {{:., _, [:erlang, ^comp_op]}, _, [{^var_name, _, ^var_ctx}, other_literal]}
-               when is_integer(literal) and is_integer(other_literal)
-               when is_float(literal) and is_float(other_literal)
-               when is_binary(literal) and is_binary(other_literal) ->
-                 true
-
-               _ ->
-                 false
-             end) do
-        []
-      else
-        _ -> right
+    [left | right] =
+      case unpack_op(call, fun, []) do
+        entries when fun == :orelse -> reconstruct_kernel_in(entries)
+        entries -> entries
       end
 
     # For example, if the expected type is true for andalso, then it can
@@ -1213,6 +1216,36 @@ defmodule Module.Types.Pattern do
     [other | acc]
   end
 
+  defp reconstruct_kernel_in([head | tail]) do
+    with {{:., dot_meta, [:erlang, :"=:="]}, meta, [left, right]} <- head,
+         true <- Macro.quoted_literal?(right),
+         false <- data_size_op?(left),
+         {[_ | _] = entries, tail} <- reconstruct_kernel_in(tail, left, []) do
+      in_args = [left, [right | entries]]
+      [{{:., dot_meta, [Kernel, :in]}, meta, in_args} | reconstruct_kernel_in(tail)]
+    else
+      _ -> [head | reconstruct_kernel_in(tail)]
+    end
+  end
+
+  defp reconstruct_kernel_in([]), do: []
+
+  defp reconstruct_kernel_in(list, left, acc) do
+    with [{{:., _, [:erlang, :"=:="]}, _, [^left, right]} | tail] <- list,
+         true <- Macro.quoted_literal?(right) do
+      reconstruct_kernel_in(tail, left, [right | acc])
+    else
+      _ -> {Enum.reverse(acc), list}
+    end
+  end
+
+  defp data_size_op?({{:., _, [:erlang, op]}, _, [_]})
+       when op in [:length, :tuple_size, :map_size],
+       do: true
+
+  defp data_size_op?(_),
+    do: false
+
   defp of_logical_all([head], disjoint?, expected, to_abort, stack, context) do
     {type, context} = of_guard(head, expected, head, stack, context)