Add map_fetch_and_put

Author: josevalim

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

@@ -1429,20 +1429,46 @@ defmodule Module.Types.Descr do
   end
 
   @doc """
-  Adds a `key` of a given type, assuming that the descr is exclusively
+  Fetches and puts a `key` of a given type, assuming that the descr is exclusively
+  a map (or dynamic).
+  """
+  def map_fetch_and_put(:term, _key, _type), do: :badmap
+
+  def map_fetch_and_put(descr, key, :term) when is_atom(key),
+    do: map_fetch_and_put_shared(descr, key, :term)
+
+  def map_fetch_and_put(descr, key, type) when is_atom(key) do
+    case :maps.take(:dynamic, type) do
+      :error -> map_fetch_and_put_shared(descr, key, type)
+      {dynamic, _static} -> map_fetch_and_put_shared(dynamic(descr), key, dynamic)
+    end
+  end
+
+  defp map_fetch_and_put_shared(descr, key, type) do
+    with {value, descr} <- map_take(descr, key, none(), &map_put_static(&1, key, type)) do
+      if empty?(value) do
+        :badkey
+      else
+        {value, descr}
+      end
+    end
+  end
+
+  @doc """
+  Puts a `key` of a given type, assuming that the descr is exclusively
   a map (or dynamic).
   """
   def map_put(:term, _key, _type), do: :badmap
-  def map_put(descr, key, :term) when is_atom(key), do: map_put_static_value(descr, key, :term)
+  def map_put(descr, key, :term) when is_atom(key), do: map_put_shared(descr, key, :term)
 
   def map_put(descr, key, type) when is_atom(key) do
     case :maps.take(:dynamic, type) do
-      :error -> map_put_static_value(descr, key, type)
-      {dynamic, _static} -> map_put_static_value(dynamic(descr), key, dynamic)
+      :error -> map_put_shared(descr, key, type)
+      {dynamic, _static} -> map_put_shared(dynamic(descr), key, dynamic)
     end
   end
 
-  defp map_put_static_value(descr, key, type) do
+  defp map_put_shared(descr, key, type) do
     with {_value, descr} <- map_take(descr, key, :term, &map_put_static(&1, key, type)) do
       {:ok, descr}
     end

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

@@ -135,7 +135,8 @@ defmodule Module.Types.Expr do
       if fallback == none() do
         Enum.reduce(pairs, map_type, fn {key, type}, acc ->
           case map_put(acc, key, type) do
-            {:ok, descr} -> descr
+            {_value, descr} -> descr
+            :badkey -> throw({:badkey, map_type, key, expr, context})
             :badmap -> throw({:badmap, map_type, expr, context})
           end
         end)

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

@@ -852,13 +852,12 @@ defmodule Module.Types.DescrTest do
       assert map_fetch(term(), :a) == :badmap
       assert map_fetch(union(open_map(), integer()), :a) == :badmap
 
-      assert map_fetch(closed_map(a: integer()), :a) == {false, integer()}
-
-      assert map_fetch(union(closed_map(a: integer()), closed_map(b: atom())), :a) ==
-               :badkey
+      assert map_fetch(open_map(), :a) == :badkey
+      assert map_fetch(open_map(a: not_set()), :a) == :badkey
+      assert map_fetch(union(closed_map(a: integer()), closed_map(b: atom())), :a) == :badkey
+      assert map_fetch(difference(closed_map(a: integer()), closed_map(a: term())), :a) == :badkey
 
-      assert map_fetch(difference(closed_map(a: integer()), closed_map(a: term())), :a) ==
-               :badkey
+      assert map_fetch(closed_map(a: integer()), :a) == {false, integer()}
 
       assert map_fetch(union(closed_map(a: integer()), closed_map(a: atom())), :a) ==
                {false, union(integer(), atom())}
@@ -1043,6 +1042,11 @@ defmodule Module.Types.DescrTest do
       assert equal?(type, open_map(a: not_set()))
     end
 
+    test "map_fetch_and_put" do
+      assert map_fetch_and_put(term(), :a, integer()) == :badmap
+      assert map_fetch_and_put(open_map(), :a, integer()) == :badkey
+    end
+
     test "map_put" do
       assert map_put(term(), :a, integer()) == :badmap
       assert map_put(integer(), :a, integer()) == :badmap

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

@@ -561,6 +561,27 @@ defmodule Module.Types.ExprTest do
                  # from: types_test.ex:LINE
                  x = :foo
              """
+
+      assert typeerror!(
+               (
+                 x = %{}
+                 %{x | x: :zero}
+               )
+             ) == ~l"""
+             expected a map within map update syntax:
+
+                 %{x | x: :zero}
+
+             but got type:
+
+                 dynamic(:foo)
+
+             where "x" was given the type:
+
+                 # type: dynamic(:foo)
+                 # from: types_test.ex:LINE
+                 x = :foo
+             """
     end
 
     test "updating to open maps" do