This is an automated email from the ASF dual-hosted git repository.
Seven-Streams 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 abcc3980 [STUBGEN][RUST] Wrap generated allocator bodies the way
rustfmt does (#847)
abcc3980 is described below
commit abcc39807321d28a7ab711f9f2e5342f19de4ef7
Author: Linzhang Li <[email protected]>
AuthorDate: Wed Oct 14 12:15:48 2026 -0400
[STUBGEN][RUST] Wrap generated allocator bodies the way rustfmt does (#847)
## Summary
The Rust backend wraps a generated allocator signature at 100 columns,
but the two
statements of the allocator body did not follow rustfmt: `cargo fmt`
rewrote the
generated blocks, `tvm-ffi-stubgen --check` then reported them stale,
and
regenerating made `cargo fmt --check` fail again.
The allocator body now comes out the way rustfmt formats it:
- `Self { base, ... }` wider than 100 columns takes one field per line
with a
trailing comma. It used to stay on one line at any width.
- A `let` wider than 100 columns first moves its call to the next line,
and goes
one argument per line only when the call alone still does not fit. It
used to
go one argument per line right away.
```rust
// before
let obj = CallFFIKernelAttrObj::new(
launch_params,
launch_fields,
num_kernel_args,
kernel_attrs,
);
// after
let obj =
CallFFIKernelAttrObj::new(launch_params, launch_fields,
num_kernel_args, kernel_attrs);
```
Output that already fits in 100 columns is unchanged.
## Testing
- New `test_allocator_body_wraps_like_rustfmt` pins the three `let`
shapes and
the two `Self { .. }` shapes at the 100/101-column boundaries.
- `pytest tests/python/test_stubgen_rust.py tests/python/test_stubgen.py
tests/python/test_stub_layout.py`:
135 passed.
- Regenerated a downstream crate of TVM bindings (`ir`, `prim`,
`target`, `tirx`)
and ran rustfmt 1.9.0 with `use_small_heuristics = "Max"` on it:
`cargo fmt --check` reports no diff (8 before), `tvm-ffi-stubgen
--check`
exits 0, and `cargo build` passes.
- `pre-commit run` on the changed files.
Signed-off-by: yuchuan <[email protected]>
---
python/tvm_ffi/stub/rust_generator/codegen.py | 26 +++++++++++++++++++-------
tests/python/test_stubgen_rust.py | 19 +++++++++++++++++++
2 files changed, 38 insertions(+), 7 deletions(-)
diff --git a/python/tvm_ffi/stub/rust_generator/codegen.py
b/python/tvm_ffi/stub/rust_generator/codegen.py
index 4180b41a..f16ada42 100644
--- a/python/tvm_ffi/stub/rust_generator/codegen.py
+++ b/python/tvm_ffi/stub/rust_generator/codegen.py
@@ -70,12 +70,24 @@ if TYPE_CHECKING:
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."""
+ """``open_ + items + close`` on one line, or wrapped the way rustfmt does
when it would overflow.
+
+ A ``let`` first moves its call to the next line; anything else, and a call
that
+ still does not fit, takes one item per line.
+ """
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}"]
+ let, eq, call = open_.partition(" = ")
+ moved = f"{indent} {call}{', '.join(items)}{close}"
+ if eq and len(moved) <= C_RUST.RUST_MAX_WIDTH:
+ return [f"{let} =", moved]
+ return [
+ open_.rstrip(),
+ *[f"{indent} {item}," for item in items],
+ f"{indent}{close.lstrip()}",
+ ]
def _check_width(target: str, field: NamedTypeSchema, rust_type: str, width:
int) -> None:
@@ -436,15 +448,15 @@ class _ObjectRenderer:
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
+ self, head: str, params: list[tuple[str, str]], call: tuple[str,
list[str]], init: list[str]
) -> list[str]:
- """Render ``<head>(<params>) -> Self { let <call>; <result> }`` inside
an ``impl`` block."""
+ """Render ``<head>(<params>) -> Self { let <call>; Self { <init> } }``
in 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}",
+ *_call_lines(" Self { ", init, " }"),
" }",
]
@@ -472,7 +484,7 @@ class _ObjectRenderer:
"pub(crate) fn new",
params,
(f"base = {base}::new", to_parent),
- f"Self {{ {', '.join(['base', *own])} }}",
+ ["base", *own],
),
"}",
]
@@ -486,7 +498,7 @@ class _ObjectRenderer:
"pub fn from_complete_fields" if custom else "pub fn new",
params,
(f"obj = {self.obj_struct}::new", forward),
- "Self { base: ObjectArc::new(obj) }",
+ ["base: ObjectArc::new(obj)"],
),
"}",
]
diff --git a/tests/python/test_stubgen_rust.py
b/tests/python/test_stubgen_rust.py
index 0ccdd166..791c7f26 100644
--- a/tests/python/test_stubgen_rust.py
+++ b/tests/python/test_stubgen_rust.py
@@ -1025,6 +1025,25 @@ def test_custom_new_renames_the_wrapper_allocator() ->
None:
assert " pub fn new(" not in text
+def test_allocator_body_wraps_like_rustfmt() -> None:
+ """Past 100 columns a `let` first moves its call to the next line; the
rest go one per line."""
+
+ def render(width: int) -> tuple[str, str]:
+ name = "f" * width
+ text, _ = _render(_info("demo.W", (_field(name, "int", 24, 8),),
total_size=32))
+ return name, text
+
+ name, text = render(70) # the `let` is exactly 100 columns
+ assert f" let obj = WObj::new({name});\n" in text
+ name, text = render(76) # 106 columns; the call alone is exactly 100
+ assert f" let obj =\n WObj::new({name});\n" in text
+ name, text = render(77) # the call alone is 101 columns; `Self { .. }` is
exactly 100
+ assert f" let obj = WObj::new(\n {name},\n );\n"
in text
+ assert f" Self {{ base, {name} }}\n" in text
+ name, text = render(78)
+ assert f" Self {{\n base,\n {name},\n
}}\n" 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")
