Type Map.put_new/3 and Map.put_new_lazy/3

josevalim · josevalim · commit 0f3fdcd84ff5 · 2025-12-15T11:03:52.000+01:00
diff --git a/lib/elixir/lib/module/types/apply.ex b/lib/elixir/lib/module/types/apply.ex
@@ -251,6 +251,8 @@ defmodule Module.Types.Apply do
         {Map, :pop, [{[open_map(), term()], tuple([term(), open_map()])}]},
         {Map, :pop, [{[open_map(), term(), term()], tuple([term(), open_map()])}]},
         {Map, :pop!, [{[open_map(), term()], tuple([term(), open_map()])}]},
+        {Map, :put_new, [{[open_map(), term(), term()], open_map()}]},
+        {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(), fun(1)], open_map()}]},
@@ -421,7 +423,7 @@ defmodule Module.Types.Apply do
     case map_update(map, @struct_key, not_set(), false, true) do
       {_value, descr, _errors} -> {:ok, return(descr, args_types, stack)}
       :badmap -> {:error, badremote(Map, :from_struct, args_types)}
-      {:error, _errors} -> {:error, {:badkeydomain, map, @struct_key, "raise"}}
+      {:error, _errors} -> {:ok, map}
     end
   end
 
@@ -455,6 +457,17 @@ defmodule Module.Types.Apply do
     end
   end
 
+  defp remote_apply(Map, :put_new, _info, [map, key, value] = args_types, stack) do
+    map_put_new(map, key, value, :put_new, args_types, stack)
+  end
+
+  defp remote_apply(Map, :put_new_lazy, _info, [map, key, fun] = args_types, stack) do
+    case fun_apply(fun, []) do
+      {:ok, value} -> map_put_new(map, key, value, :put_new_lazy, args_types, stack)
+      reason -> {:error, {:badapply, fun, [], reason}}
+    end
+  end
+
   defp remote_apply(Map, :pop, _info, args_types, stack) do
     [map, key, default] =
       case args_types do
@@ -566,15 +579,15 @@ defmodule Module.Types.Apply do
     case map_update(map, key, value, false, true) do
       {_value, descr, _errors} -> {:ok, return(descr, args_types, stack)}
       :badmap -> {:error, badremote(:maps, :put, args_types)}
-      {:error, _errors} -> {:error, {:badkeydomain, map, key, "raise"}}
+      {:error, _errors} -> {:ok, map}
     end
   end
 
   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, args_types)}
-      {:error, _errors} -> {:error, {:badkeydomain, map, key, "raise"}}
+      {:error, _errors} -> {:ok, map}
     end
   end
 
@@ -969,6 +982,19 @@ defmodule Module.Types.Apply do
 
   ## Map helpers
 
+  defp map_put_new(map, key, value, name, args_types, stack) do
+    fun = fn
+      true, type -> union(type, value)
+      false, type -> if empty?(type), do: value, else: type
+    end
+
+    case map_update_fun(map, key, fun, false, true) do
+      {_value, descr, _errors} -> {:ok, return(descr, args_types, stack)}
+      :badmap -> {:error, badremote(Map, name, args_types)}
+      {:error, _errors} -> {:ok, map}
+    end
+  end
+
   def map_update_or_replace_lazy(name, [map, key, fun] = args_types, stack, error) do
     try do
       {map, fun} =
@@ -977,7 +1003,7 @@ defmodule Module.Types.Apply do
           _ -> {map, fun}
         end
 
-      fun_apply = fn arg_type ->
+      fun_apply = fn _optional?, arg_type ->
         case fun_apply(fun, [arg_type]) do
           {:ok, res} -> res
           reason -> throw({:badapply, reason, [arg_type]})
diff --git a/lib/elixir/lib/module/types/descr.ex b/lib/elixir/lib/module/types/descr.ex
@@ -2999,10 +2999,12 @@ defmodule Module.Types.Descr do
   def map_update(descr, key_descr, type, return_type?, force?) do
     case type do
       %{dynamic: dynamic} ->
-        map_update_unchecked(dynamic(descr), key_descr, fn _ -> dynamic end, return_type?, force?)
+        fun = fn _, _ -> dynamic end
+        map_update_unchecked(dynamic(descr), key_descr, fun, return_type?, force?)
 
       %{} ->
-        map_update_unchecked(descr, key_descr, fn _ -> type end, return_type?, force?)
+        fun = fn _, _ -> type end
+        map_update_unchecked(descr, key_descr, fun, return_type?, force?)
     end
   end
 
@@ -3020,12 +3022,19 @@ defmodule Module.Types.Descr do
   def map_update_fun(descr, key_descr, type_fun, return_type? \\ true, force? \\ false) do
     gradual? = gradual?(descr)
 
-    type_fun = fn value ->
-      value = remove_optional(value)
+    type_fun = fn optional?, value ->
+      if is_function(type_fun, 1) do
+        case type_fun.(if gradual?, do: dynamic(value), else: value) do
+          %{dynamic: dynamic} -> dynamic
+          descr -> descr
+        end
+      else
+        value = if gradual?, do: dynamic(value), else: value
 
-      case type_fun.(if gradual?, do: dynamic(value), else: value) do
-        %{dynamic: dynamic} -> dynamic
-        descr -> descr
+        case type_fun.(optional?, value) do
+          %{dynamic: dynamic} -> dynamic
+          descr -> descr
+        end
       end
     end
 
@@ -3176,7 +3185,7 @@ defmodule Module.Types.Descr do
         {acc_value, acc_descr, acc_errors, acc_found?}
       else
         acc_value = union(value, acc_value)
-        acc_descr = union(map_put_key_static(descr, key, type_fun.(value)), acc_descr)
+        acc_descr = union(map_put_key_static(descr, key, type_fun.(optional?, value)), acc_descr)
 
         # The field will be missing if we are not forcing,
         # we are in static mode and the value is optional.
@@ -3253,7 +3262,8 @@ defmodule Module.Types.Descr do
           if :sets.is_element(key, negated) do
             {key, value}
           else
-            {key, union(value, type_fun.(value))}
+            {optional?, call_value} = pop_optional_static(value)
+            {key, union(value, type_fun.(optional?, call_value))}
           end
         end)
 
@@ -3364,13 +3374,16 @@ defmodule Module.Types.Descr do
         :open
 
       :closed ->
-        Map.from_keys(domain_keys, if_set(type_fun.(none())))
+        Map.from_keys(domain_keys, if_set(type_fun.(true, none())))
 
       domains = %{} ->
         Enum.reduce(domain_keys, domains, fn domain_key, acc ->
           case acc do
-            %{^domain_key => value} -> %{acc | domain_key => union(value, type_fun.(value))}
-            %{} -> Map.put(acc, domain_key, if_set(type_fun.(none())))
+            %{^domain_key => value} ->
+              %{acc | domain_key => union(value, type_fun.(true, remove_optional(value)))}
+
+            %{} ->
+              Map.put(acc, domain_key, if_set(type_fun.(true, none())))
           end
         end)
     end
@@ -3438,7 +3451,7 @@ defmodule Module.Types.Descr do
     {required_keys, optional_keys, maybe_negated_set, required_domains, optional_domains} =
       split_keys
 
-    type_fun = fn _ -> type end
+    type_fun = fn _, _ -> type end
     bdd = map_update_put_negated(bdd, maybe_negated_set, type_fun)
 
     descr =
diff --git a/lib/elixir/test/elixir/module/types/map_test.exs b/lib/elixir/test/elixir/module/types/map_test.exs
@@ -756,6 +756,119 @@ defmodule Module.Types.MapTest do
     end
   end
 
+  describe "Map.put_new_lazy/3" do
+    test "checking" do
+      assert typecheck!(Map.put_new_lazy(%{}, :key, fn -> :value end)) ==
+               closed_map(key: atom([:value]))
+
+      assert typecheck!(Map.put_new_lazy(%{key: 123}, :key, fn -> :value end)) ==
+               closed_map(key: integer())
+
+      assert typecheck!([x], Map.put_new_lazy(x, :key, fn -> :value end)) ==
+               dynamic(open_map(key: term()))
+
+      # If one of them succeeds, we are still fine!
+      assert typecheck!(
+               [condition?],
+               Map.put_new_lazy(%{foo: 123}, if(condition?, do: :foo, else: :bar), fn -> "123" end)
+             ) == union(closed_map(foo: integer()), closed_map(foo: integer(), bar: binary()))
+
+      assert typecheck!([], Map.put_new_lazy(%{789 => "binary"}, 123, fn -> 456 end)) ==
+               closed_map([{domain_key(:integer), union(binary(), integer())}])
+
+      assert typecheck!([x], Map.put_new_lazy(x, 123, fn -> 456 end)) == dynamic(open_map())
+    end
+
+    test "inference" do
+      assert typecheck!(
+               [x],
+               (
+                 _ = Map.put_new_lazy(x, :key, fn -> :value end)
+                 x
+               )
+             ) == dynamic(open_map())
+    end
+
+    test "errors" do
+      assert typeerror!([x = []], Map.put_new_lazy(x, :key, fn -> :value end)) |> strip_ansi() =~
+               """
+               incompatible types given to Map.put_new_lazy/3:
+
+                   Map.put_new_lazy(x, :key, fn -> :value end)
+
+               given types:
+
+                   empty_list(), :key, (-> dynamic(:value))
+
+               but expected one of:
+
+                   map(), term(), (-> term())
+               """
+
+      assert typeerror!(Map.put_new_lazy(%{}, :foo, 123)) =~
+               """
+               expected a 0-arity function on function call within Map.put_new_lazy/3:
+
+                   Map.put_new_lazy(%{}, :foo, 123)
+
+               but got type:
+
+                   integer()
+               """
+    end
+  end
+
+  describe "Map.put_new/3" do
+    test "checking" do
+      assert typecheck!(Map.put_new(%{}, :key, :value)) ==
+               closed_map(key: atom([:value]))
+
+      assert typecheck!(Map.put_new(%{key: 123}, :key, :value)) ==
+               closed_map(key: integer())
+
+      assert typecheck!([x], Map.put_new(x, :key, :value)) ==
+               dynamic(open_map(key: term()))
+
+      # If one of them succeeds, we are still fine!
+      assert typecheck!(
+               [condition?],
+               Map.put_new(%{foo: 123}, if(condition?, do: :foo, else: :bar), "123")
+             ) == union(closed_map(foo: integer()), closed_map(foo: integer(), bar: binary()))
+
+      assert typecheck!([], Map.put_new(%{789 => "binary"}, 123, 456)) ==
+               closed_map([{domain_key(:integer), union(binary(), integer())}])
+
+      assert typecheck!([x], Map.put_new(x, 123, 456)) == dynamic(open_map())
+    end
+
+    test "inference" do
+      assert typecheck!(
+               [x],
+               (
+                 _ = Map.put_new(x, :key, :value)
+                 x
+               )
+             ) == dynamic(open_map())
+    end
+
+    test "errors" do
+      assert typeerror!([x = []], Map.put_new(x, :key, :value)) |> strip_ansi() =~
+               """
+               incompatible types given to Map.put_new/3:
+
+                   Map.put_new(x, :key, :value)
+
+               given types:
+
+                   empty_list(), :key, :value
+
+               but expected one of:
+
+                   map(), term(), term()
+               """
+    end
+  end
+
   describe "Map.replace/3" do
     test "checking" do
       assert typecheck!(Map.replace(%{key: 123}, :key, :value)) ==
