Optimize URI form and RFC 3986 query encoding with SWAR

lib/elixir/lib/uri.ex

@@ -168,8 +168,8 @@ defmodule URI do
   end
 
   defp encode_kv_pair({key, value}, :rfc3986) do
-    encode(Kernel.to_string(key), &char_unreserved?/1) <>
-      "=" <> encode(Kernel.to_string(value), &char_unreserved?/1)
+    encode_unreserved(Kernel.to_string(key), :percent) <>
+      "=" <> encode_unreserved(Kernel.to_string(value), :percent)
   end
 
   defp encode_kv_pair({key, value}, :www_form) do
@@ -340,6 +340,10 @@ defmodule URI do
     character in @reserved_characters
   end
 
+  defguardp unreserved_char?(character)
+            when character in ?0..?9 or character in ?a..?z or character in ?A..?Z or
+                   character in ~c"~_-."
+
   @doc """
   Checks if `character` is an unreserved one in a URI.
 
@@ -357,9 +361,52 @@ defmodule URI do
   """
   @spec char_unreserved?(byte) :: boolean
   def char_unreserved?(character) do
-    character in ?0..?9 or character in ?a..?z or character in ?A..?Z or character in ~c"~_-."
+    unreserved_char?(character)
   end
 
+  # A 56-bit word remains a small integer on 64-bit BEAM; using 64 bits would
+  # allocate a bignum for the range arithmetic below.
+  @swar_ones 0x01010101010101
+  @swar_mask80 0x80808080808080
+  @swar_threshold 7
+  @compile {:inline, encode_unreserved_byte: 2, hex: 1}
+
+  defmacrop swar_range(word, first, last) do
+    ge = @swar_ones * (0x80 - first)
+    gt = @swar_ones * (0x7F - last)
+
+    quote do
+      bxor(unquote(word) + unquote(ge), unquote(word) + unquote(gt))
+    end
+  end
+
+  # For a range [lo, hi], the high bit of
+  # `bxor(word + (0x80 - lo), word + (0x7F - hi))` is set in exactly the
+  # lanes inside the range. The ASCII check prevents carries between lanes.
+  # Punctuation uses ranges too, avoiding false positives from subtraction-
+  # based zero-byte detectors when a borrow crosses lanes.
+  defguardp unreserved_word?(word)
+            when band(word, @swar_mask80) == 0 and
+                   band(
+                     bor(
+                       bor(
+                         bor(
+                           swar_range(word, ?0, ?9),
+                           swar_range(word, ?A, ?Z)
+                         ),
+                         swar_range(word, ?a, ?z)
+                       ),
+                       bor(
+                         swar_range(word, ?-, ?.),
+                         bor(
+                           swar_range(word, ?_, ?_),
+                           swar_range(word, ?~, ?~)
+                         )
+                       )
+                     ),
+                     @swar_mask80
+                   ) == @swar_mask80
+
   @doc """
   Checks if `character` is allowed unescaped in a URI.
 
@@ -433,14 +480,134 @@ defmodule URI do
   """
   @spec encode_www_form(binary) :: binary
   def encode_www_form(string) when is_binary(string) do
-    for <<byte <- string>>, into: "" do
-      case percent(byte, &char_unreserved?/1) do
-        "%20" -> "+"
-        percent -> percent
-      end
+    encode_unreserved(string, :www_form)
+  end
+
+  defp encode_unreserved(string, mode) when byte_size(string) < @swar_threshold do
+    case string do
+      <<>> -> string
+      _ -> encode_unreserved_small(string, "", mode)
     end
   end
 
+  defp encode_unreserved(string, mode), do: encode_unreserved(string, "", mode)
+
+  defp encode_unreserved_small(<<?\s, rest::binary>>, acc, :www_form) do
+    encode_unreserved_small(rest, <<acc::binary, ?+>>, :www_form)
+  end
+
+  defp encode_unreserved_small(<<byte, rest::binary>>, acc, mode)
+       when not unreserved_char?(byte) do
+    encode_unreserved_small(
+      rest,
+      <<acc::binary, ?%, hex(bsr(byte, 4)), hex(band(byte, 15))>>,
+      mode
+    )
+  end
+
+  defp encode_unreserved_small(<<byte, rest::binary>>, acc, mode) do
+    encode_unreserved_small(rest, <<acc::binary, byte>>, mode)
+  end
+
+  defp encode_unreserved_small(<<>>, acc, _mode), do: acc
+
+  defp encode_unreserved(<<?\s, rest::binary>>, acc, :www_form) do
+    encode_unreserved(rest, <<acc::binary, ?+>>, :www_form)
+  end
+
+  # Avoid the word guard when its first byte already makes failure certain.
+  defp encode_unreserved(<<byte, rest::binary>>, acc, mode)
+       when not unreserved_char?(byte) do
+    encode_unreserved(
+      rest,
+      <<acc::binary, ?%, hex(bsr(byte, 4)), hex(band(byte, 15))>>,
+      mode
+    )
+  end
+
+  # Seven bytes are checked with SWAR in each stride.
+  defp encode_unreserved(<<word::56, rest::binary>>, acc, mode)
+       when unreserved_word?(word) do
+    encode_unreserved(rest, <<acc::binary, word::56>>, mode)
+  end
+
+  defp encode_unreserved(<<_::56, _::binary>> = string, acc, mode) do
+    encode_unreserved_fallback(string, acc, mode)
+  end
+
+  defp encode_unreserved(<<>>, acc, _mode), do: acc
+
+  defp encode_unreserved(rest, acc, mode) do
+    encode_unreserved_small(rest, acc, mode)
+  end
+
+  # Consume through the first disallowed byte instead of checking overlapping words.
+  defp encode_unreserved_fallback(<<byte1, byte2, rest::binary>>, acc, mode)
+       when byte_size(rest) >= 5 and not unreserved_char?(byte2) do
+    encoded = encode_unreserved_byte(byte2, mode)
+    encode_unreserved(rest, <<acc::binary, byte1, encoded::binary>>, mode)
+  end
+
+  defp encode_unreserved_fallback(<<byte1, byte2, byte3, rest::binary>>, acc, mode)
+       when byte_size(rest) >= 4 and not unreserved_char?(byte3) do
+    encode_unreserved(
+      rest,
+      <<acc::binary, byte1, byte2, encode_unreserved_byte(byte3, mode)::binary>>,
+      mode
+    )
+  end
+
+  defp encode_unreserved_fallback(<<byte1, byte2, byte3, byte4, rest::binary>>, acc, mode)
+       when byte_size(rest) >= 3 and not unreserved_char?(byte4) do
+    encode_unreserved(
+      rest,
+      <<acc::binary, byte1, byte2, byte3, encode_unreserved_byte(byte4, mode)::binary>>,
+      mode
+    )
+  end
+
+  defp encode_unreserved_fallback(<<byte1, byte2, byte3, byte4, byte5, rest::binary>>, acc, mode)
+       when byte_size(rest) >= 2 and not unreserved_char?(byte5) do
+    encode_unreserved(
+      rest,
+      <<acc::binary, byte1, byte2, byte3, byte4, encode_unreserved_byte(byte5, mode)::binary>>,
+      mode
+    )
+  end
+
+  defp encode_unreserved_fallback(
+         <<byte1, byte2, byte3, byte4, byte5, byte6, rest::binary>>,
+         acc,
+         mode
+       )
+       when byte_size(rest) >= 1 and not unreserved_char?(byte6) do
+    encode_unreserved(
+      rest,
+      <<acc::binary, byte1, byte2, byte3, byte4, byte5,
+        encode_unreserved_byte(byte6, mode)::binary>>,
+      mode
+    )
+  end
+
+  defp encode_unreserved_fallback(
+         <<byte1, byte2, byte3, byte4, byte5, byte6, byte7, rest::binary>>,
+         acc,
+         mode
+       ) do
+    encode_unreserved(
+      rest,
+      <<acc::binary, byte1, byte2, byte3, byte4, byte5, byte6,
+        encode_unreserved_byte(byte7, mode)::binary>>,
+      mode
+    )
+  end
+
+  defp encode_unreserved_byte(?\s, :www_form), do: <<?+>>
+
+  defp encode_unreserved_byte(byte, _mode) do
+    <<?%, hex(bsr(byte, 4)), hex(band(byte, 15))>>
+  end
+
   defp percent(char, predicate) do
     if predicate.(char) do
       <<char>>

lib/elixir/test/elixir/uri_test.exs

@@ -20,6 +20,46 @@ defmodule URITest do
     assert URI.encode_www_form("4test ~1.x") == "4test+~1.x"
     assert URI.encode_www_form("poll:146%") == "poll%3A146%25"
     assert URI.encode_www_form("/\n+/ゆ") == "%2F%0A%2B%2F%E3%82%86"
+
+    boundary_input = :binary.copy("a /%+~_", 10)
+
+    for length <- [0, 1, 2, 3, 4, 5, 6, 7, 8, 62, 63, 64, 65] do
+      input = binary_part(boundary_input, 0, length)
+      assert URI.encode_www_form(input) == encode_www_form_reference(input)
+    end
+
+    safe_prefix = :binary.copy("a", 63)
+    safe = :binary.copy("a", 14)
+
+    for position <- 0..6, byte <- 0..255 do
+      <<prefix::binary-size(^position), _replaced, suffix::binary>> = safe
+      input = <<safe_prefix::binary, prefix::binary, byte, suffix::binary>>
+      assert URI.encode_www_form(input) == encode_www_form_reference(input)
+    end
+  end
+
+  defp encode_unreserved_reference(input) do
+    for <<byte <- input>>, into: "" do
+      if unreserved_reference?(byte) do
+        <<byte>>
+      else
+        "%" <> Base.encode16(<<byte>>)
+      end
+    end
+  end
+
+  defp encode_www_form_reference(input) do
+    for <<byte <- input>>, into: "" do
+      cond do
+        byte == ?\s -> "+"
+        unreserved_reference?(byte) -> <<byte>>
+        true -> "%" <> Base.encode16(<<byte>>)
+      end
+    end
+  end
+
+  defp unreserved_reference?(byte) do
+    byte in ?0..?9 or byte in ?a..?z or byte in ?A..?Z or byte in [?~, ?_, ?-, ?.]
   end
 
   test "encode_query/1,2" do
@@ -39,6 +79,25 @@ defmodule URITest do
     assert URI.encode_query([{"foo[]", "+=/?&# Ñ"}], :www_form) ==
              "foo%5B%5D=%2B%3D%2F%3F%26%23+%C3%91"
 
+    boundary_input = :binary.copy("a /%+~_", 10)
+
+    for length <- [0, 1, 2, 3, 4, 5, 6, 7, 8, 62, 63, 64, 65] do
+      input = binary_part(boundary_input, 0, length)
+      expected = encode_unreserved_reference(input)
+      assert URI.encode_query([{input, input}], :rfc3986) == expected <> "=" <> expected
+    end
+
+    safe_prefix = :binary.copy("a", 63)
+    word = :binary.copy("a", 7)
+
+    for position <- 0..6, byte <- 0..255 do
+      <<word_prefix::binary-size(^position), _replaced, word_suffix::binary>> = word
+      input = <<safe_prefix::binary, word_prefix::binary, byte, word_suffix::binary>>
+
+      assert URI.encode_query([{input, ""}], :rfc3986) ==
+               encode_unreserved_reference(input) <> "="
+    end
+
     assert_raise ArgumentError, fn ->
       URI.encode_query([{"foo", ~c"bar"}])
     end