Add type inference for literal equality in guards

Author: josevalim

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

@@ -477,6 +477,7 @@ defmodule Module.Types.Apply do
     remote_domain(mod, fun, args, expected, elem(expr, 1), stack, context)
   end
 
+  @number union(integer(), float())
   @empty_list empty_list()
   @non_empty_list non_empty_list(term())
   @empty_map empty_map()
@@ -526,8 +527,41 @@ defmodule Module.Types.Apply do
     end
   end
 
-  defp custom_compare(name, left, right, _expected, expr, stack, context, of_fun) do
-    compare(name, left, right, false, expr, stack, context, of_fun)
+  defp custom_compare(name, arg, literal, expected, expr, stack, context, of_fun) do
+    {literal_type, context} = of_fun.(literal, term(), expr, stack, context)
+
+    case booleaness(expected) do
+      booleaness when booleaness in [:maybe_both, :none] ->
+        compare(name, arg, literal, false, expr, stack, context, of_fun)
+
+      booleaness ->
+        {polarity, return} =
+          case booleaness do
+            :maybe_true -> {name in [:==, :"=:="], @atom_true}
+            :maybe_false -> {name in [:"/=", :"=/="], @atom_false}
+          end
+
+        # 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)
+          {actual, context} = of_fun.(arg, expected, expr, stack, context)
+          result = if compatible?(actual, expected), do: return, else: boolean()
+          {result, 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
+              number_type?(literal_type) and name in [:==, :"/="] -> union(literal_type, @number)
+              # Otherwise we have the literal type as is
+              true -> literal_type
+            end
+
+          {_, context} = of_fun.(arg, expected, expr, stack, context)
+          {boolean(), context}
+        end
+    end
   end
 
   defp compare(name, left, right, literal?, expr, stack, context, of_fun) do

lib/elixir/test/elixir/module/types/pattern_test.exs

@@ -818,6 +818,64 @@ defmodule Module.Types.PatternTest do
     end
   end
 
+  describe "equality in guards" do
+    test "with non-singleton literals" do
+      assert typecheck!([x], x == "foo", x) == dynamic(binary())
+      assert typecheck!([x], x === "foo", x) == dynamic(binary())
+      assert typecheck!([x], not (x == "foo"), x) == dynamic()
+      assert typecheck!([x], not (x === "foo"), x) == dynamic()
+
+      assert typecheck!([x], x != "foo", x) == dynamic()
+      assert typecheck!([x], x !== "foo", x) == dynamic()
+      assert typecheck!([x], not (x != "foo"), x) == dynamic(binary())
+      assert typecheck!([x], not (x !== "foo"), x) == dynamic(binary())
+    end
+
+    test "with number literals" do
+      assert typecheck!([x], x == 1, x) == dynamic(union(integer(), float()))
+      assert typecheck!([x], x === 1, x) == dynamic(integer())
+      assert typecheck!([x], not (x == 1), x) == dynamic()
+      assert typecheck!([x], not (x === 1), x) == dynamic()
+
+      assert typecheck!([x], x != 1, x) == dynamic()
+      assert typecheck!([x], x !== 1, x) == dynamic()
+      assert typecheck!([x], not (x != 1), x) == dynamic(union(integer(), float()))
+      assert typecheck!([x], not (x !== 1), x) == dynamic(integer())
+
+      assert typecheck!([x], x == 1.0, x) == dynamic(union(integer(), float()))
+      assert typecheck!([x], x === 1.0, x) == dynamic(float())
+      assert typecheck!([x], not (x == 1.0), x) == dynamic()
+      assert typecheck!([x], not (x === 1.0), x) == dynamic()
+
+      assert typecheck!([x], x != 1.0, x) == dynamic()
+      assert typecheck!([x], x !== 1.0, x) == dynamic()
+      assert typecheck!([x], not (x != 1.0), x) == dynamic(union(integer(), float()))
+      assert typecheck!([x], not (x !== 1.0), x) == dynamic(float())
+    end
+
+    test "with singleton literals" do
+      assert typecheck!([x], x == :foo, x) == dynamic(atom([:foo]))
+      assert typecheck!([x], x === :foo, x) == dynamic(atom([:foo]))
+      assert typecheck!([x], not (x == :foo), x) == dynamic(negation(atom([:foo])))
+      assert typecheck!([x], not (x === :foo), x) == dynamic(negation(atom([:foo])))
+
+      assert typecheck!([x], x != :foo, x) == dynamic(negation(atom([:foo])))
+      assert typecheck!([x], x !== :foo, x) == dynamic(negation(atom([:foo])))
+      assert typecheck!([x], not (x != :foo), x) == dynamic(atom([:foo]))
+      assert typecheck!([x], not (x !== :foo), x) == dynamic(atom([:foo]))
+
+      assert typecheck!([x], x == [], x) == dynamic(empty_list())
+      assert typecheck!([x], x === [], x) == dynamic(empty_list())
+      assert typecheck!([x], not (x == []), x) == dynamic(negation(empty_list()))
+      assert typecheck!([x], not (x === []), x) == dynamic(negation(empty_list()))
+
+      assert typecheck!([x], x != [], x) == dynamic(negation(empty_list()))
+      assert typecheck!([x], x !== [], x) == dynamic(negation(empty_list()))
+      assert typecheck!([x], not (x != []), x) == dynamic(empty_list())
+      assert typecheck!([x], not (x !== []), x) == dynamic(empty_list())
+    end
+  end
+
   describe "comparison in guards" do
     test "length equality" do
       assert typecheck!([x], length(x) != 0, x) == dynamic(non_empty_list(term()))