[mypyc] Use faster base64 encode implementation in librt.base64 (#20237)

Vendor optimized base64 implementation from
https://github.com/aklomp/base64.
This is based on commit 9e8ed65048ff0f703fad3deb03bf66ac7f78a4d7 (May
2025).

Enable SIMD on macOS (64-bit ARM only). Other platforms probably use a
generic version. I'll look into enabling SIMD more generally in a
follow-up PR.

A `b64encode` micro-benchmark was up to 11 times faster compared to the
stdlib `base64` module (on a MacBook Pro).

Author: JukkaL

LICENSE

@@ -227,3 +227,38 @@ FOR ANY SPECIAL, INDIRECT OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT
 OF OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
+
+= = = = =
+
+Files under lib-rt/base64 are licensed under the following license.
+
+= = = = =
+
+Copyright (c) 2005-2007, Nick Galbreath
+Copyright (c) 2015-2018, Wojciech Muła
+Copyright (c) 2016-2017, Matthieu Darbois
+Copyright (c) 2013-2022, Alfred Klomp
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are
+met:
+
+- Redistributions of source code must retain the above copyright notice,
+  this list of conditions and the following disclaimer.
+
+- Redistributions in binary form must reproduce the above copyright
+  notice, this list of conditions and the following disclaimer in the
+  documentation and/or other materials provided with the distribution.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS
+IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED
+TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
+PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
+HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
+SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED
+TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
+PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
+LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
+NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
+SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

mypyc/build.py

@@ -26,7 +26,7 @@
 import sys
 import time
 from collections.abc import Iterable
-from typing import TYPE_CHECKING, Any, NoReturn, Union, cast
+from typing import TYPE_CHECKING, Any, NamedTuple, NoReturn, Union, cast
 
 from mypy.build import BuildSource
 from mypy.errors import CompileError
@@ -42,7 +42,52 @@
 from mypyc.namegen import exported_name
 from mypyc.options import CompilerOptions
 
-LIBRT_MODULES = [("librt.internal", "librt_internal.c"), ("librt.base64", "librt_base64.c")]
+
+class ModDesc(NamedTuple):
+    module: str
+    c_files: list[str]
+    other_files: list[str]
+    include_dirs: list[str]
+
+
+LIBRT_MODULES = [
+    ModDesc("librt.internal", ["librt_internal.c"], [], []),
+    ModDesc(
+        "librt.base64",
+        [
+            "librt_base64.c",
+            "base64/lib.c",
+            "base64/codec_choose.c",
+            "base64/tables/tables.c",
+            "base64/arch/generic/codec.c",
+            "base64/arch/ssse3/codec.c",
+            "base64/arch/sse41/codec.c",
+            "base64/arch/sse42/codec.c",
+            "base64/arch/avx/codec.c",
+            "base64/arch/avx2/codec.c",
+            "base64/arch/avx512/codec.c",
+            "base64/arch/neon32/codec.c",
+            "base64/arch/neon64/codec.c",
+        ],
+        [
+            "base64/arch/generic/32/enc_loop.c",
+            "base64/arch/generic/64/enc_loop.c",
+            "base64/arch/generic/32/dec_loop.c",
+            "base64/arch/generic/enc_head.c",
+            "base64/arch/generic/enc_tail.c",
+            "base64/arch/generic/dec_head.c",
+            "base64/arch/generic/dec_tail.c",
+            "base64/arch/neon64/dec_loop.c",
+            "base64/arch/neon64/enc_loop_asm.c",
+            "base64/codecs.h",
+            "base64/env.h",
+            "base64/tables/tables.h",
+            "base64/tables/table_dec_32bit.h",
+            "base64/tables/table_enc_12bit.h",
+        ],
+        ["base64"],
+    ),
+]
 
 try:
     # Import setuptools so that it monkey-patch overrides distutils
@@ -677,17 +722,19 @@ def mypycify(
             rt_file = os.path.join(build_dir, name)
             with open(os.path.join(include_dir(), name), encoding="utf-8") as f:
                 write_file(rt_file, f.read())
-        for mod, file_name in LIBRT_MODULES:
-            rt_file = os.path.join(build_dir, file_name)
-            with open(os.path.join(include_dir(), file_name), encoding="utf-8") as f:
-                write_file(rt_file, f.read())
+        for mod, file_names, addit_files, includes in LIBRT_MODULES:
+            for file_name in file_names + addit_files:
+                rt_file = os.path.join(build_dir, file_name)
+                with open(os.path.join(include_dir(), file_name), encoding="utf-8") as f:
+                    write_file(rt_file, f.read())
             extensions.append(
                 get_extension()(
                     mod,
                     sources=[
-                        os.path.join(build_dir, file) for file in [file_name] + RUNTIME_C_FILES
+                        os.path.join(build_dir, file) for file in file_names + RUNTIME_C_FILES
                     ],
-                    include_dirs=[include_dir()],
+                    include_dirs=[include_dir()]
+                    + [os.path.join(include_dir(), d) for d in includes],
                     extra_compile_args=cflags,
                 )
             )

mypyc/lib-rt/base64/arch/avx/codec.c

@@ -0,0 +1,68 @@
+#include <stdint.h>
+#include <stddef.h>
+#include <stdlib.h>
+
+#include "libbase64.h"
+#include "../../tables/tables.h"
+#include "../../codecs.h"
+#include "config.h"
+#include "../../env.h"
+
+#if HAVE_AVX
+#include <immintrin.h>
+
+// Only enable inline assembly on supported compilers and on 64-bit CPUs.
+#ifndef BASE64_AVX_USE_ASM
+# if (defined(__GNUC__) || defined(__clang__)) && BASE64_WORDSIZE == 64
+#  define BASE64_AVX_USE_ASM 1
+# else
+#  define BASE64_AVX_USE_ASM 0
+# endif
+#endif
+
+#include "../ssse3/dec_reshuffle.c"
+#include "../ssse3/dec_loop.c"
+
+#if BASE64_AVX_USE_ASM
+# include "enc_loop_asm.c"
+#else
+# include "../ssse3/enc_translate.c"
+# include "../ssse3/enc_reshuffle.c"
+# include "../ssse3/enc_loop.c"
+#endif
+
+#endif	// HAVE_AVX
+
+void
+base64_stream_encode_avx BASE64_ENC_PARAMS
+{
+#if HAVE_AVX
+	#include "../generic/enc_head.c"
+
+	// For supported compilers, use a hand-optimized inline assembly
+	// encoder. Otherwise fall back on the SSSE3 encoder, but compiled with
+	// AVX flags to generate better optimized AVX code.
+
+#if BASE64_AVX_USE_ASM
+	enc_loop_avx(&s, &slen, &o, &olen);
+#else
+	enc_loop_ssse3(&s, &slen, &o, &olen);
+#endif
+
+	#include "../generic/enc_tail.c"
+#else
+	base64_enc_stub(state, src, srclen, out, outlen);
+#endif
+}
+
+int
+base64_stream_decode_avx BASE64_DEC_PARAMS
+{
+#if HAVE_AVX
+	#include "../generic/dec_head.c"
+	dec_loop_ssse3(&s, &slen, &o, &olen);
+	#include "../generic/dec_tail.c"
+#else
+	return base64_dec_stub(state, src, srclen, out, outlen);
+#endif
+}