[mypyc] Add new type for expressing header dependencies

mypyc/codegen/emitmodule.py

@@ -56,7 +56,15 @@
     short_id_from_name,
 )
 from mypyc.errors import Errors
-from mypyc.ir.deps import LIBRT_BASE64, LIBRT_STRINGS, LIBRT_TIME, LIBRT_VECS, SourceDep
+from mypyc.ir.deps import (
+    LIBRT_BASE64,
+    LIBRT_STRINGS,
+    LIBRT_TIME,
+    LIBRT_VECS,
+    Capsule,
+    HeaderDep,
+    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
@@ -436,24 +444,31 @@ def load_scc_from_cache(
     return modules
 
 
-def collect_source_dependencies(
-    modules: dict[str, ModuleIR], *, internal: bool = True
-) -> set[SourceDep]:
-    """Collect all SourceDep dependencies from all modules.
-
-    If internal is set to False, returns only the dependencies that can be exported to C extensions
-    dependent on the one currently being compiled.
-    """
+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):
-                if internal == dep.internal:
+                if dep.internal:
                     source_deps.add(dep)
+            elif isinstance(dep, Capsule):
+                source_deps.add(dep.internal_dep())
+    return source_deps
+
+
+def collect_header_dependencies(modules: dict[str, ModuleIR], *, internal: bool) -> set[str]:
+    """Collect all header dependencies from all modules."""
+    header_deps: set[str] = set()
+    for module in modules.values():
+        for dep in module.dependencies:
+            if isinstance(dep, (SourceDep, HeaderDep)):
+                if dep.internal == internal:
+                    header_deps.add(dep.get_header())
             else:
                 capsule_dep = dep.internal_dep() if internal else dep.external_dep()
-                source_deps.add(capsule_dep)
-    return source_deps
+                header_deps.add(capsule_dep.get_header())
+    return header_deps
 
 
 def compile_modules_to_c(
@@ -652,9 +667,9 @@ def emit_dep_headers(decls: Emitter, internal: bool) -> None:
             if self.compiler_options.depends_on_librt_internal:
                 decls.emit_line(f'#include "internal/librt_internal{suffix}.h"')
             # Include headers for conditional source files
-            source_deps = collect_source_dependencies(self.modules, internal=internal)
-            for source_dep in sorted(source_deps, key=lambda d: d.path):
-                decls.emit_line(f'#include "{source_dep.get_header()}"')
+            header_deps = collect_header_dependencies(self.modules, internal=internal)
+            for header_dep in sorted(header_deps):
+                decls.emit_line(f'#include "{header_dep}"')
 
         emit_dep_headers(ext_declarations, False)
 

mypyc/common.py

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

mypyc/ir/deps.py

@@ -26,11 +26,8 @@ def internal_dep(self) -> SourceDep:
         module = self.name.split(".")[-1]
         return SourceDep(f"{module}/librt_{module}_api.c", include_dirs=[module])
 
-    # TODO: This SourceDep is really only used for its associated header so it would make more sense
-    # to add a separate type. Alternatively, see if this can be removed altogether if we move the
-    # definitions that depend on this header from the external header of the C extension.
-    def external_dep(self) -> SourceDep:
-        """External source dependency of the capsule that may be included in external headers of C
+    def external_dep(self) -> HeaderDep:
+        """External header dependency of the capsule that may be included in external headers of C
         extensions that depend on the capsule.
 
         The external headers of the C extensions are included by other C extensions that don't
@@ -42,7 +39,7 @@ def external_dep(self) -> SourceDep:
         including the internal header would result in undefined symbols.
         """
         module = self.name.split(".")[-1]
-        return SourceDep(f"{module}/librt_{module}.c", include_dirs=[module], internal=False)
+        return HeaderDep(f"{module}/librt_{module}.h", include_dirs=[module], internal=False)
 
 
 class SourceDep:
@@ -76,7 +73,36 @@ def get_header(self) -> str:
         return self.path.replace(".c", ".h")
 
 
-Dependency = Capsule | SourceDep
+class HeaderDep:
+    """Defines a C header file that a primitive may require.
+
+    The header gets explicitly #included if the dependency is used.
+    include_dirs are passed to the C compiler when the generated extension
+    is compiled separately and needs to include the header.
+    """
+
+    def __init__(
+        self, path: str, *, include_dirs: list[str] | None = None, internal: bool = True
+    ) -> None:
+        # Relative path from mypyc/lib-rt, e.g. 'strings/librt_strings.h'
+        self.path: Final = path
+        self.include_dirs: Final = include_dirs or []
+        self.internal: Final = internal
+
+    def __repr__(self) -> str:
+        return f"HeaderDep(path={self.path!r})"
+
+    def __eq__(self, other: object) -> bool:
+        return isinstance(other, HeaderDep) and self.path == other.path
+
+    def __hash__(self) -> int:
+        return hash(("HeaderDep", self.path))
+
+    def get_header(self) -> str:
+        return self.path
+
+
+Dependency = Capsule | SourceDep | HeaderDep
 
 
 LIBRT_STRINGS: Final = Capsule("librt.strings")

mypyc/ir/module_ir.py

@@ -4,7 +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.deps import Capsule, Dependency, HeaderDep, SourceDep
 from mypyc.ir.func_ir import FuncDecl, FuncIR
 from mypyc.ir.ops import DeserMaps
 from mypyc.ir.rtypes import RType, deserialize_type
@@ -48,6 +48,14 @@ def serialize(self) -> JsonDict:
                     "internal": dep.internal,
                 }
                 serialized_deps.append(source_dep)
+            elif isinstance(dep, HeaderDep):
+                header_dep: JsonDict = {
+                    "type": "HeaderDep",
+                    "path": dep.path,
+                    "include_dirs": dep.include_dirs,
+                    "internal": dep.internal,
+                }
+                serialized_deps.append(header_dep)
 
         return {
             "fullname": self.fullname,
@@ -82,6 +90,14 @@ def deserialize(cls, data: JsonDict, ctx: DeserMaps) -> ModuleIR:
                         internal=dep_dict["internal"],
                     )
                 )
+            elif dep_dict["type"] == "HeaderDep":
+                deps.add(
+                    HeaderDep(
+                        dep_dict["path"],
+                        include_dirs=dep_dict["include_dirs"],
+                        internal=dep_dict["internal"],
+                    )
+                )
         module.dependencies = deps
 
         return module

mypyc/test/test_cheader.py

@@ -7,7 +7,7 @@
 import re
 import unittest
 
-from mypyc.ir.deps import SourceDep
+from mypyc.ir.deps import HeaderDep, SourceDep
 from mypyc.ir.ops import PrimitiveDescription
 from mypyc.primitives import (
     bytearray_ops,
@@ -79,7 +79,7 @@ def check_name(name: str) -> None:
         for op in all_ops:
             if op.dependencies:
                 for dep in op.dependencies:
-                    if isinstance(dep, SourceDep):
+                    if isinstance(dep, (SourceDep, HeaderDep)):
                         header_fnam = os.path.join(base_dir, dep.get_header())
                         if os.path.isfile(header_fnam):
                             with open(os.path.join(base_dir, header_fnam)) as f: