From 9590077eaf68a1509e88bc8f881eb440431135b7 Mon Sep 17 00:00:00 2001 From: yuchuan Date: Thu, 3 Sep 2026 20:08:50 -0400 Subject: [PATCH 1/8] [STUBGEN] Generate complete field mirrors for the Rust target. Signed-off-by: yuchuan --- examples/rust_stubgen/README.md | 26 +- .../rust/src/generated/rust_stubgen/mod.rs | 32 ++ examples/rust_stubgen/rust/src/main.rs | 9 +- examples/rust_stubgen/src/int_pair.cc | 26 +- python/tvm_ffi/stub/cli.py | 49 +- python/tvm_ffi/stub/rust_generator/codegen.py | 224 ++++++-- python/tvm_ffi/stub/rust_generator/consts.py | 52 +- .../tvm_ffi/stub/rust_generator/directives.py | 28 +- python/tvm_ffi/stub/utils.py | 8 + tests/python/test_stubgen_rust.py | 518 +++++++++++++++--- 10 files changed, 789 insertions(+), 183 deletions(-) diff --git a/examples/rust_stubgen/README.md b/examples/rust_stubgen/README.md index 080b5d6a4..aa4d70573 100644 --- a/examples/rust_stubgen/README.md +++ b/examples/rust_stubgen/README.md @@ -18,15 +18,23 @@ # Rust Stub Generation `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. - -Every object is bound *opaquely*: Rust gets a `#[repr(C)]` wrapper that embeds -only the parent, a reference type, `Deref`, the upcasts along the ancestor -chain, and one accessor per reflected field that reads through the C ABI -getter. The object's bytes are never reproduced, so the binding is correct for -any registered type; construction goes through the registered global functions. +into Rust bindings. This example registers two objects in `src/int_pair.cc` +and lets CMake regenerate `rust/src/generated/` after every build. + +Every object gets a `#[repr(C)]` wrapper, a reference type, `Deref`, and the +upcasts along its ancestor chain. What the wrapper holds depends on whether the +reflected fields account for every byte of the object: + +- `rust_stubgen.IntRange` does, so its binding is *complete*: the struct mirrors + the fields at their real offsets and widths, and Rust reads `range.begin` + directly. A `const` assertion pins the struct's size and alignment to the + reflected facts. +- `rust_stubgen.IntPair` has a vtable in front of the object header, so its + binding is *opaque*: the struct embeds only the parent and every field is read + through an accessor that calls the C ABI getter. + +Construction goes through the registered global functions in both cases. + A builtin parent such as `ffi.IntEnum` has no `Obj` in the crate; the import section defines a header-only stand-in per builtin ancestor so the derived type depth matches the registry. 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 4e7f47e61..cd99f220e 100644 --- a/examples/rust_stubgen/rust/src/generated/rust_stubgen/mod.rs +++ b/examples/rust_stubgen/rust/src/generated/rust_stubgen/mod.rs @@ -61,6 +61,7 @@ impl TryFrom for PairKind { } } +/// Opaque: bytes [44, 56) of [24, 56) are not accounted for by reflected fields. Fields are read through the C ABI getters. #[repr(C)] #[derive(tvm_ffi::derive::Object)] #[type_key = "rust_stubgen.IntPair"] @@ -97,3 +98,34 @@ impl IntPairObj { } } // tvm-ffi-stubgen(end) + +// tvm-ffi-stubgen(begin): object/rust_stubgen.IntRange +/// Complete: reflected fields fill [24, 40) exactly (alignment padding [36, 40)). +#[repr(C)] +#[derive(tvm_ffi::derive::Object)] +#[type_key = "rust_stubgen.IntRange"] +#[type_final] +pub struct IntRangeObj { + base: Object, + pub begin: i64, + pub extent: i32, +} + +const _: () = { + assert!(::core::mem::size_of::() == 40); + assert!(::core::mem::align_of::() == 8); +}; + +#[repr(C)] +#[derive(tvm_ffi::derive::ObjectRef, Clone)] +pub struct IntRange { + data: ObjectArc, +} + +impl Deref for IntRange { + type Target = IntRangeObj; + fn deref(&self) -> &IntRangeObj { + &self.data + } +} +// tvm-ffi-stubgen(end) diff --git a/examples/rust_stubgen/rust/src/main.rs b/examples/rust_stubgen/rust/src/main.rs index 27afeb3a7..7b0f360df 100644 --- a/examples/rust_stubgen/rust/src/main.rs +++ b/examples/rust_stubgen/rust/src/main.rs @@ -20,7 +20,7 @@ mod generated; -use generated::rust_stubgen::{IntPair, PairKind}; +use generated::rust_stubgen::{IntPair, IntRange, PairKind}; use tvm_ffi::{Module, Result}; /// Path of the C++ shared library built by CMake into `../build`. @@ -52,5 +52,12 @@ fn main() -> Result<()> { .call_tuple((pair.clone(),))? .try_into()?; println!("sum={sum}"); + + // `IntRange` has a reproducible layout: its fields are plain struct members. + let range: IntRange = tvm_ffi::cached_global_func!("rust_stubgen.IntRange") + .call_tuple((10i64, 5i64))? + .try_into()?; + println!("begin={} extent={}", range.begin, range.extent); + assert_eq!(range.begin + i64::from(range.extent), 15); Ok(()) } diff --git a/examples/rust_stubgen/src/int_pair.cc b/examples/rust_stubgen/src/int_pair.cc index 3d7631b45..85d44ed33 100644 --- a/examples/rust_stubgen/src/int_pair.cc +++ b/examples/rust_stubgen/src/int_pair.cc @@ -52,16 +52,40 @@ class IntPair : public ffi::ObjectRef { TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(IntPair, ffi::ObjectRef, IntPairObj); }; +// A plain data object: every byte is accounted for by a reflected field, so the +// generated binding mirrors the layout and Rust reads the fields directly. +class IntRangeObj : public ffi::Object { + public: + int64_t begin; + int32_t extent; + + IntRangeObj(int64_t begin, int32_t extent) : begin(begin), extent(extent) {} + + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("rust_stubgen.IntRange", IntRangeObj, ffi::Object); +}; + +class IntRange : public ffi::ObjectRef { + public: + IntRange(int64_t begin, int32_t extent) { data_ = ffi::make_object(begin, extent); } + + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(IntRange, ffi::ObjectRef, IntRangeObj); +}; + TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; refl::ObjectDef(refl::init(false)) .def_ro("a", &IntPairObj::a, "the first operand") .def_ro("b", &IntPairObj::b, "the second operand") .def_ro("kind", &IntPairObj::kind, "0 = unordered, 1 = ordered"); + refl::ObjectDef(refl::init(false)) + .def_ro("begin", &IntRangeObj::begin, "the first value") + .def_ro("extent", &IntRangeObj::extent, "the number of values"); refl::GlobalDef() .def("rust_stubgen.IntPair", [](int64_t a, int64_t b, int32_t kind) { return IntPair(a, b, kind); }) - .def("rust_stubgen.IntPairSum", [](const IntPair& pair) { return pair->Sum(); }); + .def("rust_stubgen.IntPairSum", [](const IntPair& pair) { return pair->Sum(); }) + .def("rust_stubgen.IntRange", + [](int64_t begin, int32_t extent) { return IntRange(begin, extent); }); } // [object.end] diff --git a/python/tvm_ffi/stub/cli.py b/python/tvm_ffi/stub/cli.py index a5e0ba393..e185ae8d4 100644 --- a/python/tvm_ffi/stub/cli.py +++ b/python/tvm_ffi/stub/cli.py @@ -24,7 +24,7 @@ import sys import traceback from pathlib import Path -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Callable from . import consts as C from .file_utils import FileInfo, collect_files, syntax_for @@ -36,7 +36,7 @@ object_info_from_type_key, toposort_objects, ) -from .utils import FuncInfo, InitConfig, Options +from .utils import DirectiveError, FuncInfo, InitConfig, Options if TYPE_CHECKING: from .generator import Generator @@ -67,13 +67,9 @@ def __main__() -> int: # - defined global functions: `tvm-ffi-stubgen(begin): global/...` # - defined object types: `tvm-ffi-stubgen(begin): object/...` ty_map: dict[str, str] = generator.default_ty_map() + directive_errors = 0 for file in files: - try: - _stage_1(file, ty_map) - except Exception: - print( - f'{C.TERM_RED}[Failed] File "{file.path}": {traceback.format_exc()}{C.TERM_RESET}' - ) + directive_errors += _run_stage(file, lambda: _stage_1(file, ty_map)) # Stage 2. Generate stubs if they are not defined on the file. generated_prefixes: set[str] = set() @@ -94,18 +90,9 @@ def __main__() -> int: for file in files: if opt.verbose: print(f"{C.TERM_CYAN}[File] {file.path}{C.TERM_RESET}") - try: - _stage_3( - file, - opt, - ty_map, - global_funcs, - generator=generator, - ) - except Exception: - print( - f'{C.TERM_RED}[Failed] File "{file.path}": {traceback.format_exc()}{C.TERM_RESET}' - ) + directive_errors += _run_stage( + file, lambda: _stage_3(file, opt, ty_map, global_funcs, generator=generator) + ) # Stage 4. Let the generator stitch the generated tree together (runs after the # files are fully written, so language-specific wiring isn't clobbered). @@ -122,7 +109,23 @@ def __main__() -> int: } write_coverage_report(Path(opt.coverage_out), classify(infos)) del dlls - return 0 + return 1 if directive_errors else 0 + + +def _run_stage(file: FileInfo, stage: Callable[[], None]) -> bool: + """Run one stage over ``file``, reporting a failure without stopping the run. + + Returns whether the failure was a :class:`DirectiveError`, which makes the + whole run exit non-zero. + """ + try: + stage() + except DirectiveError as e: + print(f'{C.TERM_RED}[Failed] File "{file.path}": {e}{C.TERM_RESET}') + return True + except Exception: + print(f'{C.TERM_RED}[Failed] File "{file.path}": {traceback.format_exc()}{C.TERM_RESET}') + return False def _stage_1( @@ -135,7 +138,7 @@ def _stage_1( try: lhs, rhs = code.param[1].split("->") except ValueError as e: - raise ValueError( + raise DirectiveError( f"Invalid ty_map format at line {code.lineno_start}. Example: `A.B -> C.D`" ) from e ty_map[lhs.strip()] = rhs.strip() @@ -239,7 +242,7 @@ def _stage_3( # noqa: PLR0912 if name in C.PIPELINE_DIRECTIVE_KINDS: continue # consumed by `_stage_1` if name not in generator.directive_kinds: - raise ValueError(f"Unknown directive `{name}` at line {code.lineno_start}") + raise DirectiveError(f"Unknown directive `{name}` at line {code.lineno_start}") generator.add_directive(imports, name, payload, code.lineno_start) # Stage 2. Process `tvm-ffi-stubgen(begin): global/...` for code in file.code_blocks: diff --git a/python/tvm_ffi/stub/rust_generator/codegen.py b/python/tvm_ffi/stub/rust_generator/codegen.py index 26f2603a1..a6f1bfe4f 100644 --- a/python/tvm_ffi/stub/rust_generator/codegen.py +++ b/python/tvm_ffi/stub/rust_generator/codegen.py @@ -14,39 +14,26 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -"""Rust code generation for ``tvm-ffi-stubgen``: the opaque binding form. - -Every reflected object renders as a ``#[repr(C)]`` struct embedding only its -parent, a reference wrapper, ``Deref``, the upcasts along the ancestor chain, -and one accessor per reflected field that reads through the C ABI getter. The -object's bytes are never reproduced. For ``tirx.IterVar`` deriving from -``ir.PrimExprConvertible``:: - - #[repr(C)] - #[derive(tvm_ffi::derive::Object)] - #[type_key = "tirx.IterVar"] - #[type_final] - pub struct IterVarObj { - base: PrimExprConvertibleObj, - } - - #[repr(C)] - #[derive(tvm_ffi::derive::ObjectRef, Clone)] - pub struct IterVar { - data: ObjectArc, - } - - impl Deref for IterVar { ... } // IterVar -> IterVarObj - impl Deref for IterVarObj { ... } // IterVarObj -> PrimExprConvertibleObj - - impl IterVarObj { - pub fn dom(&self) -> Result> { - FieldGetter::new(Self::type_index(), "dom")?.get(self) - } - ... - } - - tvm_ffi::impl_object_upcast!(IterVar => PrimExprConvertible); +"""Rust code generation for ``tvm-ffi-stubgen``. + +Every reflected object gets a ``#[repr(C)]`` object struct, a reference wrapper, +read-only ``Deref``, and the upcasts along its ancestor chain. What the object +struct holds depends on the verdict of :mod:`tvm_ffi.stub.layout`: + +- *complete*: the layout is reproducible, so the struct mirrors every physical + field at its real offset and width, public, borrowed directly. A ``const`` + assertion pins the struct's ``size_of`` / ``align_of`` to the reflected facts, + so a mirror rustc lays out differently fails to compile. +- *opaque*: the struct embeds only its parent, and one accessor per reflected + field reads through the C ABI getter. The bytes are never reproduced, so the + binding is correct for every registered type. + +The two target-language rules the classifier leaves to its caller live here: a +field without a Rust mirror (``Optional``, a ``Union``, ``void*``, ...) +makes the type opaque and is read as ``Any``; an ``opaque`` directive vetoes a +reproducible layout. ``field`` / ``nullable`` / ``enum`` directives shape the +field types of both forms; where a directive names a scalar width, it is checked +against the reflected field size at generation time. Construction and behaviour go through the registered global functions, hand-written outside the markers. A builtin parent (``ffi.IntEnum``, say) has @@ -61,6 +48,9 @@ from typing import TYPE_CHECKING from .. import consts as C +from ..layout import Verdict, classify +from ..lib_state import object_info_from_type_key +from ..utils import DirectiveError from . import consts as C_RUST from .utils import RustImports, builtin_mirror_name, render_rust_type, rust_ident @@ -72,6 +62,15 @@ from .directives import EnumSpec +def _check_width(target: str, field: NamedTypeSchema, rust_type: str, width: int) -> None: + """Reject a directive whose scalar type does not match the reflected field size.""" + if field.size is not None and field.size != width: + raise DirectiveError( + f"Directive on `{target}` maps a {field.size}-byte field to `{rust_type}` " + f"({width} bytes)" + ) + + @dataclasses.dataclass class _ObjectRenderer: """Renders one ``object/`` block into Rust source lines.""" @@ -100,14 +99,17 @@ def obj_struct(self) -> str: # --- name resolution --------------------------------------------------- - def _ty_render(self, origin: str) -> str | None: + def _resolve(self, origin: str, imports: RustImports) -> str | None: """Resolve a leaf origin to its in-scope Rust name (recording its ``use``), or ``None``.""" mapped = self.ty_map.get(origin) if mapped is None: if "." not in origin or origin.startswith("ctypes."): return None mapped = self._generated_type_path(origin) - return self.imports.record(mapped) + return imports.record(mapped) + + def _ty_render(self, origin: str) -> str | None: + return self._resolve(origin, self.imports) def _generated_type_path(self, type_key: str) -> str: """Spell a generated type key from this file. @@ -143,7 +145,88 @@ def _base_type(self) -> tuple[str, bool]: assert not any(self._generated(key) for key in chain), (self.type_key, chain) return self.imports.record_builtin_base(chain), False - # --- pieces ------------------------------------------------------------ + # --- classification ---------------------------------------------------- + + def classify(self) -> Verdict: + """Classify this object with its ancestors, under the file's directives.""" + infos = {key: object_info_from_type_key(key) for key in self.info.ancestors} + infos[self.type_key] = self.info + owner_of = {id(f): key for key, owner in infos.items() for f in owner.fields} + scratch = RustImports() + + def renderable(field: NamedTypeSchema) -> bool: + return self._field_mirror(owner_of[id(field)], field, scratch) is not None + + verdicts = classify( + infos, forced_opaque=self.imports.directives.opaque, field_renderable=renderable + ) + return verdicts[self.type_key] + + # --- field types --------------------------------------------------------- + + def _field_mirror(self, owner: str, field: NamedTypeSchema, imports: RustImports) -> str | None: + """Render the type of ``field`` in a ``#[repr(C)]`` mirror; ``None`` when it has none. + + Scalars take the width the registry recorded; ``Optional`` fields take + the in-place mirror of their C++ layout; directives override the rest. + """ + directives = self.imports.directives + target = f"{owner}.{field.name}" + enum = directives.enums.get(target) + if enum is not None: + _check_width(target, field, enum.repr, C_RUST.RUST_SCALAR_WIDTHS[enum.repr]) + return enum.name + override = directives.field_types.get(target) + if override is not None: + width = C_RUST.RUST_SCALAR_WIDTHS.get(override) + if width is not None: + _check_width(target, field, override, width) + mirror: str | None = imports.record(override) if "::" in override else override + elif field.origin == "Optional": + mirror = self._optional_mirror(field, imports) + else: + narrowed = C_RUST.RUST_SCALAR_BY_SIZE.get((field.origin, field.size)) + mirror = narrowed or render_rust_type(field, lambda o: self._resolve(o, imports)) + if mirror is None: + return None + if target in directives.nullable and not mirror.startswith("Option<"): + if field.size not in (None, C_RUST.RUST_POINTER_SIZE): + raise DirectiveError( + f"`nullable` directive on `{target}`: the field is {field.size} bytes, " + "not a pointer-sized object reference" + ) + mirror = f"Option<{mirror}>" + return mirror + + def _optional_mirror(self, field: NamedTypeSchema, imports: RustImports) -> str | None: + """Mirror an ``Optional`` field in place. + + An ``ObjectRef``-derived payload is a pointer-sized nullable pointer in + C++, mirrored by Rust's niche-optimized ``Option``. Every other + payload stays a 16-byte ``TVMFFIAny`` cell, mirrored by + ``tvm_ffi::Optional``. ``Optional`` has no mirror, and neither + does a field whose size disagrees with its payload kind. + """ + (payload,) = field.args # TypeSchema's post_init enforces exactly one argument. + if payload.origin == "Any": + return None + inner = render_rust_type(payload, lambda o: self._resolve(o, imports)) + if inner is None: + return None + any_backed = ( + payload.origin in C_RUST.RUST_ANY_BACKED_OPTIONAL_PAYLOADS + or payload.origin == "Optional" + ) + expected = ( + C_RUST.RUST_OPTIONAL_FIELD_SIZE + if any_backed + else C_RUST.RUST_OBJECT_OPTIONAL_FIELD_SIZE + ) + if field.size not in (None, expected): + return None + if any_backed: + return f"{imports.record(C_RUST.RUST_OPTIONAL_PATH)}<{inner}>" + return f"Option<{inner}>" def _accessor_lines(self, field: NamedTypeSchema) -> list[str]: """One ``pub fn (&self) -> Result`` through the C ABI getter. @@ -177,7 +260,13 @@ def _accessor_lines(self, field: NamedTypeSchema) -> list[str]: ] if target in directives.nullable and not rust_type.startswith("Option<"): rust_type = f"Option<{rust_type}>" - return [f"pub fn {name}(&self) -> Result<{rust_type}> {{", f" {getter}.get(self)", "}"] + return [ + f"pub fn {name}(&self) -> Result<{rust_type}> {{", + f" {getter}.get(self)", + "}", + ] + + # --- pieces ------------------------------------------------------------ def _enum_lines(self, spec: EnumSpec) -> list[str]: """Render the open integer newtype an ``enum`` directive declares.""" @@ -232,14 +321,49 @@ def _upcast_lines(self) -> list[str]: pairs = ", ".join(f"{self.leaf} => {target}" for target in targets) return [f"tvm_ffi::impl_object_upcast!({pairs});"] + def _struct_lines(self, verdict: Verdict, base: str) -> list[str]: + """Render the object struct: every field when complete, the parent alone when opaque.""" + header = [ + "#[repr(C)]", + "#[derive(tvm_ffi::derive::Object)]", + f'#[type_key = "{self.type_key}"]', + *(["#[type_final]"] if self.info.is_final else []), + f"pub struct {self.obj_struct} {{", + f" base: {base},", + ] + if not verdict.is_complete: + return [ + f"/// Opaque: {verdict.detail}. Fields are read through the C ABI getters.", + *header, + "}", + ] + members = [] + for field in sorted(self.info.fields, key=lambda f: f.offset or 0): + mirror = self._field_mirror(self.type_key, field, self.imports) + assert mirror is not None # the verdict already ran the renderability check + members.append(f" pub {rust_ident(field.name)}: {mirror},") + return [ + f"/// Complete: {verdict.detail}.", + *header, + *members, + "}", + "", + "const _: () = {", + f" assert!(::core::mem::size_of::<{self.obj_struct}>() == {verdict.total_size});", + f" assert!(::core::mem::align_of::<{self.obj_struct}>() == {verdict.alignment});", + "};", + ] + def body(self) -> list[str]: """Build the Rust source lines for the object.""" + verdict = self.classify() # Derive macros are spelled by full path: their leaves collide with `Object` / `ObjectRef`. self.imports.record("std::ops::Deref") self.imports.record("tvm_ffi::ObjectArc") base, has_parent = self._base_type() fields = self.info.fields - if fields: + accessors = bool(fields) and not verdict.is_complete + if accessors: self.imports.record("tvm_ffi::ObjectCore") # `Self::type_index()` self.imports.record("tvm_ffi::FieldGetter") self.imports.record("tvm_ffi::Result") @@ -251,17 +375,7 @@ def body(self) -> list[str]: for f in fields if f"{self.type_key}.{f.name}" in enums ] - sections.append( - [ - "#[repr(C)]", - "#[derive(tvm_ffi::derive::Object)]", - f'#[type_key = "{self.type_key}"]', - *(["#[type_final]"] if self.info.is_final else []), - f"pub struct {self.obj_struct} {{", - f" base: {base},", - "}", - ] - ) + sections.append(self._struct_lines(verdict, base)) sections.append( [ "#[repr(C)]", @@ -274,16 +388,16 @@ def body(self) -> list[str]: sections.append(self._deref_lines(self.leaf, self.obj_struct, "data")) if has_parent: sections.append(self._deref_lines(self.obj_struct, base, "base")) - if fields: - accessors: list[str] = [] + if accessors: + lines_: list[str] = [] for i, field in enumerate(fields): if i: - accessors.append("") - accessors += self._accessor_lines(field) + lines_.append("") + lines_ += self._accessor_lines(field) sections.append( [ f"impl {self.obj_struct} {{", - *[f" {line}" if line else "" for line in accessors], + *[f" {line}" if line else "" for line in lines_], "}", ] ) @@ -306,7 +420,7 @@ def generate_rust_object( opt: Options, obj_info: ObjectInfo, ) -> None: - """Emit the opaque Rust binding of ``obj_info`` into an ``object/`` block.""" + """Emit the Rust binding of ``obj_info`` into an ``object/`` block.""" assert len(code.lines) >= 2 assert isinstance(obj_info.type_key, str) renderer = _ObjectRenderer( diff --git a/python/tvm_ffi/stub/rust_generator/consts.py b/python/tvm_ffi/stub/rust_generator/consts.py index 218def679..c8310e9ea 100644 --- a/python/tvm_ffi/stub/rust_generator/consts.py +++ b/python/tvm_ffi/stub/rust_generator/consts.py @@ -19,7 +19,7 @@ from __future__ import annotations #: One-line directives the Rust backend consumes. -RUST_DIRECTIVE_KINDS = frozenset({"import-object", "field", "nullable", "enum"}) +RUST_DIRECTIVE_KINDS = frozenset({"import-object", "field", "nullable", "enum", "opaque"}) #: Default FFI-origin -> Rust-type map; ``::`` paths get a ``use``, bare names do not. RUST_TY_MAP_DEFAULTS = { @@ -54,6 +54,56 @@ #: Origins without a crate mirror; such a field is read as ``tvm_ffi::Any``. RUST_UNSUPPORTED_ORIGINS = frozenset({"Dict", "List", "Union", "tuple"}) +#: Width-correct scalar for a ``#[repr(C)]`` struct field, keyed by +#: ``(ffi origin, sizeof(T))``: the type schema erases scalar widths, so the +#: width comes from the reflected field size. Signedness is not recorded; +#: unsigned C++ fields render as the same-width signed type. +RUST_SCALAR_BY_SIZE = { + ("int", 1): "i8", + ("int", 2): "i16", + ("int", 4): "i32", + ("int", 8): "i64", + ("float", 4): "f32", + ("float", 8): "f64", +} + +#: Byte width of the scalar Rust types a ``field`` / ``enum`` directive may name, +#: checked against the reflected field size at generation time. +RUST_SCALAR_WIDTHS = { + "i8": 1, + "u8": 1, + "bool": 1, + "i16": 2, + "u16": 2, + "i32": 4, + "u32": 4, + "f32": 4, + "i64": 8, + "u64": 8, + "f64": 8, +} + +#: Size of an object reference field; ``nullable`` may only wrap those. +RUST_POINTER_SIZE = 8 + +#: In-place mirror of a non-object ``Optional`` field: C++ ``ffi::Optional`` +#: is a single 16-byte ``TVMFFIAny`` cell (``nullopt == kTVMFFINone``) for +#: payloads that are not ``ObjectRef``-derived. Object payloads use the +#: pointer-sized form (``nullopt == nullptr``), mirrored by Rust's ``Option``. +RUST_OPTIONAL_PATH = "tvm_ffi::Optional" +RUST_OPTIONAL_FIELD_SIZE = 16 +RUST_OBJECT_OPTIONAL_FIELD_SIZE = 8 + +#: ``Optional`` payload origins whose C++ optional stays 16-byte Any-backed: +#: everything that is not a pointer-sized object reference (strings and bytes +#: are 16-byte values themselves). A nested ``Optional`` payload also stays +#: Any-backed; it is special-cased where this set is consulted. The reflected +#: field size is checked either way, so a payload this set misclassifies makes +#: the field unrenderable rather than mis-mirrored. +RUST_ANY_BACKED_OPTIONAL_PAYLOADS = frozenset( + {"int", "float", "bool", "Device", "dtype", "DataType", "str", "bytes"} +) + #: ``use``-path rewrites: builtin ``ffi.*`` type keys live at the crate root. RUST_MOD_MAP = { "ffi": "tvm_ffi", diff --git a/python/tvm_ffi/stub/rust_generator/directives.py b/python/tvm_ffi/stub/rust_generator/directives.py index 388ffb5f1..2c5ae403e 100644 --- a/python/tvm_ffi/stub/rust_generator/directives.py +++ b/python/tvm_ffi/stub/rust_generator/directives.py @@ -16,15 +16,16 @@ # under the License. """The Rust backend's one-line directives: payload grammar and per-file storage. -All three address one reflected field as ``.``:: +Three address one reflected field as ``.``, one addresses a type:: // tvm-ffi-stubgen(field): tirx.Add.a -> PrimExpr // tvm-ffi-stubgen(nullable): ir.Expr.span // tvm-ffi-stubgen(enum): tirx.For.kind -> ForKind(i32) { Serial=0, Parallel=1 } + // tvm-ffi-stubgen(opaque): ir.SourceName -``field`` sets the accessor's Rust type (a name in scope, or a ``::`` path to +``field`` sets the field's Rust type (a name in scope, or a ``::`` path to ``use``); ``nullable`` wraps it in ``Option``; ``enum`` declares an open integer -newtype the accessor returns. +newtype for it; ``opaque`` keeps a type opaque even when its layout is reproducible. """ from __future__ import annotations @@ -32,6 +33,8 @@ import dataclasses import re +from ..utils import DirectiveError + _ENUM_RE = re.compile( r"^(?P\S+)\s*->\s*(?P[A-Za-z_]\w*)\((?P[iu](?:8|16|32|64))\)" r"\s*(?:\{(?P[^{}]*)\})?$" @@ -55,9 +58,10 @@ class Directives: field_types: dict[str, str] = dataclasses.field(default_factory=dict) nullable: set[str] = dataclasses.field(default_factory=set) enums: dict[str, EnumSpec] = dataclasses.field(default_factory=dict) + opaque: set[str] = dataclasses.field(default_factory=set) def add(self, name: str, payload: str, lineno: int) -> None: - """Parse and store one directive; raise ``ValueError`` on a malformed payload.""" + """Parse and store one directive; raise :class:`DirectiveError` when malformed.""" if name == "field": target, rust_type = _split_arrow(name, payload, lineno) self.field_types[target] = rust_type @@ -66,12 +70,22 @@ def add(self, name: str, payload: str, lineno: int) -> None: elif name == "enum": target, spec = _parse_enum(payload, lineno) self.enums[target] = spec + elif name == "opaque": + self.opaque.add(_type_target(name, payload, lineno)) else: - raise ValueError(f"Unknown directive `{name}` at line {lineno}") + raise DirectiveError(f"Unknown directive `{name}` at line {lineno}") + + +def _invalid(name: str, lineno: int, expected: str) -> DirectiveError: + return DirectiveError(f"Invalid `{name}` directive at line {lineno}. Expected `{expected}`") -def _invalid(name: str, lineno: int, expected: str) -> ValueError: - return ValueError(f"Invalid `{name}` directive at line {lineno}. Expected `{expected}`") +def _type_target(name: str, text: str, lineno: int) -> str: + """Validate a ```` reference.""" + target = text.strip() + if not target or " " in target: + raise _invalid(name, lineno, "") + return target def _field_target(name: str, text: str, lineno: int) -> str: diff --git a/python/tvm_ffi/stub/utils.py b/python/tvm_ffi/stub/utils.py index 7f5ea6258..5d46e3df9 100644 --- a/python/tvm_ffi/stub/utils.py +++ b/python/tvm_ffi/stub/utils.py @@ -35,6 +35,14 @@ from tvm_ffi.core import TypeField +class DirectiveError(ValueError): + """A one-line directive is malformed, unknown, or disagrees with reflection. + + The pipeline reports it and exits non-zero: generated code must never be + shaped by a directive it could not honour. + """ + + def _parse_type_schema(raw: str | dict[str, Any]) -> TypeSchema: """Parse a type schema from either a JSON string or an already-parsed dict.""" if isinstance(raw, dict): diff --git a/tests/python/test_stubgen_rust.py b/tests/python/test_stubgen_rust.py index 9f2e0fdfb..1cda32f35 100644 --- a/tests/python/test_stubgen_rust.py +++ b/tests/python/test_stubgen_rust.py @@ -14,11 +14,12 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -"""Tests for the Rust backend of ``tvm-ffi-stubgen``: opaque bindings.""" +"""Tests for the Rust backend of ``tvm-ffi-stubgen``: complete and opaque bindings.""" from __future__ import annotations import re +from collections.abc import Iterator from pathlib import Path import pytest @@ -29,6 +30,8 @@ from tvm_ffi.stub.cli import _stage_3 from tvm_ffi.stub.file_utils import CodeBlock, FileInfo from tvm_ffi.stub.generator import get_generator +from tvm_ffi.stub.lib_state import object_info_from_type_key +from tvm_ffi.stub.rust_generator import codegen from tvm_ffi.stub.rust_generator import consts as RC from tvm_ffi.stub.rust_generator.codegen import ( finalize_rust_module_tree, @@ -38,31 +41,73 @@ ) from tvm_ffi.stub.rust_generator.directives import Directives, EnumSpec from tvm_ffi.stub.rust_generator.utils import RustImports, RustUse, render_rust_type, rust_ident -from tvm_ffi.stub.utils import InitConfig, NamedTypeSchema, ObjectInfo, Options +from tvm_ffi.stub.utils import DirectiveError, InitConfig, NamedTypeSchema, ObjectInfo, Options RUST = get_generator("rust") +HEADER = 24 # sizeof(TVMFFIObject) + + +def _field( + name: str, + schema: str | TypeSchema, + offset: int | None = None, + size: int | None = None, + alignment: int | None = None, +) -> NamedTypeSchema: + if isinstance(schema, str): + schema = TypeSchema(schema) + if alignment is None and size is not None: + alignment = min(size, 8) # a 16-byte `TVMFFIAny` cell is 8-aligned + return NamedTypeSchema(name, schema, offset=offset, size=size, alignment=alignment) def _info( type_key: str, - fields: tuple[tuple[str, TypeSchema], ...] = (), + fields: tuple[NamedTypeSchema, ...] = (), *, parent: str | None = "ffi.Object", ancestors: list[str] | None = None, is_final: bool | None = None, + total_size: int | None = None, ) -> ObjectInfo: if ancestors is None: ancestors = ["ffi.Object"] if parent in (None, "ffi.Object") else ["ffi.Object", parent] return ObjectInfo( - fields=[NamedTypeSchema(name, schema) for name, schema in fields], + fields=list(fields), methods=[], type_key=type_key, parent_type_key=parent, ancestors=ancestors, is_final=is_final, + total_size=total_size, ) +#: Ancestors the tests define with byte facts; anything else a test names but +#: does not define resolves to a type without metadata of its own. +_SYNTHETIC: dict[str, ObjectInfo] = {} + + +def _register(*infos: ObjectInfo) -> None: + for info in infos: + assert info.type_key is not None + _SYNTHETIC[info.type_key] = info + + +@pytest.fixture(autouse=True) +def _synthetic_registry(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: + def lookup(type_key: str) -> ObjectInfo: + if type_key in _SYNTHETIC: + return _SYNTHETIC[type_key] + if type_key.startswith(("ffi.", "testing.")): + return object_info_from_type_key(type_key) + return _info(type_key, total_size=None) + + monkeypatch.setattr(codegen, "object_info_from_type_key", lookup) + yield + _SYNTHETIC.clear() + + def _object_block(type_key: str) -> CodeBlock: return CodeBlock( kind="object", @@ -86,6 +131,9 @@ def _uses(imports: RustImports) -> set[str]: return {item.path for item in imports.items} +NO_METADATA = "/// Opaque: no metadata of its own: total_size is unknown. Fields are read through the C ABI getters." + + # --------------------------------------------------------------------------- # `use` modelling and type rendering # --------------------------------------------------------------------------- @@ -167,12 +215,14 @@ def test_directives_parse() -> None: directives.add("nullable", "ir.Expr.span", 2) directives.add("enum", "tirx.For.kind -> ForKind(i32) { Serial=0, Parallel = 1 }", 3) directives.add("enum", "tirx.For.mode -> Mode(u8)", 4) + directives.add("opaque", " ir.SourceName ", 5) assert directives.field_types == {"tirx.Add.a": "PrimExpr"} assert directives.nullable == {"ir.Expr.span"} assert directives.enums == { "tirx.For.kind": EnumSpec("ForKind", "i32", (("Serial", 0), ("Parallel", 1))), "tirx.For.mode": EnumSpec("Mode", "u8", ()), } + assert directives.opaque == {"ir.SourceName"} @pytest.mark.parametrize( @@ -185,17 +235,18 @@ def test_directives_parse() -> None: ("enum", "tirx.For.kind -> ForKind", "Name(i32)"), ("enum", "tirx.For.kind -> ForKind(i128)", "Name(i32)"), ("enum", "tirx.For.kind -> ForKind(i32) { Serial }", "Name(i32)"), + ("opaque", "ir.SourceName ir.Source", ""), ("upcast", "tirx.Add -> PrimExpr", "Unknown directive"), ], ) def test_directives_reject_malformed(name: str, payload: str, expected: str) -> None: - with pytest.raises(ValueError, match=re.escape(expected)) as exc: + with pytest.raises(DirectiveError, match=re.escape(expected)) as exc: Directives().add(name, payload, 7) assert "at line 7" in str(exc.value) def test_generator_declares_its_directives_and_records_imports() -> None: - assert RUST.directive_kinds == {"import-object", "field", "nullable", "enum"} + assert RUST.directive_kinds == {"import-object", "field", "nullable", "enum", "opaque"} imports = RUST.new_imports() RUST.add_directive(imports, "import-object", "tvm_ffi.libinfo.Foo;False;_Foo", 1) RUST.add_directive(imports, "nullable", "demo.Node.span", 2) @@ -204,51 +255,52 @@ def test_generator_declares_its_directives_and_records_imports() -> None: # --------------------------------------------------------------------------- -# Object rendering +# Opaque rendering (no byte facts: the layout cannot be proven) # --------------------------------------------------------------------------- -ROOT_EXPECTED = """\ +ROOT_EXPECTED = f"""\ +{NO_METADATA} #[repr(C)] #[derive(tvm_ffi::derive::Object)] #[type_key = "demo.Pair"] -pub struct PairObj { +pub struct PairObj {{ base: Object, -} +}} #[repr(C)] #[derive(tvm_ffi::derive::ObjectRef, Clone)] -pub struct Pair { +pub struct Pair {{ data: ObjectArc, -} +}} -impl Deref for Pair { +impl Deref for Pair {{ type Target = PairObj; - fn deref(&self) -> &PairObj { + fn deref(&self) -> &PairObj {{ &self.data - } -} + }} +}} -impl PairObj { - pub fn a(&self) -> Result { +impl PairObj {{ + pub fn a(&self) -> Result {{ FieldGetter::new(Self::type_index(), "a")?.get(self) - } + }} - pub fn tag(&self) -> Result> { + pub fn tag(&self) -> Result> {{ FieldGetter::new(Self::type_index(), "tag")?.get(self) - } + }} - pub fn items(&self) -> Result> { + pub fn items(&self) -> Result> {{ FieldGetter::new(Self::type_index(), "items")?.get(self) - } + }} - pub fn owner(&self) -> Result { + pub fn owner(&self) -> Result {{ FieldGetter::new(Self::type_index(), "owner")?.get(self) - } + }} - pub fn r#type(&self) -> Result { + pub fn r#type(&self) -> Result {{ FieldGetter::new(Self::type_index(), "type")?.get_any(self) - } -}""" + }} +}}""" def test_render_root_object() -> None: @@ -256,11 +308,11 @@ def test_render_root_object() -> None: info = _info( "demo.Pair", ( - ("a", TypeSchema("int")), - ("tag", TypeSchema("Optional", (TypeSchema("str"),))), - ("items", TypeSchema("Array")), - ("owner", TypeSchema("Object")), - ("type", TypeSchema("Union", (TypeSchema("int"), TypeSchema("str")))), + _field("a", "int"), + _field("tag", TypeSchema("Optional", (TypeSchema("str"),))), + _field("items", "Array"), + _field("owner", "Object"), + _field("type", TypeSchema("Union", (TypeSchema("int"), TypeSchema("str")))), ), ) text, imports = _render(info) @@ -290,7 +342,7 @@ def test_render_derived_object_same_module() -> None: """A generated parent is embedded, dereferenced to, and upcast to along the chain.""" info = _info( "demo.Add", - (("a", TypeSchema("demo.Expr")),), + (_field("a", "demo.Expr"),), parent="demo.Expr", ancestors=["ffi.Object", "demo.BaseExpr", "demo.Expr"], is_final=True, @@ -317,7 +369,7 @@ def test_render_object_under_builtin_parent() -> None: """A builtin parent is embedded via header-only stand-ins so `TYPE_DEPTH` matches the registry.""" info = _info( "demo.Color", - (("value", TypeSchema("int")),), + (_field("value", "int"),), parent="ffi.IntEnum", ancestors=["ffi.Object", "ffi.Enum", "ffi.IntEnum"], ) @@ -401,82 +453,83 @@ def test_import_section_defines_builtin_mirrors_once() -> None: assert "\n".join(block.lines[1:-1]) == BUILTIN_MIRRORS_EXPECTED -ITER_VAR_EXPECTED = """\ +ITER_VAR_EXPECTED = f"""\ #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] #[repr(transparent)] pub struct IterVarType(i32); #[allow(non_upper_case_globals)] -impl IterVarType { +impl IterVarType {{ pub const kDataPar: Self = Self(0); pub const kThreadIndex: Self = Self(1); - pub const fn from_raw(value: i32) -> Self { + pub const fn from_raw(value: i32) -> Self {{ Self(value) - } - pub const fn as_raw(self) -> i32 { + }} + pub const fn as_raw(self) -> i32 {{ self.0 - } -} + }} +}} -impl TryFrom for IterVarType { +impl TryFrom for IterVarType {{ type Error = Error; - fn try_from(value: i64) -> Result { - i32::try_from(value).map(Self).map_err(|_| { - Error::new(VALUE_ERROR, &format!("IterVarType value {value} does not fit i32"), "") - }) - } -} - + fn try_from(value: i64) -> Result {{ + i32::try_from(value).map(Self).map_err(|_| {{ + Error::new(VALUE_ERROR, &format!("IterVarType value {{value}} does not fit i32"), "") + }}) + }} +}} + +{NO_METADATA} #[repr(C)] #[derive(tvm_ffi::derive::Object)] #[type_key = "tirx.IterVar"] #[type_final] -pub struct IterVarObj { +pub struct IterVarObj {{ base: PrimExprConvertibleObj, -} +}} #[repr(C)] #[derive(tvm_ffi::derive::ObjectRef, Clone)] -pub struct IterVar { +pub struct IterVar {{ data: ObjectArc, -} +}} -impl Deref for IterVar { +impl Deref for IterVar {{ type Target = IterVarObj; - fn deref(&self) -> &IterVarObj { + fn deref(&self) -> &IterVarObj {{ &self.data - } -} + }} +}} -impl Deref for IterVarObj { +impl Deref for IterVarObj {{ type Target = PrimExprConvertibleObj; - fn deref(&self) -> &PrimExprConvertibleObj { + fn deref(&self) -> &PrimExprConvertibleObj {{ &self.base - } -} + }} +}} -impl IterVarObj { - pub fn dom(&self) -> Result> { +impl IterVarObj {{ + pub fn dom(&self) -> Result> {{ FieldGetter::new(Self::type_index(), "dom")?.get(self) - } + }} - pub fn var(&self) -> Result { + pub fn var(&self) -> Result {{ FieldGetter::new(Self::type_index(), "var")?.get(self) - } + }} - pub fn iter_type(&self) -> Result { + pub fn iter_type(&self) -> Result {{ let raw: i64 = FieldGetter::new(Self::type_index(), "iter_type")?.get(self)?; IterVarType::try_from(raw) - } + }} - pub fn thread_tag(&self) -> Result { + pub fn thread_tag(&self) -> Result {{ FieldGetter::new(Self::type_index(), "thread_tag")?.get(self) - } + }} - pub fn span(&self) -> Result> { + pub fn span(&self) -> Result> {{ FieldGetter::new(Self::type_index(), "span")?.get(self) - } -} + }} +}} tvm_ffi::impl_object_upcast!(IterVar => PrimExprConvertible);""" @@ -486,11 +539,11 @@ def test_render_iter_var_golden() -> None: info = _info( "tirx.IterVar", ( - ("dom", TypeSchema("ir.Range")), - ("var", TypeSchema("ir.Var")), - ("iter_type", TypeSchema("int")), - ("thread_tag", TypeSchema("str")), - ("span", TypeSchema("ir.Span")), + _field("dom", "ir.Range"), + _field("var", "ir.Var"), + _field("iter_type", "int"), + _field("thread_tag", "str"), + _field("span", "ir.Span"), ), parent="ir.PrimExprConvertible", is_final=True, @@ -518,7 +571,10 @@ def test_render_iter_var_golden() -> None: def test_field_directive_with_path_records_use_and_nullable_does_not_double_wrap() -> None: info = _info( "demo.Node", - (("buffer", TypeSchema("demo.Var")), ("dom", TypeSchema("Optional", (TypeSchema("int"),)))), + ( + _field("buffer", "demo.Var"), + _field("dom", TypeSchema("Optional", (TypeSchema("int"),))), + ), ) imports = RUST.new_imports() RUST.add_directive(imports, "field", "demo.Node.buffer -> crate::typed::BufferVar", 1) @@ -529,6 +585,242 @@ def test_field_directive_with_path_records_use_and_nullable_does_not_double_wrap assert "crate::typed::BufferVar" in _uses(imports) +# --------------------------------------------------------------------------- +# Complete rendering (byte facts prove the layout) +# --------------------------------------------------------------------------- + + +def _span() -> ObjectInfo: + return _info( + "ir.Span", + ( + _field("source_name", "ir.SourceName", 24, 8), + _field("line", "int", 32, 4), + _field("column", "int", 36, 4), + _field("end_line", "int", 40, 4), + _field("end_column", "int", 44, 4), + ), + total_size=48, + ) + + +def _expr() -> ObjectInfo: + return _info( + "ir.Expr", + (_field("span", "ir.Span", 24, 8), _field("ty", "ir.Type", 32, 8)), + total_size=40, + ) + + +def _add() -> ObjectInfo: + return _info( + "tirx.Add", + (_field("a", "ir.Expr", 40, 8), _field("b", "ir.Expr", 48, 8)), + parent="ir.Expr", + total_size=56, + is_final=True, + ) + + +EXPR_EXPECTED = """\ +/// Complete: reflected fields fill [24, 40) exactly. +#[repr(C)] +#[derive(tvm_ffi::derive::Object)] +#[type_key = "ir.Expr"] +pub struct ExprObj { + base: Object, + pub span: Option, + pub ty: Type, +} + +const _: () = { + assert!(::core::mem::size_of::() == 40); + assert!(::core::mem::align_of::() == 8); +}; + +#[repr(C)] +#[derive(tvm_ffi::derive::ObjectRef, Clone)] +pub struct Expr { + data: ObjectArc, +} + +impl Deref for Expr { + type Target = ExprObj; + fn deref(&self) -> &ExprObj { + &self.data + } +}""" + +ADD_EXPECTED = """\ +/// Complete: reflected fields fill [40, 56) exactly. +#[repr(C)] +#[derive(tvm_ffi::derive::Object)] +#[type_key = "tirx.Add"] +#[type_final] +pub struct AddObj { + base: ExprObj, + pub a: PrimExpr, + pub b: PrimExpr, +} + +const _: () = { + assert!(::core::mem::size_of::() == 56); + assert!(::core::mem::align_of::() == 8); +}; + +#[repr(C)] +#[derive(tvm_ffi::derive::ObjectRef, Clone)] +pub struct Add { + data: ObjectArc, +} + +impl Deref for Add { + type Target = AddObj; + fn deref(&self) -> &AddObj { + &self.data + } +} + +impl Deref for AddObj { + type Target = ExprObj; + fn deref(&self) -> &ExprObj { + &self.base + } +} + +tvm_ffi::impl_object_upcast!(Add => Expr);""" + + +def test_render_complete_expr_golden() -> None: + """`ir.Expr` as tvm-rust-ext hand-writes it: `span: Option`, `ty: Type`, no getters.""" + imports = RUST.new_imports() + RUST.add_directive(imports, "nullable", "ir.Expr.span", 1) + text, imports = _render(_expr(), imports) + assert text == EXPR_EXPECTED + assert _uses(imports) == {"std::ops::Deref", "tvm_ffi::Object", "tvm_ffi::ObjectArc"} + + +def test_render_complete_add_golden() -> None: + """`tirx.Add` on top of a complete `ir.Expr`, with `field` directives narrowing `a` / `b`.""" + _register(_expr()) + imports = RUST.new_imports() + RUST.add_directive(imports, "field", "tirx.Add.a -> PrimExpr", 1) + RUST.add_directive(imports, "field", "tirx.Add.b -> PrimExpr", 2) + text, imports = _render(_add(), imports) + assert text == ADD_EXPECTED + assert _uses(imports) == { + "std::ops::Deref", + "tvm_ffi::ObjectArc", + "super::ir::ExprObj", + "super::ir::Expr", + } + + +def test_render_complete_span_and_prim_type() -> None: + """Scalars take their reflected width; alignment padding is reported, not mirrored.""" + text, _ = _render(_span()) + assert ( + "pub struct SpanObj {\n" + " base: Object,\n" + " pub source_name: SourceName,\n" + " pub line: i32,\n" + " pub column: i32,\n" + " pub end_line: i32,\n" + " pub end_column: i32,\n" + "}" + ) in text + assert "/// Complete: reflected fields fill [24, 48) exactly." in text + + _register(_info("ir.Type", (_field("span", "ir.Span", 24, 8),), total_size=32)) + prim_type = _info( + "ir.PrimType", + (_field("dtype", "dtype", 32, 4, 2),), + parent="ir.Type", + total_size=40, + is_final=True, + ) + text, imports = _render(prim_type) + assert ( + "/// Complete: reflected fields fill [32, 40) exactly (alignment padding [36, 40))." in text + ) + assert "pub struct PrimTypeObj {\n base: TypeObj,\n pub dtype: DLDataType,\n}" in text + assert "assert!(::core::mem::size_of::() == 40);" in text + assert "tvm_ffi::DLDataType" in _uses(imports) + + +def test_complete_optional_field_mirrors() -> None: + """`Optional` fields mirror their C++ layout: a 16-byte cell or a nullable pointer.""" + info = _info( + "demo.Opt", + ( + _field("count", TypeSchema("Optional", (TypeSchema("int"),)), 24, 16), + _field("name", TypeSchema("Optional", (TypeSchema("str"),)), 40, 16), + _field("items", TypeSchema("Optional", (TypeSchema("Array"),)), 56, 8), + ), + total_size=64, + ) + text, imports = _render(info) + assert "pub count: Optional," in text + assert "pub name: Optional," in text + assert "pub items: Option>," in text + assert "tvm_ffi::Optional" in _uses(imports) + + +@pytest.mark.parametrize( + "field", + [ + _field("x", TypeSchema("Optional", (TypeSchema("Any"),)), 24, 16), + _field("x", TypeSchema("Optional", (TypeSchema("int"),)), 24, 8), # not the 16-byte cell + _field("x", TypeSchema("Union", (TypeSchema("int"), TypeSchema("str"))), 24, 16), + _field("x", "ctypes.c_void_p", 24, 8), + ], +) +def test_unrenderable_field_keeps_the_type_opaque(field: NamedTypeSchema) -> None: + """A field without a mirror demotes an otherwise complete type; it is read as `Any`.""" + assert field.size is not None + info = _info("demo.Holder", (field,), total_size=HEADER + field.size) + text, _ = _render(info) + assert "/// Opaque: field 'x'" in text + assert "has no native mirror" in text + assert "impl HolderObj {\n pub fn x(&self) -> Result<" in text + assert "const _: () =" not in text + + +def test_opaque_directive_vetoes_a_complete_type() -> None: + imports = RUST.new_imports() + RUST.add_directive(imports, "opaque", "ir.Span", 1) + text, _ = _render(_span(), imports) + assert "/// Opaque: vetoed by directive although the layout is reproducible." in text + assert "pub fn line(&self) -> Result {" in text + + +def test_opaque_parent_keeps_the_child_opaque() -> None: + """A child of an opaque type cannot embed a mirror of its parent: it stays opaque.""" + _register(_info("ir.Expr", (_field("span", "ir.Span", 24, 8),), total_size=40)) # hole + text, _ = _render(_add()) + assert "/// Opaque: parent 'ir.Expr' is opaque (uncovered-bytes)." in text + assert " base: ExprObj," in text + assert "pub fn a(&self) -> Result {" in text + + +@pytest.mark.parametrize( + ("name", "payload", "message"), + [ + ("field", "demo.Pair.count -> i64", "maps a 4-byte field to `i64` (8 bytes)"), + ("enum", "demo.Pair.count -> Kind(i64)", "maps a 4-byte field to `i64` (8 bytes)"), + ("nullable", "demo.Pair.count", "the field is 4 bytes, not a pointer-sized"), + ], +) +def test_directive_disagreeing_with_bytes_is_an_error( + name: str, payload: str, message: str +) -> None: + info = _info("demo.Pair", (_field("count", "int", 24, 4),), total_size=32) + imports = RUST.new_imports() + RUST.add_directive(imports, name, payload, 1) + with pytest.raises(DirectiveError, match=re.escape(message)): + _render(info, imports) + + # --------------------------------------------------------------------------- # File scaffolding # --------------------------------------------------------------------------- @@ -578,10 +870,44 @@ def test_finalize_module_tree(tmp_path: Path) -> None: # --------------------------------------------------------------------------- -# The pipeline end to end +# The real registry, and the pipeline end to end # --------------------------------------------------------------------------- +def test_registry_complete_chain_is_mirrored() -> None: + base, _ = _render(object_info_from_type_key("testing.TestCxxClassBase")) + assert ( + "/// Complete: reflected fields fill [24, 40) exactly (alignment padding [36, 40))." in base + ) + assert ( + "pub struct TestCxxClassBaseObj {\n base: Object,\n pub v_i64: i64,\n pub v_i32: i32,\n}" + in base + ) + assert "assert!(::core::mem::size_of::() == 40);" in base + assert "FieldGetter" not in base + + dd, imports = _render(object_info_from_type_key("testing.TestCxxClassDerivedDerived")) + assert ( + " base: TestCxxClassDerivedObj,\n pub v_str: String,\n pub v_bool: bool,\n}" in dd + ) + assert dd.endswith( + "tvm_ffi::impl_object_upcast!(TestCxxClassDerivedDerived => TestCxxClassBase, " + "TestCxxClassDerivedDerived => TestCxxClassDerived);" + ) + assert "tvm_ffi::String" in _uses(imports) + + +def test_registry_hidden_field_and_vptr_stay_opaque() -> None: + hidden, _ = _render(object_info_from_type_key("testing.TestCxxClassHiddenField")) + assert ( + "/// Opaque: bytes [32, 40) of [24, 48) are not accounted for by reflected fields. " + "Fields are read through the C ABI getters." + ) in hidden + assert "pub fn v_i32(&self) -> Result {" in hidden + poly, _ = _render(object_info_from_type_key("testing.TestCxxClassPolymorphic")) + assert "/// Opaque: bytes [32, 40) of [24, 40)" in poly + + def test_stage_3_applies_directives_to_a_registered_type(tmp_path: Path) -> None: src = tmp_path / "mod.rs" src.write_text( @@ -602,9 +928,8 @@ def test_stage_3_applies_directives_to_a_registered_type(tmp_path: Path) -> None _stage_3(info, Options(dry_run=True), RUST.default_ty_map(), {}, RUST) text = "\n".join(line for block in info.code_blocks for line in block.lines) assert "pub struct Kind(i32);" in text - assert "pub fn v_i32(&self) -> Result {" in text - assert "pub fn v_i64(&self) -> Result {" in text - assert "use tvm_ffi::FieldGetter;" in text + assert " pub v_i32: Kind,\n" in text + assert "use tvm_ffi::Error;" in text def test_cli_init_generates_a_module_tree(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: @@ -627,9 +952,9 @@ def test_cli_init_generates_a_module_tree(tmp_path: Path, monkeypatch: pytest.Mo assert (tmp_path / "mod.rs").read_text(encoding="utf-8") == "pub mod testing;\n" text = (tmp_path / "testing" / "mod.rs").read_text(encoding="utf-8") assert text.startswith("#![allow(dead_code, unused_imports)]\n") - assert "use tvm_ffi::FieldGetter;" in text + assert "use tvm_ffi::FieldGetter;" in text # the opaque fixtures read through getters assert '#[type_key = "testing.TestCxxClassDerivedDerived"]' in text - assert " base: TestCxxClassDerivedObj," in text + assert " base: TestCxxClassDerivedObj,\n pub v_str: String," in text assert ( "tvm_ffi::impl_object_upcast!(TestCxxClassDerivedDerived => TestCxxClassBase, " "TestCxxClassDerivedDerived => TestCxxClassDerived);" @@ -645,3 +970,24 @@ def test_cli_init_generates_a_module_tree(tmp_path: Path, monkeypatch: pytest.Mo # Running again over the generated tree is a no-op. assert stub_cli.__main__() == 0 assert (tmp_path / "testing" / "mod.rs").read_text(encoding="utf-8") == text + + +def test_cli_exits_non_zero_on_a_directive_error( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + src = tmp_path / "mod.rs" + src.write_text( + "\n".join( + [ + f"{C.RUST_SYNTAX.directive('field')} testing.TestCxxClassBase.v_i32 -> i64", + f"{C.RUST_SYNTAX.begin} object/testing.TestCxxClassBase", + C.RUST_SYNTAX.end, + "", + ] + ), + encoding="utf-8", + ) + monkeypatch.setattr("sys.argv", ["tvm-ffi-stubgen", "--target", "rust", str(src)]) + assert stub_cli.__main__() == 1 + src.write_text(f"{C.RUST_SYNTAX.directive('upcast')} testing.TestCxxClassBase -> X\n") + assert stub_cli.__main__() == 1 From 0a987671316e1a9c00576ce3ea72c4a8713bc560 Mon Sep 17 00:00:00 2001 From: yuchuan Date: Thu, 3 Sep 2026 14:43:07 -0400 Subject: [PATCH 2/8] [STUBGEN] Generate new allocators for complete Rust bindings Signed-off-by: yuchuan --- examples/rust_stubgen/README.md | 25 +++- .../rust/src/generated/rust_stubgen/mod.rs | 11 ++ examples/rust_stubgen/rust/src/main.rs | 33 +++-- examples/rust_stubgen/src/int_pair.cc | 6 +- python/tvm_ffi/stub/rust_generator/codegen.py | 129 ++++++++++++++++-- python/tvm_ffi/stub/rust_generator/consts.py | 8 +- .../tvm_ffi/stub/rust_generator/directives.py | 32 +++-- tests/python/test_stubgen_rust.py | 115 +++++++++++++++- 8 files changed, 312 insertions(+), 47 deletions(-) diff --git a/examples/rust_stubgen/README.md b/examples/rust_stubgen/README.md index aa4d70573..2261b7d75 100644 --- a/examples/rust_stubgen/README.md +++ b/examples/rust_stubgen/README.md @@ -28,12 +28,12 @@ reflected fields account for every byte of the object: - `rust_stubgen.IntRange` does, so its binding is *complete*: the struct mirrors the fields at their real offsets and widths, and Rust reads `range.begin` directly. A `const` assertion pins the struct's size and alignment to the - reflected facts. + reflected facts. The object is allocated in Rust, by a generated function + that takes every field by value; a C++ function then reads it back. - `rust_stubgen.IntPair` has a vtable in front of the object header, so its binding is *opaque*: the struct embeds only the parent and every field is read - through an accessor that calls the C ABI getter. - -Construction goes through the registered global functions in both cases. + through an accessor that calls the C ABI getter. It is constructed by a + registered global function. A builtin parent such as `ffi.IntEnum` has no `Obj` in the crate; the import section defines a header-only stand-in per builtin ancestor so the @@ -64,6 +64,17 @@ open newtype: // tvm-ffi-stubgen(enum): rust_stubgen.IntPair.kind -> PairKind(i32) { Unordered=0, Ordered=1 } ``` -Two more are available: `field` names the Rust type of a field's accessor -(`// tvm-ffi-stubgen(field): rust_stubgen.IntPair.a -> MyInt`) and `nullable` -wraps it in `Option` (`// tvm-ffi-stubgen(nullable): rust_stubgen.IntPair.a`). +It also marks `IntRange` as having a hand-written `new` (in `main.rs`, where it +validates the extent), so the generator emits none of its own: + +```rust +// tvm-ffi-stubgen(custom-new): rust_stubgen.IntRange +``` + +Without the marker the generator emits `IntRange::new` itself, and a +hand-written one is a duplicate definition. Three more directives 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`), +and `upcast` adds a conversion to a hand-written typed view +(`// tvm-ffi-stubgen(upcast): rust_stubgen.IntRange -> MyView`). 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 cd99f220e..e45477b7e 100644 --- a/examples/rust_stubgen/rust/src/generated/rust_stubgen/mod.rs +++ b/examples/rust_stubgen/rust/src/generated/rust_stubgen/mod.rs @@ -99,6 +99,10 @@ impl IntPairObj { } // tvm-ffi-stubgen(end) +// `IntRange::new` is hand-written in main.rs (it validates the extent), so the +// generator does not emit one; it still emits `IntRangeObj::new` for it to call. +// tvm-ffi-stubgen(custom-new): rust_stubgen.IntRange + // tvm-ffi-stubgen(begin): object/rust_stubgen.IntRange /// Complete: reflected fields fill [24, 40) exactly (alignment padding [36, 40)). #[repr(C)] @@ -128,4 +132,11 @@ impl Deref for IntRange { &self.data } } + +impl IntRangeObj { + pub(crate) fn new(begin: i64, extent: i32) -> Self { + let base = Object::new(); + Self { base, begin, extent } + } +} // tvm-ffi-stubgen(end) diff --git a/examples/rust_stubgen/rust/src/main.rs b/examples/rust_stubgen/rust/src/main.rs index 7b0f360df..748dcffb0 100644 --- a/examples/rust_stubgen/rust/src/main.rs +++ b/examples/rust_stubgen/rust/src/main.rs @@ -16,12 +16,25 @@ * specific language governing permissions and limitations * under the License. */ -//! Use the stubgen-generated `IntPair` binding (see ../../README.md). +//! Use the stubgen-generated `IntPair` and `IntRange` bindings (see ../../README.md). mod generated; -use generated::rust_stubgen::{IntPair, IntRange, PairKind}; -use tvm_ffi::{Module, Result}; +use generated::rust_stubgen::{IntPair, IntRange, IntRangeObj, PairKind}; +use tvm_ffi::{Error, Module, ObjectArc, ObjectRefCore, Result, VALUE_ERROR}; + +/// The hand-written constructor of `IntRange`: the `custom-new` directive in +/// the generated file keeps the generator from emitting one. +impl IntRange { + pub fn new(begin: i64, extent: i32) -> Result { + if extent < 0 { + return Err(Error::new(VALUE_ERROR, "IntRange extent must not be negative", "")); + } + let data = ObjectArc::new(IntRangeObj::new(begin, extent)); + // SAFETY: `data` holds a freshly allocated `IntRangeObj`, the container type of `IntRange`. + Ok(unsafe { Self::from_data(data) }) + } +} /// Path of the C++ shared library built by CMake into `../build`. fn lib_path() -> String { @@ -53,11 +66,15 @@ fn main() -> Result<()> { .try_into()?; println!("sum={sum}"); - // `IntRange` has a reproducible layout: its fields are plain struct members. - let range: IntRange = tvm_ffi::cached_global_func!("rust_stubgen.IntRange") - .call_tuple((10i64, 5i64))? - .try_into()?; + // `IntRange` has a reproducible layout: it is allocated in Rust and its + // fields are plain struct members, on both sides of the ABI. + let range = IntRange::new(10, 5)?; println!("begin={} extent={}", range.begin, range.extent); - assert_eq!(range.begin + i64::from(range.extent), 15); + let end: i64 = tvm_ffi::cached_global_func!("rust_stubgen.IntRangeEnd") + .call_tuple((range.clone(),))? + .try_into()?; + println!("end={end}"); + assert_eq!(end, 15); + assert!(IntRange::new(0, -1).is_err()); Ok(()) } diff --git a/examples/rust_stubgen/src/int_pair.cc b/examples/rust_stubgen/src/int_pair.cc index 85d44ed33..703712677 100644 --- a/examples/rust_stubgen/src/int_pair.cc +++ b/examples/rust_stubgen/src/int_pair.cc @@ -66,8 +66,6 @@ class IntRangeObj : public ffi::Object { class IntRange : public ffi::ObjectRef { public: - IntRange(int64_t begin, int32_t extent) { data_ = ffi::make_object(begin, extent); } - TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(IntRange, ffi::ObjectRef, IntRangeObj); }; @@ -84,8 +82,8 @@ TVM_FFI_STATIC_INIT_BLOCK() { .def("rust_stubgen.IntPair", [](int64_t a, int64_t b, int32_t kind) { return IntPair(a, b, kind); }) .def("rust_stubgen.IntPairSum", [](const IntPair& pair) { return pair->Sum(); }) - .def("rust_stubgen.IntRange", - [](int64_t begin, int32_t extent) { return IntRange(begin, extent); }); + .def("rust_stubgen.IntRangeEnd", + [](const IntRange& range) { return range->begin + range->extent; }); } // [object.end] diff --git a/python/tvm_ffi/stub/rust_generator/codegen.py b/python/tvm_ffi/stub/rust_generator/codegen.py index a6f1bfe4f..f671601e9 100644 --- a/python/tvm_ffi/stub/rust_generator/codegen.py +++ b/python/tvm_ffi/stub/rust_generator/codegen.py @@ -23,23 +23,29 @@ - *complete*: the layout is reproducible, so the struct mirrors every physical field at its real offset and width, public, borrowed directly. A ``const`` assertion pins the struct's ``size_of`` / ``align_of`` to the reflected facts, - so a mirror rustc lays out differently fails to compile. + so a mirror rustc lays out differently fails to compile. The type is also + allocatable from Rust: ``Obj::new`` (crate-private) takes every + physical field root to leaf so derived allocators can chain to it, and the + wrapper's public ``new`` takes the same parameters and does one + ``ObjectArc::new``. A ``custom-new`` directive leaves the wrapper's ``new`` + to hand-written code; ``Obj::new`` is still generated for it to call. - *opaque*: the struct embeds only its parent, and one accessor per reflected field reads through the C ABI getter. The bytes are never reproduced, so the - binding is correct for every registered type. + binding is correct for every registered type. Nothing allocates it. The two target-language rules the classifier leaves to its caller live here: a field without a Rust mirror (``Optional``, a ``Union``, ``void*``, ...) makes the type opaque and is read as ``Any``; an ``opaque`` directive vetoes a reproducible layout. ``field`` / ``nullable`` / ``enum`` directives shape the field types of both forms; where a directive names a scalar width, it is checked -against the reflected field size at generation time. - -Construction and behaviour go through the registered global functions, -hand-written outside the markers. A builtin parent (``ffi.IntEnum``, say) has -no ``Obj`` in the crate: the import section defines a header-only -stand-in per builtin ancestor, so ``derive(Object)`` computes the registry's -``TYPE_DEPTH``. +against the reflected field size at generation time. ``upcast`` appends typed +views to the ancestor chain. + +Nothing here calls methods or packed constructors: behaviour goes through the +registered global functions, hand-written outside the markers. A builtin parent +(``ffi.IntEnum``, say) has no ``Obj`` in the crate: the import section +defines a header-only stand-in per builtin ancestor, so ``derive(Object)`` +computes the registry's ``TYPE_DEPTH``. """ from __future__ import annotations @@ -62,6 +68,15 @@ from .directives import EnumSpec +def _call_lines(open_: str, items: list[str], close: str) -> list[str]: + """``open_ + items + close`` on one line, or one item per line when it would overflow.""" + line = f"{open_}{', '.join(items)}{close}" + if len(line) <= C_RUST.RUST_MAX_WIDTH: + return [line] + indent = open_[: len(open_) - len(open_.lstrip())] + return [open_, *[f"{indent} {item}," for item in items], f"{indent}{close}"] + + def _check_width(target: str, field: NamedTypeSchema, rust_type: str, width: int) -> None: """Reject a directive whose scalar type does not match the reflected field size.""" if field.size is not None and field.size != width: @@ -310,12 +325,14 @@ def _deref_lines(self, source: str, target: str, member: str) -> list[str]: ] def _upcast_lines(self) -> list[str]: - """``impl_object_upcast!`` from the wrapper to every ancestor's wrapper.""" + """``impl_object_upcast!`` to every ancestor's wrapper, then the ``upcast`` directives.""" targets = [ self.imports.record(self._generated_type_path(key)) for key in self.info.ancestors if self._generated(key) ] + for view in self.imports.directives.upcasts.get(self.type_key, []): + targets.append(self.imports.record(view) if "::" in view else view) if not targets: return [] pairs = ", ".join(f"{self.leaf} => {target}" for target in targets) @@ -354,6 +371,96 @@ def _struct_lines(self, verdict: Verdict, base: str) -> list[str]: "};", ] + # --- allocators -------------------------------------------------------- + + def _allocator_params(self, key: str, info: ObjectInfo) -> list[tuple[str, str]]: + """``(field, type)`` of every physical field root to leaf, as ``Obj::new`` takes them.""" + parent = info.parent_type_key + inherited: list[tuple[str, str]] = [] + if parent is not None and self._generated(parent): + inherited = self._allocator_params(parent, object_info_from_type_key(parent)) + return self._level_params(key, info, inherited) + + def _level_params( + self, key: str, info: ObjectInfo, inherited: list[tuple[str, str]] + ) -> list[tuple[str, str]]: + """Extend the parent's parameters with ``key``'s own fields by offset. + + A ``field`` directive on ``.`` for an inherited field narrows + that parameter (``tirx.Add.ty -> PrimType``); the allocator body upcasts + it with ``.into()`` when handing it to the parent. + """ + params = [(name, self._narrowed(key, name, rust_type)) for name, rust_type in inherited] + for field in sorted(info.fields, key=lambda f: f.offset or 0): + mirror = self._field_mirror(key, field, self.imports) + assert mirror is not None # complete: every field along the chain has a mirror + params.append((field.name, mirror)) + return params + + def _narrowed(self, key: str, field_name: str, rust_type: str) -> str: + override = self.imports.directives.field_types.get(f"{key}.{field_name}") + if override is None: + return rust_type + return self.imports.record(override) if "::" in override else override + + def _fn_lines( + self, head: str, params: list[tuple[str, str]], call: tuple[str, list[str]], result: str + ) -> list[str]: + """Render ``() -> Self { let ; }`` inside an ``impl`` block.""" + plist = [f"{rust_ident(name)}: {rust_type}" for name, rust_type in params] + binding, args = call + return [ + *_call_lines(f" {head}(", plist, ") -> Self {"), + *_call_lines(f" let {binding}(", args, ");"), + f" {result}", + " }", + ] + + def _allocator_sections(self, base: str, has_parent: bool) -> list[list[str]]: + """``Obj::new`` and, unless ``custom-new`` reserves it, the wrapper's ``new``. + + Both take every physical field root to leaf. + """ + inherited: list[tuple[str, str]] = [] + if has_parent: + parent = self.info.parent_type_key + assert parent is not None + inherited = self._allocator_params(parent, object_info_from_type_key(parent)) + params = self._level_params(self.type_key, self.info, inherited) + to_parent = [ + f"{rust_ident(name)}.into()" if rust_type != parent_type else rust_ident(name) + for (name, rust_type), (_, parent_type) in zip(params, inherited) + ] + own = [rust_ident(f.name) for f in sorted(self.info.fields, key=lambda f: f.offset or 0)] + forward = [rust_ident(name) for name, _ in params] + sections = [ + [ + f"impl {self.obj_struct} {{", + *self._fn_lines( + "pub(crate) fn new", + params, + (f"base = {base}::new", to_parent), + f"Self {{ {', '.join(['base', *own])} }}", + ), + "}", + ] + ] + if self.type_key not in self.imports.directives.custom_new: + sections.append( + [ + f"impl {self.leaf} {{", + " /// Lossless complete-field allocation.", + *self._fn_lines( + "pub fn new", + params, + (f"obj = {self.obj_struct}::new", forward), + "Self { data: ObjectArc::new(obj) }", + ), + "}", + ] + ) + return sections + def body(self) -> list[str]: """Build the Rust source lines for the object.""" verdict = self.classify() @@ -401,6 +508,8 @@ def body(self) -> list[str]: "}", ] ) + elif verdict.is_complete: + sections += self._allocator_sections(base, has_parent) upcasts = self._upcast_lines() if upcasts: sections.append(upcasts) diff --git a/python/tvm_ffi/stub/rust_generator/consts.py b/python/tvm_ffi/stub/rust_generator/consts.py index c8310e9ea..ee465acd1 100644 --- a/python/tvm_ffi/stub/rust_generator/consts.py +++ b/python/tvm_ffi/stub/rust_generator/consts.py @@ -19,7 +19,9 @@ from __future__ import annotations #: One-line directives the Rust backend consumes. -RUST_DIRECTIVE_KINDS = frozenset({"import-object", "field", "nullable", "enum", "opaque"}) +RUST_DIRECTIVE_KINDS = frozenset( + {"import-object", "field", "nullable", "enum", "opaque", "upcast", "custom-new"} +) #: Default FFI-origin -> Rust-type map; ``::`` paths get a ``use``, bare names do not. RUST_TY_MAP_DEFAULTS = { @@ -121,3 +123,7 @@ "unsized virtual yield".split() ) RUST_NOT_RAW_IDENTIFIERS = frozenset({"self", "Self", "super", "crate"}) + +#: ``rustfmt``'s default ``max_width``: an allocator signature that would +#: overflow it is rendered one parameter per line. +RUST_MAX_WIDTH = 100 diff --git a/python/tvm_ffi/stub/rust_generator/directives.py b/python/tvm_ffi/stub/rust_generator/directives.py index 2c5ae403e..dded0d8b4 100644 --- a/python/tvm_ffi/stub/rust_generator/directives.py +++ b/python/tvm_ffi/stub/rust_generator/directives.py @@ -16,16 +16,21 @@ # under the License. """The Rust backend's one-line directives: payload grammar and per-file storage. -Three address one reflected field as ``.``, one addresses a type:: +Three address one reflected field as ``.``, three address a type:: // tvm-ffi-stubgen(field): tirx.Add.a -> PrimExpr // tvm-ffi-stubgen(nullable): ir.Expr.span // tvm-ffi-stubgen(enum): tirx.For.kind -> ForKind(i32) { Serial=0, Parallel=1 } // tvm-ffi-stubgen(opaque): ir.SourceName + // tvm-ffi-stubgen(upcast): tirx.Add -> PrimExpr + // tvm-ffi-stubgen(custom-new): tirx.Add ``field`` sets the field's Rust type (a name in scope, or a ``::`` path to -``use``); ``nullable`` wraps it in ``Option``; ``enum`` declares an open integer -newtype for it; ``opaque`` keeps a type opaque even when its layout is reproducible. +``use``); on a field inherited from an ancestor it narrows the allocator +parameter instead. ``nullable`` wraps it in ``Option``; ``enum`` declares an open +integer newtype for it; ``opaque`` keeps a type opaque even when its layout is +reproducible. ``upcast`` adds a typed view outside the ancestor chain; +``custom-new`` says the wrapper's ``new`` is hand-written, so none is generated. """ from __future__ import annotations @@ -53,18 +58,20 @@ class EnumSpec: @dataclasses.dataclass class Directives: - """The Rust directives of one file, keyed by ``.``.""" + """The Rust directives of one file, keyed by ``.`` or ````.""" field_types: dict[str, str] = dataclasses.field(default_factory=dict) nullable: set[str] = dataclasses.field(default_factory=set) enums: dict[str, EnumSpec] = dataclasses.field(default_factory=dict) opaque: set[str] = dataclasses.field(default_factory=set) + upcasts: dict[str, list[str]] = dataclasses.field(default_factory=dict) + custom_new: set[str] = dataclasses.field(default_factory=set) def add(self, name: str, payload: str, lineno: int) -> None: """Parse and store one directive; raise :class:`DirectiveError` when malformed.""" if name == "field": - target, rust_type = _split_arrow(name, payload, lineno) - self.field_types[target] = rust_type + lhs, rust_type = _split_arrow(name, payload, lineno, ". -> ") + self.field_types[_field_target(name, lhs, lineno)] = rust_type elif name == "nullable": self.nullable.add(_field_target(name, payload, lineno)) elif name == "enum": @@ -72,6 +79,11 @@ def add(self, name: str, payload: str, lineno: int) -> None: self.enums[target] = spec elif name == "opaque": self.opaque.add(_type_target(name, payload, lineno)) + elif name == "upcast": + lhs, rust_type = _split_arrow(name, payload, lineno, " -> ") + self.upcasts.setdefault(_type_target(name, lhs, lineno), []).append(rust_type) + elif name == "custom-new": + self.custom_new.add(_type_target(name, payload, lineno)) else: raise DirectiveError(f"Unknown directive `{name}` at line {lineno}") @@ -96,12 +108,12 @@ def _field_target(name: str, text: str, lineno: int) -> str: return target -def _split_arrow(name: str, payload: str, lineno: int) -> tuple[str, str]: - """Split ``. -> ``.""" +def _split_arrow(name: str, payload: str, lineno: int, expected: str) -> tuple[str, str]: + """Split `` -> ``; the caller validates the target.""" lhs, arrow, rhs = payload.partition("->") if not arrow or not rhs.strip(): - raise _invalid(name, lineno, ". -> ") - return _field_target(name, lhs, lineno), rhs.strip() + raise _invalid(name, lineno, expected) + return lhs, rhs.strip() def _parse_enum(payload: str, lineno: int) -> tuple[str, EnumSpec]: diff --git a/tests/python/test_stubgen_rust.py b/tests/python/test_stubgen_rust.py index 1cda32f35..cf7c5b58a 100644 --- a/tests/python/test_stubgen_rust.py +++ b/tests/python/test_stubgen_rust.py @@ -216,6 +216,9 @@ def test_directives_parse() -> None: directives.add("enum", "tirx.For.kind -> ForKind(i32) { Serial=0, Parallel = 1 }", 3) directives.add("enum", "tirx.For.mode -> Mode(u8)", 4) directives.add("opaque", " ir.SourceName ", 5) + directives.add("upcast", "tirx.Add -> PrimExpr", 6) + directives.add("upcast", "tirx.Add -> crate::typed::TypedExpr", 7) + directives.add("custom-new", " tirx.Add ", 8) assert directives.field_types == {"tirx.Add.a": "PrimExpr"} assert directives.nullable == {"ir.Expr.span"} assert directives.enums == { @@ -223,6 +226,8 @@ def test_directives_parse() -> None: "tirx.For.mode": EnumSpec("Mode", "u8", ()), } assert directives.opaque == {"ir.SourceName"} + assert directives.upcasts == {"tirx.Add": ["PrimExpr", "crate::typed::TypedExpr"]} + assert directives.custom_new == {"tirx.Add"} @pytest.mark.parametrize( @@ -236,7 +241,10 @@ def test_directives_parse() -> None: ("enum", "tirx.For.kind -> ForKind(i128)", "Name(i32)"), ("enum", "tirx.For.kind -> ForKind(i32) { Serial }", "Name(i32)"), ("opaque", "ir.SourceName ir.Source", ""), - ("upcast", "tirx.Add -> PrimExpr", "Unknown directive"), + ("upcast", "tirx.Add", "-> "), + ("upcast", "tirx.Add PrimExpr -> PrimExpr", ""), + ("custom-new", "", ""), + ("typed-view", "tirx.Add -> PrimExpr", "Unknown directive"), ], ) def test_directives_reject_malformed(name: str, payload: str, expected: str) -> None: @@ -246,7 +254,15 @@ def test_directives_reject_malformed(name: str, payload: str, expected: str) -> def test_generator_declares_its_directives_and_records_imports() -> None: - assert RUST.directive_kinds == {"import-object", "field", "nullable", "enum", "opaque"} + assert RUST.directive_kinds == { + "import-object", + "field", + "nullable", + "enum", + "opaque", + "upcast", + "custom-new", + } imports = RUST.new_imports() RUST.add_directive(imports, "import-object", "tvm_ffi.libinfo.Foo;False;_Foo", 1) RUST.add_directive(imports, "nullable", "demo.Node.span", 2) @@ -649,6 +665,21 @@ def _add() -> ObjectInfo: fn deref(&self) -> &ExprObj { &self.data } +} + +impl ExprObj { + pub(crate) fn new(span: Option, ty: Type) -> Self { + let base = Object::new(); + Self { base, span, ty } + } +} + +impl Expr { + /// Lossless complete-field allocation. + pub fn new(span: Option, ty: Type) -> Self { + let obj = ExprObj::new(span, ty); + Self { data: ObjectArc::new(obj) } + } }""" ADD_EXPECTED = """\ @@ -688,7 +719,22 @@ def _add() -> ObjectInfo: } } -tvm_ffi::impl_object_upcast!(Add => Expr);""" +impl AddObj { + pub(crate) fn new(span: Option, ty: PrimType, a: PrimExpr, b: PrimExpr) -> Self { + let base = ExprObj::new(span, ty.into()); + Self { base, a, b } + } +} + +impl Add { + /// Lossless complete-field allocation. + pub fn new(span: Option, ty: PrimType, a: PrimExpr, b: PrimExpr) -> Self { + let obj = AddObj::new(span, ty, a, b); + Self { data: ObjectArc::new(obj) } + } +} + +tvm_ffi::impl_object_upcast!(Add => Expr, Add => PrimExpr);""" def test_render_complete_expr_golden() -> None: @@ -701,11 +747,19 @@ def test_render_complete_expr_golden() -> None: def test_render_complete_add_golden() -> None: - """`tirx.Add` on top of a complete `ir.Expr`, with `field` directives narrowing `a` / `b`.""" + """`tirx.Add` as tvm-rust-ext hand-writes it, on top of a complete `ir.Expr`. + + `field` directives narrow `a` / `b` and the inherited allocator parameter + `ty` (upcast with `.into()` on the way to `ExprObj::new`); `upcast` adds + the `PrimExpr` view. + """ _register(_expr()) imports = RUST.new_imports() - RUST.add_directive(imports, "field", "tirx.Add.a -> PrimExpr", 1) - RUST.add_directive(imports, "field", "tirx.Add.b -> PrimExpr", 2) + RUST.add_directive(imports, "nullable", "ir.Expr.span", 1) + RUST.add_directive(imports, "field", "tirx.Add.a -> PrimExpr", 2) + RUST.add_directive(imports, "field", "tirx.Add.b -> PrimExpr", 3) + RUST.add_directive(imports, "field", "tirx.Add.ty -> PrimType", 4) + RUST.add_directive(imports, "upcast", "tirx.Add -> PrimExpr", 5) text, imports = _render(_add(), imports) assert text == ADD_EXPECTED assert _uses(imports) == { @@ -713,6 +767,8 @@ def test_render_complete_add_golden() -> None: "tvm_ffi::ObjectArc", "super::ir::ExprObj", "super::ir::Expr", + "super::ir::Span", + "super::ir::Type", } @@ -786,6 +842,34 @@ def test_unrenderable_field_keeps_the_type_opaque(field: NamedTypeSchema) -> Non assert "const _: () =" not in text +def test_custom_new_leaves_the_wrapper_allocator_to_hand_written_code() -> None: + """`custom-new` drops `Add::new`; `AddObj::new` stays for that code and derived types to call.""" + _register(_expr()) + imports = RUST.new_imports() + RUST.add_directive(imports, "custom-new", "tirx.Add", 1) + text, _ = _render(_add(), imports) + assert ( + "impl AddObj {\n pub(crate) fn new(span: Span, ty: Type, a: Expr, b: Expr) -> Self {" + in text + ) + assert "impl Add {" not in text + assert "pub fn new(" not in text + + +def test_upcast_directive_adds_typed_views() -> None: + """`upcast` targets follow the ancestor chain; a `::` path is imported. Opaque: no allocator.""" + info = _info("demo.Leaf", parent="demo.Base") + imports = RUST.new_imports() + RUST.add_directive(imports, "upcast", "demo.Leaf -> crate::typed::LeafView", 1) + RUST.add_directive(imports, "upcast", "demo.Leaf -> Other", 2) + text, imports = _render(info, imports) + assert text.endswith( + "tvm_ffi::impl_object_upcast!(Leaf => Base, Leaf => LeafView, Leaf => Other);" + ) + assert "crate::typed::LeafView" in _uses(imports) + assert "fn new(" not in text + + def test_opaque_directive_vetoes_a_complete_type() -> None: imports = RUST.new_imports() RUST.add_directive(imports, "opaque", "ir.Span", 1) @@ -885,11 +969,28 @@ def test_registry_complete_chain_is_mirrored() -> None: ) assert "assert!(::core::mem::size_of::() == 40);" in base assert "FieldGetter" not in base + assert " pub fn new(v_i64: i64, v_i32: i32) -> Self {" in base dd, imports = _render(object_info_from_type_key("testing.TestCxxClassDerivedDerived")) assert ( " base: TestCxxClassDerivedObj,\n pub v_str: String,\n pub v_bool: bool,\n}" in dd ) + # The allocator flattens the chain; a signature over 100 columns wraps. + assert ( + "impl TestCxxClassDerivedDerivedObj {\n" + " pub(crate) fn new(\n" + " v_i64: i64,\n" + " v_i32: i32,\n" + " v_f64: f64,\n" + " v_f32: f32,\n" + " v_str: String,\n" + " v_bool: bool,\n" + " ) -> Self {\n" + " let base = TestCxxClassDerivedObj::new(v_i64, v_i32, v_f64, v_f32);\n" + " Self { base, v_str, v_bool }\n" + " }\n" + "}" + ) in dd assert dd.endswith( "tvm_ffi::impl_object_upcast!(TestCxxClassDerivedDerived => TestCxxClassBase, " "TestCxxClassDerivedDerived => TestCxxClassDerived);" @@ -989,5 +1090,5 @@ def test_cli_exits_non_zero_on_a_directive_error( ) monkeypatch.setattr("sys.argv", ["tvm-ffi-stubgen", "--target", "rust", str(src)]) assert stub_cli.__main__() == 1 - src.write_text(f"{C.RUST_SYNTAX.directive('upcast')} testing.TestCxxClassBase -> X\n") + src.write_text(f"{C.RUST_SYNTAX.directive('typed-view')} testing.TestCxxClassBase -> X\n") assert stub_cli.__main__() == 1 From b51a85ac53f54f5fbb9567b5e24f32ec54a5c865 Mon Sep 17 00:00:00 2001 From: yuchuan Date: Thu, 3 Sep 2026 22:03:06 -0400 Subject: [PATCH 3/8] [STUBGEN] Keep types under builtin parents opaque in the Rust target A type under a builtin such as `ffi.Enum` embeds a header-only stand-in for it, but the classifier judged the builtin from its registry bytes and called a fieldless `testing.TestEnumVariant` complete: the generated struct then asserted `size_of == 48` on 24 bytes and grew allocators. `layout.classify` gains an `unmirrored` set and a `no-mirror` reason for types whose bytes the target never reproduces; the Rust backend passes every builtin ancestor below `ffi.Object`, so such types and everything under them stay opaque with an explanatory verdict. Signed-off-by: yuchuan --- python/tvm_ffi/stub/layout.py | 18 ++++++++++++++- python/tvm_ffi/stub/rust_generator/codegen.py | 16 +++++++++++-- tests/python/test_stub_layout.py | 10 ++++++++ tests/python/test_stubgen_rust.py | 23 +++++++++++++++++++ 4 files changed, 64 insertions(+), 3 deletions(-) diff --git a/python/tvm_ffi/stub/layout.py b/python/tvm_ffi/stub/layout.py index 798676a45..c89d16042 100644 --- a/python/tvm_ffi/stub/layout.py +++ b/python/tvm_ffi/stub/layout.py @@ -70,6 +70,7 @@ "uncovered-bytes", "unrenderable-field", "by-directive", + "no-mirror", ] """Why a type is opaque. @@ -82,6 +83,9 @@ a vptr, ...). - ``unrenderable-field``: the caller's predicate rejected a field. - ``by-directive``: the caller vetoed a type whose layout is reproducible. +- ``no-mirror``: the caller's target never reproduces this type's bytes (a + builtin its runtime owns), whatever the layout says; everything below it + is ``parent-opaque``. """ @@ -175,6 +179,7 @@ def classify( infos: Mapping[str, ObjectInfo], *, forced_opaque: AbstractSet[str] = frozenset(), + unmirrored: AbstractSet[str] = frozenset(), field_renderable: Callable[[NamedTypeSchema], bool] | None = None, ) -> dict[str, Verdict]: """Classify every type in ``infos``, parents before children. @@ -188,6 +193,10 @@ def classify( Type keys vetoed by the caller (semantic blockers such as interned identities). The veto only demotes a type that would otherwise be complete; a type that is opaque for a layout reason keeps that reason. + unmirrored + Type keys whose bytes the target never reproduces (builtins its runtime + owns). Such a type is opaque with reason ``no-mirror`` before any layout + evidence is weighed, and so is every type below it (``parent-opaque``). field_renderable Predicate deciding whether a reflected field has a native mirror in the target language. A rejected field demotes its type to opaque with @@ -203,7 +212,9 @@ def _classify(type_key: str) -> Verdict: raise KeyError(f"Ancestor {type_key!r} is not among the types to classify") info = infos[type_key] parent = None if info.parent_type_key is None else _classify(info.parent_type_key) - verdicts[type_key] = _classify_one(info, parent, forced_opaque, field_renderable) + verdicts[type_key] = _classify_one( + info, parent, forced_opaque, unmirrored, field_renderable + ) return verdicts[type_key] for type_key in infos: @@ -215,6 +226,7 @@ def _classify_one( info: ObjectInfo, parent: Verdict | None, forced_opaque: AbstractSet[str], + unmirrored: AbstractSet[str], field_renderable: Callable[[NamedTypeSchema], bool] | None, ) -> Verdict: assert info.type_key is not None, "cannot classify an ObjectInfo without a type key" @@ -228,6 +240,10 @@ def _classify_one( total_size=info.total_size, is_final=info.is_final, ) + if info.type_key in unmirrored: + verdict.reason = "no-mirror" + verdict.detail = "its bytes are owned by the target's runtime and never mirrored" + return verdict outcome = _prove_layout(info, parent, verdict) if outcome is None: outcome = _apply_target_rules(info, forced_opaque, field_renderable) diff --git a/python/tvm_ffi/stub/rust_generator/codegen.py b/python/tvm_ffi/stub/rust_generator/codegen.py index f671601e9..ef945d898 100644 --- a/python/tvm_ffi/stub/rust_generator/codegen.py +++ b/python/tvm_ffi/stub/rust_generator/codegen.py @@ -45,7 +45,9 @@ registered global functions, hand-written outside the markers. A builtin parent (``ffi.IntEnum``, say) has no ``Obj`` in the crate: the import section defines a header-only stand-in per builtin ancestor, so ``derive(Object)`` -computes the registry's ``TYPE_DEPTH``. +computes the registry's ``TYPE_DEPTH``. A stand-in carries no bytes beyond the +header, so a type under such a parent is always opaque (``no-mirror``), as is +everything below it. """ from __future__ import annotations @@ -172,8 +174,18 @@ def classify(self) -> Verdict: def renderable(field: NamedTypeSchema) -> bool: return self._field_mirror(owner_of[id(field)], field, scratch) is not None + # The crate never reproduces a builtin's bytes: a type under `ffi.IntEnum` embeds a + # header-only stand-in (see `_base_type`), so nothing below such a parent is complete. + unmirrored = { + key + for key in self.info.ancestors + if key != C_RUST.RUST_ROOT_TYPE_KEY and not self._generated(key) + } verdicts = classify( - infos, forced_opaque=self.imports.directives.opaque, field_renderable=renderable + infos, + forced_opaque=self.imports.directives.opaque, + unmirrored=unmirrored, + field_renderable=renderable, ) return verdicts[self.type_key] diff --git a/tests/python/test_stub_layout.py b/tests/python/test_stub_layout.py index e0a6edcb6..7e4379f5f 100644 --- a/tests/python/test_stub_layout.py +++ b/tests/python/test_stub_layout.py @@ -239,6 +239,16 @@ def test_children_of_any_opaque_parent_are_parent_opaque() -> None: assert "by-directive" in verdicts["t.Child"].detail +def test_unmirrored_type_and_its_descendants_are_opaque() -> None: + child = _info("t.Child", parent="t.Base", total_size=48, fields=(_field("x", 40, 8),)) + infos = {"ffi.Object": ROOT, "t.Base": BASE, "t.Child": child} + verdicts = classify(infos, unmirrored={"t.Base"}) + assert verdicts["t.Base"].reason == "no-mirror" + assert verdicts["t.Base"].own_bytes is None # unconditional: no layout evidence is weighed + assert verdicts["t.Child"].reason == "parent-opaque" + assert "no-mirror" in verdicts["t.Child"].detail + + def test_order_does_not_matter_but_ancestors_must_be_present() -> None: child = _info("t.Child", parent="t.Base", total_size=48, fields=(_field("x", 40, 8),)) verdicts = classify({"t.Child": child, "t.Base": BASE, "ffi.Object": ROOT}) diff --git a/tests/python/test_stubgen_rust.py b/tests/python/test_stubgen_rust.py index cf7c5b58a..aa79b00e8 100644 --- a/tests/python/test_stubgen_rust.py +++ b/tests/python/test_stubgen_rust.py @@ -878,6 +878,29 @@ def test_opaque_directive_vetoes_a_complete_type() -> None: assert "pub fn line(&self) -> Result {" in text +def test_builtin_parent_keeps_the_type_opaque() -> None: + """The crate never mirrors a builtin's bytes: the header-only stand-in cannot be complete.""" + # Fieldless under `ffi.Enum` (48 bytes): the fill criterion alone would call this complete. + info = _info( + "demo.Flag", parent="ffi.Enum", ancestors=["ffi.Object", "ffi.Enum"], total_size=48 + ) + text, _ = _render(info) + assert "/// Opaque: parent 'ffi.Enum' is opaque (no-mirror)." in text + assert " base: FfiEnumObj," in text + assert "const _: () =" not in text + assert "fn new(" not in text + # The registry fixtures under `ffi.Enum` / `ffi.IntEnum` / `ffi.StrEnum` follow the same rule. + for type_key, parent in ( + ("testing.TestEnumVariant", "ffi.Enum"), + ("testing.TestCxxIntEnum", "ffi.IntEnum"), + ("testing.TestCxxStrEnum", "ffi.StrEnum"), + ): + text, _ = _render(object_info_from_type_key(type_key)) + assert f"/// Opaque: parent '{parent}' is opaque (no-mirror)." in text + assert "const _: () =" not in text + assert "fn new(" not in text + + def test_opaque_parent_keeps_the_child_opaque() -> None: """A child of an opaque type cannot embed a mirror of its parent: it stays opaque.""" _register(_info("ir.Expr", (_field("span", "ir.Span", 24, 8),), total_size=40)) # hole From f3c682dbfd0924613d79cacb5dcc53e29107cdb7 Mon Sep 17 00:00:00 2001 From: yuchuan Date: Fri, 4 Sep 2026 11:49:54 -0400 Subject: [PATCH 4/8] simplify. Signed-off-by: yuchuan --- examples/rust_stubgen/README.md | 12 ++- python/tvm_ffi/stub/cli.py | 3 +- python/tvm_ffi/stub/layout.py | 8 +- python/tvm_ffi/stub/rust_generator/codegen.py | 74 +++++++------------ python/tvm_ffi/stub/rust_generator/consts.py | 23 ++---- python/tvm_ffi/stub/utils.py | 6 +- 6 files changed, 41 insertions(+), 85 deletions(-) diff --git a/examples/rust_stubgen/README.md b/examples/rust_stubgen/README.md index 2261b7d75..76c35146d 100644 --- a/examples/rust_stubgen/README.md +++ b/examples/rust_stubgen/README.md @@ -28,8 +28,8 @@ reflected fields account for every byte of the object: - `rust_stubgen.IntRange` does, so its binding is *complete*: the struct mirrors the fields at their real offsets and widths, and Rust reads `range.begin` directly. A `const` assertion pins the struct's size and alignment to the - reflected facts. The object is allocated in Rust, by a generated function - that takes every field by value; a C++ function then reads it back. + reflected facts. A generated `new` allocates it in Rust from every field; a + C++ function then reads it back. - `rust_stubgen.IntPair` has a vtable in front of the object header, so its binding is *opaque*: the struct embeds only the parent and every field is read through an accessor that calls the C ABI getter. It is constructed by a @@ -64,16 +64,14 @@ open newtype: // tvm-ffi-stubgen(enum): rust_stubgen.IntPair.kind -> PairKind(i32) { Unordered=0, Ordered=1 } ``` -It also marks `IntRange` as having a hand-written `new` (in `main.rs`, where it -validates the extent), so the generator emits none of its own: +It also marks `IntRange`'s `new` as hand-written (in `main.rs`, where it +validates the extent), so the generator does not emit one: ```rust // tvm-ffi-stubgen(custom-new): rust_stubgen.IntRange ``` -Without the marker the generator emits `IntRange::new` itself, and a -hand-written one is a duplicate definition. Three more directives are -available: `field` names the Rust type of a field +Three more directives 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`), and `upcast` adds a conversion to a hand-written typed view diff --git a/python/tvm_ffi/stub/cli.py b/python/tvm_ffi/stub/cli.py index e185ae8d4..6bea12001 100644 --- a/python/tvm_ffi/stub/cli.py +++ b/python/tvm_ffi/stub/cli.py @@ -115,8 +115,7 @@ def __main__() -> int: def _run_stage(file: FileInfo, stage: Callable[[], None]) -> bool: """Run one stage over ``file``, reporting a failure without stopping the run. - Returns whether the failure was a :class:`DirectiveError`, which makes the - whole run exit non-zero. + Returns whether it was a :class:`DirectiveError` (the run then exits non-zero). """ try: stage() diff --git a/python/tvm_ffi/stub/layout.py b/python/tvm_ffi/stub/layout.py index c89d16042..73c68bed1 100644 --- a/python/tvm_ffi/stub/layout.py +++ b/python/tvm_ffi/stub/layout.py @@ -83,9 +83,8 @@ a vptr, ...). - ``unrenderable-field``: the caller's predicate rejected a field. - ``by-directive``: the caller vetoed a type whose layout is reproducible. -- ``no-mirror``: the caller's target never reproduces this type's bytes (a - builtin its runtime owns), whatever the layout says; everything below it - is ``parent-opaque``. +- ``no-mirror``: the caller's target never reproduces this type's bytes + (a builtin its runtime owns); everything below it is ``parent-opaque``. """ @@ -195,8 +194,7 @@ def classify( complete; a type that is opaque for a layout reason keeps that reason. unmirrored Type keys whose bytes the target never reproduces (builtins its runtime - owns). Such a type is opaque with reason ``no-mirror`` before any layout - evidence is weighed, and so is every type below it (``parent-opaque``). + owns): opaque with reason ``no-mirror`` before any layout evidence is weighed. field_renderable Predicate deciding whether a reflected field has a native mirror in the target language. A rejected field demotes its type to opaque with diff --git a/python/tvm_ffi/stub/rust_generator/codegen.py b/python/tvm_ffi/stub/rust_generator/codegen.py index ef945d898..5a81405a5 100644 --- a/python/tvm_ffi/stub/rust_generator/codegen.py +++ b/python/tvm_ffi/stub/rust_generator/codegen.py @@ -17,37 +17,23 @@ """Rust code generation for ``tvm-ffi-stubgen``. Every reflected object gets a ``#[repr(C)]`` object struct, a reference wrapper, -read-only ``Deref``, and the upcasts along its ancestor chain. What the object -struct holds depends on the verdict of :mod:`tvm_ffi.stub.layout`: - -- *complete*: the layout is reproducible, so the struct mirrors every physical - field at its real offset and width, public, borrowed directly. A ``const`` - assertion pins the struct's ``size_of`` / ``align_of`` to the reflected facts, - so a mirror rustc lays out differently fails to compile. The type is also - allocatable from Rust: ``Obj::new`` (crate-private) takes every - physical field root to leaf so derived allocators can chain to it, and the - wrapper's public ``new`` takes the same parameters and does one - ``ObjectArc::new``. A ``custom-new`` directive leaves the wrapper's ``new`` - to hand-written code; ``Obj::new`` is still generated for it to call. -- *opaque*: the struct embeds only its parent, and one accessor per reflected - field reads through the C ABI getter. The bytes are never reproduced, so the - binding is correct for every registered type. Nothing allocates it. - -The two target-language rules the classifier leaves to its caller live here: a -field without a Rust mirror (``Optional``, a ``Union``, ``void*``, ...) +``Deref``, and the upcasts along its ancestor chain. The struct's contents follow +the verdict of :mod:`tvm_ffi.stub.layout`: + +- *complete*: the struct mirrors every physical field at its real offset and + width, pinned by a ``const`` size/alignment assertion. ``Obj::new`` + (crate-private) and the wrapper's ``new`` take every field root to leaf; + ``custom-new`` leaves the wrapper's ``new`` to hand-written code. +- *opaque*: the struct embeds only its parent, and each field is read through + the C ABI getter. Nothing allocates it. + +A field without a Rust mirror (``Optional``, ``Union``, ``void*``, ...) makes the type opaque and is read as ``Any``; an ``opaque`` directive vetoes a -reproducible layout. ``field`` / ``nullable`` / ``enum`` directives shape the -field types of both forms; where a directive names a scalar width, it is checked -against the reflected field size at generation time. ``upcast`` appends typed -views to the ancestor chain. - -Nothing here calls methods or packed constructors: behaviour goes through the -registered global functions, hand-written outside the markers. A builtin parent -(``ffi.IntEnum``, say) has no ``Obj`` in the crate: the import section -defines a header-only stand-in per builtin ancestor, so ``derive(Object)`` -computes the registry's ``TYPE_DEPTH``. A stand-in carries no bytes beyond the -header, so a type under such a parent is always opaque (``no-mirror``), as is -everything below it. +reproducible layout; a scalar width named by a directive is checked against the +reflected field size. A builtin parent (``ffi.IntEnum``, say) has no +``Obj`` in the crate: the import section defines a header-only stand-in +per builtin ancestor so ``derive(Object)`` computes the registry's +``TYPE_DEPTH``, and everything under such a parent stays opaque (``no-mirror``). """ from __future__ import annotations @@ -174,8 +160,7 @@ def classify(self) -> Verdict: def renderable(field: NamedTypeSchema) -> bool: return self._field_mirror(owner_of[id(field)], field, scratch) is not None - # The crate never reproduces a builtin's bytes: a type under `ffi.IntEnum` embeds a - # header-only stand-in (see `_base_type`), so nothing below such a parent is complete. + # Builtin ancestors are header-only stand-ins (`_base_type`): none below is complete. unmirrored = { key for key in self.info.ancestors @@ -192,10 +177,9 @@ def renderable(field: NamedTypeSchema) -> bool: # --- field types --------------------------------------------------------- def _field_mirror(self, owner: str, field: NamedTypeSchema, imports: RustImports) -> str | None: - """Render the type of ``field`` in a ``#[repr(C)]`` mirror; ``None`` when it has none. + """Render the ``#[repr(C)]`` mirror type of ``field``, or ``None``. - Scalars take the width the registry recorded; ``Optional`` fields take - the in-place mirror of their C++ layout; directives override the rest. + Directives win; then scalars by reflected width, ``Optional`` by C++ layout, else schema. """ directives = self.imports.directives target = f"{owner}.{field.name}" @@ -228,11 +212,9 @@ def _field_mirror(self, owner: str, field: NamedTypeSchema, imports: RustImports def _optional_mirror(self, field: NamedTypeSchema, imports: RustImports) -> str | None: """Mirror an ``Optional`` field in place. - An ``ObjectRef``-derived payload is a pointer-sized nullable pointer in - C++, mirrored by Rust's niche-optimized ``Option``. Every other - payload stays a 16-byte ``TVMFFIAny`` cell, mirrored by - ``tvm_ffi::Optional``. ``Optional`` has no mirror, and neither - does a field whose size disagrees with its payload kind. + An object payload is a nullable pointer (``Option``); any other payload + is a 16-byte ``TVMFFIAny`` cell (``tvm_ffi::Optional``). ``Optional`` + and a size mismatch have no mirror. """ (payload,) = field.args # TypeSchema's post_init enforces exactly one argument. if payload.origin == "Any": @@ -386,7 +368,7 @@ def _struct_lines(self, verdict: Verdict, base: str) -> list[str]: # --- allocators -------------------------------------------------------- def _allocator_params(self, key: str, info: ObjectInfo) -> list[tuple[str, str]]: - """``(field, type)`` of every physical field root to leaf, as ``Obj::new`` takes them.""" + """``(field, type)`` of every physical field root to leaf, as ``Obj::new`` takes.""" parent = info.parent_type_key inherited: list[tuple[str, str]] = [] if parent is not None and self._generated(parent): @@ -398,9 +380,8 @@ def _level_params( ) -> list[tuple[str, str]]: """Extend the parent's parameters with ``key``'s own fields by offset. - A ``field`` directive on ``.`` for an inherited field narrows - that parameter (``tirx.Add.ty -> PrimType``); the allocator body upcasts - it with ``.into()`` when handing it to the parent. + A ``field`` directive on an inherited field narrows that parameter; the + body hands it to the parent with ``.into()``. """ params = [(name, self._narrowed(key, name, rust_type)) for name, rust_type in inherited] for field in sorted(info.fields, key=lambda f: f.offset or 0): @@ -429,10 +410,7 @@ def _fn_lines( ] def _allocator_sections(self, base: str, has_parent: bool) -> list[list[str]]: - """``Obj::new`` and, unless ``custom-new`` reserves it, the wrapper's ``new``. - - Both take every physical field root to leaf. - """ + """``Obj::new`` and, unless ``custom-new`` reserves it, the wrapper's ``new``.""" inherited: list[tuple[str, str]] = [] if has_parent: parent = self.info.parent_type_key diff --git a/python/tvm_ffi/stub/rust_generator/consts.py b/python/tvm_ffi/stub/rust_generator/consts.py index ee465acd1..e377c3be1 100644 --- a/python/tvm_ffi/stub/rust_generator/consts.py +++ b/python/tvm_ffi/stub/rust_generator/consts.py @@ -56,10 +56,7 @@ #: Origins without a crate mirror; such a field is read as ``tvm_ffi::Any``. RUST_UNSUPPORTED_ORIGINS = frozenset({"Dict", "List", "Union", "tuple"}) -#: Width-correct scalar for a ``#[repr(C)]`` struct field, keyed by -#: ``(ffi origin, sizeof(T))``: the type schema erases scalar widths, so the -#: width comes from the reflected field size. Signedness is not recorded; -#: unsigned C++ fields render as the same-width signed type. +#: Mirror scalar by ``(origin, reflected size)``: the schema erases widths and signedness. RUST_SCALAR_BY_SIZE = { ("int", 1): "i8", ("int", 2): "i16", @@ -69,8 +66,7 @@ ("float", 8): "f64", } -#: Byte width of the scalar Rust types a ``field`` / ``enum`` directive may name, -#: checked against the reflected field size at generation time. +#: Byte width of the scalars a ``field`` / ``enum`` directive may name; checked against the field. RUST_SCALAR_WIDTHS = { "i8": 1, "u8": 1, @@ -88,20 +84,12 @@ #: Size of an object reference field; ``nullable`` may only wrap those. RUST_POINTER_SIZE = 8 -#: In-place mirror of a non-object ``Optional`` field: C++ ``ffi::Optional`` -#: is a single 16-byte ``TVMFFIAny`` cell (``nullopt == kTVMFFINone``) for -#: payloads that are not ``ObjectRef``-derived. Object payloads use the -#: pointer-sized form (``nullopt == nullptr``), mirrored by Rust's ``Option``. +#: C++ ``Optional`` is a 16-byte ``TVMFFIAny`` cell, or a nullable pointer for object payloads. RUST_OPTIONAL_PATH = "tvm_ffi::Optional" RUST_OPTIONAL_FIELD_SIZE = 16 RUST_OBJECT_OPTIONAL_FIELD_SIZE = 8 -#: ``Optional`` payload origins whose C++ optional stays 16-byte Any-backed: -#: everything that is not a pointer-sized object reference (strings and bytes -#: are 16-byte values themselves). A nested ``Optional`` payload also stays -#: Any-backed; it is special-cased where this set is consulted. The reflected -#: field size is checked either way, so a payload this set misclassifies makes -#: the field unrenderable rather than mis-mirrored. +#: ``Optional`` payloads kept as the 16-byte cell (a nested ``Optional`` too); the size is checked. RUST_ANY_BACKED_OPTIONAL_PAYLOADS = frozenset( {"int", "float", "bool", "Device", "dtype", "DataType", "str", "bytes"} ) @@ -124,6 +112,5 @@ ) RUST_NOT_RAW_IDENTIFIERS = frozenset({"self", "Self", "super", "crate"}) -#: ``rustfmt``'s default ``max_width``: an allocator signature that would -#: overflow it is rendered one parameter per line. +#: ``rustfmt``'s default ``max_width``; a wider allocator signature wraps one parameter per line. RUST_MAX_WIDTH = 100 diff --git a/python/tvm_ffi/stub/utils.py b/python/tvm_ffi/stub/utils.py index 5d46e3df9..c74f29309 100644 --- a/python/tvm_ffi/stub/utils.py +++ b/python/tvm_ffi/stub/utils.py @@ -36,11 +36,7 @@ class DirectiveError(ValueError): - """A one-line directive is malformed, unknown, or disagrees with reflection. - - The pipeline reports it and exits non-zero: generated code must never be - shaped by a directive it could not honour. - """ + """A directive is malformed, unknown, or contradicts reflection; the run exits non-zero.""" def _parse_type_schema(raw: str | dict[str, Any]) -> TypeSchema: From cdc95b7c9f4f6df874374e214253cd46dd9a6062 Mon Sep 17 00:00:00 2001 From: yuchuan Date: Fri, 4 Sep 2026 11:56:33 -0400 Subject: [PATCH 5/8] change the example. Signed-off-by: yuchuan --- examples/rust_stubgen/README.md | 24 ++++++------- .../rust/src/generated/rust_stubgen/mod.rs | 34 +++++++++++-------- examples/rust_stubgen/rust/src/main.rs | 20 +++++------ examples/rust_stubgen/src/int_pair.cc | 17 +++------- python/tvm_ffi/stub/rust_generator/codegen.py | 3 +- tests/python/test_stubgen_rust.py | 17 ++++++++++ 6 files changed, 65 insertions(+), 50 deletions(-) diff --git a/examples/rust_stubgen/README.md b/examples/rust_stubgen/README.md index 76c35146d..8ab17308a 100644 --- a/examples/rust_stubgen/README.md +++ b/examples/rust_stubgen/README.md @@ -22,18 +22,18 @@ into Rust bindings. This example registers two objects in `src/int_pair.cc` and lets CMake regenerate `rust/src/generated/` after every build. Every object gets a `#[repr(C)]` wrapper, a reference type, `Deref`, and the -upcasts along its ancestor chain. What the wrapper holds depends on whether the -reflected fields account for every byte of the object: - -- `rust_stubgen.IntRange` does, so its binding is *complete*: the struct mirrors - the fields at their real offsets and widths, and Rust reads `range.begin` - directly. A `const` assertion pins the struct's size and alignment to the - reflected facts. A generated `new` allocates it in Rust from every field; a - C++ function then reads it back. -- `rust_stubgen.IntPair` has a vtable in front of the object header, so its - binding is *opaque*: the struct embeds only the parent and every field is read - through an accessor that calls the C ABI getter. It is constructed by a - registered global function. +upcasts along its ancestor chain. Both objects here are plain data, so their +reflected fields account for every byte and the bindings are *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 an `IntPair` and an `IntRange` +that way, reads `pair.a` directly, and hands them to C++ functions that read +them 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 +embeds only the parent, every field is read through an accessor that calls the +C ABI getter, and construction stays on the C++ side. A builtin parent such as `ffi.IntEnum` has no `Obj` in the crate; the import section defines a header-only stand-in per builtin ancestor so the 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 e45477b7e..423802afc 100644 --- a/examples/rust_stubgen/rust/src/generated/rust_stubgen/mod.rs +++ b/examples/rust_stubgen/rust/src/generated/rust_stubgen/mod.rs @@ -23,16 +23,14 @@ // tvm-ffi-stubgen(begin): import-section use std::ops::Deref; use tvm_ffi::Error; -use tvm_ffi::FieldGetter; use tvm_ffi::Object; use tvm_ffi::ObjectArc; -use tvm_ffi::ObjectCore; use tvm_ffi::Result; use tvm_ffi::VALUE_ERROR; // tvm-ffi-stubgen(end) -// The `kind` field is an integer on the C++ side; this directive gives it an -// open integer newtype in Rust and makes the accessor return it. +// 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 } // tvm-ffi-stubgen(begin): object/rust_stubgen.IntPair @@ -61,15 +59,23 @@ impl TryFrom for PairKind { } } -/// Opaque: bytes [44, 56) of [24, 56) are not accounted for by reflected fields. Fields are read through the C ABI getters. +/// Complete: reflected fields fill [24, 48) exactly (alignment padding [44, 48)). #[repr(C)] #[derive(tvm_ffi::derive::Object)] #[type_key = "rust_stubgen.IntPair"] #[type_final] pub struct IntPairObj { base: Object, + pub a: i64, + pub b: i64, + pub kind: PairKind, } +const _: () = { + assert!(::core::mem::size_of::() == 48); + assert!(::core::mem::align_of::() == 8); +}; + #[repr(C)] #[derive(tvm_ffi::derive::ObjectRef, Clone)] pub struct IntPair { @@ -84,17 +90,17 @@ impl Deref for IntPair { } impl IntPairObj { - pub fn a(&self) -> Result { - FieldGetter::new(Self::type_index(), "a")?.get(self) - } - - pub fn b(&self) -> Result { - FieldGetter::new(Self::type_index(), "b")?.get(self) + pub(crate) fn new(a: i64, b: i64, kind: PairKind) -> Self { + let base = Object::new(); + Self { base, a, b, kind } } +} - pub fn kind(&self) -> Result { - let raw: i64 = FieldGetter::new(Self::type_index(), "kind")?.get(self)?; - PairKind::try_from(raw) +impl IntPair { + /// Lossless complete-field allocation. + pub fn new(a: i64, b: i64, kind: PairKind) -> Self { + let obj = IntPairObj::new(a, b, kind); + Self { data: ObjectArc::new(obj) } } } // tvm-ffi-stubgen(end) diff --git a/examples/rust_stubgen/rust/src/main.rs b/examples/rust_stubgen/rust/src/main.rs index 748dcffb0..8ef80c0e5 100644 --- a/examples/rust_stubgen/rust/src/main.rs +++ b/examples/rust_stubgen/rust/src/main.rs @@ -49,25 +49,23 @@ fn lib_path() -> String { } fn main() -> Result<()> { - // Load the C++ library so `IntPair` is registered with the FFI type registry. - // Keep it alive for as long as the bindings are used. + // Load the C++ library so the objects are registered with the FFI type + // registry. Keep it alive for as long as the bindings are used. let _lib = Module::load_from_file(lib_path())?; - // The object is opaque to Rust: it is constructed by the registered C++ - // function and its fields are read through the reflection getters. - let pair: IntPair = tvm_ffi::cached_global_func!("rust_stubgen.IntPair") - .call_tuple((1i64, 2i64, i64::from(PairKind::Ordered.as_raw())))? - .try_into()?; - println!("a={} b={} kind={:?}", pair.a()?, pair.b()?, pair.kind()?); - assert_eq!(pair.kind()?, PairKind::Ordered); + // Both objects have a reproducible layout: they are allocated in Rust and + // their fields are plain struct members, on both sides of the ABI. + let pair = IntPair::new(1, 2, PairKind::Ordered); + println!("a={} b={} kind={:?}", pair.a, pair.b, pair.kind); + assert_eq!(pair.kind, PairKind::Ordered); let sum: i64 = tvm_ffi::cached_global_func!("rust_stubgen.IntPairSum") .call_tuple((pair.clone(),))? .try_into()?; println!("sum={sum}"); + assert_eq!(sum, 3); - // `IntRange` has a reproducible layout: it is allocated in Rust and its - // fields are plain struct members, on both sides of the ABI. + // `IntRange::new` is hand-written above (`custom-new`), so it can validate. let range = IntRange::new(10, 5)?; println!("begin={} extent={}", range.begin, range.extent); let end: i64 = tvm_ffi::cached_global_func!("rust_stubgen.IntRangeEnd") diff --git a/examples/rust_stubgen/src/int_pair.cc b/examples/rust_stubgen/src/int_pair.cc index 703712677..93d1347ff 100644 --- a/examples/rust_stubgen/src/int_pair.cc +++ b/examples/rust_stubgen/src/int_pair.cc @@ -18,7 +18,7 @@ */ /*! * \file int_pair.cc - * \brief A tvm-ffi library that registers one object for the Rust stub generator. + * \brief A tvm-ffi library that registers two objects for the Rust stub generator. */ #include @@ -29,9 +29,9 @@ namespace rust_stubgen { namespace ffi = tvm::ffi; // [object.begin] -// A polymorphic object: the vtable in front of the object header means Rust -// cannot mirror its bytes, so the generated binding reads every field through -// the reflection getters and construction stays on the C++ side. +// Plain data objects: every byte is accounted for by a reflected field, so the +// generated bindings mirror the layout; Rust allocates them and reads the fields +// directly, and the registered functions below read them back. class IntPairObj : public ffi::Object { public: int64_t a; @@ -39,21 +39,16 @@ class IntPairObj : public ffi::Object { int32_t kind; IntPairObj(int64_t a, int64_t b, int32_t kind) : a(a), b(b), kind(kind) {} - virtual ~IntPairObj() = default; - virtual int64_t Sum() const { return a + b; } + int64_t Sum() const { return a + b; } TVM_FFI_DECLARE_OBJECT_INFO_FINAL("rust_stubgen.IntPair", IntPairObj, ffi::Object); }; class IntPair : public ffi::ObjectRef { public: - IntPair(int64_t a, int64_t b, int32_t kind) { data_ = ffi::make_object(a, b, kind); } - TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(IntPair, ffi::ObjectRef, IntPairObj); }; -// A plain data object: every byte is accounted for by a reflected field, so the -// generated binding mirrors the layout and Rust reads the fields directly. class IntRangeObj : public ffi::Object { public: int64_t begin; @@ -79,8 +74,6 @@ TVM_FFI_STATIC_INIT_BLOCK() { .def_ro("begin", &IntRangeObj::begin, "the first value") .def_ro("extent", &IntRangeObj::extent, "the number of values"); refl::GlobalDef() - .def("rust_stubgen.IntPair", - [](int64_t a, int64_t b, int32_t kind) { return IntPair(a, b, kind); }) .def("rust_stubgen.IntPairSum", [](const IntPair& pair) { return pair->Sum(); }) .def("rust_stubgen.IntRangeEnd", [](const IntRange& range) { return range->begin + range->extent; }); diff --git a/python/tvm_ffi/stub/rust_generator/codegen.py b/python/tvm_ffi/stub/rust_generator/codegen.py index 5a81405a5..906a0102f 100644 --- a/python/tvm_ffi/stub/rust_generator/codegen.py +++ b/python/tvm_ffi/stub/rust_generator/codegen.py @@ -281,6 +281,7 @@ def _enum_lines(self, spec: EnumSpec) -> list[str]: """Render the open integer newtype an ``enum`` directive declares.""" error = self.imports.record("tvm_ffi::Error") value_error = self.imports.record("tvm_ffi::VALUE_ERROR") + result = self.imports.record("tvm_ffi::Result") return [ "#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]", "#[repr(transparent)]", @@ -299,7 +300,7 @@ def _enum_lines(self, spec: EnumSpec) -> list[str]: "", f"impl TryFrom for {spec.name} {{", f" type Error = {error};", - " fn try_from(value: i64) -> Result {", + f" fn try_from(value: i64) -> {result} {{", f" {spec.repr}::try_from(value).map(Self).map_err(|_| {{", f' {error}::new({value_error}, &format!("{spec.name} value {{value}} does not fit ' f'{spec.repr}"), "")', diff --git a/tests/python/test_stubgen_rust.py b/tests/python/test_stubgen_rust.py index aa79b00e8..b97a8cdba 100644 --- a/tests/python/test_stubgen_rust.py +++ b/tests/python/test_stubgen_rust.py @@ -804,6 +804,23 @@ def test_render_complete_span_and_prim_type() -> None: assert "tvm_ffi::DLDataType" in _uses(imports) +def test_render_complete_enum_field() -> None: + """An `enum` directive types the mirrored field; the newtype brings its own `Result` import.""" + info = _info( + "demo.Pair", + (_field("a", "int", 24, 8), _field("kind", "int", 32, 4)), + total_size=40, + is_final=True, + ) + imports = RUST.new_imports() + RUST.add_directive(imports, "enum", "demo.Pair.kind -> Kind(i32) { A=0, B=1 }", 1) + text, imports = _render(info, imports) + assert " pub kind: Kind,\n" in text + assert " pub fn new(a: i64, kind: Kind) -> Self {" in text + assert {"tvm_ffi::Result", "tvm_ffi::Error", "tvm_ffi::VALUE_ERROR"} <= _uses(imports) + assert "tvm_ffi::FieldGetter" not in _uses(imports) + + def test_complete_optional_field_mirrors() -> None: """`Optional` fields mirror their C++ layout: a 16-byte cell or a nullable pointer.""" info = _info( From 94569593017437d866a0f060de0d17f7af5cfd59 Mon Sep 17 00:00:00 2001 From: yuchuan Date: Fri, 4 Sep 2026 12:05:08 -0400 Subject: [PATCH 6/8] simplify. Signed-off-by: yuchuan --- examples/rust_stubgen/README.md | 33 +++++++-------- .../rust/src/generated/rust_stubgen/mod.rs | 42 ------------------- examples/rust_stubgen/rust/src/main.rs | 37 ++++------------ examples/rust_stubgen/src/int_pair.cc | 31 +++----------- 4 files changed, 26 insertions(+), 117 deletions(-) diff --git a/examples/rust_stubgen/README.md b/examples/rust_stubgen/README.md index 8ab17308a..cad8f240a 100644 --- a/examples/rust_stubgen/README.md +++ b/examples/rust_stubgen/README.md @@ -18,17 +18,17 @@ # Rust Stub Generation `tvm-ffi-stubgen --target rust` turns the reflection metadata of a C++ library -into Rust bindings. This example registers two objects in `src/int_pair.cc` -and lets CMake regenerate `rust/src/generated/` after every build. +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. Every object gets a `#[repr(C)]` wrapper, a reference type, `Deref`, and the -upcasts along its ancestor chain. Both objects here are plain data, so their -reflected fields account for every byte and the bindings are *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 an `IntPair` and an `IntRange` -that way, reads `pair.a` directly, and hands them to C++ functions that read -them back. +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. 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 @@ -64,15 +64,10 @@ open newtype: // tvm-ffi-stubgen(enum): rust_stubgen.IntPair.kind -> PairKind(i32) { Unordered=0, Ordered=1 } ``` -It also marks `IntRange`'s `new` as hand-written (in `main.rs`, where it -validates the extent), so the generator does not emit one: - -```rust -// tvm-ffi-stubgen(custom-new): rust_stubgen.IntRange -``` - -Three more directives are available: `field` names the Rust type of a field +Four more directives 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`), -and `upcast` adds a conversion to a hand-written typed view -(`// tvm-ffi-stubgen(upcast): rust_stubgen.IntRange -> MyView`). +`upcast` adds a conversion to a hand-written typed view +(`// tvm-ffi-stubgen(upcast): rust_stubgen.IntPair -> MyView`), and +`custom-new` keeps the generator from emitting the wrapper's `new` when it is +hand-written (`// tvm-ffi-stubgen(custom-new): rust_stubgen.IntPair`). 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 423802afc..d1085be2e 100644 --- a/examples/rust_stubgen/rust/src/generated/rust_stubgen/mod.rs +++ b/examples/rust_stubgen/rust/src/generated/rust_stubgen/mod.rs @@ -104,45 +104,3 @@ impl IntPair { } } // tvm-ffi-stubgen(end) - -// `IntRange::new` is hand-written in main.rs (it validates the extent), so the -// generator does not emit one; it still emits `IntRangeObj::new` for it to call. -// tvm-ffi-stubgen(custom-new): rust_stubgen.IntRange - -// tvm-ffi-stubgen(begin): object/rust_stubgen.IntRange -/// Complete: reflected fields fill [24, 40) exactly (alignment padding [36, 40)). -#[repr(C)] -#[derive(tvm_ffi::derive::Object)] -#[type_key = "rust_stubgen.IntRange"] -#[type_final] -pub struct IntRangeObj { - base: Object, - pub begin: i64, - pub extent: i32, -} - -const _: () = { - assert!(::core::mem::size_of::() == 40); - assert!(::core::mem::align_of::() == 8); -}; - -#[repr(C)] -#[derive(tvm_ffi::derive::ObjectRef, Clone)] -pub struct IntRange { - data: ObjectArc, -} - -impl Deref for IntRange { - type Target = IntRangeObj; - fn deref(&self) -> &IntRangeObj { - &self.data - } -} - -impl IntRangeObj { - pub(crate) fn new(begin: i64, extent: i32) -> Self { - let base = Object::new(); - Self { base, begin, extent } - } -} -// tvm-ffi-stubgen(end) diff --git a/examples/rust_stubgen/rust/src/main.rs b/examples/rust_stubgen/rust/src/main.rs index 8ef80c0e5..f49fadbef 100644 --- a/examples/rust_stubgen/rust/src/main.rs +++ b/examples/rust_stubgen/rust/src/main.rs @@ -16,25 +16,12 @@ * specific language governing permissions and limitations * under the License. */ -//! Use the stubgen-generated `IntPair` and `IntRange` bindings (see ../../README.md). +//! Use the stubgen-generated `IntPair` binding (see ../../README.md). mod generated; -use generated::rust_stubgen::{IntPair, IntRange, IntRangeObj, PairKind}; -use tvm_ffi::{Error, Module, ObjectArc, ObjectRefCore, Result, VALUE_ERROR}; - -/// The hand-written constructor of `IntRange`: the `custom-new` directive in -/// the generated file keeps the generator from emitting one. -impl IntRange { - pub fn new(begin: i64, extent: i32) -> Result { - if extent < 0 { - return Err(Error::new(VALUE_ERROR, "IntRange extent must not be negative", "")); - } - let data = ObjectArc::new(IntRangeObj::new(begin, extent)); - // SAFETY: `data` holds a freshly allocated `IntRangeObj`, the container type of `IntRange`. - Ok(unsafe { Self::from_data(data) }) - } -} +use generated::rust_stubgen::{IntPair, PairKind}; +use tvm_ffi::{Module, Result}; /// Path of the C++ shared library built by CMake into `../build`. fn lib_path() -> String { @@ -49,12 +36,12 @@ fn lib_path() -> String { } fn main() -> Result<()> { - // Load the C++ library so the objects are registered with the FFI type - // registry. Keep it alive for as long as the bindings are used. + // Load the C++ library so `IntPair` is registered with the FFI type registry. + // Keep it alive for as long as the binding is used. let _lib = Module::load_from_file(lib_path())?; - // Both objects have a reproducible layout: they are allocated in Rust and - // their fields are plain struct members, on both sides of the ABI. + // 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); println!("a={} b={} kind={:?}", pair.a, pair.b, pair.kind); assert_eq!(pair.kind, PairKind::Ordered); @@ -64,15 +51,5 @@ fn main() -> Result<()> { .try_into()?; println!("sum={sum}"); assert_eq!(sum, 3); - - // `IntRange::new` is hand-written above (`custom-new`), so it can validate. - let range = IntRange::new(10, 5)?; - println!("begin={} extent={}", range.begin, range.extent); - let end: i64 = tvm_ffi::cached_global_func!("rust_stubgen.IntRangeEnd") - .call_tuple((range.clone(),))? - .try_into()?; - println!("end={end}"); - assert_eq!(end, 15); - assert!(IntRange::new(0, -1).is_err()); Ok(()) } diff --git a/examples/rust_stubgen/src/int_pair.cc b/examples/rust_stubgen/src/int_pair.cc index 93d1347ff..b758333dd 100644 --- a/examples/rust_stubgen/src/int_pair.cc +++ b/examples/rust_stubgen/src/int_pair.cc @@ -18,7 +18,7 @@ */ /*! * \file int_pair.cc - * \brief A tvm-ffi library that registers two objects for the Rust stub generator. + * \brief A tvm-ffi library that registers one object for the Rust stub generator. */ #include @@ -29,9 +29,9 @@ namespace rust_stubgen { namespace ffi = tvm::ffi; // [object.begin] -// Plain data objects: every byte is accounted for by a reflected field, so the -// generated bindings mirror the layout; Rust allocates them and reads the fields -// directly, and the registered functions below read them back. +// A plain data object: every byte is accounted for by a reflected field, so the +// generated binding mirrors the layout; Rust allocates it and reads the fields +// directly, and the registered function below reads it back. class IntPairObj : public ffi::Object { public: int64_t a; @@ -49,34 +49,13 @@ class IntPair : public ffi::ObjectRef { TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(IntPair, ffi::ObjectRef, IntPairObj); }; -class IntRangeObj : public ffi::Object { - public: - int64_t begin; - int32_t extent; - - IntRangeObj(int64_t begin, int32_t extent) : begin(begin), extent(extent) {} - - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("rust_stubgen.IntRange", IntRangeObj, ffi::Object); -}; - -class IntRange : public ffi::ObjectRef { - public: - TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(IntRange, ffi::ObjectRef, IntRangeObj); -}; - TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; refl::ObjectDef(refl::init(false)) .def_ro("a", &IntPairObj::a, "the first operand") .def_ro("b", &IntPairObj::b, "the second operand") .def_ro("kind", &IntPairObj::kind, "0 = unordered, 1 = ordered"); - refl::ObjectDef(refl::init(false)) - .def_ro("begin", &IntRangeObj::begin, "the first value") - .def_ro("extent", &IntRangeObj::extent, "the number of values"); - refl::GlobalDef() - .def("rust_stubgen.IntPairSum", [](const IntPair& pair) { return pair->Sum(); }) - .def("rust_stubgen.IntRangeEnd", - [](const IntRange& range) { return range->begin + range->extent; }); + refl::GlobalDef().def("rust_stubgen.IntPairSum", [](const IntPair& pair) { return pair->Sum(); }); } // [object.end] From 84859938a0b5cf8f5adbf2b020d24ed1a4dbbee1 Mon Sep 17 00:00:00 2001 From: yuchuan Date: Fri, 4 Sep 2026 12:27:58 -0400 Subject: [PATCH 7/8] simplify. Signed-off-by: yuchuan --- python/tvm_ffi/stub/cli.py | 48 +++++++++---------- python/tvm_ffi/stub/rust_generator/codegen.py | 31 +++++------- .../tvm_ffi/stub/rust_generator/directives.py | 10 ++-- python/tvm_ffi/stub/utils.py | 4 -- tests/python/test_stubgen_rust.py | 27 ++--------- 5 files changed, 43 insertions(+), 77 deletions(-) diff --git a/python/tvm_ffi/stub/cli.py b/python/tvm_ffi/stub/cli.py index 6bea12001..a5e0ba393 100644 --- a/python/tvm_ffi/stub/cli.py +++ b/python/tvm_ffi/stub/cli.py @@ -24,7 +24,7 @@ import sys import traceback from pathlib import Path -from typing import TYPE_CHECKING, Callable +from typing import TYPE_CHECKING from . import consts as C from .file_utils import FileInfo, collect_files, syntax_for @@ -36,7 +36,7 @@ object_info_from_type_key, toposort_objects, ) -from .utils import DirectiveError, FuncInfo, InitConfig, Options +from .utils import FuncInfo, InitConfig, Options if TYPE_CHECKING: from .generator import Generator @@ -67,9 +67,13 @@ def __main__() -> int: # - defined global functions: `tvm-ffi-stubgen(begin): global/...` # - defined object types: `tvm-ffi-stubgen(begin): object/...` ty_map: dict[str, str] = generator.default_ty_map() - directive_errors = 0 for file in files: - directive_errors += _run_stage(file, lambda: _stage_1(file, ty_map)) + try: + _stage_1(file, ty_map) + except Exception: + print( + f'{C.TERM_RED}[Failed] File "{file.path}": {traceback.format_exc()}{C.TERM_RESET}' + ) # Stage 2. Generate stubs if they are not defined on the file. generated_prefixes: set[str] = set() @@ -90,9 +94,18 @@ def __main__() -> int: for file in files: if opt.verbose: print(f"{C.TERM_CYAN}[File] {file.path}{C.TERM_RESET}") - directive_errors += _run_stage( - file, lambda: _stage_3(file, opt, ty_map, global_funcs, generator=generator) - ) + try: + _stage_3( + file, + opt, + ty_map, + global_funcs, + generator=generator, + ) + except Exception: + print( + f'{C.TERM_RED}[Failed] File "{file.path}": {traceback.format_exc()}{C.TERM_RESET}' + ) # Stage 4. Let the generator stitch the generated tree together (runs after the # files are fully written, so language-specific wiring isn't clobbered). @@ -109,22 +122,7 @@ def __main__() -> int: } write_coverage_report(Path(opt.coverage_out), classify(infos)) del dlls - return 1 if directive_errors else 0 - - -def _run_stage(file: FileInfo, stage: Callable[[], None]) -> bool: - """Run one stage over ``file``, reporting a failure without stopping the run. - - Returns whether it was a :class:`DirectiveError` (the run then exits non-zero). - """ - try: - stage() - except DirectiveError as e: - print(f'{C.TERM_RED}[Failed] File "{file.path}": {e}{C.TERM_RESET}') - return True - except Exception: - print(f'{C.TERM_RED}[Failed] File "{file.path}": {traceback.format_exc()}{C.TERM_RESET}') - return False + return 0 def _stage_1( @@ -137,7 +135,7 @@ def _stage_1( try: lhs, rhs = code.param[1].split("->") except ValueError as e: - raise DirectiveError( + raise ValueError( f"Invalid ty_map format at line {code.lineno_start}. Example: `A.B -> C.D`" ) from e ty_map[lhs.strip()] = rhs.strip() @@ -241,7 +239,7 @@ def _stage_3( # noqa: PLR0912 if name in C.PIPELINE_DIRECTIVE_KINDS: continue # consumed by `_stage_1` if name not in generator.directive_kinds: - raise DirectiveError(f"Unknown directive `{name}` at line {code.lineno_start}") + raise ValueError(f"Unknown directive `{name}` at line {code.lineno_start}") generator.add_directive(imports, name, payload, code.lineno_start) # Stage 2. Process `tvm-ffi-stubgen(begin): global/...` for code in file.code_blocks: diff --git a/python/tvm_ffi/stub/rust_generator/codegen.py b/python/tvm_ffi/stub/rust_generator/codegen.py index 906a0102f..8819dbc9a 100644 --- a/python/tvm_ffi/stub/rust_generator/codegen.py +++ b/python/tvm_ffi/stub/rust_generator/codegen.py @@ -44,7 +44,6 @@ from .. import consts as C from ..layout import Verdict, classify from ..lib_state import object_info_from_type_key -from ..utils import DirectiveError from . import consts as C_RUST from .utils import RustImports, builtin_mirror_name, render_rust_type, rust_ident @@ -68,7 +67,7 @@ def _call_lines(open_: str, items: list[str], close: str) -> list[str]: def _check_width(target: str, field: NamedTypeSchema, rust_type: str, width: int) -> None: """Reject a directive whose scalar type does not match the reflected field size.""" if field.size is not None and field.size != width: - raise DirectiveError( + raise ValueError( f"Directive on `{target}` maps a {field.size}-byte field to `{rust_type}` " f"({width} bytes)" ) @@ -174,7 +173,7 @@ def renderable(field: NamedTypeSchema) -> bool: ) return verdicts[self.type_key] - # --- field types --------------------------------------------------------- + # --- field types ------------------------------------------------------- def _field_mirror(self, owner: str, field: NamedTypeSchema, imports: RustImports) -> str | None: """Render the ``#[repr(C)]`` mirror type of ``field``, or ``None``. @@ -202,7 +201,7 @@ def _field_mirror(self, owner: str, field: NamedTypeSchema, imports: RustImports return None if target in directives.nullable and not mirror.startswith("Option<"): if field.size not in (None, C_RUST.RUST_POINTER_SIZE): - raise DirectiveError( + raise ValueError( f"`nullable` directive on `{target}`: the field is {field.size} bytes, " "not a pointer-sized object reference" ) @@ -237,6 +236,8 @@ def _optional_mirror(self, field: NamedTypeSchema, imports: RustImports) -> str return f"{imports.record(C_RUST.RUST_OPTIONAL_PATH)}<{inner}>" return f"Option<{inner}>" + # --- pieces ------------------------------------------------------------ + def _accessor_lines(self, field: NamedTypeSchema) -> list[str]: """One ``pub fn (&self) -> Result`` through the C ABI getter. @@ -269,13 +270,7 @@ def _accessor_lines(self, field: NamedTypeSchema) -> list[str]: ] if target in directives.nullable and not rust_type.startswith("Option<"): rust_type = f"Option<{rust_type}>" - return [ - f"pub fn {name}(&self) -> Result<{rust_type}> {{", - f" {getter}.get(self)", - "}", - ] - - # --- pieces ------------------------------------------------------------ + return [f"pub fn {name}(&self) -> Result<{rust_type}> {{", f" {getter}.get(self)", "}"] def _enum_lines(self, spec: EnumSpec) -> list[str]: """Render the open integer newtype an ``enum`` directive declares.""" @@ -460,8 +455,8 @@ def body(self) -> list[str]: self.imports.record("tvm_ffi::ObjectArc") base, has_parent = self._base_type() fields = self.info.fields - accessors = bool(fields) and not verdict.is_complete - if accessors: + has_accessors = bool(fields) and not verdict.is_complete + if has_accessors: self.imports.record("tvm_ffi::ObjectCore") # `Self::type_index()` self.imports.record("tvm_ffi::FieldGetter") self.imports.record("tvm_ffi::Result") @@ -486,16 +481,16 @@ def body(self) -> list[str]: sections.append(self._deref_lines(self.leaf, self.obj_struct, "data")) if has_parent: sections.append(self._deref_lines(self.obj_struct, base, "base")) - if accessors: - lines_: list[str] = [] + if has_accessors: + accessors: list[str] = [] for i, field in enumerate(fields): if i: - lines_.append("") - lines_ += self._accessor_lines(field) + accessors.append("") + accessors += self._accessor_lines(field) sections.append( [ f"impl {self.obj_struct} {{", - *[f" {line}" if line else "" for line in lines_], + *[f" {line}" if line else "" for line in accessors], "}", ] ) diff --git a/python/tvm_ffi/stub/rust_generator/directives.py b/python/tvm_ffi/stub/rust_generator/directives.py index dded0d8b4..90a01c42d 100644 --- a/python/tvm_ffi/stub/rust_generator/directives.py +++ b/python/tvm_ffi/stub/rust_generator/directives.py @@ -38,8 +38,6 @@ import dataclasses import re -from ..utils import DirectiveError - _ENUM_RE = re.compile( r"^(?P\S+)\s*->\s*(?P[A-Za-z_]\w*)\((?P[iu](?:8|16|32|64))\)" r"\s*(?:\{(?P[^{}]*)\})?$" @@ -68,7 +66,7 @@ class Directives: custom_new: set[str] = dataclasses.field(default_factory=set) def add(self, name: str, payload: str, lineno: int) -> None: - """Parse and store one directive; raise :class:`DirectiveError` when malformed.""" + """Parse and store one directive; raise ``ValueError`` on a malformed payload.""" if name == "field": lhs, rust_type = _split_arrow(name, payload, lineno, ". -> ") self.field_types[_field_target(name, lhs, lineno)] = rust_type @@ -85,11 +83,11 @@ def add(self, name: str, payload: str, lineno: int) -> None: elif name == "custom-new": self.custom_new.add(_type_target(name, payload, lineno)) else: - raise DirectiveError(f"Unknown directive `{name}` at line {lineno}") + raise ValueError(f"Unknown directive `{name}` at line {lineno}") -def _invalid(name: str, lineno: int, expected: str) -> DirectiveError: - return DirectiveError(f"Invalid `{name}` directive at line {lineno}. Expected `{expected}`") +def _invalid(name: str, lineno: int, expected: str) -> ValueError: + return ValueError(f"Invalid `{name}` directive at line {lineno}. Expected `{expected}`") def _type_target(name: str, text: str, lineno: int) -> str: diff --git a/python/tvm_ffi/stub/utils.py b/python/tvm_ffi/stub/utils.py index c74f29309..7f5ea6258 100644 --- a/python/tvm_ffi/stub/utils.py +++ b/python/tvm_ffi/stub/utils.py @@ -35,10 +35,6 @@ from tvm_ffi.core import TypeField -class DirectiveError(ValueError): - """A directive is malformed, unknown, or contradicts reflection; the run exits non-zero.""" - - def _parse_type_schema(raw: str | dict[str, Any]) -> TypeSchema: """Parse a type schema from either a JSON string or an already-parsed dict.""" if isinstance(raw, dict): diff --git a/tests/python/test_stubgen_rust.py b/tests/python/test_stubgen_rust.py index b97a8cdba..d5db77eff 100644 --- a/tests/python/test_stubgen_rust.py +++ b/tests/python/test_stubgen_rust.py @@ -41,7 +41,7 @@ ) from tvm_ffi.stub.rust_generator.directives import Directives, EnumSpec from tvm_ffi.stub.rust_generator.utils import RustImports, RustUse, render_rust_type, rust_ident -from tvm_ffi.stub.utils import DirectiveError, InitConfig, NamedTypeSchema, ObjectInfo, Options +from tvm_ffi.stub.utils import InitConfig, NamedTypeSchema, ObjectInfo, Options RUST = get_generator("rust") HEADER = 24 # sizeof(TVMFFIObject) @@ -248,7 +248,7 @@ def test_directives_parse() -> None: ], ) def test_directives_reject_malformed(name: str, payload: str, expected: str) -> None: - with pytest.raises(DirectiveError, match=re.escape(expected)) as exc: + with pytest.raises(ValueError, match=re.escape(expected)) as exc: Directives().add(name, payload, 7) assert "at line 7" in str(exc.value) @@ -941,7 +941,7 @@ def test_directive_disagreeing_with_bytes_is_an_error( info = _info("demo.Pair", (_field("count", "int", 24, 4),), total_size=32) imports = RUST.new_imports() RUST.add_directive(imports, name, payload, 1) - with pytest.raises(DirectiveError, match=re.escape(message)): + with pytest.raises(ValueError, match=re.escape(message)): _render(info, imports) @@ -1111,24 +1111,3 @@ def test_cli_init_generates_a_module_tree(tmp_path: Path, monkeypatch: pytest.Mo # Running again over the generated tree is a no-op. assert stub_cli.__main__() == 0 assert (tmp_path / "testing" / "mod.rs").read_text(encoding="utf-8") == text - - -def test_cli_exits_non_zero_on_a_directive_error( - tmp_path: Path, monkeypatch: pytest.MonkeyPatch -) -> None: - src = tmp_path / "mod.rs" - src.write_text( - "\n".join( - [ - f"{C.RUST_SYNTAX.directive('field')} testing.TestCxxClassBase.v_i32 -> i64", - f"{C.RUST_SYNTAX.begin} object/testing.TestCxxClassBase", - C.RUST_SYNTAX.end, - "", - ] - ), - encoding="utf-8", - ) - monkeypatch.setattr("sys.argv", ["tvm-ffi-stubgen", "--target", "rust", str(src)]) - assert stub_cli.__main__() == 1 - src.write_text(f"{C.RUST_SYNTAX.directive('typed-view')} testing.TestCxxClassBase -> X\n") - assert stub_cli.__main__() == 1 From ef00323f134de967e7bbbf5693967950cfd6c9ce Mon Sep 17 00:00:00 2001 From: yuchuan Date: Fri, 4 Sep 2026 14:10:46 -0400 Subject: [PATCH 8/8] add from_complete_fields and base/data protection. Signed-off-by: yuchuan --- examples/rust_stubgen/README.md | 4 +- python/tvm_ffi/stub/rust_generator/codegen.py | 39 +++++++++++-------- python/tvm_ffi/stub/rust_generator/consts.py | 5 +++ .../tvm_ffi/stub/rust_generator/directives.py | 3 +- python/tvm_ffi/stub/rust_generator/utils.py | 8 +++- tests/python/test_stubgen_rust.py | 36 +++++++++++++++-- 6 files changed, 69 insertions(+), 26 deletions(-) diff --git a/examples/rust_stubgen/README.md b/examples/rust_stubgen/README.md index cad8f240a..42b2758c1 100644 --- a/examples/rust_stubgen/README.md +++ b/examples/rust_stubgen/README.md @@ -69,5 +69,5 @@ Four more directives are available: `field` names the Rust type of a field 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` keeps the generator from emitting the wrapper's `new` when it is -hand-written (`// tvm-ffi-stubgen(custom-new): rust_stubgen.IntPair`). +`custom-new` names the generated allocator `from_complete_fields` when `new` +is hand-written (`// tvm-ffi-stubgen(custom-new): rust_stubgen.IntPair`). diff --git a/python/tvm_ffi/stub/rust_generator/codegen.py b/python/tvm_ffi/stub/rust_generator/codegen.py index 8819dbc9a..78d778fb0 100644 --- a/python/tvm_ffi/stub/rust_generator/codegen.py +++ b/python/tvm_ffi/stub/rust_generator/codegen.py @@ -23,7 +23,8 @@ - *complete*: the struct mirrors every physical field at its real offset and width, pinned by a ``const`` size/alignment assertion. ``Obj::new`` (crate-private) and the wrapper's ``new`` take every field root to leaf; - ``custom-new`` leaves the wrapper's ``new`` to hand-written code. + ``custom-new`` leaves the wrapper's ``new`` to hand-written code and names + the generated one ``from_complete_fields``. - *opaque*: the struct embeds only its parent, and each field is read through the C ABI getter. Nothing allocates it. @@ -336,7 +337,7 @@ def _struct_lines(self, verdict: Verdict, base: str) -> list[str]: f'#[type_key = "{self.type_key}"]', *(["#[type_final]"] if self.info.is_final else []), f"pub struct {self.obj_struct} {{", - f" base: {base},", + f" base: {base},", # a reflected `base` field becomes `base_` (`rust_ident`) ] if not verdict.is_complete: return [ @@ -406,7 +407,10 @@ def _fn_lines( ] def _allocator_sections(self, base: str, has_parent: bool) -> list[list[str]]: - """``Obj::new`` and, unless ``custom-new`` reserves it, the wrapper's ``new``.""" + """``Obj::new`` and the wrapper's ``new``. + + ``custom-new`` names the wrapper's allocator ``from_complete_fields`` instead. + """ inherited: list[tuple[str, str]] = [] if has_parent: parent = self.info.parent_type_key @@ -431,20 +435,20 @@ def _allocator_sections(self, base: str, has_parent: bool) -> list[list[str]]: "}", ] ] - if self.type_key not in self.imports.directives.custom_new: - sections.append( - [ - f"impl {self.leaf} {{", - " /// Lossless complete-field allocation.", - *self._fn_lines( - "pub fn new", - params, - (f"obj = {self.obj_struct}::new", forward), - "Self { data: ObjectArc::new(obj) }", - ), - "}", - ] - ) + custom = self.type_key in self.imports.directives.custom_new + sections.append( + [ + f"impl {self.leaf} {{", + " /// Lossless complete-field allocation.", + *self._fn_lines( + "pub fn from_complete_fields" if custom else "pub fn new", + params, + (f"obj = {self.obj_struct}::new", forward), + "Self { data: ObjectArc::new(obj) }", + ), + "}", + ] + ) return sections def body(self) -> list[str]: @@ -474,6 +478,7 @@ def body(self) -> list[str]: "#[repr(C)]", "#[derive(tvm_ffi::derive::ObjectRef, Clone)]", f"pub struct {self.leaf} {{", + # a reflected `data` field becomes `data_` (`rust_ident`) f" data: ObjectArc<{self.obj_struct}>,", "}", ] diff --git a/python/tvm_ffi/stub/rust_generator/consts.py b/python/tvm_ffi/stub/rust_generator/consts.py index e377c3be1..0636cbfb2 100644 --- a/python/tvm_ffi/stub/rust_generator/consts.py +++ b/python/tvm_ffi/stub/rust_generator/consts.py @@ -112,5 +112,10 @@ ) RUST_NOT_RAW_IDENTIFIERS = frozenset({"self", "Self", "super", "crate"}) +#: Member names the generated structs use themselves: ``base`` is the parent slot of every object +#: struct, ``data`` the wrapper's ``ObjectArc``. A reflected field with one of these names would +#: collide (or shadow through ``Deref``), so ``rust_ident`` spells it ``base_`` / ``data_``. +RUST_RESERVED_MEMBERS = frozenset({"base", "data"}) + #: ``rustfmt``'s default ``max_width``; a wider allocator signature wraps one parameter per line. RUST_MAX_WIDTH = 100 diff --git a/python/tvm_ffi/stub/rust_generator/directives.py b/python/tvm_ffi/stub/rust_generator/directives.py index 90a01c42d..185b90652 100644 --- a/python/tvm_ffi/stub/rust_generator/directives.py +++ b/python/tvm_ffi/stub/rust_generator/directives.py @@ -30,7 +30,8 @@ parameter instead. ``nullable`` wraps it in ``Option``; ``enum`` declares an open integer newtype for it; ``opaque`` keeps a type opaque even when its layout is reproducible. ``upcast`` adds a typed view outside the ancestor chain; -``custom-new`` says the wrapper's ``new`` is hand-written, so none is generated. +``custom-new`` says the wrapper's ``new`` is hand-written; the generated one is +named ``from_complete_fields`` instead. """ from __future__ import annotations diff --git a/python/tvm_ffi/stub/rust_generator/utils.py b/python/tvm_ffi/stub/rust_generator/utils.py index c60f52d74..27510efd9 100644 --- a/python/tvm_ffi/stub/rust_generator/utils.py +++ b/python/tvm_ffi/stub/rust_generator/utils.py @@ -133,9 +133,13 @@ def _generic(base: str | None, *params: str | None) -> str | None: def rust_ident(name: str) -> str: - """Spell a reflected field name in Rust: drop the C++ trailing underscore, escape keywords.""" + """Spell a reflected field name in Rust: drop the C++ trailing underscore, escape collisions. + + Keywords become raw identifiers; the four that cannot, and the names the + generated structs use themselves (``base``, ``data``), get a trailing underscore. + """ name = name.rstrip("_") or name - if name in C.RUST_NOT_RAW_IDENTIFIERS: + if name in C.RUST_NOT_RAW_IDENTIFIERS or name in C.RUST_RESERVED_MEMBERS: return f"{name}_" if name in C.RUST_KEYWORDS: return f"r#{name}" diff --git a/tests/python/test_stubgen_rust.py b/tests/python/test_stubgen_rust.py index d5db77eff..446e1ade0 100644 --- a/tests/python/test_stubgen_rust.py +++ b/tests/python/test_stubgen_rust.py @@ -202,6 +202,8 @@ def test_rust_ident() -> None: assert rust_ident("type") == "r#type" assert rust_ident("self") == "self_" assert rust_ident("crate") == "crate_" + assert rust_ident("base") == "base_" # the parent slot of every generated object struct + assert rust_ident("data") == "data_" # the wrapper's `ObjectArc` member # --------------------------------------------------------------------------- @@ -804,6 +806,26 @@ def test_render_complete_span_and_prim_type() -> None: assert "tvm_ffi::DLDataType" in _uses(imports) +def test_reserved_member_names_get_a_trailing_underscore() -> None: + """A reflected `base` or `data` field must not collide with the generated members.""" + info = _info( + "demo.Ramp", + (_field("base", "int", 24, 8), _field("data", "int", 32, 8)), + total_size=40, + is_final=True, + ) + text, _ = _render(info) + assert " base: Object,\n pub base_: i64,\n pub data_: i64,\n" in text + assert " pub fn new(base_: i64, data_: i64) -> Self {" in text + assert " Self { base, base_, data_ }" in text + # The opaque form keeps the reflected name on the C ABI side. + text, _ = _render(_info("demo.Node", (_field("base", "int"),))) + assert ( + 'pub fn base_(&self) -> Result {\n FieldGetter::new(Self::type_index(), "base")' + in text + ) + + def test_render_complete_enum_field() -> None: """An `enum` directive types the mirrored field; the newtype brings its own `Result` import.""" info = _info( @@ -859,8 +881,8 @@ def test_unrenderable_field_keeps_the_type_opaque(field: NamedTypeSchema) -> Non assert "const _: () =" not in text -def test_custom_new_leaves_the_wrapper_allocator_to_hand_written_code() -> None: - """`custom-new` drops `Add::new`; `AddObj::new` stays for that code and derived types to call.""" +def test_custom_new_renames_the_wrapper_allocator() -> None: + """`custom-new`: `Add::new` stays hand-written, the allocator is `from_complete_fields`.""" _register(_expr()) imports = RUST.new_imports() RUST.add_directive(imports, "custom-new", "tirx.Add", 1) @@ -869,8 +891,14 @@ def test_custom_new_leaves_the_wrapper_allocator_to_hand_written_code() -> None: "impl AddObj {\n pub(crate) fn new(span: Span, ty: Type, a: Expr, b: Expr) -> Self {" in text ) - assert "impl Add {" not in text - assert "pub fn new(" not in text + assert ( + "impl Add {\n /// Lossless complete-field allocation.\n" + " pub fn from_complete_fields(span: Span, ty: Type, a: Expr, b: Expr) -> Self {\n" + " let obj = AddObj::new(span, ty, a, b);\n" + " Self { data: ObjectArc::new(obj) }\n" + " }\n}" + ) in text + assert " pub fn new(" not in text def test_upcast_directive_adds_typed_views() -> None: