Type check Map.put/3 and Map.delete/2

Author: josevalim

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

@@ -244,12 +244,14 @@ defmodule Module.Types.Apply do
         {:erlang, :tuple_to_list, [{[open_tuple([])], dynamic(list(term()))}]},
 
         ## Map
+        {:maps, :remove, [{[term(), open_map()], open_map()}]},
         {:maps, :from_keys, [{[list(term()), term()], open_map()}]},
         {:maps, :find,
          [{[term(), open_map()], tuple([atom([:ok]), term()]) |> union(atom([:error]))}]},
         {:maps, :get, [{[term(), open_map()], term()}]},
         {:maps, :is_key, [{[term(), open_map()], boolean()}]},
         {:maps, :keys, [{[open_map()], dynamic(list(term()))}]},
+        {:maps, :put, [{[term(), term(), open_map()], open_map()}]},
         {:maps, :take,
          [{[term(), open_map()], tuple([term(), open_map()]) |> union(atom([:error]))}]},
         {:maps, :to_list, [{[open_map()], dynamic(list(tuple([term(), term()])))}]},
@@ -425,17 +427,11 @@ defmodule Module.Types.Apply do
     end
   end
 
-  defp remote_apply(:maps, :keys, _info, [map], stack) do
-    case map_to_list(map, fn key, _value -> key end) do
-      {:ok, list_type} -> {:ok, return(list_type, [map], stack)}
-      :badmap -> {:error, badremote(:maps, :keys, 1)}
-    end
-  end
-
-  defp remote_apply(:maps, :values, _info, [map], stack) do
-    case map_to_list(map, fn _key, value -> value end) do
-      {:ok, list_type} -> {:ok, return(list_type, [map], stack)}
-      :badmap -> {:error, badremote(:maps, :keys, 1)}
+  defp remote_apply(:maps, :remove, _info, [key, map] = args_types, stack) do
+    case map_update(map, key, not_set(), false, true) do
+      {_value, descr, _errors} -> {:ok, return(descr, args_types, stack)}
+      :badmap -> {:error, badremote(:maps, :remove, 2)}
+      {:error, _errors} -> {:error, {:badkeydomain, map, key, nil}}
     end
   end
 
@@ -461,10 +457,17 @@ defmodule Module.Types.Apply do
     end
   end
 
-  defp remote_apply(:maps, :update, _info, [key, value, map] = args_types, stack) do
-    case map_update(map, key, value, false, false) do
+  defp remote_apply(:maps, :keys, _info, [map], stack) do
+    case map_to_list(map, fn key, _value -> key end) do
+      {:ok, list_type} -> {:ok, return(list_type, [map], stack)}
+      :badmap -> {:error, badremote(:maps, :keys, 1)}
+    end
+  end
+
+  defp remote_apply(:maps, :put, _info, [key, value, map] = args_types, stack) do
+    case map_update(map, key, value, false, true) do
       {_value, descr, _errors} -> {:ok, return(descr, args_types, stack)}
-      :badmap -> {:error, badremote(:maps, :update, 3)}
+      :badmap -> {:error, badremote(:maps, :put, 3)}
       {:error, _errors} -> {:error, {:badkeydomain, map, key, nil}}
     end
   end
@@ -500,6 +503,21 @@ defmodule Module.Types.Apply do
     end
   end
 
+  defp remote_apply(:maps, :update, _info, [key, value, map] = args_types, stack) do
+    case map_update(map, key, value, false, false) do
+      {_value, descr, _errors} -> {:ok, return(descr, args_types, stack)}
+      :badmap -> {:error, badremote(:maps, :update, 3)}
+      {:error, _errors} -> {:error, {:badkeydomain, map, key, nil}}
+    end
+  end
+
+  defp remote_apply(:maps, :values, _info, [map], stack) do
+    case map_to_list(map, fn _key, value -> value end) do
+      {:ok, list_type} -> {:ok, return(list_type, [map], stack)}
+      :badmap -> {:error, badremote(:maps, :keys, 1)}
+    end
+  end
+
   defp remote_apply(_mod, _fun, info, args_types, stack) do
     remote_apply(info, args_types, stack)
   end

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

@@ -80,6 +80,54 @@ defmodule Module.Types.MapTest do
     end
   end
 
+  describe "Map.delete/2" do
+    test "checking" do
+      assert typecheck!(Map.delete(%{}, :key)) ==
+               closed_map(key: not_set())
+
+      assert typecheck!(Map.delete(%{key: 123}, :key)) ==
+               closed_map(key: not_set())
+
+      assert typecheck!([x], Map.delete(x, :key)) ==
+               dynamic(open_map(key: not_set()))
+
+      # If one of them succeeds, we are still fine!
+      assert typecheck!(
+               [condition?],
+               Map.delete(%{foo: 123}, if(condition?, do: :foo, else: :bar))
+             ) ==
+               union(
+                 closed_map(foo: not_set()),
+                 closed_map(foo: integer(), bar: not_set())
+               )
+
+      assert typecheck!([x], Map.delete(x, 123)) == dynamic(open_map())
+    end
+
+    test "inference" do
+      assert typecheck!(
+               [x],
+               (
+                 _ = Map.delete(x, :key)
+                 x
+               )
+             ) == dynamic(open_map())
+    end
+
+    test "errors" do
+      assert typeerror!([x = []], Map.delete(x, :key)) =~
+               "incompatible types given to Map.delete/2"
+    end
+
+    test "combined with put" do
+      assert typecheck!([x], x |> Map.delete(:key) |> Map.put(:key, "123")) ==
+               dynamic(open_map(key: binary()))
+
+      assert typecheck!([x, y], x |> Map.delete(:key) |> Map.put(String.to_atom(y), "123")) ==
+               dynamic(open_map(key: if_set(binary())))
+    end
+  end
+
   describe "Map.fetch/2" do
     test "checking" do
       assert typecheck!(Map.fetch(%{key: 123}, :key)) ==
@@ -190,6 +238,46 @@ defmodule Module.Types.MapTest do
     end
   end
 
+  describe "Map.put/3" do
+    test "checking" do
+      assert typecheck!(Map.put(%{}, :key, :value)) ==
+               closed_map(key: atom([:value]))
+
+      assert typecheck!(Map.put(%{key: 123}, :key, :value)) ==
+               closed_map(key: atom([:value]))
+
+      assert typecheck!([x], Map.put(x, :key, :value)) ==
+               dynamic(open_map(key: atom([:value])))
+
+      # If one of them succeeds, we are still fine!
+      assert typecheck!(
+               [condition?],
+               Map.put(%{foo: 123}, if(condition?, do: :foo, else: :bar), "123")
+             ) ==
+               union(
+                 closed_map(foo: binary()),
+                 closed_map(foo: integer(), bar: binary())
+               )
+
+      assert typecheck!([x], Map.put(x, 123, 456)) == dynamic(open_map())
+    end
+
+    test "inference" do
+      assert typecheck!(
+               [x],
+               (
+                 _ = Map.put(x, :key, :value)
+                 x
+               )
+             ) == dynamic(open_map())
+    end
+
+    test "errors" do
+      assert typeerror!([x = []], Map.put(x, :key, :value)) =~
+               "incompatible types given to Map.put/3"
+    end
+  end
+
   describe "Map.replace!/3" do
     test "checking" do
       assert typecheck!(Map.replace!(%{key: 123}, :key, :value)) ==
@@ -201,8 +289,8 @@ defmodule Module.Types.MapTest do
       # If one of them succeeds, we are still fine!
       assert typecheck!(
                [condition?],
-               Map.replace!(%{foo: 123}, if(condition?, do: :foo, else: :bar), :value)
-             ) == closed_map(foo: atom([:value]))
+               Map.replace!(%{foo: 123}, if(condition?, do: :foo, else: :bar), "123")
+             ) == closed_map(foo: binary())
 
       assert typecheck!([x], Map.replace!(x, 123, 456)) == dynamic(open_map())
     end