This is an automated email from the ASF dual-hosted git repository.
tlopex pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm-ffi.git
The following commit(s) were added to refs/heads/main by this push:
new e3831a03 [STUBGEN][RUST] Error on dependencies a partial run does not
ask for (#748)
e3831a03 is described below
commit e3831a0378819136c7e71ac233e5ca3a5f5f6b46
Author: Linzhang Li <[email protected]>
AuthorDate: Sat Sep 5 21:18:30 2026 -0400
[STUBGEN][RUST] Error on dependencies a partial run does not ask for (#748)
[STUBGEN][RUST] Error on dependencies a partial stubgen
## Summary
A binding refers to other type keys: its parent (`base`), its ancestors
(upcasts), and its field types. The generator spelled every non-`ffi.*`
key by module path (`super::ir::ExprObj`) whether or not anything
provided it, so a partial run silently referenced modules that did not
exist. Now a referenced key must come from the crate (`ffi.*`), a
`ty-map` in the file (hand-written), or an `object/<key>` block in any
file of the run; otherwise the object fails with every missing key
listed:
```text
`tirx.Add` (line 3) depends on `ir.Expr`, `ir.Span`, `ir.Type`, which this
run does not ask for:
add an `object/<key>` block or a `ty-map` directive for each
```
`ty-map` now covers the parent too: with `ir.Expr -> crate::ir::Expr`
the child embeds `crate::ir::ExprObj` as `base`, derefs and upcasts to
it, and its allocator calls `ExprObj::new(...)` with the parent's
fields, which is the contract a hand-written parent must honour. Before,
only field types were mapped and `base` still pointed at
`super::ir::ExprObj`. Dependencies are not closed over automatically, as
discussed: the error is the to-do list.
## Changes
- `rust_generator/codegen.py`: `_generated()` becomes `_provider()`
(crate / mapped / generated / missing) and the base slot, upcasts,
allocator recursion, classification, and field resolution branch on it;
`generate_rust_object` takes `declared` and raises once per object.
- `stub/cli.py`: computes `declared` (every `object/` key of the run)
before stage 3 and passes it through `_stage_3`; `generate_object_block`
gains the parameter in the protocol and both backends (Python ignores
it).
- `examples/rust_stubgen/README.md`: "Partial generation" section.
- `tests/python/test_stubgen_rust.py`: missing keys reported together, a
key declared in another file resolves, a `ty-map`'d parent redirects
base / `Deref` / upcast, the registry case through `_stage_3`.
## Testing
`test_stubgen_rust.py` + `test_stubgen.py`: 105 passed; pre-commit lint
clean. Against `libtvm_compiler.so`, a skeleton with only
`object/tirx.Add` fails with the message above (exit 2); with three
`ty-map` lines it generates `base: ExprObj` from `crate::ir`.
Signed-off-by: yuchuan <[email protected]>
---
examples/rust_stubgen/README.md | 11 +++
python/tvm_ffi/stub/cli.py | 13 +++-
python/tvm_ffi/stub/generator.py | 8 ++-
python/tvm_ffi/stub/python_generator/generator.py | 2 +
python/tvm_ffi/stub/rust_generator/codegen.py | 70 +++++++++++++++----
python/tvm_ffi/stub/rust_generator/generator.py | 4 +-
tests/python/test_stubgen_rust.py | 83 +++++++++++++++++++++--
7 files changed, 169 insertions(+), 22 deletions(-)
diff --git a/examples/rust_stubgen/README.md b/examples/rust_stubgen/README.md
index 42b2758c..f4d257e3 100644
--- a/examples/rust_stubgen/README.md
+++ b/examples/rust_stubgen/README.md
@@ -71,3 +71,14 @@ wraps it in `Option` (`// tvm-ffi-stubgen(nullable):
rust_stubgen.IntPair.a`),
(`// tvm-ffi-stubgen(upcast): rust_stubgen.IntPair -> MyView`), and
`custom-new` names the generated allocator `from_complete_fields` when `new`
is hand-written (`// tvm-ffi-stubgen(custom-new): rust_stubgen.IntPair`).
+
+## Partial generation
+
+Every type an object refers to (its parent, its ancestors, the types of its
+fields) has to be provided in the same `tvm-ffi-stubgen` run: `ffi.*` types
+come from the `tvm-ffi` crate, an `object/<key>` block in any processed file
+generates it, and a `ty-map` points at a hand-written binding whose object
+struct is named `<Name>Obj`
+(`// tvm-ffi-stubgen(ty-map): rust_stubgen.IntPair -> crate::hand::IntPair`).
+Anything else is an error listing the missing keys, so a partial binding never
+references a module that does not exist.
diff --git a/python/tvm_ffi/stub/cli.py b/python/tvm_ffi/stub/cli.py
index a1afe73f..3fdffa4b 100644
--- a/python/tvm_ffi/stub/cli.py
+++ b/python/tvm_ffi/stub/cli.py
@@ -39,6 +39,8 @@ from .lib_state import (
from .utils import FuncInfo, InitConfig, Options
if TYPE_CHECKING:
+ from collections.abc import Container
+
from .generator import Generator
@@ -97,6 +99,13 @@ def __main__() -> int:
# Stage 3: Process
# - `tvm-ffi-stubgen(begin): global/...`
# - `tvm-ffi-stubgen(begin): object/...`
+ # Every `object/` block of the run: what a generated binding may refer to.
+ declared = frozenset(
+ code.param
+ for file in files
+ for code in file.code_blocks
+ if code.kind == "object" and isinstance(code.param, str)
+ )
for file in files:
if opt.verbose:
print(f"{C.TERM_CYAN}[File] {file.path}{C.TERM_RESET}")
@@ -107,6 +116,7 @@ def __main__() -> int:
ty_map,
global_funcs,
generator=generator,
+ declared=declared,
)
except Exception:
failed += 1
@@ -238,6 +248,7 @@ def _stage_3( # noqa: PLR0912
ty_map: dict[str, str],
global_funcs: dict[str, list[FuncInfo]],
generator: Generator,
+ declared: Container[str] = frozenset(),
) -> bool:
"""Process one file's blocks; return whether its content is (or would be)
changed."""
defined_funcs: set[str] = set()
@@ -268,7 +279,7 @@ def _stage_3( # noqa: PLR0912
obj_info = object_info_from_type_key(type_key)
type_key = ty_map.get(type_key, type_key)
defined_types.add(generator.canonical_type_name(type_key))
- generator.generate_object_block(code, ty_map, imports, opt,
obj_info)
+ generator.generate_object_block(code, ty_map, imports, opt,
obj_info, declared)
# Stage 4. Add imports for used types.
for code in file.code_blocks:
if code.kind == "import-section":
diff --git a/python/tvm_ffi/stub/generator.py b/python/tvm_ffi/stub/generator.py
index 19322274..add7503a 100644
--- a/python/tvm_ffi/stub/generator.py
+++ b/python/tvm_ffi/stub/generator.py
@@ -45,6 +45,7 @@ from .python_generator import PythonGenerator
from .rust_generator import RustGenerator
if TYPE_CHECKING:
+ from collections.abc import Container
from pathlib import Path
from .file_utils import CodeBlock
@@ -133,8 +134,13 @@ class Generator(Protocol):
imports: Any,
opt: Options,
obj_info: ObjectInfo,
+ declared: Container[str] = frozenset(),
) -> None:
- """Emit a type definition (fields + methods + init) for an
``object/<key>`` block."""
+ """Emit a type definition (fields + methods + init) for an
``object/<key>`` block.
+
+ ``declared`` lists the type keys that have an ``object/`` block
anywhere in
+ this run, for generators that must know what a binding may refer to.
+ """
...
def generate_import_section_block(
diff --git a/python/tvm_ffi/stub/python_generator/generator.py
b/python/tvm_ffi/stub/python_generator/generator.py
index 906c9a8f..cd81821d 100644
--- a/python/tvm_ffi/stub/python_generator/generator.py
+++ b/python/tvm_ffi/stub/python_generator/generator.py
@@ -32,6 +32,7 @@ from . import consts as PC
from .utils import ImportItem, PythonImports
if TYPE_CHECKING:
+ from collections.abc import Container
from pathlib import Path
from ..file_utils import CodeBlock
@@ -94,6 +95,7 @@ class PythonGenerator:
imports: PythonImports,
opt: Options,
obj_info: ObjectInfo,
+ declared: Container[str] = frozenset(),
) -> None:
"""Emit a Python class definition for an ``object/<key>`` block."""
G.generate_python_object(code, ty_map, imports.items, opt, obj_info)
diff --git a/python/tvm_ffi/stub/rust_generator/codegen.py
b/python/tvm_ffi/stub/rust_generator/codegen.py
index bed9af62..c5da9b4f 100644
--- a/python/tvm_ffi/stub/rust_generator/codegen.py
+++ b/python/tvm_ffi/stub/rust_generator/codegen.py
@@ -35,6 +35,12 @@ reflected field size. A builtin parent (``ffi.IntEnum``,
say) has no
``<Leaf>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``).
+
+Every type key an object refers to (parent, ancestors, field types) must be
+provided in the same run: ``ffi.*`` by the crate, a ``ty-map`` by a
hand-written
+binding whose object struct is ``<Name>Obj``, anything else by an ``object/``
+block in one of the processed files. Otherwise the block is an error naming the
+missing keys, so a partial binding never references a module that does not
exist.
"""
from __future__ import annotations
@@ -49,6 +55,7 @@ from . import consts as C_RUST
from .utils import RustImports, builtin_mirror_name, render_rust_type,
rust_ident
if TYPE_CHECKING:
+ from collections.abc import Container
from pathlib import Path
from ..file_utils import CodeBlock
@@ -83,6 +90,10 @@ class _ObjectRenderer:
ty_map: dict[str, str]
#: Module segments of the file this object lands in (``tirx.transform.X``
-> ``("tirx", "transform")``).
mod_segments: tuple[str, ...]
+ #: Type keys with an ``object/`` block somewhere in this run (see
:meth:`_provider`).
+ declared: Container[str]
+ #: Referenced type keys nobody provides; reported together once the body
is built.
+ missing: set[str] = dataclasses.field(default_factory=set)
@property
def type_key(self) -> str:
@@ -108,6 +119,7 @@ class _ObjectRenderer:
if mapped is None:
if "." not in origin or origin.startswith("ctypes."):
return None
+ self._provider(origin) # the crate or this run; a key nobody
provides is recorded
mapped = self._generated_type_path(origin)
return imports.record(mapped)
@@ -129,23 +141,38 @@ class _ObjectRenderer:
supers = "super::" * len(self.mod_segments)
return f"{supers or 'self::'}{type_key.replace('.', '::')}"
- def _generated(self, type_key: str) -> bool:
- """Whether ``type_key`` has a generated binding (builtin ``ffi.*``
types live in the crate)."""
- return type_key.partition(".")[0] not in C_RUST.RUST_MOD_MAP
+ def _provider(self, type_key: str) -> str | None:
+ """Who provides the binding of ``type_key``; a key nobody does is
recorded in :attr:`missing`.
+
+ ``"crate"`` for builtin ``ffi.*`` types, ``"mapped"`` for a ``ty-map``
to a
+ hand-written binding, ``"generated"`` for an ``object/`` block in this
run.
+ """
+ if type_key.partition(".")[0] in C_RUST.RUST_MOD_MAP:
+ return "crate"
+ if type_key in self.ty_map:
+ return "mapped"
+ if type_key in self.declared:
+ return "generated"
+ self.missing.add(type_key)
+ return None
def _base_type(self) -> tuple[str, bool]:
- """Resolve the ``base`` struct and whether it is a generated parent.
+ """Resolve the ``base`` struct and whether it is a generated or
``ty-map``'d parent.
A builtin parent below ``ffi.Object`` is embedded as its header-only
stand-in (see :meth:`RustImports.record_builtin_base`).
"""
parent = self.info.parent_type_key
- if parent is not None and self._generated(parent):
- return self.imports.record(self._generated_type_path(parent) +
"Obj"), True
+ if parent is not None:
+ provider = self._provider(parent)
+ if provider == "mapped":
+ return self.imports.record(self.ty_map[parent] + "Obj"), True
+ if provider != "crate": # generated in this run, or missing
(reported after the body)
+ return self.imports.record(self._generated_type_path(parent) +
"Obj"), True
chain = [key for key in self.info.ancestors if key !=
C_RUST.RUST_ROOT_TYPE_KEY]
if parent not in (None, C_RUST.RUST_ROOT_TYPE_KEY, *chain):
chain.append(parent)
- assert not any(self._generated(key) for key in chain), (self.type_key,
chain)
+ assert all(self._provider(key) == "crate" for key in chain),
(self.type_key, chain)
return self.imports.record_builtin_base(chain), False
# --- classification ----------------------------------------------------
@@ -164,7 +191,7 @@ class _ObjectRenderer:
unmirrored = {
key
for key in self.info.ancestors
- if key != C_RUST.RUST_ROOT_TYPE_KEY and not self._generated(key)
+ if key != C_RUST.RUST_ROOT_TYPE_KEY and self._provider(key) ==
"crate"
}
verdicts = classify(
infos,
@@ -317,11 +344,13 @@ class _ObjectRenderer:
def _upcast_lines(self) -> list[str]:
"""``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)
- ]
+ targets: list[str] = []
+ for key in self.info.ancestors:
+ provider = self._provider(key)
+ if provider == "mapped":
+ targets.append(self.imports.record(self.ty_map[key]))
+ elif provider != "crate":
+
targets.append(self.imports.record(self._generated_type_path(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:
@@ -368,7 +397,7 @@ class _ObjectRenderer:
"""``(field, type)`` of every physical field root to leaf, as
``<key>Obj::new`` takes."""
parent = info.parent_type_key
inherited: list[tuple[str, str]] = []
- if parent is not None and self._generated(parent):
+ if parent is not None and self._provider(parent) != "crate":
inherited = self._allocator_params(parent,
object_info_from_type_key(parent))
return self._level_params(key, info, inherited)
@@ -518,8 +547,12 @@ def generate_rust_object(
imports: RustImports,
opt: Options,
obj_info: ObjectInfo,
+ declared: Container[str],
) -> None:
- """Emit the Rust binding of ``obj_info`` into an ``object/<key>`` block."""
+ """Emit the Rust binding of ``obj_info`` into an ``object/<key>`` block.
+
+ Raises ``ValueError`` when it refers to a type key this run does not
provide.
+ """
assert len(code.lines) >= 2
assert isinstance(obj_info.type_key, str)
renderer = _ObjectRenderer(
@@ -527,8 +560,15 @@ def generate_rust_object(
imports=imports,
ty_map=ty_map,
mod_segments=tuple(obj_info.type_key.split(".")[:-1]),
+ declared=declared,
)
body = renderer.body()
+ if renderer.missing:
+ keys = ", ".join(f"`{key}`" for key in sorted(renderer.missing))
+ raise ValueError(
+ f"`{obj_info.type_key}` (line {code.lineno_start}) depends on
{keys}, which this run "
+ "does not ask for: add an `object/<key>` block or a `ty-map`
directive for each"
+ )
indent = " " * code.indent
code.lines = [
code.lines[0],
diff --git a/python/tvm_ffi/stub/rust_generator/generator.py
b/python/tvm_ffi/stub/rust_generator/generator.py
index 85b7e398..30b44dbd 100644
--- a/python/tvm_ffi/stub/rust_generator/generator.py
+++ b/python/tvm_ffi/stub/rust_generator/generator.py
@@ -35,6 +35,7 @@ from .codegen import (
from .utils import RustImports, RustUse
if TYPE_CHECKING:
+ from collections.abc import Container
from pathlib import Path
from ..file_utils import CodeBlock
@@ -93,9 +94,10 @@ class RustGenerator:
imports: RustImports,
opt: Options,
obj_info: ObjectInfo,
+ declared: Container[str] = frozenset(),
) -> None:
"""Emit the opaque Rust binding for an ``object/<key>`` block."""
- generate_rust_object(code, ty_map, imports, opt, obj_info)
+ generate_rust_object(code, ty_map, imports, opt, obj_info, declared)
def generate_import_section_block(
self, code: CodeBlock, imports: RustImports, opt: Options,
defined_types: set[str]
diff --git a/tests/python/test_stubgen_rust.py
b/tests/python/test_stubgen_rust.py
index 438a1131..39493400 100644
--- a/tests/python/test_stubgen_rust.py
+++ b/tests/python/test_stubgen_rust.py
@@ -19,7 +19,7 @@
from __future__ import annotations
import re
-from collections.abc import Iterator
+from collections.abc import Container, Iterator
from pathlib import Path
import pytest
@@ -27,7 +27,7 @@ import tvm_ffi.stub.cli as stub_cli
import tvm_ffi.testing # noqa: F401 (loads the `testing.*` fixture types)
from tvm_ffi.core import TypeSchema
from tvm_ffi.stub import consts as C
-from tvm_ffi.stub.cli import _stage_3
+from tvm_ffi.stub.cli import _stage_1, _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
@@ -118,12 +118,26 @@ def _object_block(type_key: str) -> CodeBlock:
)
-def _render(info: ObjectInfo, imports: RustImports | None = None) ->
tuple[str, RustImports]:
+class _Everything:
+ """Stands in for the run's declared type keys when a test does not care:
all are asked for."""
+
+ def __contains__(self, key: object) -> bool:
+ return True
+
+
+ALL_DECLARED = _Everything()
+
+
+def _render(
+ info: ObjectInfo,
+ imports: RustImports | None = None,
+ declared: Container[str] = ALL_DECLARED,
+) -> tuple[str, RustImports]:
"""Render ``info`` into a fresh object block; return the body text and the
collector."""
imports = RustImports() if imports is None else imports
assert info.type_key is not None
block = _object_block(info.type_key)
- generate_rust_object(block, RUST.default_ty_map(), imports, Options(),
info)
+ generate_rust_object(block, RUST.default_ty_map(), imports, Options(),
info, declared)
return "\n".join(block.lines[1:-1]), imports
@@ -1227,3 +1241,64 @@ def test_cli_check_rejects_init_flags(tmp_path: Path,
monkeypatch: pytest.Monkey
with pytest.raises(SystemExit) as excinfo:
stub_cli.__main__()
assert excinfo.value.code == 2
+
+
+# ---------------------------------------------------------------------------
+# Partial generation: every referenced type key must be provided in the run
+# ---------------------------------------------------------------------------
+
+
+def test_unasked_dependencies_are_an_error() -> None:
+ info = _info("tirx.Add", (_field("a", "tirx.Var"),), parent="ir.Expr")
+ expected = (
+ r"`tirx.Add` \(line 1\) depends on `ir.Expr`, `tirx.Var`, which this
run does not ask"
+ )
+ with pytest.raises(ValueError, match=expected):
+ _render(info, declared=frozenset({"tirx.Add"}))
+ # Declared in another file of the run: reached by module path, as before.
+ text, imports = _render(info, declared=frozenset({"tirx.Add", "ir.Expr",
"tirx.Var"}))
+ assert " base: ExprObj," in text
+ assert "pub fn a(&self) -> Result<Var> {" in text # same module: a bare
name
+ assert {"super::ir::ExprObj", "super::ir::Expr"} <= _uses(imports)
+
+
+def test_ty_map_parent_redirects_base_deref_and_upcast() -> None:
+ info = _info("tirx.Add", parent="ir.Expr", ancestors=["ffi.Object",
"ir.Expr"])
+ ty_map = {**RUST.default_ty_map(), "ir.Expr": "crate::ir::Expr"}
+ block = _object_block("tirx.Add")
+ imports = RustImports()
+ generate_rust_object(block, ty_map, imports, Options(), info,
frozenset({"tirx.Add"}))
+ text = "\n".join(block.lines[1:-1])
+ assert " base: ExprObj," in text
+ assert "impl Deref for AddObj {\n type Target = ExprObj;" in text
+ assert text.endswith("tvm_ffi::impl_object_upcast!(Add => Expr);")
+ assert {"crate::ir::ExprObj", "crate::ir::Expr"} <= _uses(imports)
+ assert not any(path.startswith("super::") for path in _uses(imports))
+
+
+def test_stage_3_checks_dependencies_across_the_run(tmp_path: Path) -> None:
+ src = tmp_path / "mod.rs"
+ blocks = [
+ f"{C.RUST_SYNTAX.begin} import-section",
+ C.RUST_SYNTAX.end,
+ f"{C.RUST_SYNTAX.begin} object/testing.TestCxxClassDerived",
+ C.RUST_SYNTAX.end,
+ "",
+ ]
+ declared = frozenset({"testing.TestCxxClassDerived"})
+ src.write_text("\n".join(blocks), encoding="utf-8")
+ info = FileInfo.from_file(src)
+ assert info is not None
+ with pytest.raises(ValueError, match=r"depends on
`testing\.TestCxxClassBase`"):
+ _stage_3(info, Options(dry_run=True), RUST.default_ty_map(), {}, RUST,
declared)
+ # A `ty-map` names a hand-written parent; its `<Name>Obj` becomes the base
slot.
+ mapped = f"{C.RUST_SYNTAX.ty_map} testing.TestCxxClassBase ->
crate::hand::TestCxxClassBase"
+ src.write_text("\n".join([mapped, *blocks]), encoding="utf-8")
+ info = FileInfo.from_file(src)
+ assert info is not None
+ ty_map = RUST.default_ty_map()
+ _stage_1(info, ty_map)
+ _stage_3(info, Options(dry_run=True), ty_map, {}, RUST, declared)
+ text = "\n".join(line for block in info.code_blocks for line in
block.lines)
+ assert " base: TestCxxClassBaseObj," in text
+ assert "use crate::hand::TestCxxClassBaseObj;" in text
