This is an automated email from the ASF dual-hosted git repository.
tlopex pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm-ffi.git
The following commit(s) were added to refs/heads/main by this push:
new 85e72e78 [STUBGEN][RUST] Skip crate-provided types and add a custom
constructor example (#755)
85e72e78 is described below
commit 85e72e7853e2fec081f427aa42c10e262ead1145
Author: Linzhang Li <[email protected]>
AuthorDate: Sun Sep 6 12:03:44 2026 -0400
[STUBGEN][RUST] Skip crate-provided types and add a custom constructor
example (#755)
## Summary
`prefix: ffi` could generate a non-generic `Array` that shadowed the
crate's
`Array<T>`, causing E0107. Prefix expansion and `--init` now use
`Generator.is_builtin` to skip runtime-provided types. Rust shares
`_provider`'s
type-key rule; Python continues to use `BUILTIN_TYPE_KEYS`.
`examples/rust_stubgen` combines the demand-driven layout from #752 with
a
hand-written `IntPair::new(a, b)` that derives `kind` from its operands.
`custom-new` renames the generated allocator to `from_complete_fields`,
and
`STUB_INIT OFF` keeps regeneration in place. Linux CI builds and runs
the example
and checks for stale generated bindings.
## Testing
- Stubgen tests: **109 passed**, including `ffi.Array` regression
coverage.
- `ruff`, `ty`, and `yamllint`: passed.
- Example builds and runs; regeneration leaves bindings unchanged and
`--check` passes.
- Verified duplicate-`new` errors, restoration of deleted blocks, and
generation
of blocks for new types.
---------
Signed-off-by: yuchuan <[email protected]>
---
.github/workflows/ci_test.yml | 13 ++++
docs/packaging/stubgen.rst | 3 +-
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/cli.py | 25 +++---
python/tvm_ffi/stub/generator.py | 8 ++
python/tvm_ffi/stub/python_generator/generator.py | 4 +
python/tvm_ffi/stub/rust_generator/codegen.py | 10 ++-
python/tvm_ffi/stub/rust_generator/generator.py | 6 +-
python/tvm_ffi/stub/rust_generator/utils.py | 5 ++
tests/python/test_stubgen_rust.py | 9 ++-
13 files changed, 161 insertions(+), 36 deletions(-)
diff --git a/.github/workflows/ci_test.yml b/.github/workflows/ci_test.yml
index ab0486de..7d3d6a2e 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/docs/packaging/stubgen.rst b/docs/packaging/stubgen.rst
index 98087e2e..7b299d46 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/<type_key>`` 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/examples/rust_stubgen/CMakeLists.txt
b/examples/rust_stubgen/CMakeLists.txt
index 970ad3ee..ec22b32c 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 f4d257e3..35906b31 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/<type_key>` 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 `<Name>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 01afb51b..cda7a4d6 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 f49fadbe..e3ca7d89 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/cli.py b/python/tvm_ffi/stub/cli.py
index 4f96e9ae..af8f8c24 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/<key>`` 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 _stage_2(
}
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 _stage_2(
[] 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/generator.py b/python/tvm_ffi/stub/generator.py
index add7503a..c22522a4 100644
--- a/python/tvm_ffi/stub/generator.py
+++ b/python/tvm_ffi/stub/generator.py
@@ -114,6 +114,14 @@ class Generator(Protocol):
"""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 cd81821d..5d784818 100644
--- a/python/tvm_ffi/stub/python_generator/generator.py
+++ b/python/tvm_ffi/stub/python_generator/generator.py
@@ -75,6 +75,10 @@ class PythonGenerator:
"""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 c5da9b4f..7385dab7 100644
--- a/python/tvm_ffi/stub/rust_generator/codegen.py
+++ b/python/tvm_ffi/stub/rust_generator/codegen.py
@@ -52,7 +52,13 @@ from .. import consts as C
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
@@ -147,7 +153,7 @@ class _ObjectRenderer:
``"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 30b44dbd..edfc7c07 100644
--- a/python/tvm_ffi/stub/rust_generator/generator.py
+++ b/python/tvm_ffi/stub/rust_generator/generator.py
@@ -32,7 +32,7 @@ from .codegen import (
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 @@ class RustGenerator:
"""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 d38e35bc..dcd68e3d 100644
--- a/python/tvm_ffi/stub/rust_generator/utils.py
+++ b/python/tvm_ffi/stub/rust_generator/utils.py
@@ -58,6 +58,11 @@ class RustUse:
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 97a446d2..05449b8c 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,11 +1405,12 @@ 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<T>`.
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)
