From 1416681537f456a687b3e80a36901589c9efbbce Mon Sep 17 00:00:00 2001 From: yuchuan Date: Sun, 6 Sep 2026 11:00:54 -0400 Subject: [PATCH 1/2] unify the check of ffi types. Signed-off-by: yuchuan --- docs/packaging/stubgen.rst | 3 ++- python/tvm_ffi/stub/cli.py | 25 +++++++++++-------- python/tvm_ffi/stub/consts.py | 1 + python/tvm_ffi/stub/generator.py | 8 ++++++ .../stub/python_generator/generator.py | 4 +++ python/tvm_ffi/stub/rust_generator/codegen.py | 13 +++++++--- .../tvm_ffi/stub/rust_generator/generator.py | 6 ++++- python/tvm_ffi/stub/rust_generator/utils.py | 5 ++++ tests/python/test_stubgen_rust.py | 22 +++++++++++++--- 9 files changed, 67 insertions(+), 20 deletions(-) diff --git a/docs/packaging/stubgen.rst b/docs/packaging/stubgen.rst index 98087e2e9..7b299d461 100644 --- a/docs/packaging/stubgen.rst +++ b/docs/packaging/stubgen.rst @@ -390,7 +390,8 @@ When you run the tool, it: ``prefix`` - Demand a Namespace Rust target only. Adds an ``object/`` block to the file for every object registered directly under the prefix that no processed file defines yet; ``skip`` - leaves one out. Code outside the blocks is preserved. + leaves one out, and a key the ``tvm_ffi`` crate binds itself (every ``ffi.*`` key) + is never added. Code outside the blocks is preserved. .. code-block:: rust diff --git a/python/tvm_ffi/stub/cli.py b/python/tvm_ffi/stub/cli.py index 4f96e9aed..af8f8c247 100644 --- a/python/tvm_ffi/stub/cli.py +++ b/python/tvm_ffi/stub/cli.py @@ -98,7 +98,7 @@ def __main__() -> int: # Stage 2b. Add the object blocks a `tvm-ffi-stubgen(prefix)` file asks for. This runs # after `--init`, which rewrites files on disk and reloads them. - failed += _roll_out_prefixes(files) + failed += _roll_out_prefixes(files, generator) # Stage 3: Process # - `tvm-ffi-stubgen(begin): global/...` @@ -166,15 +166,14 @@ def _stage_1( ty_map[lhs.strip()] = rhs.strip() -def _roll_out_prefixes(files: list[FileInfo]) -> int: +def _roll_out_prefixes(files: list[FileInfo], generator: Generator) -> int: """Append an ``object/`` block for each registered object under a file's ``prefix``. - Keys with a block in any file of the run, or named by ``skip``, are left alone; an - ``import-section`` is added when the file has none. Returns the number of bad files. + Keys with a block in any file of the run, named by ``skip``, or bound by the target's + runtime are left alone; an ``import-section`` is added when the file has none. Returns + the number of bad files. """ - defined = { - code.param for file in files for code in file.code_blocks if code.kind == "object" - } | C.BUILTIN_TYPE_KEYS + defined = {code.param for file in files for code in file.code_blocks if code.kind == "object"} registry = collect_type_keys() owners: dict[str, Path] = {} failed = 0 @@ -200,7 +199,11 @@ def _roll_out_prefixes(files: list[FileInfo]) -> int: ) continue skipped = {c.param[1].strip() for c in directives if c.param[0] == "skip"} - keys = [key for key in registry[prefix] if key not in defined and key not in skipped] + keys = [ + key + for key in registry[prefix] + if key not in defined and key not in skipped and not generator.is_builtin(key) + ] blocks = file.code_blocks if not any(c.kind == "import-section" for c in blocks): at = blocks.index(head) + 1 @@ -251,7 +254,7 @@ def _find_or_insert_file(path: Path) -> FileInfo: } defined_objs: set[str] = { # ty: ignore[invalid-assignment] code.param for file in files for code in file.code_blocks if code.kind == "object" - } | C.BUILTIN_TYPE_KEYS + } skipped: set[str] = { code.param[1].strip() for file in files @@ -275,7 +278,9 @@ def _find_or_insert_file(path: Path) -> FileInfo: [] if prefix in defined_func_prefixes else global_funcs.get(prefix, []), key=lambda f: f.schema.name, ) - objs = sorted(set(obj_names) - defined_objs - skipped) + objs = sorted( + key for key in set(obj_names) - defined_objs - skipped if not generator.is_builtin(key) + ) object_infos = toposort_objects(objs) if not funcs and not object_infos: continue diff --git a/python/tvm_ffi/stub/consts.py b/python/tvm_ffi/stub/consts.py index 5df28d5f8..164322677 100644 --- a/python/tvm_ffi/stub/consts.py +++ b/python/tvm_ffi/stub/consts.py @@ -115,6 +115,7 @@ def directive(self, name: str) -> str: # the FFI reflection registry (see `utils.ObjectInfo.from_type_info`). FN_NAME_MAP: dict[str, str] = {} +#: Type keys the Python package binds itself (see ``PythonGenerator.is_builtin``). BUILTIN_TYPE_KEYS = { "ffi.Bytes", "ffi.Error", diff --git a/python/tvm_ffi/stub/generator.py b/python/tvm_ffi/stub/generator.py index add7503a5..c22522a42 100644 --- a/python/tvm_ffi/stub/generator.py +++ b/python/tvm_ffi/stub/generator.py @@ -114,6 +114,14 @@ def extra_export_names(self, imports: Any) -> set[str]: """Return extra public-export names implied by the collected imports.""" ... + def is_builtin(self, type_key: str) -> bool: + """Whether the target's runtime binds ``type_key`` itself. + + Such a key gets no ``object/`` block: neither ``--init`` scaffolds one nor a + ``prefix`` directive rolls one out. + """ + ... + # --- per-block generation (mutates `code.lines`) ------------------------ def generate_global_funcs_block( diff --git a/python/tvm_ffi/stub/python_generator/generator.py b/python/tvm_ffi/stub/python_generator/generator.py index cd81821d7..5d784818c 100644 --- a/python/tvm_ffi/stub/python_generator/generator.py +++ b/python/tvm_ffi/stub/python_generator/generator.py @@ -75,6 +75,10 @@ def extra_export_names(self, imports: PythonImports) -> set[str]: """Return extra ``__all__`` names implied by the collected imports.""" return {"LIB"} if imports.has_lib_load else set() + def is_builtin(self, type_key: str) -> bool: + """Whether the ``tvm_ffi`` package binds ``type_key`` itself.""" + return type_key in C.BUILTIN_TYPE_KEYS + # --- per-block generation (mutates `code.lines`) ------------------------ def generate_global_funcs_block( diff --git a/python/tvm_ffi/stub/rust_generator/codegen.py b/python/tvm_ffi/stub/rust_generator/codegen.py index c5da9b4f4..d7db7a6b7 100644 --- a/python/tvm_ffi/stub/rust_generator/codegen.py +++ b/python/tvm_ffi/stub/rust_generator/codegen.py @@ -52,7 +52,13 @@ from ..layout import Verdict, classify from ..lib_state import object_info_from_type_key from . import consts as C_RUST -from .utils import RustImports, builtin_mirror_name, render_rust_type, rust_ident +from .utils import ( + RustImports, + builtin_mirror_name, + is_crate_type_key, + render_rust_type, + rust_ident, +) if TYPE_CHECKING: from collections.abc import Container @@ -132,8 +138,7 @@ def _generated_type_path(self, type_key: str) -> str: Same module: the bare leaf. Elsewhere: ``super::`` per segment of this file's module, then the full path (edition 2021 rejects ``use ir::Expr``). """ - head, _, _ = type_key.partition(".") - if head in C_RUST.RUST_MOD_MAP: + if is_crate_type_key(type_key): return type_key mod, _, type_leaf = type_key.rpartition(".") if tuple(mod.split(".")) == self.mod_segments: @@ -147,7 +152,7 @@ def _provider(self, type_key: str) -> str | None: ``"crate"`` for builtin ``ffi.*`` types, ``"mapped"`` for a ``ty-map`` to a hand-written binding, ``"generated"`` for an ``object/`` block in this run. """ - if type_key.partition(".")[0] in C_RUST.RUST_MOD_MAP: + if is_crate_type_key(type_key): return "crate" if type_key in self.ty_map: return "mapped" diff --git a/python/tvm_ffi/stub/rust_generator/generator.py b/python/tvm_ffi/stub/rust_generator/generator.py index 30b44dbdc..edfc7c078 100644 --- a/python/tvm_ffi/stub/rust_generator/generator.py +++ b/python/tvm_ffi/stub/rust_generator/generator.py @@ -32,7 +32,7 @@ generate_rust_import_section, generate_rust_object, ) -from .utils import RustImports, RustUse +from .utils import RustImports, RustUse, is_crate_type_key if TYPE_CHECKING: from collections.abc import Container @@ -75,6 +75,10 @@ def extra_export_names(self, imports: RustImports) -> set[str]: """No extra export names for Rust.""" return set() + def is_builtin(self, type_key: str) -> bool: + """Every ``ffi.*`` key is bound by the crate (the rule ``codegen`` resolves types with).""" + return is_crate_type_key(type_key) + # --- per-block generation ------------------------------------------------- def generate_global_funcs_block( diff --git a/python/tvm_ffi/stub/rust_generator/utils.py b/python/tvm_ffi/stub/rust_generator/utils.py index d38e35bcc..dcd68e3dd 100644 --- a/python/tvm_ffi/stub/rust_generator/utils.py +++ b/python/tvm_ffi/stub/rust_generator/utils.py @@ -58,6 +58,11 @@ def as_use_line(self) -> str: return f"use {self.path};" +def is_crate_type_key(type_key: str) -> bool: + """Whether the ``tvm_ffi`` crate binds ``type_key`` itself (every ``ffi.*`` key).""" + return type_key.partition(".")[0] in C.RUST_MOD_MAP + + def builtin_mirror_name(type_key: str) -> str: """Name of the header-only stand-in for a builtin type (``ffi.IntEnum -> FfiIntEnumObj``). diff --git a/tests/python/test_stubgen_rust.py b/tests/python/test_stubgen_rust.py index 97a446d21..825559542 100644 --- a/tests/python/test_stubgen_rust.py +++ b/tests/python/test_stubgen_rust.py @@ -1369,7 +1369,7 @@ def test_roll_out_matches_the_prefix_exactly( monkeypatch.setattr( stub_cli, "collect_type_keys", - lambda: {"a": ["a.Y", "a.X"], "a.b": ["a.b.Z"], "ffi": ["ffi.Object"]}, + lambda: {"a": ["a.Y", "a.X"], "a.b": ["a.b.Z"], "ffi": ["ffi.Array", "ffi.Object"]}, ) monkeypatch.setattr( stub_cli, @@ -1392,7 +1392,7 @@ def test_roll_out_matches_the_prefix_exactly( ) info = FileInfo.from_file(src) assert info is not None - assert stub_cli._roll_out_prefixes([info]) == 0 + assert stub_cli._roll_out_prefixes([info], RUST) == 0 assert [line for block in info.code_blocks for line in block.lines] == [ f"{C.RUST_SYNTAX.directive('prefix')} a.", f"{begin} import-section", @@ -1405,14 +1405,28 @@ def test_roll_out_matches_the_prefix_exactly( end, "mod tail {}", ] - # Builtin type keys live in the crate and are never rolled out. + # Every `ffi.*` key lives in the crate and is never rolled out, `ffi.Array` included: a + # generated non-generic `Array` would shadow the crate's `Array`. src.write_text(f"{C.RUST_SYNTAX.directive('prefix')} ffi\n", encoding="utf-8") info = FileInfo.from_file(src) assert info is not None - assert stub_cli._roll_out_prefixes([info]) == 0 + assert stub_cli._roll_out_prefixes([info], RUST) == 0 assert not any(block.kind == "object" for block in info.code_blocks) +def test_prefix_ffi_rolls_out_nothing_from_the_registry( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """The real registry has `ffi.Array`, `ffi.Map`, `ffi.Module`...; none of them gets a block.""" + src = tmp_path / "ffi.rs" + src.write_text(f"{C.RUST_SYNTAX.directive('prefix')} ffi\n", encoding="utf-8") + monkeypatch.setattr("sys.argv", ["tvm-ffi-stubgen", "--target", "rust", str(tmp_path)]) + assert stub_cli.__main__() == 0 + text = src.read_text(encoding="utf-8") + assert f"{C.RUST_SYNTAX.begin} object/" not in text + assert "pub struct Array" not in text + + def test_prefix_survives_init(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: """`--init` rewrites the file on disk and reloads it; the roll-out must come after that.""" (tmp_path / "testing").mkdir() From b26aa5fd7fc9cf0fa84725ca558c2d12af093869 Mon Sep 17 00:00:00 2001 From: yuchuan Date: Sun, 6 Sep 2026 11:28:52 -0400 Subject: [PATCH 2/2] fix the doc. Signed-off-by: yuchuan --- .github/workflows/ci_test.yml | 13 +++ examples/rust_stubgen/CMakeLists.txt | 2 +- examples/rust_stubgen/README.md | 88 ++++++++++++++++--- .../rust/src/generated/rust_stubgen/mod.rs | 17 +++- examples/rust_stubgen/rust/src/main.rs | 7 +- python/tvm_ffi/stub/consts.py | 1 - python/tvm_ffi/stub/rust_generator/codegen.py | 3 +- tests/python/test_stubgen_rust.py | 13 --- 8 files changed, 111 insertions(+), 33 deletions(-) diff --git a/.github/workflows/ci_test.yml b/.github/workflows/ci_test.yml index ab0486de6..7d3d6a2e0 100644 --- a/.github/workflows/ci_test.yml +++ b/.github/workflows/ci_test.yml @@ -209,6 +209,19 @@ jobs: run: | cargo test + # Build the Rust stubgen example against the installed package. Its post-build + # step refreshes the bindings in place, so a diff means the checked-in file was stale. + - name: Run rust stubgen example + if: ${{ matrix.os == 'ubuntu-latest' }} + working-directory: examples/rust_stubgen + env: + CMAKE_BUILD_PARALLEL_LEVEL: ${{ steps.env_vars.outputs.cpu_count }} + run: | + cmake -B build -DCMAKE_BUILD_TYPE=Release + cmake --build build + git diff --exit-code -- rust/src/generated + cd rust && cargo run + # Build + test the orcjit wheel. On Linux the action downloads a prebuilt LLVM # tarball from a GitHub Release (built from source under the manylinux # gcc-toolset so the ABI stays under the floor); macOS/Windows install LLVM via diff --git a/examples/rust_stubgen/CMakeLists.txt b/examples/rust_stubgen/CMakeLists.txt index 970ad3ee1..ec22b32c5 100644 --- a/examples/rust_stubgen/CMakeLists.txt +++ b/examples/rust_stubgen/CMakeLists.txt @@ -36,7 +36,7 @@ tvm_ffi_configure_target( STUB_TARGET rust STUB_INIT - ON + OFF STUB_DIR "./rust/src/generated" ) diff --git a/examples/rust_stubgen/README.md b/examples/rust_stubgen/README.md index f4d257e3d..35906b312 100644 --- a/examples/rust_stubgen/README.md +++ b/examples/rust_stubgen/README.md @@ -19,16 +19,17 @@ `tvm-ffi-stubgen --target rust` turns the reflection metadata of a C++ library into Rust bindings. This example registers one object, `rust_stubgen.IntPair` -(`src/int_pair.cc`), and lets CMake regenerate `rust/src/generated/` after -every build. +(`src/int_pair.cc`), and binds it in `rust/src/generated/rust_stubgen/mod.rs`: +a file that mixes generated blocks with hand-written code, and that CMake +refreshes in place after every build. Every object gets a `#[repr(C)]` wrapper, a reference type, `Deref`, and the upcasts along its ancestor chain. `IntPair` is plain data, so its reflected fields account for every byte and the binding is *complete*: the struct mirrors the fields at their real offsets and widths, a `const` assertion pins its size -and alignment to the reflected facts, and a generated `new` allocates the object -in Rust. `main.rs` builds one that way, reads `pair.a` directly, and hands it to -a C++ function that reads it back. +and alignment to the reflected facts, and a generated allocator builds the +object in Rust. `main.rs` builds one that way, reads `pair.a` directly, and +hands it to a C++ function that reads it back. An object whose layout cannot be reproduced (a polymorphic one, say, with a vtable in front of the object header) is bound *opaquely* instead: the struct @@ -42,7 +43,7 @@ derived type depth matches the registry. ## Build and run ```bash -# 1. Build the C++ library; the post-build step runs the stub generator. +# 1. Build the C++ library; the post-build step refreshes the bindings in place. cmake -B build -DCMAKE_BUILD_TYPE=Release cmake --build build @@ -50,27 +51,71 @@ cmake --build build cd rust && cargo run ``` +The program prints: + +```text +a=2 b=1 kind=PairKind(0) +sum=3 +``` + The Rust crate depends on the `tvm-ffi` crate of this repository and needs `tvm-ffi-config` on `PATH` (activate the virtual environment where the `apache-tvm-ffi` package is installed). +## Demand-driven layout + +The binding file opens with a `prefix` line: + +```rust +// tvm-ffi-stubgen(prefix): rust_stubgen +``` + +It makes the file own the `rust_stubgen` namespace: every object registered +directly under it gets an `object/` block in this file on the next +run, and a `skip` line (`// tvm-ffi-stubgen(skip): rust_stubgen.Internal`) +leaves one out. Add a second object to `int_pair.cc` and rebuild, and its block +appears after `IntPair`; delete a block and it comes back, so dropping a binding +means writing `skip`. Code outside the blocks is never touched. + +Nothing requires a `generated/` directory: the blocks can live in any `.rs` +file the command line points at. This example keeps them under +`rust/src/generated/` only to make the split visible. + +## Hand-written constructors + +`IntPair::new` is hand-written. The `custom-new` directive makes the generator +name its lossless allocator `from_complete_fields` instead of `new`, and the +hand-written constructor calls it after deriving `kind` from the operands: + +```rust +// tvm-ffi-stubgen(custom-new): rust_stubgen.IntPair + +impl IntPair { + pub fn new(a: i64, b: i64) -> Self { + let kind = if a <= b { PairKind::Ordered } else { PairKind::Unordered }; + Self::from_complete_fields(a, b, kind) + } +} +``` + +Without the directive the block would define `IntPair::new` too and the crate +would not compile. + ## Directives The generated file keeps one-line directives the generator reads on every run. -`rust/src/generated/rust_stubgen/mod.rs` declares the integer field `kind` as an -open newtype: +Besides `prefix` and `custom-new`, this example declares the integer field +`kind` as an open newtype: ```rust // tvm-ffi-stubgen(enum): rust_stubgen.IntPair.kind -> PairKind(i32) { Unordered=0, Ordered=1 } ``` -Four more directives are available: `field` names the Rust type of a field +Three more are available: `field` names the Rust type of a field (`// tvm-ffi-stubgen(field): rust_stubgen.IntPair.a -> MyInt`), `nullable` wraps it in `Option` (`// tvm-ffi-stubgen(nullable): rust_stubgen.IntPair.a`), -`upcast` adds a conversion to a hand-written typed view -(`// tvm-ffi-stubgen(upcast): rust_stubgen.IntPair -> MyView`), and -`custom-new` names the generated allocator `from_complete_fields` when `new` -is hand-written (`// tvm-ffi-stubgen(custom-new): rust_stubgen.IntPair`). +and `upcast` adds a conversion to a hand-written typed view +(`// tvm-ffi-stubgen(upcast): rust_stubgen.IntPair -> MyView`). ## Partial generation @@ -82,3 +127,20 @@ struct is named `Obj` (`// tvm-ffi-stubgen(ty-map): rust_stubgen.IntPair -> crate::hand::IntPair`). Anything else is an error listing the missing keys, so a partial binding never references a module that does not exist. + +## Workflow + +1. Write the skeleton: `rust/src/generated/mod.rs` with `pub mod rust_stubgen;` + and `rust_stubgen/mod.rs` holding nothing but the `prefix` line. +2. Build. The post-build step adds the import section and one block per + object. Read the generated code, then add directives and hand-written code + outside the blocks. +3. Rebuild whenever the C++ side changes, or run the post-build command + directly; a `repo: local` pre-commit hook can do the same: + + ```bash + tvm-ffi-stubgen rust/src/generated --target rust --dlls build/librust_stubgen.so + ``` + +4. In CI, run the same command with `--check`. It writes nothing and exits with + status 1 when a block is out of date (2 when a file fails to process). diff --git a/examples/rust_stubgen/rust/src/generated/rust_stubgen/mod.rs b/examples/rust_stubgen/rust/src/generated/rust_stubgen/mod.rs index 01afb51b3..cda7a4d6a 100644 --- a/examples/rust_stubgen/rust/src/generated/rust_stubgen/mod.rs +++ b/examples/rust_stubgen/rust/src/generated/rust_stubgen/mod.rs @@ -20,6 +20,9 @@ //! FFI bindings for `rust_stubgen` (generated by tvm-ffi-stubgen). +// Every object registered under `rust_stubgen` gets its block in this file; `skip` leaves one out. +// tvm-ffi-stubgen(prefix): rust_stubgen + // tvm-ffi-stubgen(begin): import-section use std::ops::Deref; use tvm_ffi::Error; @@ -32,6 +35,8 @@ use tvm_ffi::VALUE_ERROR; // The `kind` field is an integer on the C++ side; this directive types it as an // open integer newtype in Rust. // tvm-ffi-stubgen(enum): rust_stubgen.IntPair.kind -> PairKind(i32) { Unordered=0, Ordered=1 } +// `IntPair::new` is hand-written below, so the generated allocator is `from_complete_fields`. +// tvm-ffi-stubgen(custom-new): rust_stubgen.IntPair // tvm-ffi-stubgen(begin): object/rust_stubgen.IntPair #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] @@ -98,9 +103,19 @@ impl IntPairObj { impl IntPair { /// Lossless complete-field allocation. - pub fn new(a: i64, b: i64, kind: PairKind) -> Self { + pub fn from_complete_fields(a: i64, b: i64, kind: PairKind) -> Self { let obj = IntPairObj::new(a, b, kind); Self { base: ObjectArc::new(obj) } } } // tvm-ffi-stubgen(end) + +// Hand-written code lives outside the blocks; the generator never touches it. +impl IntPair { + /// Allocate a pair whose `kind` follows from the operands: ordered when + /// `a <= b`, unordered otherwise. + pub fn new(a: i64, b: i64) -> Self { + let kind = if a <= b { PairKind::Ordered } else { PairKind::Unordered }; + Self::from_complete_fields(a, b, kind) + } +} diff --git a/examples/rust_stubgen/rust/src/main.rs b/examples/rust_stubgen/rust/src/main.rs index f49fadbef..e3ca7d898 100644 --- a/examples/rust_stubgen/rust/src/main.rs +++ b/examples/rust_stubgen/rust/src/main.rs @@ -41,10 +41,11 @@ fn main() -> Result<()> { let _lib = Module::load_from_file(lib_path())?; // The object has a reproducible layout: it is allocated in Rust and its - // fields are plain struct members, on both sides of the ABI. - let pair = IntPair::new(1, 2, PairKind::Ordered); + // fields are plain struct members, on both sides of the ABI. `IntPair::new` + // is hand-written next to the generated block and derives `kind` itself. + let pair = IntPair::new(2, 1); println!("a={} b={} kind={:?}", pair.a, pair.b, pair.kind); - assert_eq!(pair.kind, PairKind::Ordered); + assert_eq!(pair.kind, PairKind::Unordered); let sum: i64 = tvm_ffi::cached_global_func!("rust_stubgen.IntPairSum") .call_tuple((pair.clone(),))? diff --git a/python/tvm_ffi/stub/consts.py b/python/tvm_ffi/stub/consts.py index 164322677..5df28d5f8 100644 --- a/python/tvm_ffi/stub/consts.py +++ b/python/tvm_ffi/stub/consts.py @@ -115,7 +115,6 @@ def directive(self, name: str) -> str: # the FFI reflection registry (see `utils.ObjectInfo.from_type_info`). FN_NAME_MAP: dict[str, str] = {} -#: Type keys the Python package binds itself (see ``PythonGenerator.is_builtin``). BUILTIN_TYPE_KEYS = { "ffi.Bytes", "ffi.Error", diff --git a/python/tvm_ffi/stub/rust_generator/codegen.py b/python/tvm_ffi/stub/rust_generator/codegen.py index d7db7a6b7..7385dab7f 100644 --- a/python/tvm_ffi/stub/rust_generator/codegen.py +++ b/python/tvm_ffi/stub/rust_generator/codegen.py @@ -138,7 +138,8 @@ def _generated_type_path(self, type_key: str) -> str: Same module: the bare leaf. Elsewhere: ``super::`` per segment of this file's module, then the full path (edition 2021 rejects ``use ir::Expr``). """ - if is_crate_type_key(type_key): + head, _, _ = type_key.partition(".") + if head in C_RUST.RUST_MOD_MAP: return type_key mod, _, type_leaf = type_key.rpartition(".") if tuple(mod.split(".")) == self.mod_segments: diff --git a/tests/python/test_stubgen_rust.py b/tests/python/test_stubgen_rust.py index 825559542..05449b8c3 100644 --- a/tests/python/test_stubgen_rust.py +++ b/tests/python/test_stubgen_rust.py @@ -1414,19 +1414,6 @@ def test_roll_out_matches_the_prefix_exactly( assert not any(block.kind == "object" for block in info.code_blocks) -def test_prefix_ffi_rolls_out_nothing_from_the_registry( - tmp_path: Path, monkeypatch: pytest.MonkeyPatch -) -> None: - """The real registry has `ffi.Array`, `ffi.Map`, `ffi.Module`...; none of them gets a block.""" - src = tmp_path / "ffi.rs" - src.write_text(f"{C.RUST_SYNTAX.directive('prefix')} ffi\n", encoding="utf-8") - monkeypatch.setattr("sys.argv", ["tvm-ffi-stubgen", "--target", "rust", str(tmp_path)]) - assert stub_cli.__main__() == 0 - text = src.read_text(encoding="utf-8") - assert f"{C.RUST_SYNTAX.begin} object/" not in text - assert "pub struct Array" not in text - - def test_prefix_survives_init(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: """`--init` rewrites the file on disk and reloads it; the roll-out must come after that.""" (tmp_path / "testing").mkdir()