This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm-ffi.git
The following commit(s) were added to refs/heads/main by this push:
new ec4dc714 [Stubgen] Rust directives for allow-missing-docs and closed
enums (#950)
ec4dc714 is described below
commit ec4dc7143feff7a12c06db39068899fd9349d466
Author: Tianqi Chen <[email protected]>
AuthorDate: Fri Nov 6 19:26:28 2026 -0500
[Stubgen] Rust directives for allow-missing-docs and closed enums (#950)
Two opt-in Rust stubgen directives:
- `tvm-ffi-stubgen(missing-docs): allow` puts `#[allow(missing_docs)]`
on every generated type and `impl` of the file, so a crate can enable
`#![warn(missing_docs)]` and document the generated types by hand.
- `closed` before the member list of an `enum` directive emits a
`#[repr(<int>)]` Rust enum of exactly those members, for value sets the
ABI fixes, instead of the open integer newtype.
Without either, the generated code is unchanged.
---
examples/rust_stubgen/README.md | 14 ++++++++
python/tvm_ffi/stub/rust_generator/codegen.py | 38 ++++++++++++++++++++-
python/tvm_ffi/stub/rust_generator/consts.py | 11 +++++-
python/tvm_ffi/stub/rust_generator/directives.py | 43 ++++++++++++++++++------
tests/python/test_stubgen_rust.py | 37 ++++++++++++++++++++
5 files changed, 131 insertions(+), 12 deletions(-)
diff --git a/examples/rust_stubgen/README.md b/examples/rust_stubgen/README.md
index b089b79a..ed67c858 100644
--- a/examples/rust_stubgen/README.md
+++ b/examples/rust_stubgen/README.md
@@ -118,6 +118,20 @@ 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.IntPair -> MyView`).
+When the C++ side fixes the value set in its ABI, `closed` before the member
+list makes `enum` emit a `#[repr(i32)]` Rust enum instead. Reading a value
+outside the set is undefined behavior, so the set should include 0, the value
of
+a field never set:
+
+```rust
+// tvm-ffi-stubgen(enum): rust_stubgen.IntPair.kind -> PairKind(i32) closed {
Unordered=0, Ordered=1 }
+```
+
+A crate built with `#![warn(missing_docs)]` adds
+`// tvm-ffi-stubgen(missing-docs): allow` to a file: every generated type and
+`impl` in it gets `#[allow(missing_docs)]`, and docs written above a block
+still document its first item.
+
## Partial generation
Every type an object refers to (its parent, its ancestors, the types of its
diff --git a/python/tvm_ffi/stub/rust_generator/codegen.py
b/python/tvm_ffi/stub/rust_generator/codegen.py
index e4378814..fbd58c4b 100644
--- a/python/tvm_ffi/stub/rust_generator/codegen.py
+++ b/python/tvm_ffi/stub/rust_generator/codegen.py
@@ -330,6 +330,10 @@ class _ObjectRenderer:
# --- pieces ------------------------------------------------------------
+ def _allow_missing_docs(self) -> list[str]:
+ """``#[allow(missing_docs)]`` under ``missing-docs: allow``, for a
crate that warns on it."""
+ return ["#[allow(missing_docs)]"] if
self.imports.directives.allow_missing_docs else []
+
def _accessor_lines(self, field: NamedTypeSchema) -> list[str]:
"""One ``pub fn <field>(&self) -> Result<T>`` through the C ABI getter.
@@ -375,15 +379,41 @@ class _ObjectRenderer:
]
def _enum_lines(self, spec: EnumSpec) -> list[str]:
- """Render the open integer newtype an ``enum`` directive declares."""
+ """Render the type an ``enum`` directive declares: an open integer
newtype, or a closed enum."""
error = self.imports.record("tvm_ffi::Error")
value_error = self.imports.record("tvm_ffi::VALUE_ERROR")
result = self.imports.record("tvm_ffi::Result")
+ allow = self._allow_missing_docs()
+ if spec.closed:
+ message = f'&format!("{spec.name} has no value {{value}}")'
+ rejected = _call_lines(
+ f" _ => Err({error}::new(", [value_error, message,
'""'], ")),"
+ )
+ return [
+ *allow,
+ "#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]",
+ f"#[repr({spec.repr})]",
+ f"pub enum {spec.name} {{",
+ *[f" {member} = {value}," for member, value in
spec.members],
+ "}",
+ "",
+ f"impl TryFrom<i64> for {spec.name} {{",
+ f" type Error = {error};",
+ f" fn try_from(value: i64) -> {result}<Self> {{",
+ " match value {",
+ *[f" {value} => Ok(Self::{member})," for member,
value in spec.members],
+ *rejected,
+ " }",
+ " }",
+ "}",
+ ]
return [
+ *allow,
"#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]",
"#[repr(transparent)]",
f"pub struct {spec.name}({spec.repr});",
"",
+ *allow,
"#[allow(non_upper_case_globals)]",
f"impl {spec.name} {{",
*[f" pub const {member}: Self = Self({value});" for member,
value in spec.members],
@@ -435,6 +465,7 @@ class _ObjectRenderer:
def _struct_lines(self, verdict: Verdict, base: str) -> list[str]:
"""Render the object struct: every field when complete, the parent
alone when opaque."""
header = [
+ *self._allow_missing_docs(),
"#[repr(C)]",
"#[derive(tvm_ffi::derive::Object)]",
f'#[type_key = "{self.type_key}"]',
@@ -546,8 +577,10 @@ class _ObjectRenderer:
]
own = [rust_ident(name) for name, _ in params[len(inherited) :]]
forward = [_forward(rust_ident(name), rust_type) for name, rust_type
in params]
+ allow = self._allow_missing_docs()
sections = [
[
+ *allow,
f"impl {self.obj_struct} {{",
*self._fn_lines(
"pub(crate) fn new" if self.info.is_final else "pub fn
new",
@@ -561,6 +594,7 @@ class _ObjectRenderer:
custom = self.type_key in self.imports.directives.custom_new
sections.append(
[
+ *allow,
f"impl {self.leaf} {{",
" /// Lossless complete-field allocation.",
*self._fn_lines(
@@ -677,6 +711,7 @@ class _ObjectRenderer:
sections.append(self._struct_lines(verdict, base))
sections.append(
[
+ *self._allow_missing_docs(),
"#[repr(C)]",
"#[derive(tvm_ffi::derive::ObjectRef, Clone)]",
f"pub struct {self.leaf} {{",
@@ -694,6 +729,7 @@ class _ObjectRenderer:
accessors += self._accessor_lines(field)
sections.append(
[
+ *self._allow_missing_docs(),
f"impl {self.obj_struct} {{",
*[f" {line}" if line else "" for line in accessors],
"}",
diff --git a/python/tvm_ffi/stub/rust_generator/consts.py
b/python/tvm_ffi/stub/rust_generator/consts.py
index 5db26ab8..6dab7590 100644
--- a/python/tvm_ffi/stub/rust_generator/consts.py
+++ b/python/tvm_ffi/stub/rust_generator/consts.py
@@ -20,7 +20,16 @@ from __future__ import annotations
#: One-line directives the Rust backend consumes.
RUST_DIRECTIVE_KINDS = frozenset(
- {"import-object", "field", "nullable", "enum", "opaque", "upcast",
"custom-new"}
+ {
+ "import-object",
+ "field",
+ "nullable",
+ "enum",
+ "opaque",
+ "upcast",
+ "custom-new",
+ "missing-docs",
+ }
)
#: Default FFI-origin -> Rust-type map; ``::`` paths get a ``use``, bare names
do not.
diff --git a/python/tvm_ffi/stub/rust_generator/directives.py
b/python/tvm_ffi/stub/rust_generator/directives.py
index ef2037ac..c322bbc7 100644
--- a/python/tvm_ffi/stub/rust_generator/directives.py
+++ b/python/tvm_ffi/stub/rust_generator/directives.py
@@ -16,7 +16,8 @@
# under the License.
"""The Rust backend's one-line directives: payload grammar and per-file
storage.
-Three address one reflected field as ``<type_key>.<field>``, three address a
type::
+Three address one reflected field as ``<type_key>.<field>``, three address a
type,
+and one applies to the whole file::
// tvm-ffi-stubgen(field): tirx.Add.a -> PrimExpr
// tvm-ffi-stubgen(nullable): ir.Expr.span
@@ -24,14 +25,18 @@ Three address one reflected field as
``<type_key>.<field>``, three address a typ
// tvm-ffi-stubgen(opaque): ir.SourceName
// tvm-ffi-stubgen(upcast): tirx.Add -> PrimExpr
// tvm-ffi-stubgen(custom-new): tirx.Add
+ // tvm-ffi-stubgen(missing-docs): allow
``field`` sets the field's Rust type (a name in scope, or a ``::`` path to
``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; the generated one is
-named ``from_fields`` instead.
+integer newtype for it, or with ``closed`` before the member list a
+``#[repr(<int>)]`` enum of exactly those members; ``opaque`` keeps a type
opaque
+even when its layout is reproducible. ``upcast`` adds an upcast to a reference
+outside the ancestor chain; ``custom-new`` says the wrapper's ``new`` is
+hand-written; the generated one is named ``from_fields`` instead.
+``missing-docs: allow`` puts ``#[allow(missing_docs)]`` on the file's generated
+types and ``impl`` blocks.
"""
from __future__ import annotations
@@ -41,18 +46,24 @@ import re
_ENUM_RE = re.compile(
r"^(?P<target>\S+)\s*->\s*(?P<name>[A-Za-z_]\w*)\((?P<repr>[iu](?:8|16|32|64))\)"
- r"\s*(?:\{(?P<body>[^{}]*)\})?$"
+ r"\s*(?:(?P<closed>closed\s*)?\{(?P<body>[^{}]*)\})?$"
)
_MEMBER_RE = re.compile(r"^(?P<name>[A-Za-z_]\w*)\s*=\s*(?P<value>-?\d+)$")
@dataclasses.dataclass(frozen=True)
class EnumSpec:
- """An ``enum`` directive: the newtype's name, its integer repr, and its
members."""
+ """An ``enum`` directive: the type's name, its integer repr, and its
members.
+
+ ``closed`` emits a Rust ``enum``: reading any other value from a mirrored
field is
+ undefined behavior, so it is for value sets fixed by the ABI that include
0, the
+ value of a field an object zero-filled at creation never set.
+ """
name: str
repr: str
members: tuple[tuple[str, int], ...]
+ closed: bool = False
@dataclasses.dataclass
@@ -65,6 +76,7 @@ class Directives:
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)
+ allow_missing_docs: bool = False
def add(self, name: str, payload: str, lineno: int) -> None:
"""Parse and store one directive; raise ``ValueError`` on a malformed
payload."""
@@ -83,6 +95,10 @@ class Directives:
self.upcasts.setdefault(_type_target(name, lhs, lineno),
[]).append(rust_type)
elif name == "custom-new":
self.custom_new.add(_type_target(name, payload, lineno))
+ elif name == "missing-docs":
+ if payload.strip() != "allow": # the one level so far
+ raise _invalid(name, lineno, "allow")
+ self.allow_missing_docs = True
else:
raise ValueError(f"Unknown directive `{name}` at line {lineno}")
@@ -116,8 +132,11 @@ def _split_arrow(name: str, payload: str, lineno: int,
expected: str) -> tuple[s
def _parse_enum(payload: str, lineno: int) -> tuple[str, EnumSpec]:
- """Parse ``<type_key>.<field> -> Name(i32) { A=0, B=1 }`` (the member list
is optional)."""
- expected = "<type_key>.<field> -> Name(i32) { A=0, B=1 }"
+ """Parse ``<type_key>.<field> -> Name(i32) [closed] { A=0, B=1 }``.
+
+ The member list is optional unless ``closed``.
+ """
+ expected = "<type_key>.<field> -> Name(i32) [closed] { A=0, B=1 }"
match = _ENUM_RE.match(payload.strip())
if match is None:
raise _invalid("enum", lineno, expected)
@@ -129,5 +148,9 @@ def _parse_enum(payload: str, lineno: int) -> tuple[str,
EnumSpec]:
if member is None:
raise _invalid("enum", lineno, expected)
members.append((member.group("name"), int(member.group("value"))))
+ closed = match.group("closed") is not None
+ if closed and not members:
+ raise _invalid("enum", lineno, expected)
target = _field_target("enum", match.group("target"), lineno)
- return target, EnumSpec(match.group("name"), match.group("repr"),
tuple(members))
+ spec = EnumSpec(match.group("name"), match.group("repr"), tuple(members),
closed)
+ return target, spec
diff --git a/tests/python/test_stubgen_rust.py
b/tests/python/test_stubgen_rust.py
index 445ba40a..0cbcfa83 100644
--- a/tests/python/test_stubgen_rust.py
+++ b/tests/python/test_stubgen_rust.py
@@ -353,6 +353,7 @@ def
test_generator_declares_its_directives_and_records_imports() -> None:
"opaque",
"upcast",
"custom-new",
+ "missing-docs",
}
imports = RUST.new_imports()
RUST.add_directive(imports, "import-object",
"tvm_ffi.libinfo.Foo;False;_Foo", 1)
@@ -955,6 +956,42 @@ def test_render_complete_enum_field() -> None:
assert "tvm_ffi::FieldGetter" not in _uses(imports)
+def test_allow_missing_docs_and_closed_enum() -> None:
+ """`missing-docs: allow` marks every generated type and `impl`; `closed`
emits a Rust enum."""
+ 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, "missing-docs", " allow ", 1)
+ RUST.add_directive(imports, "enum", "demo.Pair.kind -> Kind(i32) closed {
A=0, B=-1 }", 2)
+ text, _ = _render(info, imports)
+ assert text.startswith(
+ "#[allow(missing_docs)]\n"
+ "#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]\n"
+ "#[repr(i32)]\n"
+ "pub enum Kind {\n A = 0,\n B = -1,\n}\n"
+ )
+ assert (
+ " match value {\n"
+ " 0 => Ok(Self::A),\n"
+ " -1 => Ok(Self::B),\n"
+ ' _ => Err(Error::new(VALUE_ERROR, &format!("Kind has no
value {value}"), "")),\n'
+ " }\n"
+ ) in text
+ # The enum, the object struct, the wrapper and both allocator `impl`s.
+ assert text.count("#[allow(missing_docs)]\n") == 5
+ with pytest.raises(ValueError, match="closed"):
+ Directives().add("enum", "demo.Pair.kind -> Kind(i32) closed {}", 3)
+ for payload in ("", "warn"):
+ with pytest.raises(
+ ValueError, match=r"`missing-docs` directive at line 4\. Expected
`allow`"
+ ):
+ Directives().add("missing-docs", payload, 4)
+
+
def test_complete_optional_field_mirrors() -> None:
"""`Optional<T>` fields mirror their C++ layout: a 16-byte cell or a
nullable pointer."""
info = _info(
