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 94133a59 [STUBGEN][RUST] Add the Rust backend with opaque object
bindings (#738)
94133a59 is described below
commit 94133a5976008c155028ae06f10cba57dc4bd6de
Author: Linzhang Li <[email protected]>
AuthorDate: Fri Sep 4 01:01:26 2026 -0400
[STUBGEN][RUST] Add the Rust backend with opaque object bindings (#738)
## Summary
Register `--target rust` in `tvm-ffi-stubgen`. Every reflected object
becomes an *opaque* Rust binding: a `#[repr(C)]` struct that embeds only
its parent, a reference wrapper, read-only `Deref`,
`impl_object_upcast!` along the ancestor chain, and one accessor per
reflected field that reads through the C ABI getter (`FieldGetter`). The
object's bytes are never reproduced, so the binding is correct for every
registered type. Three directives (`field`, `nullable`, `enum`) shape
the accessor types. A CMake `STUB_TARGET` option and an example project
with its generated module checked in round it out.
Rebased on `main` after #736; depends on nothing else.
## Motivation
This is the first Rust step of the series that replaces #609 (after the
layout classifier in #730 and the directive channel in #736). Starting
with the opaque form keeps "never generate a wrong layout" true from the
first Rust PR. The follow-ups attach the classifier and upgrade the
types whose layout the registry can prove to full field mirrors, then
add allocators.
The opaque form is not a stopgap. Polymorphic types, types with
unreflected bytes, and types the registry has no layout for stay opaque
forever, and reading their fields through the reflected getter is the
only correct option. It is exactly what tvm-rust-ext hand-writes today
for `IterVar`, `Axis`, `TileLayout`, `BufferRegion`, `SourceName`, and
`Source`; this PR generates that shape from directives alone.
## Changes
- `python/tvm_ffi/stub/rust_generator/` (new)
- `consts.py`: FFI-origin to Rust-type map, `ffi.*` to crate-root
rewrite, declared directive names, keyword table.
- `utils.py`: `RustUse` / `RustImports` (per-file `use` collector plus
directives), `render_rust_type` (returns `None` where the crate has no
mirror), `rust_ident`.
- `directives.py`: grammar of `field: K.f -> T`, `nullable: K.f`, `enum:
K.f -> Name(i32) { A=0, ... }`; malformed payloads are rejected with the
line number.
- `codegen.py`: `_ObjectRenderer` (module-tree name resolution,
accessors, the `#[repr(transparent)]` enum newtype with `TryFrom<i64>`,
`Deref`, upcasts), plus the import section, `--init` scaffold and `pub
mod` stitching carried over from #609.
- `generator.py`: the `Generator` protocol adapter, declaring
`{import-object, field, nullable, enum}`.
- `stub/generator.py`: registers `"rust"`.
- `cmake/Utils/Library.cmake`: `STUB_TARGET`, passed as `--target`.
- `examples/rust_stubgen/`: a polymorphic C++ object with two registered
global functions, a Rust program that constructs it through the FFI and
reads it through the generated accessors, and the generated `mod.rs`
with an `enum` directive.
Decisions worth a look:
- A field without a Rust mirror (`Union`, `Dict`, `List`, `tuple`,
`void*`, an unmapped bare origin) gets an accessor returning `Any`.
Nothing is skipped; there is no `UnsupportedTypeError`.
- A builtin `ffi.*` parent has no generated `<Leaf>Obj`, so such a type
embeds the object header and upcasts only to generated ancestors.
- A `use` whose leaf is already taken is spelled in full at the use site
instead of failing.
- No allocator is generated: construction goes through the registered
global functions, as in the example.
Not carried over from #609: the builder, `DerefMut`, `same_as` /
`downcast`, method generation, the offset warnings, and every Rust
runtime change.
## Testing
`tests/python/test_stubgen_rust.py` (31 cases): `use` modelling and
collisions; value-type rendering including origins without a mirror; the
three directive grammars and their rejections; goldens for a root object
(header base, `Any` fallback, `r#type`), same-module and cross-module
derived objects, a builtin-parent object, and the hand-written `IterVar`
of tvm-rust-ext reproduced from directives; import section, `--init`
scaffold and module-tree stitching; `_stage_3` over a registered type
with an `enum` directive; the CLI in `--init` mode over the `testing.`
prefix, run twice to confirm idempotence. Full Python suite passes on
the rebased branch.
The example builds end to end: `cmake --build` regenerates
`rust/src/generated/rust_stubgen/mod.rs`, and `cargo run` prints `a=1
b=2 kind=PairKind(1)` and `sum=3` against the C++ library.
---------
Signed-off-by: yuchuan <[email protected]>
---
cmake/Utils/Library.cmake | 8 +-
examples/rust_stubgen/CMakeLists.txt | 43 ++
examples/rust_stubgen/README.md | 61 ++
examples/rust_stubgen/rust/Cargo.toml | 26 +
examples/rust_stubgen/rust/build.rs | 47 ++
examples/rust_stubgen/rust/src/generated/mod.rs | 20 +
.../rust/src/generated/rust_stubgen/mod.rs | 99 ++++
examples/rust_stubgen/rust/src/main.rs | 56 ++
examples/rust_stubgen/src/int_pair.cc | 68 +++
python/tvm_ffi/stub/generator.py | 2 +
python/tvm_ffi/stub/rust_generator/__init__.py | 23 +
python/tvm_ffi/stub/rust_generator/codegen.py | 425 ++++++++++++++
python/tvm_ffi/stub/rust_generator/consts.py | 73 +++
python/tvm_ffi/stub/rust_generator/directives.py | 108 ++++
python/tvm_ffi/stub/rust_generator/generator.py | 147 +++++
python/tvm_ffi/stub/rust_generator/utils.py | 142 +++++
tests/python/test_stubgen.py | 6 +-
tests/python/test_stubgen_rust.py | 647 +++++++++++++++++++++
18 files changed, 1997 insertions(+), 4 deletions(-)
diff --git a/cmake/Utils/Library.cmake b/cmake/Utils/Library.cmake
index 5c062125..6f263b58 100644
--- a/cmake/Utils/Library.cmake
+++ b/cmake/Utils/Library.cmake
@@ -172,6 +172,7 @@ endfunction ()
# target_name
# [LINK_SHARED ON|OFF] [LINK_HEADER ON|OFF] [DEBUG_SYMBOL ON|OFF]
[MSVC_FLAGS ON|OFF]
# [STUB_INIT ON|OFF] [STUB_DIR <dir>] [STUB_PKG <pkg>] [STUB_PREFIX <prefix>]
+# [STUB_TARGET python|rust]
# )
# Configure a target to integrate with TVM-FFI CMake utilities:
# - Link against tvm_ffi::header and/or tvm_ffi::shared
@@ -194,13 +195,14 @@ endfunction ()
# STUB_INIT: Whether to allow generating new directives. Default: OFF
(ON/OFF-style)
# STUB_PKG: Package name passed to stub generator (requires STUB_DIR and
STUB_INIT=ON; default: ${SKBUILD_PROJECT_NAME} if set, otherwise target name)
# STUB_PREFIX: Module prefix passed to stub generator (requires STUB_DIR
and STUB_INIT=ON; default: "<STUB_PKG>.")
+# STUB_TARGET: Code generator backend passed to the stub generator as
--target (default: python)
# ~~~
function (tvm_ffi_configure_target target)
if (NOT target)
message(
FATAL_ERROR
"tvm_ffi_configure_target: missing target name. "
- "Usage: tvm_ffi_configure_target(<target> [LINK_SHARED ON|OFF]
[LINK_HEADER ON|OFF] [DEBUG_SYMBOL ON|OFF] [MSVC_FLAGS ON|OFF] [STUB_INIT
ON|OFF] [STUB_DIR <dir>] [STUB_PKG <pkg>] [STUB_PREFIX <prefix>])"
+ "Usage: tvm_ffi_configure_target(<target> [LINK_SHARED ON|OFF]
[LINK_HEADER ON|OFF] [DEBUG_SYMBOL ON|OFF] [MSVC_FLAGS ON|OFF] [STUB_INIT
ON|OFF] [STUB_DIR <dir>] [STUB_PKG <pkg>] [STUB_PREFIX <prefix>] [STUB_TARGET
python|rust])"
)
endif ()
@@ -219,6 +221,7 @@ function (tvm_ffi_configure_target target)
STUB_DIR
STUB_PKG
STUB_PREFIX
+ STUB_TARGET
)
set(tvm_ffi_arg_multiValueArgs)
@@ -347,6 +350,9 @@ function (tvm_ffi_configure_target target)
REQUIRED
)
set(tvm_ffi_stub_cli_args "${tvm_ffi_arg__STUB_DIR_ABS}" --dlls
$<TARGET_FILE:${target}>)
+ if (DEFINED tvm_ffi_arg__STUB_TARGET AND tvm_ffi_arg__STUB_TARGET)
+ list(APPEND tvm_ffi_stub_cli_args --target "${tvm_ffi_arg__STUB_TARGET}")
+ endif ()
if (tvm_ffi_arg__STUB_INIT)
list(
APPEND
diff --git a/examples/rust_stubgen/CMakeLists.txt
b/examples/rust_stubgen/CMakeLists.txt
new file mode 100644
index 00000000..970ad3ee
--- /dev/null
+++ b/examples/rust_stubgen/CMakeLists.txt
@@ -0,0 +1,43 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+cmake_minimum_required(VERSION 3.18)
+project(rust_stubgen LANGUAGES CXX)
+
+find_package(
+ Python
+ COMPONENTS Interpreter
+ REQUIRED
+)
+execute_process(
+ COMMAND "${Python_EXECUTABLE}" -m tvm_ffi.config --cmakedir
+ OUTPUT_STRIP_TRAILING_WHITESPACE
+ OUTPUT_VARIABLE tvm_ffi_ROOT COMMAND_ERROR_IS_FATAL ANY
+)
+find_package(tvm_ffi CONFIG REQUIRED)
+
+# [example.cmake.begin]
+add_library(rust_stubgen SHARED src/int_pair.cc)
+tvm_ffi_configure_target(
+ rust_stubgen
+ STUB_TARGET
+ rust
+ STUB_INIT
+ ON
+ STUB_DIR
+ "./rust/src/generated"
+)
+# [example.cmake.end]
diff --git a/examples/rust_stubgen/README.md b/examples/rust_stubgen/README.md
new file mode 100644
index 00000000..080b5d6a
--- /dev/null
+++ b/examples/rust_stubgen/README.md
@@ -0,0 +1,61 @@
+<!--- Licensed to the Apache Software Foundation (ASF) under one -->
+<!--- or more contributor license agreements. See the NOTICE file -->
+<!--- distributed with this work for additional information -->
+<!--- regarding copyright ownership. The ASF licenses this file -->
+<!--- to you under the Apache License, Version 2.0 (the -->
+<!--- "License"); you may not use this file except in compliance -->
+<!--- with the License. You may obtain a copy of the License at -->
+
+<!--- http://www.apache.org/licenses/LICENSE-2.0 -->
+
+<!--- Unless required by applicable law or agreed to in writing, -->
+<!--- software distributed under the License is distributed on an -->
+<!--- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -->
+<!--- KIND, either express or implied. See the License for the -->
+<!--- specific language governing permissions and limitations -->
+<!--- under the License. -->
+
+# 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.
+A builtin parent such as `ffi.IntEnum` has no `<Leaf>Obj` in the crate; the
+import section defines a header-only stand-in per builtin ancestor so the
+derived type depth matches the registry.
+
+## Build and run
+
+```bash
+# 1. Build the C++ library; the post-build step runs the stub generator.
+cmake -B build -DCMAKE_BUILD_TYPE=Release
+cmake --build build
+
+# 2. Build and run the Rust program against it.
+cd rust && cargo run
+```
+
+The Rust crate depends on the `tvm-ffi` crate of this repository and needs
+`tvm-ffi-config` on `PATH` (activate the virtual environment where the
+`apache-tvm-ffi` package is installed).
+
+## Directives
+
+The generated file keeps one-line directives the generator reads on every run.
+`rust/src/generated/rust_stubgen/mod.rs` declares the integer field `kind` as
an
+open newtype:
+
+```rust
+// 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`).
diff --git a/examples/rust_stubgen/rust/Cargo.toml
b/examples/rust_stubgen/rust/Cargo.toml
new file mode 100644
index 00000000..776b10d1
--- /dev/null
+++ b/examples/rust_stubgen/rust/Cargo.toml
@@ -0,0 +1,26 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+[package]
+name = "rust_stubgen_example"
+version = "0.1.0"
+edition = "2021"
+license = "Apache-2.0"
+publish = false
+
+[dependencies]
+tvm-ffi = { path = "../../../rust/tvm-ffi" }
diff --git a/examples/rust_stubgen/rust/build.rs
b/examples/rust_stubgen/rust/build.rs
new file mode 100644
index 00000000..b6245607
--- /dev/null
+++ b/examples/rust_stubgen/rust/build.rs
@@ -0,0 +1,47 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+//! `tvm-ffi-sys`'s build script links libtvm_ffi, but the loader-path
+//! environment it emits only applies to its own package; re-emit it here so a
+//! plain `cargo run` finds libtvm_ffi at startup.
+
+use std::env;
+use std::process::Command;
+
+fn main() {
+ let output = Command::new("tvm-ffi-config")
+ .arg("--libdir")
+ .output()
+ .expect("failed to run tvm-ffi-config; install tvm-ffi and activate
the virtualenv");
+ assert!(output.status.success(), "tvm-ffi-config --libdir failed");
+ let lib_dir = String::from_utf8(output.stdout).unwrap().trim().to_string();
+
+ let loader_var = match env::var("CARGO_CFG_TARGET_OS").as_deref() {
+ Ok("windows") => "PATH",
+ Ok("macos") => "DYLD_LIBRARY_PATH",
+ _ => "LD_LIBRARY_PATH",
+ };
+ let sep = if loader_var == "PATH" { ";" } else { ":" };
+ let prev = env::var(loader_var).unwrap_or_default();
+ let val = if prev.is_empty() {
+ lib_dir
+ } else {
+ format!("{prev}{sep}{lib_dir}")
+ };
+ println!("cargo:rustc-env={loader_var}={val}");
+}
diff --git a/examples/rust_stubgen/rust/src/generated/mod.rs
b/examples/rust_stubgen/rust/src/generated/mod.rs
new file mode 100644
index 00000000..cf99c109
--- /dev/null
+++ b/examples/rust_stubgen/rust/src/generated/mod.rs
@@ -0,0 +1,20 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+pub mod rust_stubgen;
diff --git a/examples/rust_stubgen/rust/src/generated/rust_stubgen/mod.rs
b/examples/rust_stubgen/rust/src/generated/rust_stubgen/mod.rs
new file mode 100644
index 00000000..4e7f47e6
--- /dev/null
+++ b/examples/rust_stubgen/rust/src/generated/rust_stubgen/mod.rs
@@ -0,0 +1,99 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+#![allow(dead_code, unused_imports)]
+
+//! FFI bindings for `rust_stubgen` (generated by tvm-ffi-stubgen).
+
+// 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.
+// tvm-ffi-stubgen(enum): rust_stubgen.IntPair.kind -> PairKind(i32) {
Unordered=0, Ordered=1 }
+
+// tvm-ffi-stubgen(begin): object/rust_stubgen.IntPair
+#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
+#[repr(transparent)]
+pub struct PairKind(i32);
+
+#[allow(non_upper_case_globals)]
+impl PairKind {
+ pub const Unordered: Self = Self(0);
+ pub const Ordered: Self = Self(1);
+ pub const fn from_raw(value: i32) -> Self {
+ Self(value)
+ }
+ pub const fn as_raw(self) -> i32 {
+ self.0
+ }
+}
+
+impl TryFrom<i64> for PairKind {
+ type Error = Error;
+ fn try_from(value: i64) -> Result<Self> {
+ i32::try_from(value).map(Self).map_err(|_| {
+ Error::new(VALUE_ERROR, &format!("PairKind value {value} does not
fit i32"), "")
+ })
+ }
+}
+
+#[repr(C)]
+#[derive(tvm_ffi::derive::Object)]
+#[type_key = "rust_stubgen.IntPair"]
+#[type_final]
+pub struct IntPairObj {
+ base: Object,
+}
+
+#[repr(C)]
+#[derive(tvm_ffi::derive::ObjectRef, Clone)]
+pub struct IntPair {
+ data: ObjectArc<IntPairObj>,
+}
+
+impl Deref for IntPair {
+ type Target = IntPairObj;
+ fn deref(&self) -> &IntPairObj {
+ &self.data
+ }
+}
+
+impl IntPairObj {
+ pub fn a(&self) -> Result<i64> {
+ FieldGetter::new(Self::type_index(), "a")?.get(self)
+ }
+
+ pub fn b(&self) -> Result<i64> {
+ FieldGetter::new(Self::type_index(), "b")?.get(self)
+ }
+
+ pub fn kind(&self) -> Result<PairKind> {
+ let raw: i64 = FieldGetter::new(Self::type_index(),
"kind")?.get(self)?;
+ PairKind::try_from(raw)
+ }
+}
+// tvm-ffi-stubgen(end)
diff --git a/examples/rust_stubgen/rust/src/main.rs
b/examples/rust_stubgen/rust/src/main.rs
new file mode 100644
index 00000000..27afeb3a
--- /dev/null
+++ b/examples/rust_stubgen/rust/src/main.rs
@@ -0,0 +1,56 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+//! Use the stubgen-generated `IntPair` binding (see ../../README.md).
+
+mod generated;
+
+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 {
+ let name = if cfg!(target_os = "windows") {
+ "rust_stubgen.dll"
+ } else if cfg!(target_os = "macos") {
+ "librust_stubgen.dylib"
+ } else {
+ "librust_stubgen.so"
+ };
+ format!("{}/../build/{}", env!("CARGO_MANIFEST_DIR"), name)
+}
+
+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.
+ 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);
+
+ let sum: i64 = tvm_ffi::cached_global_func!("rust_stubgen.IntPairSum")
+ .call_tuple((pair.clone(),))?
+ .try_into()?;
+ println!("sum={sum}");
+ Ok(())
+}
diff --git a/examples/rust_stubgen/src/int_pair.cc
b/examples/rust_stubgen/src/int_pair.cc
new file mode 100644
index 00000000..3d7631b4
--- /dev/null
+++ b/examples/rust_stubgen/src/int_pair.cc
@@ -0,0 +1,68 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+/*!
+ * \file int_pair.cc
+ * \brief A tvm-ffi library that registers one object for the Rust stub
generator.
+ */
+#include <tvm/ffi/tvm_ffi.h>
+
+#include <cstdint>
+
+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.
+class IntPairObj : public ffi::Object {
+ public:
+ int64_t a;
+ int64_t b;
+ 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; }
+
+ 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<IntPairObj>(a, b, kind); }
+
+ TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(IntPair, ffi::ObjectRef,
IntPairObj);
+};
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+ namespace refl = tvm::ffi::reflection;
+ refl::ObjectDef<IntPairObj>(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::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(); });
+}
+// [object.end]
+
+} // namespace rust_stubgen
diff --git a/python/tvm_ffi/stub/generator.py b/python/tvm_ffi/stub/generator.py
index 3cd896a3..19322274 100644
--- a/python/tvm_ffi/stub/generator.py
+++ b/python/tvm_ffi/stub/generator.py
@@ -42,6 +42,7 @@ from typing import TYPE_CHECKING, Any, Protocol
from . import consts as C
from .python_generator import PythonGenerator
+from .rust_generator import RustGenerator
if TYPE_CHECKING:
from pathlib import Path
@@ -193,6 +194,7 @@ class Generator(Protocol):
_GENERATORS: dict[str, Generator] = {
"python": PythonGenerator(),
+ "rust": RustGenerator(),
}
diff --git a/python/tvm_ffi/stub/rust_generator/__init__.py
b/python/tvm_ffi/stub/rust_generator/__init__.py
new file mode 100644
index 00000000..ea520cbe
--- /dev/null
+++ b/python/tvm_ffi/stub/rust_generator/__init__.py
@@ -0,0 +1,23 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Rust code generator for ``tvm-ffi-stubgen``."""
+
+from __future__ import annotations
+
+from .generator import RustGenerator
+
+__all__ = ["RustGenerator"]
diff --git a/python/tvm_ffi/stub/rust_generator/codegen.py
b/python/tvm_ffi/stub/rust_generator/codegen.py
new file mode 100644
index 00000000..26f2603a
--- /dev/null
+++ b/python/tvm_ffi/stub/rust_generator/codegen.py
@@ -0,0 +1,425 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# 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<IterVarObj>,
+ }
+
+ impl Deref for IterVar { ... } // IterVar -> IterVarObj
+ impl Deref for IterVarObj { ... } // IterVarObj ->
PrimExprConvertibleObj
+
+ impl IterVarObj {
+ pub fn dom(&self) -> Result<Option<Range>> {
+ FieldGetter::new(Self::type_index(), "dom")?.get(self)
+ }
+ ...
+ }
+
+ tvm_ffi::impl_object_upcast!(IterVar => PrimExprConvertible);
+
+Construction and behaviour go through the registered global functions,
+hand-written outside the markers. 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``.
+"""
+
+from __future__ import annotations
+
+import dataclasses
+from typing import TYPE_CHECKING
+
+from .. import consts as C
+from . import consts as C_RUST
+from .utils import RustImports, builtin_mirror_name, render_rust_type,
rust_ident
+
+if TYPE_CHECKING:
+ from pathlib import Path
+
+ from ..file_utils import CodeBlock
+ from ..utils import InitConfig, NamedTypeSchema, ObjectInfo, Options
+ from .directives import EnumSpec
+
+
[email protected]
+class _ObjectRenderer:
+ """Renders one ``object/<key>`` block into Rust source lines."""
+
+ info: ObjectInfo
+ imports: RustImports
+ ty_map: dict[str, str]
+ #: Module segments of the file this object lands in (``tirx.transform.X``
-> ``("tirx", "transform")``).
+ mod_segments: tuple[str, ...]
+
+ @property
+ def type_key(self) -> str:
+ """The object's type key."""
+ assert self.info.type_key is not None
+ return self.info.type_key
+
+ @property
+ def leaf(self) -> str:
+ """The reference wrapper's name (``IterVar``)."""
+ return self.type_key.rsplit(".", 1)[-1]
+
+ @property
+ def obj_struct(self) -> str:
+ """The object struct's name (``IterVarObj``)."""
+ return f"{self.leaf}Obj"
+
+ # --- name resolution ---------------------------------------------------
+
+ def _ty_render(self, origin: str) -> 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)
+
+ def _generated_type_path(self, type_key: str) -> str:
+ """Spell a generated type key from this file.
+
+ Same module: the bare leaf. Elsewhere: ``super::`` per segment of this
+ file's module, then the full path (edition 2021 rejects ``use
ir::Expr``).
+ """
+ head, _, _ = type_key.partition(".")
+ if head in C_RUST.RUST_MOD_MAP:
+ return type_key
+ mod, _, type_leaf = type_key.rpartition(".")
+ if tuple(mod.split(".")) == self.mod_segments:
+ return type_leaf
+ 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 _base_type(self) -> tuple[str, bool]:
+ """Resolve the ``base`` struct and whether it is a generated 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
+ 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)
+ return self.imports.record_builtin_base(chain), False
+
+ # --- pieces ------------------------------------------------------------
+
+ def _accessor_lines(self, field: NamedTypeSchema) -> list[str]:
+ """One ``pub fn <field>(&self) -> Result<T>`` through the C ABI getter.
+
+ ``T`` comes from the directives, else the schema; without a Rust type,
``Any``.
+ """
+ directives = self.imports.directives
+ target = f"{self.type_key}.{field.name}"
+ name = rust_ident(field.name)
+ getter = f'FieldGetter::new(Self::type_index(), "{field.name}")?'
+
+ enum = directives.enums.get(target)
+ if enum is not None:
+ return [
+ f"pub fn {name}(&self) -> Result<{enum.name}> {{",
+ f" let raw: i64 = {getter}.get(self)?;",
+ f" {enum.name}::try_from(raw)",
+ "}",
+ ]
+ override = directives.field_types.get(target)
+ if override is not None:
+ rust_type = self.imports.record(override) if "::" in override else
override
+ else:
+ rust_type = render_rust_type(field, self._ty_render)
+ if rust_type is None:
+ any_type = self.imports.record("tvm_ffi::Any")
+ return [
+ f"pub fn {name}(&self) -> Result<{any_type}> {{",
+ f" {getter}.get_any(self)",
+ "}",
+ ]
+ 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)", "}"]
+
+ 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")
+ return [
+ "#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]",
+ "#[repr(transparent)]",
+ f"pub struct {spec.name}({spec.repr});",
+ "",
+ "#[allow(non_upper_case_globals)]",
+ f"impl {spec.name} {{",
+ *[f" pub const {member}: Self = Self({value});" for member,
value in spec.members],
+ f" pub const fn from_raw(value: {spec.repr}) -> Self {{",
+ " Self(value)",
+ " }",
+ f" pub const fn as_raw(self) -> {spec.repr} {{",
+ " self.0",
+ " }",
+ "}",
+ "",
+ f"impl TryFrom<i64> for {spec.name} {{",
+ f" type Error = {error};",
+ " fn try_from(value: i64) -> Result<Self> {",
+ 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}"), "")',
+ " })",
+ " }",
+ "}",
+ ]
+
+ def _deref_lines(self, source: str, target: str, member: str) -> list[str]:
+ return [
+ f"impl Deref for {source} {{",
+ f" type Target = {target};",
+ f" fn deref(&self) -> &{target} {{",
+ f" &self.{member}",
+ " }",
+ "}",
+ ]
+
+ def _upcast_lines(self) -> list[str]:
+ """``impl_object_upcast!`` from the wrapper to every ancestor's
wrapper."""
+ targets = [
+ self.imports.record(self._generated_type_path(key))
+ for key in self.info.ancestors
+ if self._generated(key)
+ ]
+ if not targets:
+ return []
+ pairs = ", ".join(f"{self.leaf} => {target}" for target in targets)
+ return [f"tvm_ffi::impl_object_upcast!({pairs});"]
+
+ def body(self) -> list[str]:
+ """Build the Rust source lines for the object."""
+ # 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:
+ self.imports.record("tvm_ffi::ObjectCore") # `Self::type_index()`
+ self.imports.record("tvm_ffi::FieldGetter")
+ self.imports.record("tvm_ffi::Result")
+
+ sections: list[list[str]] = []
+ enums = self.imports.directives.enums
+ sections += [
+ self._enum_lines(enums[f"{self.type_key}.{f.name}"])
+ 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(
+ [
+ "#[repr(C)]",
+ "#[derive(tvm_ffi::derive::ObjectRef, Clone)]",
+ f"pub struct {self.leaf} {{",
+ f" data: ObjectArc<{self.obj_struct}>,",
+ "}",
+ ]
+ )
+ 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] = []
+ for i, field in enumerate(fields):
+ if i:
+ accessors.append("")
+ accessors += self._accessor_lines(field)
+ sections.append(
+ [
+ f"impl {self.obj_struct} {{",
+ *[f" {line}" if line else "" for line in accessors],
+ "}",
+ ]
+ )
+ upcasts = self._upcast_lines()
+ if upcasts:
+ sections.append(upcasts)
+
+ lines: list[str] = []
+ for i, section in enumerate(sections):
+ if i:
+ lines.append("")
+ lines += section
+ return lines
+
+
+def generate_rust_object(
+ code: CodeBlock,
+ ty_map: dict[str, str],
+ imports: RustImports,
+ opt: Options,
+ obj_info: ObjectInfo,
+) -> None:
+ """Emit the opaque Rust binding of ``obj_info`` into an ``object/<key>``
block."""
+ assert len(code.lines) >= 2
+ assert isinstance(obj_info.type_key, str)
+ renderer = _ObjectRenderer(
+ info=obj_info,
+ imports=imports,
+ ty_map=ty_map,
+ mod_segments=tuple(obj_info.type_key.split(".")[:-1]),
+ )
+ body = renderer.body()
+ indent = " " * code.indent
+ code.lines = [
+ code.lines[0],
+ *[(indent + line) if line else "" for line in body],
+ code.lines[-1],
+ ]
+ _ = opt # accepted for protocol parity
+
+
+# --- import section (`use` statements) --------------------------------------
+
+
+def _builtin_mirror_lines(type_key: str, base: str) -> list[str]:
+ """Render the header-only stand-in for one builtin ancestor."""
+ return [
+ f"/// Header-only stand-in for the builtin `{type_key}`; it only
carries the ancestor depth.",
+ "#[allow(dead_code)]",
+ "#[repr(C)]",
+ "#[derive(tvm_ffi::derive::Object)]",
+ f'#[type_key = "{type_key}"]',
+ f"struct {builtin_mirror_name(type_key)} {{",
+ f" base: {base},",
+ "}",
+ ]
+
+
+def generate_rust_import_section(
+ code: CodeBlock,
+ imports: RustImports,
+ opt: Options,
+ defined_types: set[str],
+) -> None:
+ """Render the ``use`` lines, then the builtin stand-ins, into an
``import-section`` block.
+
+ Imports of types defined in this file are dropped; the rest are deduped
and sorted.
+ """
+ assert len(code.lines) >= 2
+ body = sorted({item.as_use_line() for item in imports.items if item.path
not in defined_types})
+ for type_key, base in imports.builtin_mirrors.items():
+ body += ["", *_builtin_mirror_lines(type_key, base)]
+ indent = " " * code.indent
+ code.lines = [
+ code.lines[0],
+ *[(indent + line) if line else "" for line in body],
+ code.lines[-1],
+ ]
+ _ = opt # accepted for protocol parity
+
+
+# --- whole-file scaffolding (`--init` mode) ---------------------------------
+
+
+def generate_rust_api_file(
+ code_blocks: list[CodeBlock],
+ ty_map: dict[str, str],
+ module_name: str,
+ object_infos: list[ObjectInfo],
+ init_cfg: InitConfig,
+ is_root: bool,
+ syntax: C.MarkerSyntax,
+) -> str:
+ """Scaffold a single Rust binding file (one file per module prefix)."""
+ append = ""
+ if not code_blocks:
+ append += "#![allow(dead_code, unused_imports)]\n"
+ append += f"\n//! FFI bindings for `{module_name}` (generated by
tvm-ffi-stubgen).\n\n"
+ if not any(c.kind == "import-section" for c in code_blocks):
+ append += f"{syntax.begin} import-section\n{syntax.end}\n\n"
+ defined = {c.param for c in code_blocks if c.kind == "object"}
+ for info in object_infos:
+ type_key = info.type_key
+ if type_key is None or type_key in defined:
+ continue
+ append += f"{syntax.begin} object/{type_key}\n{syntax.end}\n\n"
+ _ = (ty_map, init_cfg, is_root) # unused for the Rust single-file layout
+ return append
+
+
+# --- module-tree stitching (auto-form `pub mod` declarations) ----------------
+
+
+def finalize_rust_module_tree(init_path: Path, prefixes: set[str]) -> None:
+ """Declare each generated prefix with ``pub mod`` in its parent's
``mod.rs``.
+
+ Missing ``mod.rs`` files are created; the user mounts ``init_path`` with
one ``mod`` line.
+ """
+ children: dict[Path, set[str]] = {}
+ for prefix in prefixes:
+ segs = [s for s in prefix.split(".") if s]
+ for i, seg in enumerate(segs):
+ parent = init_path.joinpath(*segs[:i])
+ children.setdefault(parent, set()).add(seg)
+
+ for parent, names in children.items():
+ parent.mkdir(parents=True, exist_ok=True)
+ mod_rs = parent / "mod.rs"
+ existing = mod_rs.read_text(encoding="utf-8") if mod_rs.exists() else
""
+ to_add = [f"pub mod {n};" for n in sorted(names) if f"pub mod {n};"
not in existing]
+ if not to_add:
+ continue
+ text = existing
+ if text and not text.endswith("\n"):
+ text += "\n"
+ if text.strip(): # separate from any existing bindings
+ text += "\n"
+ text += "\n".join(to_add) + "\n"
+ mod_rs.write_text(text, encoding="utf-8")
diff --git a/python/tvm_ffi/stub/rust_generator/consts.py
b/python/tvm_ffi/stub/rust_generator/consts.py
new file mode 100644
index 00000000..218def67
--- /dev/null
+++ b/python/tvm_ffi/stub/rust_generator/consts.py
@@ -0,0 +1,73 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Rust-specific constants for the ``tvm-ffi-stubgen`` Rust backend."""
+
+from __future__ import annotations
+
+#: One-line directives the Rust backend consumes.
+RUST_DIRECTIVE_KINDS = frozenset({"import-object", "field", "nullable",
"enum"})
+
+#: Default FFI-origin -> Rust-type map; ``::`` paths get a ``use``, bare names
do not.
+RUST_TY_MAP_DEFAULTS = {
+ "int": "i64",
+ "float": "f64",
+ "bool": "bool",
+ "None": "()",
+ "str": "tvm_ffi::String",
+ "bytes": "tvm_ffi::Bytes",
+ "Any": "tvm_ffi::Any",
+ "Callable": "tvm_ffi::Function",
+ "Array": "tvm_ffi::Array", # the crate's own Array<T>, NOT Vec
+ "Map": "tvm_ffi::Map", # the crate's own Map<K, V>, NOT HashMap
+ # An object value is the `ObjectRef` handle; `Object` only appears as an
embedded `base`.
+ "Object": "tvm_ffi::object::ObjectRef",
+ "Tensor": "tvm_ffi::Tensor",
+ "Shape": "tvm_ffi::Shape",
+ "Device": "tvm_ffi::DLDevice",
+ "dtype": "tvm_ffi::DLDataType",
+ "DataType": "tvm_ffi::DLDataType",
+ # --- builtin object type keys (ffi.*) ---
+ "ffi.String": "tvm_ffi::String",
+ "ffi.Bytes": "tvm_ffi::Bytes",
+ "ffi.Module": "tvm_ffi::Module",
+ "ffi.Error": "tvm_ffi::Error",
+ "ffi.Object": "tvm_ffi::object::ObjectRef",
+ "ffi.Tensor": "tvm_ffi::Tensor",
+ "ffi.Shape": "tvm_ffi::Shape",
+ "ffi.Function": "tvm_ffi::Function",
+}
+
+#: Origins without a crate mirror; such a field is read as ``tvm_ffi::Any``.
+RUST_UNSUPPORTED_ORIGINS = frozenset({"Dict", "List", "Union", "tuple"})
+
+#: ``use``-path rewrites: builtin ``ffi.*`` type keys live at the crate root.
+RUST_MOD_MAP = {
+ "ffi": "tvm_ffi",
+}
+
+#: Root of the object hierarchy.
+RUST_ROOT_TYPE_KEY = "ffi.Object"
+
+#: Keywords a field name may collide with: spelled ``r#name``, or ``name_``
for the four
+#: that cannot be raw identifiers.
+RUST_KEYWORDS = frozenset(
+ "as async await break const continue crate dyn else enum extern false fn
for if impl in "
+ "let loop match mod move mut pub ref return self Self static struct super
trait true type "
+ "unsafe use where while abstract become box do final gen macro override
priv try typeof "
+ "unsized virtual yield".split()
+)
+RUST_NOT_RAW_IDENTIFIERS = frozenset({"self", "Self", "super", "crate"})
diff --git a/python/tvm_ffi/stub/rust_generator/directives.py
b/python/tvm_ffi/stub/rust_generator/directives.py
new file mode 100644
index 00000000..388ffb5f
--- /dev/null
+++ b/python/tvm_ffi/stub/rust_generator/directives.py
@@ -0,0 +1,108 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""The Rust backend's one-line directives: payload grammar and per-file
storage.
+
+All three address one reflected field as ``<type_key>.<field>``::
+
+ // 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 }
+
+``field`` sets the accessor'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.
+"""
+
+from __future__ import annotations
+
+import dataclasses
+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>[^{}]*)\})?$"
+)
+_MEMBER_RE = re.compile(r"^(?P<name>[A-Za-z_]\w*)\s*=\s*(?P<value>-?\d+)$")
+
+
[email protected](frozen=True)
+class EnumSpec:
+ """An ``enum`` directive: the newtype's name, its integer repr, and its
members."""
+
+ name: str
+ repr: str
+ members: tuple[tuple[str, int], ...]
+
+
[email protected]
+class Directives:
+ """The Rust directives of one file, keyed by ``<type_key>.<field>``."""
+
+ 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)
+
+ def add(self, name: str, payload: str, lineno: int) -> None:
+ """Parse and store one directive; raise ``ValueError`` on a malformed
payload."""
+ if name == "field":
+ target, rust_type = _split_arrow(name, payload, lineno)
+ self.field_types[target] = rust_type
+ elif name == "nullable":
+ self.nullable.add(_field_target(name, payload, lineno))
+ elif name == "enum":
+ target, spec = _parse_enum(payload, lineno)
+ self.enums[target] = spec
+ else:
+ raise ValueError(f"Unknown directive `{name}` at line {lineno}")
+
+
+def _invalid(name: str, lineno: int, expected: str) -> ValueError:
+ return ValueError(f"Invalid `{name}` directive at line {lineno}. Expected
`{expected}`")
+
+
+def _field_target(name: str, text: str, lineno: int) -> str:
+ """Validate a ``<type_key>.<field>`` reference."""
+ target = text.strip()
+ if not target or " " in target or "." not in target.strip("."):
+ raise _invalid(name, lineno, "<type_key>.<field>")
+ return target
+
+
+def _split_arrow(name: str, payload: str, lineno: int) -> tuple[str, str]:
+ """Split ``<type_key>.<field> -> <rust type>``."""
+ lhs, arrow, rhs = payload.partition("->")
+ if not arrow or not rhs.strip():
+ raise _invalid(name, lineno, "<type_key>.<field> -> <RustType>")
+ return _field_target(name, lhs, lineno), rhs.strip()
+
+
+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 }"
+ match = _ENUM_RE.match(payload.strip())
+ if match is None:
+ raise _invalid("enum", lineno, expected)
+ members: list[tuple[str, int]] = []
+ for item in (match.group("body") or "").split(","):
+ if not item.strip():
+ continue
+ member = _MEMBER_RE.match(item.strip())
+ if member is None:
+ raise _invalid("enum", lineno, expected)
+ members.append((member.group("name"), int(member.group("value"))))
+ target = _field_target("enum", match.group("target"), lineno)
+ return target, EnumSpec(match.group("name"), match.group("repr"),
tuple(members))
diff --git a/python/tvm_ffi/stub/rust_generator/generator.py
b/python/tvm_ffi/stub/rust_generator/generator.py
new file mode 100644
index 00000000..85b7e398
--- /dev/null
+++ b/python/tvm_ffi/stub/rust_generator/generator.py
@@ -0,0 +1,147 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""The Rust code generator for ``tvm-ffi-stubgen``.
+
+:class:`RustGenerator` implements the :class:`tvm_ffi.stub.generator.Generator`
+protocol, delegating the actual rendering to ``rust_generator.codegen``.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from .. import consts as C
+from . import consts as C_RUST
+from .codegen import (
+ finalize_rust_module_tree,
+ generate_rust_api_file,
+ generate_rust_import_section,
+ generate_rust_object,
+)
+from .utils import RustImports, RustUse
+
+if TYPE_CHECKING:
+ from pathlib import Path
+
+ from ..file_utils import CodeBlock
+ from ..utils import FuncInfo, InitConfig, ObjectInfo, Options
+
+
+class RustGenerator:
+ """Generator that emits opaque Rust bindings for reflected objects (see
:mod:`.codegen`)."""
+
+ name = "rust"
+ syntax = C.RUST_SYNTAX
+ source_exts = frozenset({".rs"})
+ directive_kinds = C_RUST.RUST_DIRECTIVE_KINDS
+
+ def default_ty_map(self) -> dict[str, str]:
+ """Return the default FFI-origin -> Rust-type name map."""
+ return C_RUST.RUST_TY_MAP_DEFAULTS.copy()
+
+ # --- per-file collection
--------------------------------------------------
+
+ def new_imports(self) -> RustImports:
+ """Create an empty per-file collector (``use`` items and
directives)."""
+ return RustImports()
+
+ def add_directive(self, imports: RustImports, name: str, payload: str,
lineno: int) -> None:
+ """Record ``import-object`` as a ``use`` (its path is the first ``;``
field); parse the rest."""
+ if name == "import-object":
+ imports.record(payload.split(";", 1)[0].strip())
+ else:
+ imports.directives.add(name, payload, lineno)
+
+ def canonical_type_name(self, type_key: str) -> str:
+ """Return the Rust path for a defined type key (matches
:attr:`RustUse.path`)."""
+ return RustUse(type_key).path
+
+ def extra_export_names(self, imports: RustImports) -> set[str]:
+ """No extra export names for Rust."""
+ return set()
+
+ # --- per-block generation
-------------------------------------------------
+
+ def generate_global_funcs_block(
+ self,
+ code: CodeBlock,
+ global_funcs: list[FuncInfo],
+ ty_map: dict[str, str],
+ imports: RustImports,
+ opt: Options,
+ ) -> None:
+ """No-op: Rust reaches global functions through
``cached_global_func!``."""
+
+ def generate_object_block(
+ self,
+ code: CodeBlock,
+ ty_map: dict[str, str],
+ imports: RustImports,
+ opt: Options,
+ obj_info: ObjectInfo,
+ ) -> None:
+ """Emit the opaque Rust binding for an ``object/<key>`` block."""
+ generate_rust_object(code, ty_map, imports, opt, obj_info)
+
+ def generate_import_section_block(
+ self, code: CodeBlock, imports: RustImports, opt: Options,
defined_types: set[str]
+ ) -> None:
+ """Emit Rust ``use`` statements for the collected imports."""
+ generate_rust_import_section(code, imports, opt, defined_types)
+
+ def generate_all_block(self, code: CodeBlock, names: set[str], opt:
Options) -> None:
+ """No-op: Rust re-exports are not generated."""
+
+ def generate_export_block(self, code: CodeBlock) -> None:
+ """No-op: submodule re-export is not generated."""
+
+ def generate_helpers_block(self, code: CodeBlock, opt: Options) -> None:
+ """No-op: the generated bindings need no per-file support code."""
+
+ # --- whole-file scaffolding (`--init` mode)
--------------------------------
+
+ def api_filename(self) -> str:
+ """One Rust file per module prefix."""
+ return "mod.rs"
+
+ def init_filename(self) -> str:
+ """No separate entry file for Rust; the API file is the module."""
+ return "mod.rs"
+
+ def generate_api_file(
+ self,
+ code_blocks: list[CodeBlock],
+ ty_map: dict[str, str],
+ module_name: str,
+ object_infos: list[ObjectInfo],
+ init_cfg: InitConfig,
+ is_root: bool,
+ ) -> str:
+ """Scaffold a Rust binding file: header + import/object markers."""
+ return generate_rust_api_file(
+ code_blocks, ty_map, module_name, object_infos, init_cfg, is_root,
self.syntax
+ )
+
+ def generate_init_file(
+ self, code_blocks: list[CodeBlock], module_name: str, submodule: str
+ ) -> str:
+ """No-op: the API file is the module entry."""
+ return ""
+
+ def finalize_init(self, init_path: Path, generated_prefixes: set[str]) ->
None:
+ """Auto-form the module tree: write ``pub mod <child>;``
declarations."""
+ finalize_rust_module_tree(init_path, generated_prefixes)
diff --git a/python/tvm_ffi/stub/rust_generator/utils.py
b/python/tvm_ffi/stub/rust_generator/utils.py
new file mode 100644
index 00000000..c60f52d7
--- /dev/null
+++ b/python/tvm_ffi/stub/rust_generator/utils.py
@@ -0,0 +1,142 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Rust generator helpers: ``use`` modelling, type rendering, identifier
spelling."""
+
+from __future__ import annotations
+
+import dataclasses
+from typing import TYPE_CHECKING, Callable
+
+from . import consts as C
+from .directives import Directives
+
+if TYPE_CHECKING:
+ from tvm_ffi.core import TypeSchema
+
+
[email protected](frozen=True, eq=True)
+class RustUse:
+ """A Rust ``use`` item: ``use <path>;``.
+
+ Dotted FFI names become ``::`` paths via :data:`~.consts.RUST_MOD_MAP`
+ (``ffi.String -> tvm_ffi::String``); bare names (``i64``) need no ``use``.
+ """
+
+ path: str
+
+ def __init__(self, name: str) -> None:
+ """Normalize ``name`` into a Rust ``use`` path and store it."""
+ if "::" not in name and "." in name:
+ head, _, tail = name.partition(".")
+ head = C.RUST_MOD_MAP.get(head, head)
+ name = f"{head}.{tail}"
+ object.__setattr__(self, "path", name.replace(".", "::"))
+
+ @property
+ def leaf(self) -> str:
+ """The final path segment (the in-scope name), e.g. ``Array`` for
``tvm_ffi::Array``."""
+ return self.path.rsplit("::", 1)[-1]
+
+ def as_use_line(self) -> str:
+ """Render the ``use`` statement, or ``""`` for a bare
prelude/primitive type."""
+ if "::" not in self.path:
+ return ""
+ return f"use {self.path};"
+
+
+def builtin_mirror_name(type_key: str) -> str:
+ """Name of the header-only stand-in for a builtin type (``ffi.IntEnum ->
FfiIntEnumObj``).
+
+ ``derive(Object)`` takes ``TYPE_DEPTH`` from the embedded base, so a type
+ under a builtin needs a base at the builtin's depth; the crate has none.
+ """
+ head, _, leaf = type_key.rpartition(".")
+ return f"{head.replace('.', '_').capitalize()}{leaf}Obj"
+
+
[email protected]
+class RustImports:
+ """The per-file collector: the ``use`` items and the Rust directives of
one file."""
+
+ items: list[RustUse] = dataclasses.field(default_factory=list)
+ directives: Directives = dataclasses.field(default_factory=Directives)
+ builtin_mirrors: dict[str, str] = dataclasses.field(default_factory=dict)
+ """Builtin ancestors mirrored in this file, root first: type key -> the
base each embeds."""
+
+ def record_builtin_base(self, chain: list[str]) -> str:
+ """Record the stand-ins for the builtin ``chain`` (root side first);
return the last name.
+
+ An empty chain (parent ``ffi.Object``) yields the crate's ``Object``.
+ """
+ base = self.record("tvm_ffi::Object")
+ for key in chain:
+ self.builtin_mirrors.setdefault(key, base)
+ base = builtin_mirror_name(key)
+ return base
+
+ def record(self, name: str) -> str:
+ """Record a ``use`` (deduped) and return the name to spell in code.
+
+ The leaf, or the full path when another path already claimed that leaf.
+ """
+ probe = RustUse(name)
+ if not probe.as_use_line():
+ return probe.leaf
+ for item in self.items:
+ if item.path == probe.path:
+ return item.leaf
+ if any(item.leaf == probe.leaf for item in self.items):
+ return probe.path
+ self.items.append(probe)
+ return probe.leaf
+
+
+def render_rust_type(schema: TypeSchema, ty_render: Callable[[str], str |
None]) -> str | None:
+ """Render ``schema`` as a Rust value type via ``ty_render`` (leaf origin
-> name), or ``None``."""
+ origin, args = schema.origin, schema.args
+ if origin in C.RUST_UNSUPPORTED_ORIGINS:
+ return None
+ if origin == "Array":
+ assert args # TypeSchema's post_init fills a missing element type.
+ return _generic(ty_render("Array"), render_rust_type(args[0],
ty_render))
+ if origin == "Map":
+ assert len(args) == 2 # TypeSchema's post_init fills a bare Map to
(Any, Any).
+ key = render_rust_type(args[0], ty_render)
+ value = render_rust_type(args[1], ty_render)
+ return _generic(ty_render("Map"), key, value)
+ if origin == "Optional":
+ (payload,) = args # TypeSchema's post_init enforces exactly one
argument.
+ return _generic("Option", render_rust_type(payload, ty_render))
+ if origin == "Callable":
+ return ty_render("Callable") # the crate's Function is type-erased
+ return ty_render(origin)
+
+
+def _generic(base: str | None, *params: str | None) -> str | None:
+ if base is None or any(p is None for p in params):
+ return None
+ return f"{base}<{', '.join(p for p in params if p is not None)}>"
+
+
+def rust_ident(name: str) -> str:
+ """Spell a reflected field name in Rust: drop the C++ trailing underscore,
escape keywords."""
+ name = name.rstrip("_") or name
+ if name in C.RUST_NOT_RAW_IDENTIFIERS:
+ return f"{name}_"
+ if name in C.RUST_KEYWORDS:
+ return f"r#{name}"
+ return name
diff --git a/tests/python/test_stubgen.py b/tests/python/test_stubgen.py
index c2606e46..712835cf 100644
--- a/tests/python/test_stubgen.py
+++ b/tests/python/test_stubgen.py
@@ -1105,9 +1105,9 @@ def
test_collect_files_filters_by_generator_exts(tmp_path: Path) -> None:
def test_generator_registry_names() -> None:
"""``--target`` choices follow the registered generators."""
- assert generator_names() == ["python"]
- with pytest.raises(ValueError, match="Known generators: python"):
- get_generator("rust")
+ assert generator_names() == ["python", "rust"]
+ with pytest.raises(ValueError, match="Known generators: python, rust"):
+ get_generator("zig")
def test_codeblock_from_begin_line_directive() -> None:
diff --git a/tests/python/test_stubgen_rust.py
b/tests/python/test_stubgen_rust.py
new file mode 100644
index 00000000..9f2e0fdf
--- /dev/null
+++ b/tests/python/test_stubgen_rust.py
@@ -0,0 +1,647 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# 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."""
+
+from __future__ import annotations
+
+import re
+from pathlib import Path
+
+import pytest
+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.file_utils import CodeBlock, FileInfo
+from tvm_ffi.stub.generator import get_generator
+from tvm_ffi.stub.rust_generator import consts as RC
+from tvm_ffi.stub.rust_generator.codegen import (
+ finalize_rust_module_tree,
+ generate_rust_api_file,
+ generate_rust_import_section,
+ generate_rust_object,
+)
+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
+
+RUST = get_generator("rust")
+
+
+def _info(
+ type_key: str,
+ fields: tuple[tuple[str, TypeSchema], ...] = (),
+ *,
+ parent: str | None = "ffi.Object",
+ ancestors: list[str] | None = None,
+ is_final: bool | 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],
+ methods=[],
+ type_key=type_key,
+ parent_type_key=parent,
+ ancestors=ancestors,
+ is_final=is_final,
+ )
+
+
+def _object_block(type_key: str) -> CodeBlock:
+ return CodeBlock(
+ kind="object",
+ param=type_key,
+ lineno_start=1,
+ lineno_end=2,
+ lines=[f"{C.RUST_SYNTAX.begin} object/{type_key}", C.RUST_SYNTAX.end],
+ )
+
+
+def _render(info: ObjectInfo, imports: RustImports | None = None) ->
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)
+ return "\n".join(block.lines[1:-1]), imports
+
+
+def _uses(imports: RustImports) -> set[str]:
+ return {item.path for item in imports.items}
+
+
+# ---------------------------------------------------------------------------
+# `use` modelling and type rendering
+# ---------------------------------------------------------------------------
+
+
+def test_rustuse_paths() -> None:
+ assert RustUse("tvm_ffi::Array").path == "tvm_ffi::Array"
+ assert RustUse("tvm_ffi::Array").leaf == "Array"
+ assert RustUse("tvm_ffi::Array").as_use_line() == "use tvm_ffi::Array;"
+ # A dotted FFI name becomes a path; the builtin `ffi` module maps to the
crate root.
+ assert RustUse("ffi.String").path == "tvm_ffi::String"
+ assert RustUse("my_pkg.sub.Foo").as_use_line() == "use my_pkg::sub::Foo;"
+ # Bare names need no `use`.
+ assert RustUse("i64").as_use_line() == ""
+
+
+def test_rustimports_record_dedups_and_resolves_collisions() -> None:
+ imports = RustImports()
+ assert imports.record("i64") == "i64"
+ assert imports.record("tvm_ffi::Array") == "Array"
+ assert imports.record("tvm_ffi::Array") == "Array"
+ assert imports.items == [RustUse("tvm_ffi::Array")]
+ # A second path wanting an already-claimed leaf is spelled in full, and
not recorded.
+ assert imports.record("other::Array") == "other::Array"
+ assert imports.items == [RustUse("tvm_ffi::Array")]
+
+
+def _render_type(schema: TypeSchema) -> tuple[str | None, RustImports]:
+ imports = RustImports()
+ ty_map = RC.RUST_TY_MAP_DEFAULTS
+
+ def ty_render(origin: str) -> str | None:
+ return imports.record(ty_map[origin]) if origin in ty_map else None
+
+ return render_rust_type(schema, ty_render), imports
+
+
+def test_render_rust_type_value_positions() -> None:
+ assert _render_type(TypeSchema("int"))[0] == "i64"
+ assert _render_type(TypeSchema("str"))[0] == "String"
+ assert _render_type(TypeSchema("Any"))[0] == "Any"
+ assert _render_type(TypeSchema("Callable", (TypeSchema("int"),)))[0] ==
"Function"
+ assert _render_type(TypeSchema("Optional", (TypeSchema("str"),)))[0] ==
"Option<String>"
+ text, imports = _render_type(TypeSchema("Map", (TypeSchema("str"),
TypeSchema("Array"))))
+ assert text == "Map<String, Array<Any>>"
+ assert _uses(imports) == {"tvm_ffi::Map", "tvm_ffi::String",
"tvm_ffi::Array", "tvm_ffi::Any"}
+
+
[email protected](
+ "schema",
+ [
+ TypeSchema("Union", (TypeSchema("int"), TypeSchema("str"))),
+ TypeSchema("Dict", (TypeSchema("str"), TypeSchema("int"))),
+ TypeSchema("tuple"),
+ TypeSchema("Array", (TypeSchema("List", (TypeSchema("int"),)),)),
+ TypeSchema("Optional", (TypeSchema("ctypes.c_void_p"),)),
+ ],
+)
+def test_render_rust_type_without_mirror(schema: TypeSchema) -> None:
+ assert _render_type(schema)[0] is None
+
+
+def test_rust_ident() -> None:
+ assert rust_ident("value") == "value"
+ assert rust_ident("imports_") == "imports"
+ assert rust_ident("type") == "r#type"
+ assert rust_ident("self") == "self_"
+ assert rust_ident("crate") == "crate_"
+
+
+# ---------------------------------------------------------------------------
+# Directives
+# ---------------------------------------------------------------------------
+
+
+def test_directives_parse() -> None:
+ directives = Directives()
+ directives.add("field", " tirx.Add.a -> PrimExpr ", 1)
+ 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)
+ 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", ()),
+ }
+
+
[email protected](
+ ("name", "payload", "expected"),
+ [
+ ("field", "tirx.Add.a", "-> <RustType>"),
+ ("field", "tirx.Add.a ->", "-> <RustType>"),
+ ("field", "Add -> PrimExpr", "<type_key>.<field>"),
+ ("nullable", "ir.Expr.span extra", "<type_key>.<field>"),
+ ("enum", "tirx.For.kind -> ForKind", "Name(i32)"),
+ ("enum", "tirx.For.kind -> ForKind(i128)", "Name(i32)"),
+ ("enum", "tirx.For.kind -> ForKind(i32) { Serial }", "Name(i32)"),
+ ("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:
+ 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"}
+ 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)
+ assert imports.items == [RustUse("tvm_ffi::libinfo::Foo")]
+ assert imports.directives.nullable == {"demo.Node.span"}
+
+
+# ---------------------------------------------------------------------------
+# Object rendering
+# ---------------------------------------------------------------------------
+
+ROOT_EXPECTED = """\
+#[repr(C)]
+#[derive(tvm_ffi::derive::Object)]
+#[type_key = "demo.Pair"]
+pub struct PairObj {
+ base: Object,
+}
+
+#[repr(C)]
+#[derive(tvm_ffi::derive::ObjectRef, Clone)]
+pub struct Pair {
+ data: ObjectArc<PairObj>,
+}
+
+impl Deref for Pair {
+ type Target = PairObj;
+ fn deref(&self) -> &PairObj {
+ &self.data
+ }
+}
+
+impl PairObj {
+ pub fn a(&self) -> Result<i64> {
+ FieldGetter::new(Self::type_index(), "a")?.get(self)
+ }
+
+ pub fn tag(&self) -> Result<Option<String>> {
+ FieldGetter::new(Self::type_index(), "tag")?.get(self)
+ }
+
+ pub fn items(&self) -> Result<Array<Any>> {
+ FieldGetter::new(Self::type_index(), "items")?.get(self)
+ }
+
+ pub fn owner(&self) -> Result<ObjectRef> {
+ FieldGetter::new(Self::type_index(), "owner")?.get(self)
+ }
+
+ pub fn r#type(&self) -> Result<Any> {
+ FieldGetter::new(Self::type_index(), "type")?.get_any(self)
+ }
+}"""
+
+
+def test_render_root_object() -> None:
+ """A root type embeds the header, gets no parent `Deref` and no upcast."""
+ 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")))),
+ ),
+ )
+ text, imports = _render(info)
+ assert text == ROOT_EXPECTED
+ assert _uses(imports) == {
+ "std::ops::Deref",
+ "tvm_ffi::Object",
+ "tvm_ffi::ObjectArc",
+ "tvm_ffi::ObjectCore",
+ "tvm_ffi::FieldGetter",
+ "tvm_ffi::Result",
+ "tvm_ffi::String",
+ "tvm_ffi::Array",
+ "tvm_ffi::Any",
+ "tvm_ffi::object::ObjectRef",
+ }
+
+
+def test_render_object_without_fields_has_no_impl() -> None:
+ text, imports = _render(_info("demo.Marker", is_final=True))
+ assert "#[type_final]" in text
+ assert "impl MarkerObj" not in text
+ assert "FieldGetter" not in _uses(imports)
+
+
+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")),),
+ parent="demo.Expr",
+ ancestors=["ffi.Object", "demo.BaseExpr", "demo.Expr"],
+ is_final=True,
+ )
+ text, imports = _render(info)
+ assert "#[type_final]\npub struct AddObj {\n base: ExprObj,\n}" in text
+ assert "impl Deref for AddObj {\n type Target = ExprObj;" in text
+ assert "pub fn a(&self) -> Result<Expr> {" in text
+ assert text.endswith("tvm_ffi::impl_object_upcast!(Add => BaseExpr, Add =>
Expr);")
+ # Same-module names are local items: nothing to `use`.
+ assert not any(path.startswith("demo") for path in _uses(imports))
+
+
+def test_render_derived_object_cross_module() -> None:
+ """A parent in another module is reached through the generated root."""
+ info = _info("tirx.Add", parent="ir.Expr", ancestors=["ffi.Object",
"ir.Expr"])
+ text, imports = _render(info)
+ assert " base: ExprObj," in text
+ assert text.endswith("tvm_ffi::impl_object_upcast!(Add => Expr);")
+ assert {"super::ir::ExprObj", "super::ir::Expr"} <= _uses(imports)
+
+
+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")),),
+ parent="ffi.IntEnum",
+ ancestors=["ffi.Object", "ffi.Enum", "ffi.IntEnum"],
+ )
+ text, imports = _render(info)
+ assert " base: FfiIntEnumObj," in text
+ # One stand-in per builtin ancestor below `ffi.Object`, chained through
`base`.
+ assert imports.builtin_mirrors == {"ffi.Enum": "Object", "ffi.IntEnum":
"FfiEnumObj"}
+ assert "tvm_ffi::Object" in _uses(imports)
+ # No Deref, upcast or crate import for the stand-ins.
+ assert "impl Deref for ColorObj" not in text
+ assert "impl_object_upcast" not in text
+ assert not any("Enum" in path for path in _uses(imports))
+ # A child of `ffi.Object` embeds the crate's header.
+ text, imports = _render(_info("demo.Root"))
+ assert " base: Object," in text
+ assert imports.builtin_mirrors == {}
+ # A generated parent is embedded by name.
+ red = _info(
+ "demo.Red",
+ parent="demo.Color",
+ ancestors=["ffi.Object", "ffi.Enum", "ffi.IntEnum", "demo.Color"],
+ )
+ text, imports = _render(red)
+ assert " base: ColorObj," in text
+ assert imports.builtin_mirrors == {}
+
+
+BUILTIN_MIRRORS_EXPECTED = """\
+use std::ops::Deref;
+use tvm_ffi::Object;
+use tvm_ffi::ObjectArc;
+
+/// Header-only stand-in for the builtin `ffi.Enum`; it only carries the
ancestor depth.
+#[allow(dead_code)]
+#[repr(C)]
+#[derive(tvm_ffi::derive::Object)]
+#[type_key = "ffi.Enum"]
+struct FfiEnumObj {
+ base: Object,
+}
+
+/// Header-only stand-in for the builtin `ffi.IntEnum`; it only carries the
ancestor depth.
+#[allow(dead_code)]
+#[repr(C)]
+#[derive(tvm_ffi::derive::Object)]
+#[type_key = "ffi.IntEnum"]
+struct FfiIntEnumObj {
+ base: FfiEnumObj,
+}
+
+/// Header-only stand-in for the builtin `ffi.StrEnum`; it only carries the
ancestor depth.
+#[allow(dead_code)]
+#[repr(C)]
+#[derive(tvm_ffi::derive::Object)]
+#[type_key = "ffi.StrEnum"]
+struct FfiStrEnumObj {
+ base: FfiEnumObj,
+}"""
+
+
+def test_import_section_defines_builtin_mirrors_once() -> None:
+ """Objects sharing builtin ancestors share one mirror chain, rendered
after the `use`s."""
+ imports = RustImports()
+ enum_chain = ["ffi.Object", "ffi.Enum"]
+ for type_key, parent in (
+ ("demo.Color", "ffi.IntEnum"),
+ ("demo.Mode", "ffi.IntEnum"),
+ ("demo.Flag", "ffi.Enum"),
+ ("demo.Op", "ffi.StrEnum"),
+ ):
+ ancestors = enum_chain if parent == "ffi.Enum" else [*enum_chain,
parent]
+ _render(_info(type_key, parent=parent, ancestors=ancestors), imports)
+ block = CodeBlock(
+ kind="import-section",
+ param="",
+ lineno_start=1,
+ lineno_end=2,
+ lines=[f"{C.RUST_SYNTAX.begin} import-section", C.RUST_SYNTAX.end],
+ )
+ generate_rust_import_section(block, imports, Options(),
defined_types=set())
+ assert "\n".join(block.lines[1:-1]) == BUILTIN_MIRRORS_EXPECTED
+
+
+ITER_VAR_EXPECTED = """\
+#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
+#[repr(transparent)]
+pub struct IterVarType(i32);
+
+#[allow(non_upper_case_globals)]
+impl IterVarType {
+ pub const kDataPar: Self = Self(0);
+ pub const kThreadIndex: Self = Self(1);
+ pub const fn from_raw(value: i32) -> Self {
+ Self(value)
+ }
+ pub const fn as_raw(self) -> i32 {
+ self.0
+ }
+}
+
+impl TryFrom<i64> for IterVarType {
+ type Error = Error;
+ fn try_from(value: i64) -> Result<Self> {
+ i32::try_from(value).map(Self).map_err(|_| {
+ Error::new(VALUE_ERROR, &format!("IterVarType value {value} does
not fit i32"), "")
+ })
+ }
+}
+
+#[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<IterVarObj>,
+}
+
+impl Deref for IterVar {
+ type Target = IterVarObj;
+ fn deref(&self) -> &IterVarObj {
+ &self.data
+ }
+}
+
+impl Deref for IterVarObj {
+ type Target = PrimExprConvertibleObj;
+ fn deref(&self) -> &PrimExprConvertibleObj {
+ &self.base
+ }
+}
+
+impl IterVarObj {
+ pub fn dom(&self) -> Result<Option<Range>> {
+ FieldGetter::new(Self::type_index(), "dom")?.get(self)
+ }
+
+ pub fn var(&self) -> Result<PrimVar> {
+ FieldGetter::new(Self::type_index(), "var")?.get(self)
+ }
+
+ pub fn iter_type(&self) -> Result<IterVarType> {
+ let raw: i64 = FieldGetter::new(Self::type_index(),
"iter_type")?.get(self)?;
+ IterVarType::try_from(raw)
+ }
+
+ pub fn thread_tag(&self) -> Result<String> {
+ FieldGetter::new(Self::type_index(), "thread_tag")?.get(self)
+ }
+
+ pub fn span(&self) -> Result<Option<Span>> {
+ FieldGetter::new(Self::type_index(), "span")?.get(self)
+ }
+}
+
+tvm_ffi::impl_object_upcast!(IterVar => PrimExprConvertible);"""
+
+
+def test_render_iter_var_golden() -> None:
+ """The shape tvm-rust-ext hand-writes for its polymorphic `IterVar`,
driven by directives."""
+ info = _info(
+ "tirx.IterVar",
+ (
+ ("dom", TypeSchema("ir.Range")),
+ ("var", TypeSchema("ir.Var")),
+ ("iter_type", TypeSchema("int")),
+ ("thread_tag", TypeSchema("str")),
+ ("span", TypeSchema("ir.Span")),
+ ),
+ parent="ir.PrimExprConvertible",
+ is_final=True,
+ )
+ imports = RUST.new_imports()
+ RUST.add_directive(imports, "nullable", "tirx.IterVar.dom", 1)
+ RUST.add_directive(imports, "field", "tirx.IterVar.var -> PrimVar", 2)
+ RUST.add_directive(imports, "nullable", "tirx.IterVar.span", 3)
+ RUST.add_directive(
+ imports,
+ "enum",
+ "tirx.IterVar.iter_type -> IterVarType(i32) { kDataPar=0,
kThreadIndex=1 }",
+ 4,
+ )
+ text, imports = _render(info, imports)
+ assert text == ITER_VAR_EXPECTED
+ assert {
+ "tvm_ffi::Error",
+ "tvm_ffi::VALUE_ERROR",
+ "super::ir::Range",
+ "super::ir::Span",
+ } <= _uses(imports)
+
+
+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"),)))),
+ )
+ imports = RUST.new_imports()
+ RUST.add_directive(imports, "field", "demo.Node.buffer ->
crate::typed::BufferVar", 1)
+ RUST.add_directive(imports, "nullable", "demo.Node.dom", 2)
+ text, imports = _render(info, imports)
+ assert "pub fn buffer(&self) -> Result<BufferVar> {" in text
+ assert "pub fn dom(&self) -> Result<Option<i64>> {" in text
+ assert "crate::typed::BufferVar" in _uses(imports)
+
+
+# ---------------------------------------------------------------------------
+# File scaffolding
+# ---------------------------------------------------------------------------
+
+
+def test_import_section_dedups_sorts_and_filters_defined_types() -> None:
+ imports = RustImports()
+ for name in ("tvm_ffi::ObjectArc", "std::ops::Deref",
"tvm_ffi::ObjectArc", "super::ir::Expr"):
+ imports.record(name)
+ block = CodeBlock(
+ kind="import-section",
+ param="",
+ lineno_start=1,
+ lineno_end=2,
+ lines=[f"{C.RUST_SYNTAX.begin} import-section", C.RUST_SYNTAX.end],
+ )
+ generate_rust_import_section(block, imports, Options(),
defined_types={"super::ir::Expr"})
+ assert block.lines[1:-1] == ["use std::ops::Deref;", "use
tvm_ffi::ObjectArc;"]
+
+
+def test_api_file_scaffold() -> None:
+ infos = [_info("demo.A"), _info("demo.B")]
+ cfg = InitConfig(pkg="demo", shared_target="demo_shared", prefix="demo.")
+ text = generate_rust_api_file([], {}, "demo", infos, cfg, True,
C.RUST_SYNTAX)
+ assert text.startswith("#![allow(dead_code, unused_imports)]\n")
+ assert "//! FFI bindings for `demo` (generated by tvm-ffi-stubgen)." in
text
+ assert f"{C.RUST_SYNTAX.begin} import-section\n{C.RUST_SYNTAX.end}" in text
+ assert text.count(f"{C.RUST_SYNTAX.begin} object/demo.") == 2
+ # Existing blocks are not re-scaffolded.
+ again = generate_rust_api_file(
+ [_object_block("demo.A")], {}, "demo", infos, cfg, True, C.RUST_SYNTAX
+ )
+ assert "object/demo.A" not in again and "object/demo.B" in again
+ assert RUST.api_filename() == RUST.init_filename() == "mod.rs"
+ assert RUST.generate_init_file([], "demo", "mod") == ""
+
+
+def test_finalize_module_tree(tmp_path: Path) -> None:
+ (tmp_path / "ir").mkdir()
+ (tmp_path / "ir" / "mod.rs").write_text("pub struct Existing;\n",
encoding="utf-8")
+ finalize_rust_module_tree(tmp_path, {"ir", "tirx.transform"})
+ assert (tmp_path / "mod.rs").read_text(encoding="utf-8") == "pub mod
ir;\npub mod tirx;\n"
+ assert (tmp_path / "tirx" / "mod.rs").read_text(encoding="utf-8") == "pub
mod transform;\n"
+ assert (tmp_path / "ir" / "mod.rs").read_text(encoding="utf-8") == "pub
struct Existing;\n"
+ finalize_rust_module_tree(tmp_path, {"ir", "tirx.transform"}) # idempotent
+ assert (tmp_path / "mod.rs").read_text(encoding="utf-8") == "pub mod
ir;\npub mod tirx;\n"
+
+
+# ---------------------------------------------------------------------------
+# The pipeline end to end
+# ---------------------------------------------------------------------------
+
+
+def test_stage_3_applies_directives_to_a_registered_type(tmp_path: Path) ->
None:
+ src = tmp_path / "mod.rs"
+ src.write_text(
+ "\n".join(
+ [
+ f"{C.RUST_SYNTAX.begin} import-section",
+ C.RUST_SYNTAX.end,
+ f"{C.RUST_SYNTAX.directive('enum')}
testing.TestCxxClassBase.v_i32 -> Kind(i32) {{ A=0 }}",
+ f"{C.RUST_SYNTAX.begin} object/testing.TestCxxClassBase",
+ C.RUST_SYNTAX.end,
+ "",
+ ]
+ ),
+ encoding="utf-8",
+ )
+ info = FileInfo.from_file(src)
+ assert info is not 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<Kind> {" in text
+ assert "pub fn v_i64(&self) -> Result<i64> {" in text
+ assert "use tvm_ffi::FieldGetter;" in text
+
+
+def test_cli_init_generates_a_module_tree(tmp_path: Path, monkeypatch:
pytest.MonkeyPatch) -> None:
+ monkeypatch.setattr(
+ "sys.argv",
+ [
+ "tvm-ffi-stubgen",
+ "--target",
+ "rust",
+ "--init-pypkg",
+ "demo",
+ "--init-lib",
+ "demo_shared",
+ "--init-prefix",
+ "testing.",
+ str(tmp_path),
+ ],
+ )
+ assert stub_cli.__main__() == 0
+ 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 '#[type_key = "testing.TestCxxClassDerivedDerived"]' in text
+ assert " base: TestCxxClassDerivedObj," in text
+ assert (
+ "tvm_ffi::impl_object_upcast!(TestCxxClassDerivedDerived =>
TestCxxClassBase, "
+ "TestCxxClassDerivedDerived => TestCxxClassDerived);"
+ ) in text
+ # Builtin ancestors are mirrored once, in the import section; enum
fixtures embed the last one.
+ assert text.count("struct FfiEnumObj {\n base: Object,\n}") == 1
+ assert text.count("struct FfiIntEnumObj {\n base: FfiEnumObj,\n}") == 1
+ assert text.count("struct FfiStrEnumObj {\n base: FfiEnumObj,\n}") == 1
+ assert text.index("struct FfiIntEnumObj") <
text.index(f"{C.RUST_SYNTAX.begin} object/")
+ assert "pub struct TestEnumVariantObj {\n base: FfiEnumObj,\n}" in text
+ assert "pub struct TestCxxIntEnumObj {\n base: FfiIntEnumObj,\n}" in
text
+ assert "pub struct TestCxxStrEnumObj {\n base: FfiStrEnumObj,\n}" in
text
+ # 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
