Fix inferred map requirements for mixed update keys

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

@@ -182,11 +182,14 @@ defmodule Module.Types.Expr do
       end)
 
     # The only information we can attach to the expected types is that
-    # certain keys are expected.
+    # certain keys are expected. The whole key type must be a singleton atom,
+    # as atom_fetch/1 may return only the atom subset of a gradual type.
     expected_pairs =
       Enum.flat_map(pairs_types, fn {key_type, _value_type} ->
-        case atom_fetch(key_type) do
-          {:finite, [key]} -> [{key, {term(), false}}]
+        with true <- subtype?(key_type, atom()),
+             {:finite, [key]} <- atom_fetch(key_type) do
+          [{key, {term(), false}}]
+        else
           _ -> []
         end
       end)

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

@@ -292,6 +292,35 @@ defmodule Module.Types.InferTest do
              )
   end
 
+  test "from map updates with mixed atom and non-atom keys", config do
+    types =
+      infer config do
+        def replace(map, flag) do
+          key = if flag, do: :a, else: 1
+          %{map | key => :new}
+        end
+      end
+
+    assert {:infer, _, [{[map, flag], _}]} = types[{:replace, 2}]
+    assert map == open_map()
+    assert equal?(flag, term())
+  end
+
+  test "from map updates with singleton atom keys", config do
+    types =
+      infer config do
+        def replace(map) do
+          key = key()
+          %{map | key => :new}
+        end
+
+        defp key, do: :a
+      end
+
+    assert {:infer, _, [{[map], _}]} = types[{:replace, 1}]
+    assert map == open_map(a: {term(), false})
+  end
+
   test "from captures", config do
     types =
       infer config do

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

@@ -209,6 +209,23 @@ defmodule Module.Types.IntegrationTest do
   end
 
   describe "type checking" do
+    test "map updates with mixed atom and non-atom keys" do
+      files = %{
+        "a.ex" => """
+        defmodule A do
+          def replace(map, flag) do
+            key = if flag, do: :a, else: 1
+            %{map | key => :new}
+          end
+
+          def run, do: replace(%{1 => :old}, false)
+        end
+        """
+      }
+
+      assert_no_warnings(files)
+    end
+
     test "inferred remote calls" do
       files = %{
         "a.ex" => """