Return if the key is optional on map_get

Author: josevalim

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

@@ -568,61 +568,56 @@ defmodule Module.Types.Descr do
   defp optional?(%{bitmap: bitmap}) when (bitmap &&& @bit_optional) != 0, do: true
   defp optional?(_), do: false
 
-  defp remove_not_set(type) do
+  defp pop_not_set(type) do
     case type do
-      %{bitmap: @bit_optional} -> Map.delete(type, :bitmap)
-      %{bitmap: bitmap} -> %{type | bitmap: bitmap &&& ~~~@bit_optional}
-      _ -> type
+      %{bitmap: @bit_optional} ->
+        {true, Map.delete(type, :bitmap)}
+
+      %{bitmap: bitmap} when (bitmap &&& @bit_optional) != 0 ->
+        {true, %{type | bitmap: bitmap - @bit_optional}}
+
+      _ ->
+        {false, type}
     end
   end
 
   defp map_tag_to_type(:open), do: term_or_not_set()
   defp map_tag_to_type(:closed), do: not_set()
+
   defp map_descr(tag, fields), do: %{map: [{tag, fields, []}]}
 
   @doc """
   Gets the type of the value returned by accessing `key` on `map`.
-  Does not guarantee the key exists. To do that, use `map_has_key?`.
-  """
-  def map_get!(%{} = descr, key) do
-    if not gradual?(descr) do
-      map_get_static(descr, key)
-    else
-      {dynamic, static} = Map.pop(descr, :dynamic)
-      dynamic_value_type = map_get_static(dynamic, key)
-      static_value_type = map_get_static(static, key)
-      union(intersection(dynamic(), dynamic_value_type), static_value_type)
-    end
-  end
 
-  @doc """
-  Check that a key is present.
+  It returns a two element tuple. The first element says if the type
+  is optional or not, the second element is the type. In static mode,
+  we likely want to raise if `map.field` (or pattern matching?) is
+  called on an optional key.
   """
-  def map_has_key?(descr, key) do
-    subtype?(descr, open_map([{key, term()}]))
-  end
+  def map_get(%{} = descr, key) do
+    case :maps.take(:dynamic, descr) do
+      :error ->
+        map_get_static(descr, key)
 
-  # Assuming `descr` is a static type. Accessing `key` will, if it succeeds,
-  # return a value of the type returned. To guarantee that the key is always
-  # present, use `map_has_key?`. To guarantee that the key may be present
-  # use `map_may_have_key?`. If key is never present, result will be `none()`.
-  defp map_get_static(descr, key) when is_atom(key) do
-    descr_map = intersection(descr, open_map())
+      {dynamic, static} ->
+        {dynamic_optional?, dynamic_type} = map_get_static(dynamic, key)
+        {static_optional?, static_type} = map_get_static(static, key)
 
-    if empty?(descr_map) do
-      none()
-    else
-      map_split_on_key(descr_map.map, key)
-      |> Enum.reduce(none(), fn typeof_key, union -> union(typeof_key, union) end)
-      |> remove_not_set()
+        {dynamic_optional? or static_optional?,
+         union(intersection(dynamic(), dynamic_type), static_type)}
     end
   end
 
-  @doc """
-  Check that a key may be present.
-  """
-  def map_may_have_key?(descr, key) do
-    compatible?(descr, open_map([{key, term()}]))
+  defp map_get_static(descr, key) when is_atom(key) do
+    case descr do
+      %{map: map} ->
+        map_split_on_key(map, key)
+        |> Enum.reduce(none(), &union/2)
+        |> pop_not_set()
+
+      %{} ->
+        {false, none()}
+    end
   end
 
   # Union is list concatenation

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

@@ -246,95 +246,77 @@ defmodule Module.Types.DescrTest do
 
   describe "map operations" do
     test "get field" do
-      assert map_get!(closed_map(a: integer()), :a) == integer()
-      assert map_get!(dynamic(), :a) == dynamic()
+      assert map_get(closed_map(a: integer()), :a) == {false, integer()}
+      assert map_get(dynamic(), :a) == {true, dynamic()}
 
       assert intersection(dynamic(), open_map(a: integer()))
-             |> map_get!(:a) == intersection(integer(), dynamic())
+             |> map_get(:a) == {false, intersection(integer(), dynamic())}
 
-      assert open_map(my_map: open_map(foo: integer()))
-             |> intersection(open_map(my_map: open_map(bar: boolean())))
-             |> map_get!(:my_map)
-             |> equal?(open_map(foo: integer(), bar: boolean()))
+      {false, value_type} =
+        open_map(my_map: open_map(foo: integer()))
+        |> intersection(open_map(my_map: open_map(bar: boolean())))
+        |> map_get(:my_map)
 
-      assert map_get!(union(closed_map(a: integer()), closed_map(a: atom())), :a) ==
-               union(integer(), atom())
+      assert equal?(value_type, open_map(foo: integer(), bar: boolean()))
 
-      assert map_get!(union(closed_map(a: integer()), closed_map(b: atom())), :a) == integer()
-      assert map_get!(term(), :a) == term()
+      assert map_get(union(closed_map(a: integer()), closed_map(a: atom())), :a) ==
+               {false, union(integer(), atom())}
 
-      assert closed_map(a: union(integer(), atom()))
-             |> difference(open_map(a: integer()))
-             |> map_get!(:a)
-             |> equal?(atom())
+      assert map_get(union(closed_map(a: integer()), closed_map(b: atom())), :a) ==
+               {true, integer()}
 
-      assert closed_map(a: integer(), b: atom())
-             |> difference(closed_map(a: integer(), b: atom([:foo])))
-             |> map_get!(:a)
-             |> equal?(integer())
+      assert map_get(term(), :a) == {true, term()}
 
-      assert closed_map(a: integer())
-             |> difference(closed_map(a: atom()))
-             |> map_get!(:a)
-             |> equal?(integer())
+      {false, value_type} =
+        closed_map(a: union(integer(), atom()))
+        |> difference(open_map(a: integer()))
+        |> map_get(:a)
 
-      assert open_map(a: integer(), b: atom())
-             |> union(closed_map(a: tuple()))
-             |> map_get!(:a)
-             |> equal?(union(integer(), tuple()))
+      assert equal?(value_type, atom())
+
+      {false, value_type} =
+        closed_map(a: integer(), b: atom())
+        |> difference(closed_map(a: integer(), b: atom([:foo])))
+        |> map_get(:a)
+
+      assert equal?(value_type, integer())
+
+      {false, value_type} =
+        closed_map(a: integer())
+        |> difference(closed_map(a: atom()))
+        |> map_get(:a)
+
+      assert equal?(value_type, integer())
+
+      {false, value_type} =
+        open_map(a: integer(), b: atom())
+        |> union(closed_map(a: tuple()))
+        |> map_get(:a)
+
+      assert equal?(value_type, union(integer(), tuple()))
+
+      {false, value_type} =
+        closed_map(a: atom())
+        |> difference(closed_map(a: atom([:foo, :bar])))
+        |> difference(closed_map(a: atom([:bar])))
+        |> map_get(:a)
 
-      assert closed_map(a: atom())
-             |> difference(closed_map(a: atom([:foo, :bar])))
-             |> difference(closed_map(a: atom([:bar])))
-             |> map_get!(:a)
-             |> equal?(intersection(atom(), negation(atom([:foo, :bar]))))
+      assert equal?(value_type, intersection(atom(), negation(atom([:foo, :bar]))))
 
       assert closed_map(a: union(atom(), pid()), b: integer(), c: tuple())
              |> difference(open_map(a: atom(), b: integer()))
              |> difference(open_map(a: atom(), c: tuple()))
-             |> map_get!(:a) == pid()
+             |> map_get(:a) == {false, pid()}
 
       assert closed_map(a: union(atom([:foo]), pid()), b: integer(), c: tuple())
              |> difference(open_map(a: atom([:foo]), b: integer()))
              |> difference(open_map(a: atom(), c: tuple()))
-             |> map_get!(:a) == pid()
+             |> map_get(:a) == {false, pid()}
 
       assert closed_map(a: union(atom([:foo, :bar, :baz]), integer()))
              |> difference(open_map(a: atom([:foo, :bar])))
              |> difference(open_map(a: atom([:foo, :baz])))
-             |> map_get!(:a) == integer()
-    end
-
-    test "key presence" do
-      assert map_has_key?(closed_map(a: integer()), :a)
-      refute map_has_key?(closed_map(a: integer()), :b)
-      refute map_has_key?(open_map(), :a)
-      refute map_has_key?(closed_map(a: union(integer(), not_set())), :a)
-      refute map_has_key?(union(closed_map(a: integer()), closed_map(b: atom())), :a)
-      assert map_has_key?(union(closed_map(a: integer()), closed_map(a: atom())), :a)
-      assert map_has_key?(intersection(dynamic(), closed_map(a: integer())), :a)
-      refute map_has_key?(intersection(dynamic(), closed_map(a: integer())), :b)
-
-      refute map_may_have_key?(closed_map(foo: integer()), :bar)
-      assert map_may_have_key?(closed_map(foo: integer()), :foo)
-      assert map_may_have_key?(dynamic(), :foo)
-      refute map_may_have_key?(intersection(dynamic(), open_map(foo: not_set())), :foo)
-    end
-
-    test "type-checking map access" do
-      # dynamic() and %{..., :a => integer(), b: not_set()}
-      t = intersection(dynamic(), open_map(a: integer(), c: not_set()))
-
-      assert subtype?(map_get!(t, :a), integer())
-      assert map_get!(t, :b) == dynamic()
-
-      assert map_has_key?(t, :a)
-      refute map_has_key?(t, :b)
-      refute map_has_key?(t, :c)
-
-      assert map_may_have_key?(t, :a)
-      assert map_may_have_key?(t, :b)
-      refute map_may_have_key?(t, :c)
+             |> map_get(:a) == {false, integer()}
     end
   end