[mypyc] Allow primitives to depend on optional C source files (#20379)

This way we don't need to compile all support rt library code every time
mypyc is run. Some primitives are not used very often, but may have
relatively complex implementations. Update `bytes.translate` as an
example.

Author: JukkaL

mypyc/analysis/capsule_deps.py

@@ -1,27 +1,29 @@
 from __future__ import annotations
 
+from mypyc.ir.deps import Dependency
 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.
+def find_implicit_op_dependencies(fn: FuncIR) -> set[Dependency] | None:
+    """Find implicit dependencies that need to be imported.
 
     Using primitives or types defined in librt submodules such as "librt.base64"
-    requires a capsule import.
+    requires dependency imports (e.g., capsule imports).
 
     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
+    deps: set[Dependency] | 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)
+            if isinstance(op, CallC) and op.dependencies is not None:
+                for dep in op.dependencies:
+                    if deps is None:
+                        deps = set()
+                    deps.add(dep)
             else:
                 assert not isinstance(op, PrimitiveOp), "Lowered IR is expected"
     return deps

mypyc/build.py

@@ -39,6 +39,7 @@
 from mypyc.codegen import emitmodule
 from mypyc.common import IS_FREE_THREADED, RUNTIME_C_FILES, shared_lib_name
 from mypyc.errors import Errors
+from mypyc.ir.deps import SourceDep
 from mypyc.ir.pprint import format_modules
 from mypyc.namegen import exported_name
 from mypyc.options import CompilerOptions
@@ -282,14 +283,14 @@ def generate_c(
     groups: emitmodule.Groups,
     fscache: FileSystemCache,
     compiler_options: CompilerOptions,
-) -> tuple[list[list[tuple[str, str]]], str]:
+) -> tuple[list[list[tuple[str, str]]], str, list[SourceDep]]:
     """Drive the actual core compilation step.
 
     The groups argument describes how modules are assigned to C
     extension modules. See the comments on the Groups type in
     mypyc.emitmodule for details.
 
-    Returns the C source code and (for debugging) the pretty printed IR.
+    Returns the C source code, (for debugging) the pretty printed IR, and list of SourceDeps.
     """
     t0 = time.time()
 
@@ -325,7 +326,10 @@ def generate_c(
     if options.mypyc_annotation_file:
         generate_annotated_html(options.mypyc_annotation_file, result, modules, mapper)
 
-    return ctext, "\n".join(format_modules(modules))
+    # Collect SourceDep dependencies
+    source_deps = sorted(emitmodule.collect_source_dependencies(modules), key=lambda d: d.path)
+
+    return ctext, "\n".join(format_modules(modules)), source_deps
 
 
 def build_using_shared_lib(
@@ -486,9 +490,9 @@ def mypyc_build(
     *,
     separate: bool | list[tuple[list[str], str | None]] = False,
     only_compile_paths: Iterable[str] | None = None,
-    skip_cgen_input: Any | None = None,
+    skip_cgen_input: tuple[list[list[tuple[str, str]]], list[str]] | None = None,
     always_use_shared_lib: bool = False,
-) -> tuple[emitmodule.Groups, list[tuple[list[str], list[str]]]]:
+) -> tuple[emitmodule.Groups, list[tuple[list[str], list[str]]], list[SourceDep]]:
     """Do the front and middle end of mypyc building, producing and writing out C source."""
     fscache = FileSystemCache()
     mypyc_sources, all_sources, options = get_mypy_config(
@@ -511,14 +515,16 @@ def mypyc_build(
 
     # We let the test harness just pass in the c file contents instead
     # so that it can do a corner-cutting version without full stubs.
+    source_deps: list[SourceDep] = []
     if not skip_cgen_input:
-        group_cfiles, ops_text = generate_c(
+        group_cfiles, ops_text, source_deps = generate_c(
             all_sources, options, groups, fscache, compiler_options=compiler_options
         )
         # TODO: unique names?
         write_file(os.path.join(compiler_options.target_dir, "ops.txt"), ops_text)
     else:
-        group_cfiles = skip_cgen_input
+        group_cfiles = skip_cgen_input[0]
+        source_deps = [SourceDep(d) for d in skip_cgen_input[1]]
 
     # Write out the generated C and collect the files for each group
     # Should this be here??
@@ -535,7 +541,7 @@ def mypyc_build(
         deps = [os.path.join(compiler_options.target_dir, dep) for dep in get_header_deps(cfiles)]
         group_cfilenames.append((cfilenames, deps))
 
-    return groups, group_cfilenames
+    return groups, group_cfilenames, source_deps
 
 
 def mypycify(
@@ -548,7 +554,7 @@ def mypycify(
     strip_asserts: bool = False,
     multi_file: bool = False,
     separate: bool | list[tuple[list[str], str | None]] = False,
-    skip_cgen_input: Any | None = None,
+    skip_cgen_input: tuple[list[list[tuple[str, str]]], list[str]] | None = None,
     target_dir: str | None = None,
     include_runtime_files: bool | None = None,
     strict_dunder_typing: bool = False,
@@ -633,7 +639,7 @@ def mypycify(
     )
 
     # Generate all the actual important C code
-    groups, group_cfilenames = mypyc_build(
+    groups, group_cfilenames, source_deps = mypyc_build(
         paths,
         only_compile_paths=only_compile_paths,
         compiler_options=compiler_options,
@@ -708,11 +714,19 @@ def mypycify(
     # compiler invocations.
     shared_cfilenames = []
     if not compiler_options.include_runtime_files:
-        for name in RUNTIME_C_FILES:
+        # Collect all files to copy: runtime files + conditional source files
+        files_to_copy = list(RUNTIME_C_FILES)
+        for source_dep in source_deps:
+            files_to_copy.append(source_dep.path)
+            files_to_copy.append(source_dep.get_header())
+
+        # Copy all files
+        for name in files_to_copy:
             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())
-            shared_cfilenames.append(rt_file)
+            if name.endswith(".c"):
+                shared_cfilenames.append(rt_file)
 
     extensions = []
     for (group_sources, lib_name), (cfilenames, deps) in zip(groups, group_cfilenames):

mypyc/codegen/emitmodule.py

@@ -27,7 +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.analysis.capsule_deps import find_implicit_op_dependencies
 from mypyc.codegen.cstring import c_string_initializer
 from mypyc.codegen.emit import (
     Emitter,
@@ -56,6 +56,7 @@
     short_id_from_name,
 )
 from mypyc.errors import Errors
+from mypyc.ir.deps import LIBRT_BASE64, LIBRT_STRINGS, SourceDep
 from mypyc.ir.func_ir import FuncIR
 from mypyc.ir.module_ir import ModuleIR, ModuleIRs, deserialize_modules
 from mypyc.ir.ops import DeserMaps, LoadLiteral
@@ -263,9 +264,9 @@ 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)
+            deps = find_implicit_op_dependencies(fn)
+            if deps is not None:
+                module.dependencies.update(deps)
             # Perform optimizations.
             do_copy_propagation(fn, compiler_options)
             do_flag_elimination(fn, compiler_options)
@@ -427,6 +428,16 @@ def load_scc_from_cache(
     return modules
 
 
+def collect_source_dependencies(modules: dict[str, ModuleIR]) -> set[SourceDep]:
+    """Collect all SourceDep dependencies from all modules."""
+    source_deps: set[SourceDep] = set()
+    for module in modules.values():
+        for dep in module.dependencies:
+            if isinstance(dep, SourceDep):
+                source_deps.add(dep)
+    return source_deps
+
+
 def compile_modules_to_c(
     result: BuildResult, compiler_options: CompilerOptions, errors: Errors, groups: Groups
 ) -> tuple[ModuleIRs, list[FileContents], Mapper]:
@@ -560,6 +571,10 @@ def generate_c_for_modules(self) -> list[tuple[str, str]]:
         if self.compiler_options.include_runtime_files:
             for name in RUNTIME_C_FILES:
                 base_emitter.emit_line(f'#include "{name}"')
+            # Include conditional source files
+            source_deps = collect_source_dependencies(self.modules)
+            for source_dep in sorted(source_deps, key=lambda d: d.path):
+                base_emitter.emit_line(f'#include "{source_dep.path}"')
         base_emitter.emit_line(f'#include "__native{self.short_group_suffix}.h"')
         base_emitter.emit_line(f'#include "__native_internal{self.short_group_suffix}.h"')
         emitter = base_emitter
@@ -611,10 +626,14 @@ 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 any("librt.base64" in mod.capsules for mod in self.modules.values()):
+        if any(LIBRT_BASE64 in mod.dependencies for mod in self.modules.values()):
             ext_declarations.emit_line("#include <librt_base64.h>")
-        if any("librt.strings" in mod.capsules for mod in self.modules.values()):
+        if any(LIBRT_STRINGS in mod.dependencies for mod in self.modules.values()):
             ext_declarations.emit_line("#include <librt_strings.h>")
+        # Include headers for conditional source files
+        source_deps = collect_source_dependencies(self.modules)
+        for source_dep in sorted(source_deps, key=lambda d: d.path):
+            ext_declarations.emit_line(f'#include "{source_dep.get_header()}"')
 
         declarations = Emitter(self.context)
         declarations.emit_line(f"#ifndef MYPYC_LIBRT_INTERNAL{self.group_suffix}_H")
@@ -1072,11 +1091,11 @@ def emit_module_exec_func(
             emitter.emit_line("if (import_librt_internal() < 0) {")
             emitter.emit_line("return -1;")
             emitter.emit_line("}")
-        if "librt.base64" in module.capsules:
+        if LIBRT_BASE64 in module.dependencies:
             emitter.emit_line("if (import_librt_base64() < 0) {")
             emitter.emit_line("return -1;")
             emitter.emit_line("}")
-        if "librt.strings" in module.capsules:
+        if LIBRT_STRINGS in module.dependencies:
             emitter.emit_line("if (import_librt_strings() < 0) {")
             emitter.emit_line("return -1;")
             emitter.emit_line("}")

mypyc/common.py

@@ -65,7 +65,8 @@
 BITMAP_TYPE: Final = "uint32_t"
 BITMAP_BITS: Final = 32
 
-# Runtime C library files
+# Runtime C library files that are always included (some ops may bring
+# extra dependencies via mypyc.ir.SourceDep)
 RUNTIME_C_FILES: Final = [
     "init.c",
     "getargs.c",

mypyc/ir/deps.py

@@ -0,0 +1,52 @@
+from typing import Final
+
+
+class Capsule:
+    """Defines a C extension capsule that a primitive may require."""
+
+    def __init__(self, name: str) -> None:
+        # Module fullname, e.g. 'librt.base64'
+        self.name: Final = name
+
+    def __repr__(self) -> str:
+        return f"Capsule(name={self.name!r})"
+
+    def __eq__(self, other: object) -> bool:
+        return isinstance(other, Capsule) and self.name == other.name
+
+    def __hash__(self) -> int:
+        return hash(("Capsule", self.name))
+
+
+class SourceDep:
+    """Defines a C source file that a primitive may require.
+
+    Each source file must also have a corresponding .h file (replace .c with .h)
+    that gets implicitly #included if the source is used.
+    """
+
+    def __init__(self, path: str) -> None:
+        # Relative path from mypyc/lib-rt, e.g. 'bytes_extra_ops.c'
+        self.path: Final = path
+
+    def __repr__(self) -> str:
+        return f"SourceDep(path={self.path!r})"
+
+    def __eq__(self, other: object) -> bool:
+        return isinstance(other, SourceDep) and self.path == other.path
+
+    def __hash__(self) -> int:
+        return hash(("SourceDep", self.path))
+
+    def get_header(self) -> str:
+        """Get the header file path by replacing .c with .h"""
+        return self.path.replace(".c", ".h")
+
+
+Dependency = Capsule | SourceDep
+
+
+LIBRT_STRINGS: Final = Capsule("librt.strings")
+LIBRT_BASE64: Final = Capsule("librt.base64")
+
+BYTES_EXTRA_OPS: Final = SourceDep("bytes_extra_ops.c")

mypyc/ir/module_ir.py

@@ -4,6 +4,7 @@
 
 from mypyc.common import JsonDict
 from mypyc.ir.class_ir import ClassIR
+from mypyc.ir.deps import Capsule, Dependency, SourceDep
 from mypyc.ir.func_ir import FuncDecl, FuncIR
 from mypyc.ir.ops import DeserMaps
 from mypyc.ir.rtypes import RType, deserialize_type
@@ -30,17 +31,25 @@ 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()
+        # Dependencies needed by the module (such as capsules or source files)
+        self.dependencies: set[Dependency] = set()
 
     def serialize(self) -> JsonDict:
+        # Serialize dependencies as a list of dicts with type information
+        serialized_deps = []
+        for dep in sorted(self.dependencies, key=lambda d: (type(d).__name__, str(d))):
+            if isinstance(dep, Capsule):
+                serialized_deps.append({"type": "Capsule", "name": dep.name})
+            elif isinstance(dep, SourceDep):
+                serialized_deps.append({"type": "SourceDep", "path": dep.path})
+
         return {
             "fullname": self.fullname,
             "imports": self.imports,
             "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),
+            "dependencies": serialized_deps,
         }
 
     @classmethod
@@ -53,7 +62,16 @@ def deserialize(cls, data: JsonDict, ctx: DeserMaps) -> ModuleIR:
             [(k, deserialize_type(t, ctx)) for k, t in data["final_names"]],
             [],
         )
-        module.capsules = set(data["capsules"])
+
+        # Deserialize dependencies
+        deps: set[Dependency] = set()
+        for dep_dict in data["dependencies"]:
+            if dep_dict["type"] == "Capsule":
+                deps.add(Capsule(dep_dict["name"]))
+            elif dep_dict["type"] == "SourceDep":
+                deps.add(SourceDep(dep_dict["path"]))
+        module.dependencies = deps
+
         return module