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 f0520693 [Rust] Stubgen: cache each accessor's FieldGetter (#961)
f0520693 is described below
commit f052069318165cb60f3fc7dd0d3252db9eddffee
Author: Tianqi Chen <[email protected]>
AuthorDate: Sat Nov 7 16:54:52 2026 -0500
[Rust] Stubgen: cache each accessor's FieldGetter (#961)
A generated Rust field accessor resolved its field on every call: a
registry check plus a scan of the type's field names. The stub generator
now keeps each accessor's `FieldGetter` in a fn-local
`ColdInitLazyLock`, so the lookup runs once and later calls go straight
to the C ABI getter:
```rust
pub fn dom(&self) -> Result<Option<Range>> {
static GETTER: ColdInitLazyLock<Result<FieldGetter>> =
ColdInitLazyLock::new(|| FieldGetter::new(IterVarObj::type_index(),
"dom"));
GETTER.clone()?.get(self)
}
```
- Applies to every accessor form (typed, `enum` directive, `Any`).
- A fn-local static cannot name `Self`, so the generator spells the
owning object type.
- The emitted text wraps the way rustfmt does, so formatted crates stay
byte-identical to the generator output.
- The stubgen expectations that pin accessor text are updated.
---
python/tvm_ffi/stub/rust_generator/codegen.py | 32 +++++++++++++++--
tests/python/test_stubgen_rust.py | 51 ++++++++++++++++++++-------
2 files changed, 68 insertions(+), 15 deletions(-)
diff --git a/python/tvm_ffi/stub/rust_generator/codegen.py
b/python/tvm_ffi/stub/rust_generator/codegen.py
index 95b9a900..d2de8b75 100644
--- a/python/tvm_ffi/stub/rust_generator/codegen.py
+++ b/python/tvm_ffi/stub/rust_generator/codegen.py
@@ -421,7 +421,8 @@ class _ObjectRenderer:
target = f"{self.type_key}.{field.name}"
name = rust_ident(field.name)
vis = self._visibility(target)
- getter = f'FieldGetter::new(Self::type_index(), "{field.name}")?'
+ static = self._getter_static_lines(field.name)
+ getter = "GETTER.clone()?"
allow = _allow_non_snake_case([name])
enum = directives.enums.get(target)
@@ -429,6 +430,7 @@ class _ObjectRenderer:
return [
*allow,
f"{vis} fn {name}(&self) -> Result<{enum.name}> {{",
+ *static,
f" let raw: i64 = {getter}.get(self)?;",
f" {enum.name}::try_from(raw)",
"}",
@@ -445,6 +447,7 @@ class _ObjectRenderer:
return [
*allow,
f"{vis} fn {name}(&self) -> Result<{any_type}> {{",
+ *static,
f" {getter}.get_any(self)",
"}",
]
@@ -453,10 +456,35 @@ class _ObjectRenderer:
return [
*allow,
f"{vis} fn {name}(&self) -> Result<{rust_type}> {{",
+ *static,
f" {getter}.get(self)",
"}",
]
+ def _getter_static_lines(self, field_name: str) -> list[str]:
+ """Declare the accessor's ``static GETTER``, which resolves the field
on first use only.
+
+ A fn-local ``static`` cannot name ``Self``, so the owner is spelled
out.
+ Wrapped as rustfmt does at the accessor's depth (``impl``, then
``fn``).
+ """
+ lazy =
self.imports.record("tvm_ffi::extra::cold_init_lazy_lock::ColdInitLazyLock")
+ static = f" static GETTER: {lazy}<Result<FieldGetter>> ="
+ owner, key = f"{self.obj_struct}::type_index()", f'"{field_name}"'
+ call = f"FieldGetter::new({owner}, {key})"
+ one_line = f"{static} {lazy}::new(|| {call});"
+ moved = f" {lazy}::new(|| {call});"
+ if len(one_line) <= C_RUST.RUST_MAX_WIDTH:
+ lines = [one_line]
+ elif len(moved) <= C_RUST.RUST_MAX_WIDTH:
+ lines = [static, moved]
+ else:
+ lines = [
+ f"{static} {lazy}::new(|| {{",
+ *_call_lines(" FieldGetter::new(", [owner, key],
")"),
+ " });",
+ ]
+ return [line[4:] for line in lines]
+
def _enum_lines(self, spec: EnumSpec) -> list[str]:
"""Render the type an ``enum`` directive declares: an open integer
newtype, or a closed enum."""
error = self.imports.record("tvm_ffi::Error")
@@ -794,7 +822,7 @@ class _ObjectRenderer:
is_enum = C_RUST.RUST_ENUM_TYPE_KEY in self.info.ancestors
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::ObjectCore") # `<Obj>::type_index()`
self.imports.record("tvm_ffi::reflection::FieldGetter")
self.imports.record("tvm_ffi::Result")
diff --git a/tests/python/test_stubgen_rust.py
b/tests/python/test_stubgen_rust.py
index 004d7203..1406698e 100644
--- a/tests/python/test_stubgen_rust.py
+++ b/tests/python/test_stubgen_rust.py
@@ -384,23 +384,33 @@ pub struct PairObj {{
impl PairObj {{
pub fn a(&self) -> Result<i64> {{
- FieldGetter::new(Self::type_index(), "a")?.get(self)
+ static GETTER: ColdInitLazyLock<Result<FieldGetter>> =
+ ColdInitLazyLock::new(|| FieldGetter::new(PairObj::type_index(),
"a"));
+ GETTER.clone()?.get(self)
}}
pub fn tag(&self) -> Result<Option<String>> {{
- FieldGetter::new(Self::type_index(), "tag")?.get(self)
+ static GETTER: ColdInitLazyLock<Result<FieldGetter>> =
+ ColdInitLazyLock::new(|| FieldGetter::new(PairObj::type_index(),
"tag"));
+ GETTER.clone()?.get(self)
}}
pub fn items(&self) -> Result<Array<Any>> {{
- FieldGetter::new(Self::type_index(), "items")?.get(self)
+ static GETTER: ColdInitLazyLock<Result<FieldGetter>> =
+ ColdInitLazyLock::new(|| FieldGetter::new(PairObj::type_index(),
"items"));
+ GETTER.clone()?.get(self)
}}
pub fn owner(&self) -> Result<ObjectRef> {{
- FieldGetter::new(Self::type_index(), "owner")?.get(self)
+ static GETTER: ColdInitLazyLock<Result<FieldGetter>> =
+ ColdInitLazyLock::new(|| FieldGetter::new(PairObj::type_index(),
"owner"));
+ GETTER.clone()?.get(self)
}}
pub fn r#type(&self) -> Result<Any> {{
- FieldGetter::new(Self::type_index(), "type")?.get_any(self)
+ static GETTER: ColdInitLazyLock<Result<FieldGetter>> =
+ ColdInitLazyLock::new(|| FieldGetter::new(PairObj::type_index(),
"type"));
+ GETTER.clone()?.get_any(self)
}}
}}"""
@@ -425,6 +435,7 @@ def test_render_root_object() -> None:
"tvm_ffi::ObjectCore",
"tvm_ffi::reflection::FieldGetter",
"tvm_ffi::Result",
+ "tvm_ffi::extra::cold_init_lazy_lock::ColdInitLazyLock",
"tvm_ffi::String",
"tvm_ffi::Array",
"tvm_ffi::Any",
@@ -605,24 +616,34 @@ impl Deref for IterVarObj {{
impl IterVarObj {{
pub fn dom(&self) -> Result<Option<Range>> {{
- FieldGetter::new(Self::type_index(), "dom")?.get(self)
+ static GETTER: ColdInitLazyLock<Result<FieldGetter>> =
+ ColdInitLazyLock::new(||
FieldGetter::new(IterVarObj::type_index(), "dom"));
+ GETTER.clone()?.get(self)
}}
pub fn var(&self) -> Result<PrimVar> {{
- FieldGetter::new(Self::type_index(), "var")?.get(self)
+ static GETTER: ColdInitLazyLock<Result<FieldGetter>> =
+ ColdInitLazyLock::new(||
FieldGetter::new(IterVarObj::type_index(), "var"));
+ GETTER.clone()?.get(self)
}}
pub fn iter_type(&self) -> Result<IterVarType> {{
- let raw: i64 = FieldGetter::new(Self::type_index(),
"iter_type")?.get(self)?;
+ static GETTER: ColdInitLazyLock<Result<FieldGetter>> =
+ ColdInitLazyLock::new(||
FieldGetter::new(IterVarObj::type_index(), "iter_type"));
+ let raw: i64 = GETTER.clone()?.get(self)?;
IterVarType::try_from(raw)
}}
pub fn thread_tag(&self) -> Result<String> {{
- FieldGetter::new(Self::type_index(), "thread_tag")?.get(self)
+ static GETTER: ColdInitLazyLock<Result<FieldGetter>> =
+ ColdInitLazyLock::new(||
FieldGetter::new(IterVarObj::type_index(), "thread_tag"));
+ GETTER.clone()?.get(self)
}}
pub fn span(&self) -> Result<Option<Span>> {{
- FieldGetter::new(Self::type_index(), "span")?.get(self)
+ static GETTER: ColdInitLazyLock<Result<FieldGetter>> =
+ ColdInitLazyLock::new(||
FieldGetter::new(IterVarObj::type_index(), "span"));
+ GETTER.clone()?.get(self)
}}
}}
@@ -919,7 +940,7 @@ def test_reflected_base_and_data_keep_their_names() -> None:
RUST.add_directive(imports, "visibility", "demo.Node.cache_ ->
pub(crate)", 1)
text, _ = _render(_info("demo.Node", (_field("base", "int"),
_field("cache_", "int"))), imports)
assert (
- 'pub fn base(&self) -> Result<i64> {\n
FieldGetter::new(Self::type_index(), "base")'
+ "pub fn base(&self) -> Result<i64> {\n static GETTER:
ColdInitLazyLock<Result<FieldGetter>> ="
in text
)
assert " pub(crate) fn cache_(&self) -> Result<i64> {" in text
@@ -1883,7 +1904,9 @@ def
test_ffi_field_type_off_the_default_ty_map_reads_as_any() -> None:
text, imports = _render(_info("demo.Doc", (_field("interrupt",
"ffi.VisitInterrupt"),)))
assert (
" pub fn interrupt(&self) -> Result<Any> {\n"
- ' FieldGetter::new(Self::type_index(),
"interrupt")?.get_any(self)\n'
+ " static GETTER: ColdInitLazyLock<Result<FieldGetter>> =\n"
+ ' ColdInitLazyLock::new(||
FieldGetter::new(DocObj::type_index(), "interrupt"));\n'
+ " GETTER.clone()?.get_any(self)\n"
) in text
assert not any("VisitInterrupt" in path for path in _uses(imports))
@@ -1903,7 +1926,9 @@ def
test_any_field_of_an_opaque_type_reads_through_get_any() -> None:
text, _ = _render(_info("demo.Literal", (_field("value", "Any"),)))
assert (
" pub fn value(&self) -> Result<Any> {\n"
- ' FieldGetter::new(Self::type_index(),
"value")?.get_any(self)\n'
+ " static GETTER: ColdInitLazyLock<Result<FieldGetter>> =\n"
+ ' ColdInitLazyLock::new(||
FieldGetter::new(LiteralObj::type_index(), "value"));\n'
+ " GETTER.clone()?.get_any(self)\n"
) in text
