Add String.to_existing_atom/2 to validate against a list of allowed atoms (#15483)

Author: sabiwara

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

@@ -288,7 +288,8 @@ defmodule Module.Types.Apply do
          [{[term(), open_map()], tuple([term(), open_map()]) |> opt_union(atom([:error]))}]},
         {:maps, :to_list, [{[open_map()], list(tuple([term(), term()]))}]},
         {:maps, :update, [{[term(), term(), open_map()], open_map()}]},
-        {:maps, :values, [{[open_map()], list(term())}]}
+        {:maps, :values, [{[open_map()], list(term())}]},
+        {String, :to_existing_atom, [{[binary(), non_empty_list(atom())], atom()}]}
       ] do
     [arity] = Enum.map(clauses, fn {args, _return} -> length(args) end) |> Enum.uniq()
 
@@ -1418,6 +1419,20 @@ defmodule Module.Types.Apply do
     end
   end
 
+  defp remote_apply(String, :to_existing_atom, info, [_string, list] = args_types, stack) do
+    # TODO remove once we add parametric types, this will just be:
+    # binary(), non_empty_list(a) -> a when a: atom()
+
+    case remote_apply(info, args_types, stack) do
+      {:ok, _} ->
+        {false, refined_atom} = list_of(list)
+        {:ok, refined_atom}
+
+      other ->
+        other
+    end
+  end
+
   defp remote_apply(_mod, _fun, info, args_types, stack) do
     remote_apply(info, args_types, stack)
   end

lib/elixir/lib/string.ex

@@ -2992,6 +2992,8 @@ defmodule String do
   Converts a string to an existing atom or raises if
   the atom does not exist.
 
+  If the list of expected atoms is known upfront, prefer `to_existing_atom/2`.
+
   The maximum atom size is of 255 Unicode code points.
   Raises an `ArgumentError` if the atom does not exist.
 
@@ -3021,6 +3023,46 @@ defmodule String do
     :erlang.binary_to_existing_atom(string, :utf8)
   end
 
+  @doc """
+  Converts a string to one of the `allowed_atoms` or raises.
+
+  Raises an `ArgumentError` if the atom either does not exist or is not within
+  the existing list.
+
+  This should be preferred to `to_existing_atom/1` if the list is known upfront,
+  since there is no risk that the atom has not been loaded.
+
+  ## Examples
+
+      iex> String.to_existing_atom("foo", [:foo, :bar])
+      :foo
+
+      iex> String.to_existing_atom("unknown", [:foo, :bar])
+      ** (ArgumentError) unexpected value: \"unknown\", the allowed atoms are: [:foo, :bar]
+
+  """
+  @spec to_existing_atom(String.t(), nonempty_list(atom)) :: atom
+  def to_existing_atom(string, [_ | _] = allowed_atoms) when is_binary(string) do
+    atom = :erlang.binary_to_existing_atom(string, :utf8)
+
+    if atom not in allowed_atoms do
+      to_existing_atom_unexpected(string, allowed_atoms)
+    end
+
+    atom
+  end
+
+  # used just to have a less cryptic stacktrace and consistent error
+  @doc false
+  def __to_existing_atom__(string, allowed_atoms) do
+    to_existing_atom_unexpected(string, allowed_atoms)
+  end
+
+  defp to_existing_atom_unexpected(string, allowed_atoms) do
+    raise ArgumentError,
+          "unexpected value: #{inspect(string)}, the allowed atoms are: #{inspect(allowed_atoms)}"
+  end
+
   @doc """
   Returns an integer whose text representation is `string`.
 

lib/elixir/src/elixir_erl_pass.erl

@@ -634,6 +634,27 @@ translate_remote(maps, merge, Meta, [Map1, Map2], S) ->
     {[TMap1, TMap2], TS} ->
       {{call, Ann, {remote, Ann, {atom, Ann, maps}, {atom, Ann, merge}}, [TMap1, TMap2]}, TS}
   end;
+translate_remote('Elixir.String', to_existing_atom, Meta, [String, List], S) ->
+  Ann = ?ann(Meta),
+  {[TString, TList], TS} = translate_args([String, List], Ann, S),
+
+  case is_list(List) andalso lists:all(fun is_atom/1, List) of
+    true ->
+      Generated = erl_anno:set_generated(true, Ann),
+      LastClause = {clause, Generated,
+        [{var, Generated, '_'}],
+        [],
+        [{call, Ann, {remote, Ann, {atom, Ann, 'Elixir.String'}, {atom, Ann, '__to_existing_atom__'}}, [TString, TList]}]},
+      Clauses = [
+        {clause, Generated,
+          [{bin, Generated, [{bin_element, Generated, {string, Generated, atom_to_list(Atom)}, default, default}]}],
+          [],
+          [{atom, Ann, Atom}]}
+      || Atom <- List] ++ [LastClause],
+      {{'case', Generated, TString, Clauses}, TS};
+    false ->
+      {{call, Ann, {remote, Ann, {atom, Ann, 'Elixir.String'}, {atom, Ann, to_existing_atom}}, [TString, TList]}, TS}
+  end;
 translate_remote(Left, Right, Meta, Args, S) ->
   Ann = ?ann(Meta),
 

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

@@ -1898,6 +1898,41 @@ defmodule Module.Types.ExprTest do
                  x in [:foo, :bar]
              """
     end
+
+    test "String.to_existing_atom/2" do
+      assert typecheck!(
+               [x],
+               String.to_existing_atom(x, [:foo, :bar])
+             ) == atom([:foo, :bar])
+
+      assert typecheck!(
+               [x],
+               (
+                 values = [:foo, :bar]
+                 String.to_existing_atom(x, values)
+               )
+             ) == atom([:foo, :bar])
+
+      assert typecheck!(
+               [x, values],
+               String.to_existing_atom(x, values)
+             ) == dynamic(atom())
+
+      assert typeerror!(
+               [x],
+               String.to_existing_atom(:not_a_string, x)
+             ) =~ "incompatible types given to String.to_existing_atom/2"
+
+      assert typeerror!(
+               [x],
+               String.to_existing_atom(x, [:foo, "not atoms"])
+             ) =~ "incompatible types given to String.to_existing_atom/2"
+
+      assert typeerror!(
+               [x],
+               String.to_existing_atom(x, [])
+             ) =~ "incompatible types given to String.to_existing_atom/2"
+    end
   end
 
   describe "case" do

lib/elixir/test/elixir/string_test.exs

@@ -1099,4 +1099,23 @@ defmodule StringTest do
     assert String.bag_distance("\r\t\xFF\v", "\xFF\r\n\xFF") == 0.25
     assert String.split("\r\t\v", "") == ["", "\r", "\t", "\v", ""]
   end
+
+  test "to_existing_atom/2" do
+    # constant
+    assert String.to_existing_atom("foo", [:foo, :bar]) == :foo
+    assert String.to_existing_atom("bar", [:foo, :bar]) == :bar
+
+    assert_raise ArgumentError, fn ->
+      String.to_existing_atom("baz", [:foo, :bar])
+    end
+
+    # variable
+    values = [:foo, :bar]
+    assert String.to_existing_atom("foo", values) == :foo
+    assert String.to_existing_atom("bar", values) == :bar
+
+    assert_raise ArgumentError, fn ->
+      String.to_existing_atom("baz", values)
+    end
+  end
 end

lib/elixir/test/erlang/control_test.erl

@@ -86,6 +86,22 @@ optimized_inspect_interpolation_test() ->
        {call, _, {remote, _,{atom, _, 'Elixir.Kernel'}, {atom, _, inspect}}, [_]},
        default, [binary]}]} = to_erl("\"#{inspect(1)}\"").
 
+optimized_string_to_existing_atom_test() ->
+  {'case', _, _,
+    [{clause, _,
+      [{bin, _, [{bin_element, _, {string, _, "foo"}, default, default}]}],
+      [],
+      [{atom, _, foo}]},
+     {clause, _,
+      [{bin, _, [{bin_element, _, {string, _, "bar"}, default, default}]}],
+      [],
+      [{atom, _, bar}]},
+     {clause, _,
+      [{var, _, '_'}],
+      [],
+      [{call, _, {remote, _, {atom, _, 'Elixir.String'}, {atom, _, '__to_existing_atom__'}}, [_, _]}]}]
+  } = to_erl("String.to_existing_atom(\"baz\", [:foo, :bar])").
+
 optimized_map_merge_test() ->
   {map, _,
     [{map_field_assoc, _, {atom, _, a}, {integer, _, 1}},