Type Map.update/4

Author: josevalim

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

@@ -256,6 +256,7 @@ defmodule Module.Types.Apply do
         {Map, :put_new_lazy, [{[open_map(), term(), fun(0)], open_map()}]},
         {Map, :replace, [{[open_map(), term(), term()], open_map()}]},
         {Map, :replace_lazy, [{[open_map(), term(), fun(1)], open_map()}]},
+        {Map, :update, [{[open_map(), term(), term(), fun(1)], open_map()}]},
         {Map, :update!, [{[open_map(), term(), fun(1)], open_map()}]},
         {:maps, :from_keys, [{[list(term()), term()], open_map()}]},
         {:maps, :find,
@@ -529,6 +530,42 @@ defmodule Module.Types.Apply do
     map_update_or_replace_lazy(:replace_lazy, args_types, stack, "do nothing")
   end
 
+  defp remote_apply(Map, :update, _info, [map, key, default, fun] = args_types, stack) do
+    try do
+      {map, default} =
+        case default do
+          %{dynamic: default} -> {dynamic(map), default}
+          _ -> {map, default}
+        end
+
+      {map, fun} =
+        case fun do
+          %{dynamic: fun} -> {dynamic(map), fun}
+          _ -> {map, fun}
+        end
+
+      fun_apply = fn optional?, arg_type ->
+        if empty?(arg_type) do
+          default
+        else
+          case fun_apply(fun, [arg_type]) do
+            {:ok, res} -> if optional?, do: union(res, default), else: res
+            reason -> throw({:badapply, reason, [arg_type]})
+          end
+        end
+      end
+
+      map_update_fun(map, key, fun_apply, false, true)
+    catch
+      {:badapply, reason, args_types} ->
+        {:error, {:badapply, fun, args_types, reason}}
+    else
+      {_value, descr, _errors} -> {:ok, return(descr, args_types, stack)}
+      :badmap -> {:error, badremote(Map, :update, args_types)}
+      {:error, _errors} -> {:ok, map}
+    end
+  end
+
   defp remote_apply(Map, :update!, _info, args_types, stack) do
     map_update_or_replace_lazy(:update!, args_types, stack, "raise")
   end

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

@@ -1186,6 +1186,96 @@ defmodule Module.Types.MapTest do
     end
   end
 
+  describe "Map.update/4" do
+    test "checking" do
+      assert typecheck!(Map.update(%{}, :key, :default, fn _ -> :value end)) ==
+               dynamic(closed_map(key: atom([:default])))
+
+      assert typecheck!(Map.update(%{key: 123}, :key, :default, fn _ -> :value end)) ==
+               dynamic(closed_map(key: atom([:value])))
+
+      assert typecheck!([x], Map.update(x, :key, :default, fn _ -> :value end)) ==
+               dynamic(open_map(key: atom([:value, :default])))
+
+      # If one of them succeeds, we are still fine!
+      assert typecheck!(
+               [condition?],
+               Map.update(%{foo: 123}, if(condition?, do: :foo, else: :bar), :default, fn _ ->
+                 "123"
+               end)
+             ) ==
+               dynamic(
+                 union(
+                   closed_map(foo: binary()),
+                   closed_map(foo: integer(), bar: atom([:default]))
+                 )
+               )
+
+      # Both succeed but different clauses
+      assert typecheck!(
+               [condition?],
+               Map.update(
+                 %{key1: :foo, key2: :bar},
+                 if(condition?, do: :key1, else: :key2),
+                 :default,
+                 fn
+                   :foo -> 123
+                   :bar -> 123.0
+                 end
+               )
+             ) ==
+               dynamic(
+                 union(
+                   closed_map(key1: atom([:foo]), key2: float()),
+                   closed_map(key1: integer(), key2: atom([:bar]))
+                 )
+               )
+
+      assert typecheck!([x], Map.update(x, 123, :default, fn _ -> 456 end)) == dynamic(open_map())
+
+      integer_to_integer_float_atom =
+        dynamic(
+          closed_map([
+            {domain_key(:integer), integer() |> union(float()) |> union(atom([:default]))}
+          ])
+        )
+
+      assert typecheck!([], Map.update(%{123 => 456}, 123, :default, fn x -> x * 1.0 end)) ==
+               integer_to_integer_float_atom
+
+      assert typecheck!([], Map.update(%{123 => 456}, 456, :default, fn x -> x * 1.0 end)) ==
+               integer_to_integer_float_atom
+    end
+
+    test "inference" do
+      assert typecheck!(
+               [x],
+               (
+                 _ = Map.update(x, :key, :default, fn _ -> :value end)
+                 x
+               )
+             ) == dynamic(open_map())
+    end
+
+    test "errors" do
+      assert typeerror!(Map.update(%{key: :foo}, :key, :default, fn :bar -> :value end))
+             |> strip_ansi() =~
+               """
+               incompatible types given on function call within Map.update/4:
+
+                   Map.update(%{key: :foo}, :key, :default, fn :bar -> :value end)
+
+               given types:
+
+                   dynamic(:foo)
+
+               but function has type:
+
+                   (:bar -> dynamic(:value))
+               """
+    end
+  end
+
   describe "Map.update!/3" do
     test "checking" do
       assert typecheck!(Map.update!(%{key: 123}, :key, fn _ -> :value end)) ==
@@ -1219,6 +1309,9 @@ defmodule Module.Types.MapTest do
 
       assert typecheck!([], Map.update!(%{123 => 456}, 123, fn x -> x * 1.0 end)) ==
                dynamic(closed_map([{domain_key(:integer), union(integer(), float())}]))
+
+      assert typecheck!([], Map.update!(%{123 => 456}, 456, fn x -> x * 1.0 end)) ==
+               dynamic(closed_map([{domain_key(:integer), union(integer(), float())}]))
     end
 
     test "inference" do