[mypyc] Allow types to have capsule/source file deps (#20667)

An RType may need primitives/definitions from a capsule for things like
runtime type checks or unboxing/boxing. Also some types might have 
custom incref/decref operations defined in a header that is not included 
by default.

Author: JukkaL

mypyc/analysis/capsule_deps.py

@@ -2,7 +2,8 @@
 
 from mypyc.ir.deps import Dependency
 from mypyc.ir.func_ir import FuncIR
-from mypyc.ir.ops import CallC, PrimitiveOp
+from mypyc.ir.ops import Assign, CallC, PrimitiveOp
+from mypyc.ir.rtypes import RStruct, RTuple, RType, RUnion
 
 
 def find_implicit_op_dependencies(fn: FuncIR) -> set[Dependency] | None:
@@ -16,14 +17,51 @@ def find_implicit_op_dependencies(fn: FuncIR) -> set[Dependency] | None:
     defined in other modules.
     """
     deps: set[Dependency] | None = None
+    # Check function signature types for dependencies
+    deps = find_type_dependencies(fn, deps)
+    # Check ops for dependencies
     for block in fn.blocks:
         for op in block.ops:
-            # TODO: Also determine implicit type object dependencies (e.g. cast targets)
+            assert not isinstance(op, PrimitiveOp), "Lowered IR is expected"
             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"
+            deps = collect_type_deps(op.type, deps)
+            if isinstance(op, Assign):
+                deps = collect_type_deps(op.dest.type, deps)
+    return deps
+
+
+def find_type_dependencies(fn: FuncIR, deps: set[Dependency] | None) -> set[Dependency] | None:
+    """Find dependencies from RTypes in function signatures.
+
+    Some RTypes (e.g., those for librt types) have associated dependencies
+    that need to be imported when the type is used.
+    """
+    # Check parameter types
+    for arg in fn.decl.sig.args:
+        deps = collect_type_deps(arg.type, deps)
+    # Check return type
+    deps = collect_type_deps(fn.decl.sig.ret_type, deps)
+    return deps
+
+
+def collect_type_deps(typ: RType, deps: set[Dependency] | None) -> set[Dependency] | None:
+    """Collect dependencies from an RType, recursively checking compound types."""
+    if typ.dependencies is not None:
+        for dep in typ.dependencies:
+            if deps is None:
+                deps = set()
+            deps.add(dep)
+    if isinstance(typ, RUnion):
+        for item in typ.items:
+            deps = collect_type_deps(item, deps)
+    elif isinstance(typ, RTuple):
+        for item in typ.types:
+            deps = collect_type_deps(item, deps)
+    elif isinstance(typ, RStruct):
+        for item in typ.types:
+            deps = collect_type_deps(item, deps)
     return deps

mypyc/ir/rtypes.py

@@ -41,6 +41,7 @@ class to enable the new behavior. In rare cases, adding a new
 from typing import TYPE_CHECKING, ClassVar, Final, Generic, TypeGuard, TypeVar, final
 
 from mypyc.common import HAVE_IMMORTAL, IS_32_BIT_PLATFORM, PLATFORM_SIZE, JsonDict, short_name
+from mypyc.ir.deps import LIBRT_STRINGS, Dependency
 from mypyc.namegen import NameGenerator
 
 if TYPE_CHECKING:
@@ -88,6 +89,7 @@ class RType:
     # we never raise an AttributeError and don't need the bitfield
     # entry.)
     error_overlap = False
+    dependencies: tuple[Dependency, ...] | None = None
 
     @abstractmethod
     def accept(self, visitor: RTypeVisitor[T]) -> T:
@@ -232,6 +234,7 @@ def __init__(
         size: int = PLATFORM_SIZE,
         error_overlap: bool = False,
         may_be_immortal: bool = True,
+        dependencies: tuple[Dependency, ...] | None = None,
     ) -> None:
         RPrimitive.primitive_map[name] = self
 
@@ -244,6 +247,7 @@ def __init__(
         self.size = size
         self.error_overlap = error_overlap
         self._may_be_immortal = may_be_immortal and HAVE_IMMORTAL
+        self.dependencies = dependencies
         if ctype == "CPyTagged":
             self.c_undefined = "CPY_INT_TAG"
         elif ctype in ("int16_t", "int32_t", "int64_t"):
@@ -517,7 +521,7 @@ def __hash__(self) -> int:
 range_rprimitive: Final = RPrimitive("builtins.range", is_unboxed=False, is_refcounted=True)
 
 KNOWN_NATIVE_TYPES: Final = {
-    name: RPrimitive(name, is_unboxed=False, is_refcounted=True)
+    name: RPrimitive(name, is_unboxed=False, is_refcounted=True, dependencies=(LIBRT_STRINGS,))
     for name in [
         "librt.internal.WriteBuffer",
         "librt.internal.ReadBuffer",

mypyc/test-data/capsule-deps.test

@@ -0,0 +1,89 @@
+[case testNoDeps]
+def f() -> None:
+    pass
+[out]
+No deps
+
+[case testStringsNoCapsuleDepWithoutExperimental]
+# Test case missing _experimental suffix -> no deps
+from librt.strings import StringWriter
+
+def f() -> None:
+    StringWriter()
+[out]
+No deps
+
+[case testStringsCapsuleDep_experimental]
+from librt.strings import StringWriter
+
+def f() -> None:
+    StringWriter()
+[out]
+Capsule(name='librt.strings')
+
+[case testStringsCapsuleDepFromParamType_experimental]
+from librt.strings import StringWriter
+
+def f(s: StringWriter) -> None:
+    pass
+[out]
+Capsule(name='librt.strings')
+
+[case testStringsCapsuleDepFromReturnType_experimental]
+from librt.strings import StringWriter
+
+def f() -> StringWriter:
+    assert False
+[out]
+Capsule(name='librt.strings')
+
+[case testStringsCapsuleDepFromUnion_experimental]
+from typing import Union
+
+from librt.strings import StringWriter
+
+def f(s: Union[StringWriter, str]) -> None:
+    pass
+[out]
+Capsule(name='librt.strings')
+
+[case testStringsCapsuleDepFromTuple_experimental]
+from librt.strings import StringWriter
+
+def f(s: tuple[int, StringWriter]) -> None:
+    pass
+[out]
+Capsule(name='librt.strings')
+
+[case testExtraFileDep_experimental]
+from librt.strings import StringWriter
+
+def f(s: StringWriter) -> int:
+    return s[0]
+[out]
+Capsule(name='librt.strings')
+SourceDep(path='stringwriter_extra_ops.c')
+
+[case testMultipleCapsuleDeps1_experimental]
+from librt.strings import StringWriter
+from librt.base64 import b64encode
+
+def f() -> None:
+    StringWriter()
+
+def g() -> bytes:
+    return b64encode(b'foo')
+[out]
+Capsule(name='librt.base64')
+Capsule(name='librt.strings')
+
+[case testMultipleCapsuleDeps2_experimental]
+from librt.strings import StringWriter
+from librt.base64 import b64encode
+
+def f() -> bytes:
+    StringWriter()
+    return b64encode(b'foo')
+[out]
+Capsule(name='librt.base64')
+Capsule(name='librt.strings')

mypyc/test/test_capsule_deps.py

@@ -0,0 +1,53 @@
+"""Test cases for capsule dependency analysis."""
+
+from __future__ import annotations
+
+import os.path
+
+from mypy.errors import CompileError
+from mypy.test.config import test_temp_dir
+from mypy.test.data import DataDrivenTestCase
+from mypyc.analysis.capsule_deps import find_implicit_op_dependencies
+from mypyc.common import TOP_LEVEL_NAME
+from mypyc.options import CompilerOptions
+from mypyc.test.testutil import (
+    ICODE_GEN_BUILTINS,
+    MypycDataSuite,
+    assert_test_output,
+    build_ir_for_single_file,
+    infer_ir_build_options_from_test_name,
+    use_custom_builtins,
+)
+from mypyc.transform.lower import lower_ir
+
+files = ["capsule-deps.test"]
+
+
+class TestCapsuleDeps(MypycDataSuite):
+    files = files
+    base_path = test_temp_dir
+
+    def run_case(self, testcase: DataDrivenTestCase) -> None:
+        options = infer_ir_build_options_from_test_name(testcase.name)
+        if options is None:
+            # Skipped test case
+            return
+        with use_custom_builtins(os.path.join(self.data_prefix, ICODE_GEN_BUILTINS), testcase):
+            try:
+                ir = build_ir_for_single_file(testcase.input, options)
+            except CompileError as e:
+                actual = e.messages
+            else:
+                all_deps: set[str] = set()
+                for fn in ir:
+                    if fn.name == TOP_LEVEL_NAME and not testcase.name.endswith("_toplevel"):
+                        continue
+                    compiler_options = CompilerOptions()
+                    lower_ir(fn, compiler_options)
+                    deps = find_implicit_op_dependencies(fn)
+                    if deps:
+                        for dep in deps:
+                            all_deps.add(repr(dep))
+                actual = sorted(all_deps) if all_deps else ["No deps"]
+
+            assert_test_output(testcase, actual, "Invalid test output", testcase.output)