Type check generators in comprehensions

Author: josevalim

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

@@ -811,14 +811,12 @@ defmodule Module.Types.Apply do
   end
 
   def format_diagnostic({:badremote, mfac, expr, args_types, domain, clauses, context}) do
-    traces = collect_traces(expr, context)
     {mod, fun, arity, converter} = mfac
     meta = elem(expr, 1)
 
-    # Protocol errors can be very verbose, so we collapse structs
-    {banner, hints, opts} =
-      cond do
-        meta[:from_interpolation] ->
+    {banner, hints, traces} =
+      case Keyword.get(meta, :type_check) do
+        :interpolation ->
           {_, _, [arg]} = expr
 
           {"""
@@ -827,34 +825,50 @@ defmodule Module.Types.Apply do
                #{expr_to_string(arg) |> indent(4)}
 
            it has type:
-           """, [:interpolation], [collapse_structs: true]}
+           """, [:interpolation], collect_traces(expr, context)}
 
-        Code.ensure_loaded?(mod) and
-            Keyword.has_key?(mod.module_info(:attributes), :__protocol__) ->
-          {nil, [{:protocol, mod}], [collapse_structs: true]}
+        :generator ->
+          {:<-, _, [_, arg]} = expr
 
-        true ->
-          {nil, [], []}
-      end
+          {"""
+           incompatible value given to for-comprehension:
 
-    explanation =
-      empty_arg_reason(converter.(args_types)) ||
-        """
-        but expected one of:
-        #{clauses_args_to_quoted_string(clauses, converter, opts)}
-        """
+               #{expr_to_string(expr) |> indent(4)}
 
-    mfa_or_fa = if mod, do: Exception.format_mfa(mod, fun, arity), else: "#{fun}/#{arity}"
+           it has type:
+           """, [:generator], collect_traces(arg, context)}
 
-    banner =
-      banner ||
-        """
-        incompatible types given to #{mfa_or_fa}:
+        _ ->
+          mfa_or_fa = if mod, do: Exception.format_mfa(mod, fun, arity), else: "#{fun}/#{arity}"
 
-            #{expr_to_string(expr) |> indent(4)}
+          {"""
+           incompatible types given to #{mfa_or_fa}:
 
-        given types:
-        """
+               #{expr_to_string(expr) |> indent(4)}
+
+           given types:
+           """, [], collect_traces(expr, context)}
+      end
+
+    explanation =
+      cond do
+        reason = empty_arg_reason(converter.(args_types)) ->
+          reason
+
+        Code.ensure_loaded?(mod) and
+            Keyword.has_key?(mod.module_info(:attributes), :__protocol__) ->
+          # Protocol errors can be very verbose, so we collapse structs
+          """
+          but expected a type that implements the #{inspect(mod)} protocol, it must be one of:
+          #{clauses_args_to_quoted_string(clauses, converter, collapse_structs: true)}
+          """
+
+        true ->
+          """
+          but expected one of:
+          #{clauses_args_to_quoted_string(clauses, converter, [])}
+          """
+      end
 
     %{
       details: %{typing_traces: traces},

lib/elixir/lib/module/types/descr.ex

@@ -31,15 +31,15 @@ defmodule Module.Types.Descr do
   @map_empty [{:closed, %{}, []}]
 
   @none %{}
-  @empty_list %{bitmap: @bit_empty_list}
-  @not_non_empty_list %{bitmap: @bit_top, atom: @atom_top, tuple: @tuple_top, map: @map_top}
   @term %{
     bitmap: @bit_top,
     atom: @atom_top,
     tuple: @tuple_top,
     map: @map_top,
     list: @non_empty_list_top
   }
+  @empty_list %{bitmap: @bit_empty_list}
+  @not_non_empty_list Map.delete(@term, :list)
 
   @empty_intersection [0, @none]
   @empty_difference [0, []]
@@ -98,6 +98,7 @@ defmodule Module.Types.Descr do
   @not_set %{optional: 1}
   @term_or_optional Map.put(@term, :optional, 1)
   @term_or_dynamic_optional Map.put(@term, :dynamic, %{optional: 1})
+  @not_atom_or_optional Map.delete(@term_or_optional, :atom)
 
   def not_set(), do: @not_set
   def if_set(:term), do: term_or_optional()
@@ -1751,10 +1752,32 @@ defmodule Module.Types.Descr do
   end
 
   # Two maps are fusible if they differ in at most one element.
+  defp non_fusible_maps?({_, fields1, []}, {_, fields2, []})
+       when map_size(fields1) > map_size(fields2) do
+    not fusible_maps?(Map.to_list(fields2), fields1, 0)
+  end
+
   defp non_fusible_maps?({_, fields1, []}, {_, fields2, []}) do
-    Enum.count_until(fields1, fn {key, value} -> Map.fetch!(fields2, key) != value end, 2) > 1
+    not fusible_maps?(Map.to_list(fields1), fields2, 0)
+  end
+
+  defp fusible_maps?([{:__struct__, value} | rest], fields, count) do
+    case Map.fetch!(fields, :__struct__) do
+      ^value -> fusible_maps?(rest, fields, count)
+      _ -> false
+    end
   end
 
+  defp fusible_maps?([{key, value} | rest], fields, count) do
+    case Map.fetch!(fields, key) do
+      ^value -> fusible_maps?(rest, fields, count)
+      _ when count == 1 -> false
+      _ when count == 0 -> fusible_maps?(rest, fields, count + 1)
+    end
+  end
+
+  defp fusible_maps?([], _fields, _count), do: true
+
   defp map_non_negated_fuse_pair({tag, fields1, []}, {_, fields2, []}) do
     fields =
       symmetrical_merge(fields1, fields2, fn _k, v1, v2 ->
@@ -1818,6 +1841,11 @@ defmodule Module.Types.Descr do
     {:empty_map, [], []}
   end
 
+  def map_literal_to_quoted({:open, %{__struct__: @not_atom_or_optional} = fields}, _opts)
+      when map_size(fields) == 1 do
+    {:non_struct_map, [], []}
+  end
+
   def map_literal_to_quoted({tag, fields}, opts) do
     case tag do
       :closed ->

lib/elixir/lib/module/types/expr.ex

@@ -471,13 +471,17 @@ defmodule Module.Types.Expr do
 
   ## Comprehensions
 
-  defp for_clause({:<-, _, [left, right]} = expr, stack, context) do
+  defp for_clause({:<-, meta, [left, right]}, stack, context) do
+    expr = {:<-, [type_check: :generator] ++ meta, [left, right]}
     {pattern, guards} = extract_head([left])
-    {_, context} = of_expr(right, stack, context)
+    {type, context} = of_expr(right, stack, context)
 
     {_type, context} =
       Pattern.of_match(pattern, guards, dynamic(), expr, :for, stack, context)
 
+    {_type, context} =
+      Apply.remote(Enumerable, :count, [right], [type], expr, stack, context)
+
     context
   end
 
@@ -511,6 +515,7 @@ defmodule Module.Types.Expr do
   end
 
   defp for_option({:uniq, _}, _stack, context) do
+    # This option is verified to be a boolean at compile-time
     context
   end
 

lib/elixir/lib/module/types/helpers.ex

@@ -90,11 +90,12 @@ defmodule Module.Types.Helpers do
         string upfront or implement the protocol accordingly
         """
 
-      {:protocol, protocol} ->
+      :generator ->
         """
 
-        #{hint()} #{inspect(protocol)} is a protocol in Elixir. Either make sure you \
-        give valid data types as arguments or implement the protocol accordingly
+        #{hint()} for-comprehensions in Elixir use the Enumerable protocol to traverse \
+        data structures. Either convert the data type into a list (or another Enumerable) \
+        or implement the protocol accordingly
         """
 
       :anonymous_rescue ->

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

@@ -96,7 +96,7 @@ defmodule Module.Types.Of do
     {Function, fun()},
     {Integer, integer()},
     {List, list(term())},
-    {Map, open_map(__struct__: not_set())},
+    {Map, open_map(__struct__: if_set(negation(atom())))},
     {Port, port()},
     {PID, pid()},
     {Reference, reference()},
@@ -339,7 +339,7 @@ defmodule Module.Types.Of do
          {{:., _, [String.Chars, :to_string]} = dot, meta, [arg]},
          {:binary, _, nil}
        ) do
-    {dot, [from_interpolation: true] ++ meta, [arg]}
+    {dot, [type_check: :interpolation] ++ meta, [arg]}
   end
 
   defp annotate_interpolation(left, _right) do

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

@@ -1494,6 +1494,15 @@ defmodule Module.Types.DescrTest do
       assert closed_map(__struct__: atom([Decimal]), coef: term(), exp: term(), sign: integer())
              |> to_quoted_string(collapse_structs: true) ==
                "%Decimal{sign: integer()}"
+
+      # Does not fuse structs
+      assert union(closed_map(__struct__: atom([Foo])), closed_map(__struct__: atom([Bar])))
+             |> to_quoted_string() ==
+               "%{__struct__: Bar} or %{__struct__: Foo}"
+
+      # Properly format non_struct_map
+      assert open_map(__struct__: if_set(negation(atom()))) |> to_quoted_string() ==
+               "non_struct_map()"
     end
   end
 

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

@@ -131,7 +131,7 @@ defmodule Module.Types.IntegrationTest do
       assert itself_arg.(Itself.Function) == dynamic(fun())
       assert itself_arg.(Itself.Integer) == dynamic(integer())
       assert itself_arg.(Itself.List) == dynamic(list(term()))
-      assert itself_arg.(Itself.Map) == dynamic(open_map(__struct__: not_set()))
+      assert itself_arg.(Itself.Map) == dynamic(open_map(__struct__: if_set(negation(atom()))))
       assert itself_arg.(Itself.Port) == dynamic(port())
       assert itself_arg.(Itself.PID) == dynamic(pid())
       assert itself_arg.(Itself.Reference) == dynamic(reference())
@@ -374,7 +374,7 @@ defmodule Module.Types.IntegrationTest do
       assert_warnings(files, warnings)
     end
 
-    test "protocol dispatch" do
+    test "String.Chars protocol dispatch" do
       files = %{
         "a.ex" => """
         defmodule FooBar do
@@ -394,7 +394,7 @@ defmodule Module.Types.IntegrationTest do
 
                 -dynamic(%Range{first: term(), last: term(), step: term()})-
 
-            but expected one of:
+            but expected a type that implements the String.Chars protocol, it must be one of:
 
                 %Date{} or %DateTime{} or %NaiveDateTime{} or %Time{} or %URI{} or %Version{} or
                   %Version.Requirement{} or atom() or binary() or float() or integer() or list(term())
@@ -416,7 +416,7 @@ defmodule Module.Types.IntegrationTest do
 
                 -dynamic(%Range{first: term(), last: term(), step: term()})-
 
-            but expected one of:
+            but expected a type that implements the String.Chars protocol, it must be one of:
 
                 %Date{} or %DateTime{} or %NaiveDateTime{} or %Time{} or %URI{} or %Version{} or
                   %Version.Requirement{} or atom() or binary() or float() or integer() or list(term())
@@ -426,8 +426,43 @@ defmodule Module.Types.IntegrationTest do
                 # type: dynamic(%Range{})
                 # from: a.ex:2:24
                 _.._//_ = data
+        """
+      ]
+
+      assert_warnings(files, warnings, consolidate_protocols: true)
+    end
+
+    test "Enumerable protocol dispatch" do
+      files = %{
+        "a.ex" => """
+        defmodule FooBar do
+          def example1(%Date{} = date), do: for(x <- date, do: x)
+        end
+        """
+      }
+
+      warnings = [
+        """
+            warning: incompatible value given to for-comprehension:
+
+                x <- date
+
+            it has type:
+
+                -dynamic(%Date{year: term(), month: term(), day: term(), calendar: term()})-
+
+            but expected a type that implements the Enumerable protocol, it must be one of:
+
+                %Date.Range{} or %File.Stream{} or %GenEvent.Stream{} or %HashDict{} or %HashSet{} or
+                  %IO.Stream{} or %MapSet{} or %Range{} or %Stream{} or fun() or list(term()) or non_struct_map()
+
+            where "date" was given the type:
+
+                # type: dynamic(%Date{})
+                # from: a.ex:2:24
+                %Date{} = date
 
-            hint: String.Chars is a protocol in Elixir. Either make sure you give valid data types as arguments or implement the protocol accordingly
+            hint: for-comprehensions in Elixir use the Enumerable protocol to traverse data structures. Either convert the data type into a list (or another Enumerable) or implement the protocol accordingly
         """
       ]