Check map updates

lib/elixir/lib/map.ex

@@ -68,7 +68,7 @@ defmodule Map do
       iex> a
       1
 
-  But this will raise a match error:
+  But this will raise a `MatchError` exception:
 
       %{:c => 3} = %{:a => 1, 2 => :b}
 
@@ -88,8 +88,10 @@ defmodule Map do
       iex> map = %{one: 1, two: 2}
       iex> %{map | one: "one"}
       %{one: "one", two: 2}
-      iex> %{map | three: 3}
-      ** (KeyError) key :three not found
+
+  When a key that does not exist in the map is updated a `KeyError` exception will be raised:
+
+      %{map | three: 3}
 
   The functions in this module that need to find a specific key work in logarithmic time.
   This means that the time it takes to find keys grows as the map grows, but it's not

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

@@ -139,28 +139,20 @@ defmodule Module.Types.Expr do
   end
 
   # %{map | ...}
-  def of_expr({:%{}, _, [{:|, _, [_map, args]}]} = expr, stack, context) do
+  def of_expr({:%{}, _, [{:|, _, [map, args]}]} = expr, stack, context) do
     stack = push_expr_stack(expr, stack)
 
-    case of_pairs(args, stack, context) do
-      {:ok, _pairs, context} -> {:ok, {:map, [{:optional, :dynamic, :dynamic}]}, context}
-      {:error, reason} -> {:error, reason}
-    end
+    additional_pairs = [{:optional, :dynamic, :dynamic}]
+    map_update(map, args, additional_pairs, stack, context)
   end
 
   # %Struct{map | ...}
-  def of_expr({:%, meta, [module, {:%{}, _, [{:|, _, [_map, args]}]}]} = expr, stack, context) do
+  def of_expr({:%, meta, [module, {:%{}, _, [{:|, _, [map, args]}]}]} = expr, stack, context) do
     context = Remote.check(module, :__struct__, 0, meta, context)
     stack = push_expr_stack(expr, stack)
 
-    case of_pairs(args, stack, context) do
-      {:ok, _pairs, context} ->
-        pairs = [{:required, {:atom, :__struct__}, {:atom, module}}]
-        {:ok, {:map, pairs}, context}
-
-      {:error, reason} ->
-        {:error, reason}
-    end
+    additional_pairs = [{:required, {:atom, :__struct__}, {:atom, module}}]
+    map_update(map, args, additional_pairs, stack, context)
   end
 
   # %{...}
@@ -395,6 +387,27 @@ defmodule Module.Types.Expr do
     end)
   end
 
+  defp map_update(map, args, additional_pairs, stack, context) do
+    with {:ok, map_type, context} <- of_expr(map, stack, context),
+         {:ok, arg_pairs, context} <- of_pairs(args, stack, context),
+         arg_pairs = pairs_to_unions(arg_pairs, context),
+         # Change value types to dynamic to reuse map unification for map updates
+         dynamic_value_pairs =
+           Enum.map(arg_pairs, fn {:required, key, _value} -> {:required, key, :dynamic} end),
+         args_type = {:map, additional_pairs ++ dynamic_value_pairs},
+         {:ok, type, context} <- unify(args_type, map_type, stack, context) do
+      # Retrieve map type and overwrite with the new value types from the map update
+      {:map, pairs} = resolve_var(type, context)
+
+      updated_pairs =
+        Enum.reduce(arg_pairs, pairs, fn {:required, key, value}, pairs ->
+          List.keyreplace(pairs, key, 1, {:required, key, value})
+        end)
+
+      {:ok, {:map, updated_pairs}, context}
+    end
+  end
+
   defp for_clause({:<-, _, [left, expr]}, stack, context) do
     {pattern, guards} = extract_head([left])
 

lib/elixir/lib/module/types/infer.ex

@@ -335,6 +335,9 @@ defmodule Module.Types.Infer do
     {type, context}
   end
 
+  def resolve_var({:var, var}, context), do: resolve_var(Map.fetch!(context.types, var), context)
+  def resolve_var(other, _context), do: other
+
   # Check unify stack to see if variable was already expanded
   defp variable_expanded?(var, stack, context) do
     Enum.any?(stack.unify_stack, &variable_same?(var, &1, context))

lib/elixir/test/elixir/map_test.exs

@@ -225,6 +225,8 @@ defmodule MapTest do
     end
   end
 
+  defp empty_map(), do: %{}
+
   test "structs" do
     assert %ExternalUser{} == %{__struct__: ExternalUser, name: "john", age: 27}
 
@@ -236,10 +238,8 @@ defmodule MapTest do
     %ExternalUser{name: name} = %ExternalUser{}
     assert name == "john"
 
-    map = %{}
-
     assert_raise BadStructError, "expected a struct named MapTest.ExternalUser, got: %{}", fn ->
-      %ExternalUser{map | name: "meg"}
+      %ExternalUser{empty_map() | name: "meg"}
     end
   end
 

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

@@ -263,6 +263,87 @@ defmodule Module.Types.ExprTest do
              )
   end
 
+  test "map update" do
+    assert quoted_expr(
+             (
+               map = %{foo: :a}
+               %{map | foo: :b}
+             )
+           ) ==
+             {:ok, {:map, [{:required, {:atom, :foo}, {:atom, :b}}]}}
+
+    assert quoted_expr([map], %{map | foo: :b}) ==
+             {:ok,
+              {:map, [{:optional, :dynamic, :dynamic}, {:required, {:atom, :foo}, {:atom, :b}}]}}
+
+    assert {:error,
+            {:unable_unify,
+             {{:map, [{:optional, :dynamic, :dynamic}, {:required, {:atom, :bar}, :dynamic}]},
+              {:map, [{:required, {:atom, :foo}, {:atom, :a}}]},
+              _}}} =
+             quoted_expr(
+               (
+                 map = %{foo: :a}
+                 %{map | bar: :b}
+               )
+             )
+  end
+
+  test "struct update" do
+    assert quoted_expr(
+             (
+               map = %Module.Types.ExprTest.Struct2{field: :a}
+               %Module.Types.ExprTest.Struct2{map | field: :b}
+             )
+           ) ==
+             {:ok,
+              {:map,
+               [
+                 {:required, {:atom, :__struct__}, {:atom, Module.Types.ExprTest.Struct2}},
+                 {:required, {:atom, :field}, {:atom, :b}}
+               ]}}
+
+    assert {:error,
+            {:unable_unify,
+             {{:map,
+               [
+                 {:required, {:atom, :__struct__}, {:atom, Module.Types.ExprTest.Struct2}},
+                 {:required, {:atom, :field}, :dynamic}
+               ]}, {:map, [{:required, {:atom, :field}, {:atom, :a}}]},
+              _}}} =
+             quoted_expr(
+               (
+                 map = %{field: :a}
+                 %Module.Types.ExprTest.Struct2{map | field: :b}
+               )
+             )
+
+    assert quoted_expr([map], %Module.Types.ExprTest.Struct2{map | field: :b}) ==
+             {:ok,
+              {:map,
+               [
+                 {:required, {:atom, :__struct__}, {:atom, Module.Types.ExprTest.Struct2}},
+                 {:required, {:atom, :field}, {:atom, :b}}
+               ]}}
+
+    assert {:error,
+            {:unable_unify,
+             {{:map,
+               [{:optional, :dynamic, :dynamic}, {:required, {:atom, :not_field}, :dynamic}]},
+              {:map,
+               [
+                 {:required, {:atom, :__struct__}, {:atom, Module.Types.ExprTest.Struct2}},
+                 {:required, {:atom, :field}, {:atom, nil}}
+               ]},
+              _}}} =
+             quoted_expr(
+               (
+                 map = %Module.Types.ExprTest.Struct2{}
+                 %{map | not_field: :b}
+               )
+             )
+  end
+
   # Use module attribute to avoid formatter adding parentheses
   @mix_module Mix