[STUBGEN][RUST] Skip crate-provided types and add a custom constructor example

.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

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
 

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"
 )

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,35 +43,79 @@ 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
 
 # 2. Build and run the Rust program against it.
 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).

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)
+    }
+}

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(),))?

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 _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

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(

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(

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
@@ -147,7 +153,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"

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(

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``).