[mypyc] Import librt.base64 capsule automatically if needed (#20233)

Allow primitives to specify the capsule they need via module name such
as `librt.base64`. This way we can import the capsule automatically only 
when there are references to the contents of the capsule in the compiled 
code.

Only make the change for `librt.base64`, but we can also do a similar
thing for `librt.internal` in a follow-up PR.

Author: JukkaL

mypyc/analysis/capsule_deps.py

@@ -0,0 +1,27 @@
+from __future__ import annotations
+
+from mypyc.ir.func_ir import FuncIR
+from mypyc.ir.ops import CallC, PrimitiveOp
+
+
+def find_implicit_capsule_dependencies(fn: FuncIR) -> set[str] | None:
+    """Find implicit dependencies on capsules that need to be imported.
+
+    Using primitives or types defined in librt submodules such as "librt.base64"
+    requires a capsule import.
+
+    Note that a module can depend on a librt module even if it doesn't explicitly
+    import it, for example via re-exported names or via return types of functions
+    defined in other modules.
+    """
+    deps: set[str] | None = None
+    for block in fn.blocks:
+        for op in block.ops:
+            # TODO: Also determine implicit type object dependencies (e.g. cast targets)
+            if isinstance(op, CallC) and op.capsule is not None:
+                if deps is None:
+                    deps = set()
+                deps.add(op.capsule)
+            else:
+                assert not isinstance(op, PrimitiveOp), "Lowered IR is expected"
+    return deps

mypyc/build.py

@@ -495,7 +495,6 @@ def mypycify(
     group_name: str | None = None,
     log_trace: bool = False,
     depends_on_librt_internal: bool = False,
-    depends_on_librt_base64: bool = False,
     install_librt: bool = False,
     experimental_features: bool = False,
 ) -> list[Extension]:
@@ -570,7 +569,6 @@ def mypycify(
         group_name=group_name,
         log_trace=log_trace,
         depends_on_librt_internal=depends_on_librt_internal,
-        depends_on_librt_base64=depends_on_librt_base64,
         experimental_features=experimental_features,
     )
 

mypyc/codegen/emitmodule.py

@@ -27,6 +27,7 @@
 from mypy.options import Options
 from mypy.plugin import Plugin, ReportConfigContext
 from mypy.util import hash_digest, json_dumps
+from mypyc.analysis.capsule_deps import find_implicit_capsule_dependencies
 from mypyc.codegen.cstring import c_string_initializer
 from mypyc.codegen.emit import Emitter, EmitterContext, HeaderDeclaration, c_array_initializer
 from mypyc.codegen.emitclass import generate_class, generate_class_reuse, generate_class_type_decl
@@ -259,6 +260,10 @@ def compile_scc_to_ir(
 
             # Switch to lower abstraction level IR.
             lower_ir(fn, compiler_options)
+            # Calculate implicit module dependencies (needed for librt)
+            capsules = find_implicit_capsule_dependencies(fn)
+            if capsules is not None:
+                module.capsules.update(capsules)
             # Perform optimizations.
             do_copy_propagation(fn, compiler_options)
             do_flag_elimination(fn, compiler_options)
@@ -604,7 +609,7 @@ def generate_c_for_modules(self) -> list[tuple[str, str]]:
         ext_declarations.emit_line("#include <CPy.h>")
         if self.compiler_options.depends_on_librt_internal:
             ext_declarations.emit_line("#include <librt_internal.h>")
-        if self.compiler_options.depends_on_librt_base64:
+        if any("librt.base64" in mod.capsules for mod in self.modules.values()):
             ext_declarations.emit_line("#include <librt_base64.h>")
 
         declarations = Emitter(self.context)
@@ -1036,7 +1041,7 @@ def emit_module_exec_func(
             emitter.emit_line("if (import_librt_internal() < 0) {")
             emitter.emit_line("return -1;")
             emitter.emit_line("}")
-        if self.compiler_options.depends_on_librt_base64:
+        if "librt.base64" in module.capsules:
             emitter.emit_line("if (import_librt_base64() < 0) {")
             emitter.emit_line("return -1;")
             emitter.emit_line("}")

mypyc/ir/module_ir.py

@@ -30,6 +30,8 @@ def __init__(
         # These are only visible in the module that defined them, so no need
         # to serialize.
         self.type_var_names = type_var_names
+        # Capsules needed by the module, specified via module names such as "librt.base64"
+        self.capsules: set[str] = set()
 
     def serialize(self) -> JsonDict:
         return {
@@ -38,18 +40,21 @@ def serialize(self) -> JsonDict:
             "functions": [f.serialize() for f in self.functions],
             "classes": [c.serialize() for c in self.classes],
             "final_names": [(k, t.serialize()) for k, t in self.final_names],
+            "capsules": sorted(self.capsules),
         }
 
     @classmethod
     def deserialize(cls, data: JsonDict, ctx: DeserMaps) -> ModuleIR:
-        return ModuleIR(
+        module = ModuleIR(
             data["fullname"],
             data["imports"],
             [ctx.functions[FuncDecl.get_id_from_json(f)] for f in data["functions"]],
             [ClassIR.deserialize(c, ctx) for c in data["classes"]],
             [(k, deserialize_type(t, ctx)) for k, t in data["final_names"]],
             [],
         )
+        module.capsules = set(data["capsules"])
+        return module
 
 
 def deserialize_modules(data: dict[str, JsonDict], ctx: DeserMaps) -> dict[str, ModuleIR]:

mypyc/ir/ops.py

@@ -708,6 +708,7 @@ def __init__(
         priority: int,
         is_pure: bool,
         experimental: bool,
+        capsule: str | None,
     ) -> None:
         # Each primitive much have a distinct name, but otherwise they are arbitrary.
         self.name: Final = name
@@ -733,6 +734,9 @@ def __init__(
         # Experimental primitives are not used unless mypyc experimental features are
         # explicitly enabled
         self.experimental = experimental
+        # Capsule that needs to imported and configured to call the primitive
+        # (name of the target module, e.g. "librt.base64").
+        self.capsule = capsule
 
     def __repr__(self) -> str:
         return f"<PrimitiveDescription {self.name!r}: {self.arg_types}>"
@@ -1233,6 +1237,7 @@ def __init__(
         *,
         is_pure: bool = False,
         returns_null: bool = False,
+        capsule: str | None = None,
     ) -> None:
         self.error_kind = error_kind
         super().__init__(line)
@@ -1250,6 +1255,9 @@ def __init__(
         # The function might return a null value that does not indicate
         # an error.
         self.returns_null = returns_null
+        # A capsule from this module must be imported and initialized before calling this
+        # function (used for C functions exported from librt). Example value: "librt.base64"
+        self.capsule = capsule
         if is_pure or returns_null:
             assert error_kind == ERR_NEVER
 

mypyc/irbuild/ll_builder.py

@@ -2075,6 +2075,7 @@ def call_c(
                 var_arg_idx,
                 is_pure=desc.is_pure,
                 returns_null=desc.returns_null,
+                capsule=desc.capsule,
             )
         )
         if desc.is_borrowed:
@@ -2159,6 +2160,7 @@ def primitive_op(
                 desc.priority,
                 is_pure=desc.is_pure,
                 returns_null=False,
+                capsule=desc.capsule,
             )
             return self.call_c(c_desc, args, line, result_type=result_type)
 

mypyc/options.py

@@ -18,7 +18,6 @@ def __init__(
         group_name: str | None = None,
         log_trace: bool = False,
         depends_on_librt_internal: bool = False,
-        depends_on_librt_base64: bool = False,
         experimental_features: bool = False,
     ) -> None:
         self.strip_asserts = strip_asserts
@@ -57,7 +56,6 @@ def __init__(
         # only for mypy itself, third-party code compiled with mypyc should not use
         # librt.internal.
         self.depends_on_librt_internal = depends_on_librt_internal
-        self.depends_on_librt_base64 = depends_on_librt_base64
         # Some experimental features are only available when building librt in
         # experimental mode (e.g. use _experimental suffix in librt run test).
         # These can't be used with a librt wheel installed from PyPI.

mypyc/primitives/misc_ops.py

@@ -473,4 +473,5 @@
     c_function_name="LibRTBase64_b64encode_internal",
     error_kind=ERR_MAGIC,
     experimental=True,
+    capsule="librt.base64",
 )

mypyc/primitives/registry.py

@@ -62,6 +62,7 @@ class CFunctionDescription(NamedTuple):
     priority: int
     is_pure: bool
     returns_null: bool
+    capsule: str | None
 
 
 # A description for C load operations including LoadGlobal and LoadAddress
@@ -100,6 +101,7 @@ def method_op(
     is_borrowed: bool = False,
     priority: int = 1,
     is_pure: bool = False,
+    capsule: str | None = None,
 ) -> PrimitiveDescription:
     """Define a c function call op that replaces a method call.
 
@@ -145,6 +147,7 @@ def method_op(
         priority,
         is_pure=is_pure,
         experimental=False,
+        capsule=capsule,
     )
     ops.append(desc)
     return desc
@@ -164,6 +167,7 @@ def function_op(
     is_borrowed: bool = False,
     priority: int = 1,
     experimental: bool = False,
+    capsule: str | None = None,
 ) -> PrimitiveDescription:
     """Define a C function call op that replaces a function call.
 
@@ -193,6 +197,7 @@ def function_op(
         priority=priority,
         is_pure=False,
         experimental=experimental,
+        capsule=capsule,
     )
     ops.append(desc)
     return desc
@@ -212,6 +217,7 @@ def binary_op(
     steals: StealsDescription = False,
     is_borrowed: bool = False,
     priority: int = 1,
+    capsule: str | None = None,
 ) -> PrimitiveDescription:
     """Define a c function call op for a binary operation.
 
@@ -240,6 +246,7 @@ def binary_op(
         priority=priority,
         is_pure=False,
         experimental=False,
+        capsule=capsule,
     )
     ops.append(desc)
     return desc
@@ -281,6 +288,7 @@ def custom_op(
         0,
         is_pure=is_pure,
         returns_null=returns_null,
+        capsule=None,
     )
 
 
@@ -297,6 +305,7 @@ def custom_primitive_op(
     steals: StealsDescription = False,
     is_borrowed: bool = False,
     is_pure: bool = False,
+    capsule: str | None = None,
 ) -> PrimitiveDescription:
     """Define a primitive op that can't be automatically generated based on the AST.
 
@@ -319,6 +328,7 @@ def custom_primitive_op(
         priority=0,
         is_pure=is_pure,
         experimental=False,
+        capsule=capsule,
     )
 
 
@@ -335,6 +345,7 @@ def unary_op(
     is_borrowed: bool = False,
     priority: int = 1,
     is_pure: bool = False,
+    capsule: str | None = None,
 ) -> PrimitiveDescription:
     """Define a primitive op for an unary operation.
 
@@ -361,6 +372,7 @@ def unary_op(
         priority=priority,
         is_pure=is_pure,
         experimental=False,
+        capsule=capsule,
     )
     ops.append(desc)
     return desc

mypyc/test-data/irbuild-base64.test

@@ -0,0 +1,37 @@
+[case testBase64_experimental]
+from librt.base64 import b64encode
+
+def enc(b: bytes) -> bytes:
+    return b64encode(b)
+[out]
+def enc(b):
+    b, r0 :: bytes
+L0:
+    r0 = LibRTBase64_b64encode_internal(b)
+    return r0
+
+[case testBase64ExperimentalDisabled]
+from librt.base64 import b64encode
+
+def enc(b: bytes) -> bytes:
+    return b64encode(b)
+[out]
+def enc(b):
+    b :: bytes
+    r0 :: dict
+    r1 :: str
+    r2 :: object
+    r3 :: object[1]
+    r4 :: object_ptr
+    r5 :: object
+    r6 :: bytes
+L0:
+    r0 = __main__.globals :: static
+    r1 = 'b64encode'
+    r2 = CPyDict_GetItem(r0, r1)
+    r3 = [b]
+    r4 = load_address r3
+    r5 = PyObject_Vectorcall(r2, r4, 1, 0)
+    keep_alive b
+    r6 = cast(bytes, r5)
+    return r6