Macro.classify_atom - binary fast path

lib/elixir/lib/macro.ex

@@ -207,6 +207,16 @@ defmodule Macro do
           escape: (binary(), char() -> binary())
         ]
 
+  # Not a defguardp because Macro is compiled before Enum during bootstrap,
+  # and Kernel.defguard/1 relies on Enum.
+  defmacrop is_ascii_identifier_char(char) do
+    quote do
+      (unquote(char) >= ?a and unquote(char) <= ?z) or
+        (unquote(char) >= ?A and unquote(char) <= ?Z) or
+        (unquote(char) >= ?0 and unquote(char) <= ?9) or unquote(char) == ?_
+    end
+  end
+
   @doc """
   Breaks a pipeline expression into a list.
 
@@ -2670,67 +2680,77 @@ defmodule Macro do
   #   * `:other` - any other atom (these are usually escaped when inspected, like
   #     `:"foo and bar"`)
   #
-  defp inner_classify(atom) when is_atom(atom) do
-    cond do
-      atom in [:%, :%{}, :{}, :<<>>, :..., :.., :., :..//, :->] ->
-        :not_callable
+  defp inner_classify(atom)
+       when atom in [:%, :%{}, :{}, :<<>>, :..., :.., :., :..//, :->],
+       do: :not_callable
 
-      # <|>, ^^^, and ~~~ are deprecated
-      atom in [:"::", :"^^^", :"~~~", :"<|>"] ->
-        :quoted_operator
+  # <|>, ^^^, and ~~~ are deprecated
+  defp inner_classify(atom) when atom in [:"::", :"^^^", :"~~~", :"<|>"],
+    do: :quoted_operator
 
-      operator?(atom, 1) or operator?(atom, 2) ->
-        :unquoted_operator
+  defp inner_classify(atom) when is_atom(atom) do
+    if operator?(atom, 1) or operator?(atom, 2) do
+      :unquoted_operator
+    else
+      # ASCII identifiers and aliases are recognized on the binary to avoid building a
+      # charlist and running the (unicode aware) tokenizer, which dominates the cost of
+      # classifying keyword list, map and struct keys. Both are matched at once because
+      # the first byte tells them apart and a second match would build a second context.
+      case Atom.to_string(atom) do
+        <<char, rest::binary>> when (char >= ?a and char <= ?z) or char == ?_ ->
+          if ascii_identifier_rest?(rest), do: :identifier, else: tokenizer_classify(atom)
+
+        "Elixir" <> rest ->
+          if valid_alias_piece?(rest), do: :alias, else: tokenizer_classify(atom)
+
+        _other ->
+          tokenizer_classify(atom)
+      end
+    end
+  end
 
-      true ->
-        charlist = Atom.to_charlist(atom)
+  defp tokenizer_classify(atom) do
+    case :elixir_config.identifier_tokenizer().tokenize(Atom.to_charlist(atom)) do
+      {kind, _acc, [], _, _, special} ->
+        cond do
+          kind != :identifier or :lists.member(:at, special) ->
+            :not_callable
 
-        if valid_alias?(charlist) do
-          :alias
-        else
-          case :elixir_config.identifier_tokenizer().tokenize(charlist) do
-            {kind, _acc, [], _, _, special} ->
-              cond do
-                kind != :identifier or :lists.member(:at, special) ->
-                  :not_callable
-
-                # identifier_tokenizer used to return errors for non-nfc, but
-                # now it nfc-normalizes everything. However, lack of nfc is
-                # still a good reason to quote an atom when printing.
-                :lists.member(:nfkc, special) ->
-                  :other
-
-                true ->
-                  :identifier
-              end
-
-            _ ->
-              :other
-          end
+          # identifier_tokenizer used to return errors for non-nfc, but
+          # now it nfc-normalizes everything. However, lack of nfc is
+          # still a good reason to quote an atom when printing.
+          :lists.member(:nfkc, special) ->
+            :other
+
+          true ->
+            :identifier
         end
+
+      _ ->
+        :other
     end
   end
 
-  defp valid_alias?([?E, ?l, ?i, ?x, ?i, ?r] ++ rest), do: valid_alias_piece?(rest)
-  defp valid_alias?(_other), do: false
+  defp ascii_identifier_rest?(<<char, rest::binary>>) when is_ascii_identifier_char(char),
+    do: ascii_identifier_rest?(rest)
 
-  defp valid_alias_piece?([?., char | rest]) when char >= ?A and char <= ?Z,
-    do: valid_alias_piece?(trim_leading_while_valid_identifier(rest))
+  defp ascii_identifier_rest?(<<char>>) when char == ?? or char == ?!, do: true
+  defp ascii_identifier_rest?(<<>>), do: true
+  defp ascii_identifier_rest?(_binary), do: false
+
+  defp valid_alias_piece?(<<>>), do: true
+
+  defp valid_alias_piece?(<<?., char, rest::binary>>) when char >= ?A and char <= ?Z,
+    do: valid_alias_piece_rest?(rest)
 
-  defp valid_alias_piece?([]), do: true
   defp valid_alias_piece?(_other), do: false
 
-  defp trim_leading_while_valid_identifier([char | rest])
-       when char >= ?a and char <= ?z
-       when char >= ?A and char <= ?Z
-       when char >= ?0 and char <= ?9
-       when char == ?_ do
-    trim_leading_while_valid_identifier(rest)
-  end
+  # Inside a piece. A helper returning the rest of the binary would build a sub binary
+  # per piece, so branch back into valid_alias_piece?/1 to keep the match context.
+  defp valid_alias_piece_rest?(<<char, rest::binary>>) when is_ascii_identifier_char(char),
+    do: valid_alias_piece_rest?(rest)
 
-  defp trim_leading_while_valid_identifier(other) do
-    other
-  end
+  defp valid_alias_piece_rest?(rest), do: valid_alias_piece?(rest)
 
   @doc """
   Default backend for `Kernel.dbg/2`.

lib/elixir/test/elixir/macro_test.exs

@@ -1883,6 +1883,126 @@ defmodule MacroTest do
     end
   end
 
+  describe "classify_atom/1" do
+    test "identifiers" do
+      assert Macro.classify_atom(:foo) == :identifier
+      assert Macro.classify_atom(:foo_bar) == :identifier
+      assert Macro.classify_atom(:foo123) == :identifier
+      assert Macro.classify_atom(:fOO) == :identifier
+      assert Macro.classify_atom(:_) == :identifier
+      assert Macro.classify_atom(:_foo) == :identifier
+      assert Macro.classify_atom(:__foo__) == :identifier
+      assert Macro.classify_atom(:foo?) == :identifier
+      assert Macro.classify_atom(:foo!) == :identifier
+      assert Macro.classify_atom(:_foo?) == :identifier
+      assert Macro.classify_atom(:olá) == :identifier
+      assert Macro.classify_atom(:こんにちは世界) == :identifier
+      assert Macro.classify_atom(:olá?) == :identifier
+      assert Macro.classify_atom(:olá!) == :identifier
+      assert Macro.classify_atom(:日本語) == :identifier
+    end
+
+    test "reserved words are identifiers" do
+      assert Macro.classify_atom(true) == :identifier
+      assert Macro.classify_atom(false) == :identifier
+      assert Macro.classify_atom(nil) == :identifier
+      assert Macro.classify_atom(:do) == :identifier
+      assert Macro.classify_atom(:end) == :identifier
+      assert Macro.classify_atom(:fn) == :identifier
+    end
+
+    test "identifiers with ? or ! in invalid positions" do
+      assert Macro.classify_atom(:"foo!!") == :quoted
+      assert Macro.classify_atom(:"foo??") == :quoted
+      assert Macro.classify_atom(:"foo?!") == :quoted
+      assert Macro.classify_atom(:"foo?bar") == :quoted
+      assert Macro.classify_atom(:"foo!bar") == :quoted
+      assert Macro.classify_atom(:"!foo") == :quoted
+      assert Macro.classify_atom(:"?foo") == :quoted
+    end
+
+    test "aliases" do
+      assert Macro.classify_atom(Foo) == :alias
+      assert Macro.classify_atom(Foo.Bar) == :alias
+      assert Macro.classify_atom(Foo.Bar.Baz) == :alias
+      assert Macro.classify_atom(:"Elixir") == :alias
+      assert Macro.classify_atom(Elixir.Elixir) == :alias
+      assert Macro.classify_atom(:"Elixir.A1_b.C2") == :alias
+    end
+
+    test "atoms prefixed by Elixir that are not aliases" do
+      assert Macro.classify_atom(:Elixirfoo) == :unquoted
+      assert Macro.classify_atom(:ElixirFoo) == :unquoted
+      assert Macro.classify_atom(:"Elixir.") == :quoted
+      assert Macro.classify_atom(:"Elixir.Foo.") == :quoted
+      assert Macro.classify_atom(:"Elixir.foo") == :quoted
+      assert Macro.classify_atom(:"Elixir.foo.Bar") == :quoted
+      assert Macro.classify_atom(:"Elixir.Foo.bar") == :quoted
+      assert Macro.classify_atom(:"Elixir.Foo..Bar") == :quoted
+      assert Macro.classify_atom(:"Elixir.1Foo") == :quoted
+    end
+
+    test "unquoted" do
+      assert Macro.classify_atom(:Foo) == :unquoted
+      assert Macro.classify_atom(:FOO) == :unquoted
+      assert Macro.classify_atom(:foo@bar) == :unquoted
+      assert Macro.classify_atom(:+) == :unquoted
+      assert Macro.classify_atom(:...) == :unquoted
+      assert Macro.classify_atom(:..) == :unquoted
+      assert Macro.classify_atom(:foo@) == :unquoted
+      assert Macro.classify_atom(:_foo@bar) == :unquoted
+      assert Macro.classify_atom(:olá@bar) == :unquoted
+      assert Macro.classify_atom(:+) == :unquoted
+      assert Macro.classify_atom(:...) == :unquoted
+      assert Macro.classify_atom(:..) == :unquoted
+
+      # atoms starting with a non-ASCII uppercase letter are not aliases
+      assert Macro.classify_atom(:Ólá) == :unquoted
+      assert Macro.classify_atom(:ΑΒΓ) == :unquoted
+
+      # mixed scripts separated by an underscore tokenize as an alias
+      assert Macro.classify_atom(:T_シャツ) == :unquoted
+    end
+
+    test "textual operators are unquoted, not identifiers" do
+      assert Macro.classify_atom(:when) == :unquoted
+      assert Macro.classify_atom(:and) == :unquoted
+      assert Macro.classify_atom(:or) == :unquoted
+      assert Macro.classify_atom(:not) == :unquoted
+      assert Macro.classify_atom(:in) == :unquoted
+    end
+
+    test "special forms are unquoted, not callable" do
+      assert Macro.classify_atom(:%) == :unquoted
+      assert Macro.classify_atom(:%{}) == :unquoted
+      assert Macro.classify_atom(:{}) == :unquoted
+      assert Macro.classify_atom(:<<>>) == :unquoted
+      assert Macro.classify_atom(:.) == :unquoted
+      assert Macro.classify_atom(:->) == :unquoted
+      assert Macro.classify_atom(:..//) == :unquoted
+    end
+
+    test "quoted" do
+      assert Macro.classify_atom(:"::") == :quoted
+      assert Macro.classify_atom(:"^^^") == :quoted
+      assert Macro.classify_atom(:"~~~") == :quoted
+      assert Macro.classify_atom(:"<|>") == :quoted
+      assert Macro.classify_atom(:"@foo") == :quoted
+      assert Macro.classify_atom(:"foo bar") == :quoted
+      assert Macro.classify_atom(:"1foo") == :quoted
+      assert Macro.classify_atom(:"foo\n") == :quoted
+      assert Macro.classify_atom(:"") == :quoted
+      assert Macro.classify_atom(:" ") == :quoted
+      assert Macro.classify_atom(:"🌢") == :quoted
+
+      # mixed-script identifiers cannot be tokenized
+      assert Macro.classify_atom(String.to_atom("аdmin")) == :quoted
+
+      nfd = :unicode.characters_to_nfd_binary("olá")
+      assert Macro.classify_atom(String.to_unsafe_atom(nfd)) == :quoted
+    end
+  end
+
   test "operator?/2" do
     assert Macro.operator?(:+, 2)
     assert Macro.operator?(:+, 1)