perf(util): encode Base64 in three-byte blocks (#1045)

Author: fallintoplace

src/iceberg/test/base64_test.cc

@@ -19,6 +19,9 @@
 
 #include "iceberg/util/base64.h"
 
+#include <cstddef>
+#include <string>
+
 #include <gtest/gtest.h>
 
 #include "iceberg/test/matchers.h"
@@ -50,6 +53,27 @@ TEST(Base64Test, Encode) {
   EXPECT_EQ("AA==", Base64::Encode({"\x00", 1}));
 }
 
+TEST(Base64Test, EncodeBinaryAcrossBlockBoundaries) {
+  for (size_t length = 0; length <= 256; ++length) {
+    std::string data(length, '\0');
+    for (size_t i = 0; i < length; ++i) {
+      data[i] = static_cast<char>((i * 73 + length) & 0xFF);
+    }
+
+    const auto encoded = Base64::Encode(data);
+    EXPECT_EQ(((length + 2) / 3) * 4, encoded.size()) << length;
+    ICEBERG_UNWRAP_OR_FAIL(auto decoded, Base64::Decode(encoded));
+    EXPECT_EQ(data, decoded) << length;
+
+    const auto url_encoded = Base64::UrlEncode(data);
+    EXPECT_EQ((length / 3) * 4 + (length % 3 == 0 ? 0 : length % 3 + 1),
+              url_encoded.size())
+        << length;
+    ICEBERG_UNWRAP_OR_FAIL(auto url_decoded, Base64::UrlDecode(url_encoded));
+    EXPECT_EQ(data, url_decoded) << length;
+  }
+}
+
 TEST(Base64Test, Decode) {
   // Empty string
   ICEBERG_UNWRAP_OR_FAIL(auto empty, Base64::Decode(""));

src/iceberg/util/base64.cc

@@ -30,28 +30,40 @@ namespace {
 // When `pad` is true, the output is padded with '=' to a multiple of 4 characters.
 std::string Base64EncodeWithAlphabet(std::string_view input, std::string_view alphabet,
                                      bool pad) {
-  std::string output;
-  output.reserve((input.size() + 2) / 3 * 4);
+  const size_t full_blocks = input.size() / 3;
+  const size_t remainder = input.size() % 3;
+  const size_t output_size =
+      full_blocks * 4 + (remainder == 0 ? 0 : (pad ? 4 : remainder + 1));
+  std::string output(output_size, '\0');
+
+  const auto* bytes = reinterpret_cast<const uint8_t*>(input.data());
+  size_t out = 0;
+  for (size_t i = 0; i < full_blocks * 3; i += 3) {
+    const uint32_t block = (static_cast<uint32_t>(bytes[i]) << 16) |
+                           (static_cast<uint32_t>(bytes[i + 1]) << 8) |
+                           static_cast<uint32_t>(bytes[i + 2]);
+    output[out++] = alphabet[(block >> 18) & 0x3F];
+    output[out++] = alphabet[(block >> 12) & 0x3F];
+    output[out++] = alphabet[(block >> 6) & 0x3F];
+    output[out++] = alphabet[block & 0x3F];
+  }
 
-  uint32_t buffer = 0;
-  int bits_collected = 0;
-  for (unsigned char byte : input) {
-    buffer = (buffer << 8) | byte;
-    bits_collected += 8;
-    while (bits_collected >= 6) {
-      bits_collected -= 6;
-      output.push_back(alphabet[(buffer >> bits_collected) & 0x3F]);
+  if (remainder > 0) {
+    const size_t i = full_blocks * 3;
+    uint32_t block = static_cast<uint32_t>(bytes[i]) << 16;
+    if (remainder == 2) {
+      block |= static_cast<uint32_t>(bytes[i + 1]) << 8;
     }
-  }
-  if (bits_collected > 0) {
-    // Pad the remaining bits on the right to form the final 6-bit group.
-    output.push_back(alphabet[(buffer << (6 - bits_collected)) & 0x3F]);
-  }
-  if (pad) {
-    while (output.size() % 4 != 0) {
-      output.push_back('=');
+    output[out++] = alphabet[(block >> 18) & 0x3F];
+    output[out++] = alphabet[(block >> 12) & 0x3F];
+    if (remainder == 2) {
+      output[out++] = alphabet[(block >> 6) & 0x3F];
+    }
+    while (out < output_size) {
+      output[out++] = '=';
     }
   }
+
   return output;
 }