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 ad6fc947 [FEAT][Rust] Record type schemas and doc strings of exported
functions (#832)
ad6fc947 is described below
commit ad6fc94765b39d2e27b360d0cd0a1cac23231fdf
Author: Gabriel Wu <[email protected]>
AuthorDate: Tue Oct 27 03:45:50 2026 +0800
[FEAT][Rust] Record type schemas and doc strings of exported functions
(#832)
- **Type schemas.** `AnyCompatible::type_schema()` mirrors C++
`TypeTraits<T>::TypeSchema`. Its default is `{"type":"<type_str>"}`,
which is right for object references and for most builtins, whose type
string is their type key. `Option`, `Array`, `Map`, `BigInt`,
`Unchanged` and `UnchangedOr` override it to write theirs as C++ does.
`RValueRef`, `Any` and `AnyView` have their schemas as parameters or
return types.
- **Function schemas.** `AsPackedCallable::type_schema()` gives a typed
function's schema,
`{"type":"ffi.Function","named_args":{"return":[...],"params":[...]}}`,
as C++ `FunctionInfo::TypeSchema` writes it. A callable whose signature
the crate cannot see, such as a downstream `AsPackedCallable` impl for a
packed function, keeps the default method, which returns the untyped
schema that `def_packed` records. Existing impls keep compiling.
- **Exports, behind a feature.** The new `export-metadata` cargo feature
is the counterpart of C++ `TVM_FFI_DLL_EXPORT_INCLUDE_METADATA`, and
like it is off by default. With the feature,
`tvm_ffi_dll_export_typed_func!` also exports
`__tvm_ffi__metadata_<name>`, and the new
`tvm_ffi_dll_export_typed_func_doc!(name, "doc")` exports
`__tvm_ffi__doc_<name>`, like `TVM_FFI_DLL_EXPORT_TYPED_FUNC_DOC`.
Without the feature, the export is unchanged and the doc macro exports
nothing, as in C++; it still checks that the doc string is a string
constant. A `#[cfg]` emitted by a `macro_rules!` expansion would test
the features of the calling crate, so the export macros call hidden
helper macros that tvm-ffi defines twice, once per state of its own
feature. As in C++, both getters return a string that the runtime
allocates (`TVMFFIStringFromByteArray`), so it outlives the library.
- **Globals, always.** `Function::register_global_typed(name, func,
doc)` registers a typed function with its schema and doc string, like
`GlobalDef().def`. `Function::register_global` now records the untyped
schema, like `def_packed`. Both register through
`TVMFFIFunctionSetGlobalFromMethodInfo`, which tvm-ffi-sys now binds.
- **Byte-identical JSON.** The metadata escapes the schema string with
the runtime's JSON writer (`ffi.json.Stringify`), which uses the same
`EscapeStringJSON` as the C++ macros. The Rust metadata therefore
matches the C++ metadata byte for byte, and Rust needs no JSON escaper
of its own.
- **Docs and CI.** The Rust guide documents typed globals, exports and
the feature. CI's Rust step also runs `cargo test -p tvm-ffi --features
export-metadata`, so the export tests run.
**API note:** a typed closure passed to `Function::from_typed` (and now
to the export macro) needs a return type with a schema. That means an
`AnyCompatible` type, `Any` or `AnyView`, which covers every return type
the crate converts today. Before, any type with `From<Out> for Any` was
accepted.
---------
Co-authored-by: Yuchuan <[email protected]>
---
.github/workflows/ci_test.yml | 1 +
docs/guides/rust_lang_guide.md | 41 ++++
rust/tvm-ffi-sys/src/c_api.rs | 4 +
rust/tvm-ffi/Cargo.toml | 6 +-
rust/tvm-ffi/src/any.rs | 2 +-
rust/tvm-ffi/src/big_int/mod.rs | 4 +
rust/tvm-ffi/src/collections/array.rs | 7 +
rust/tvm-ffi/src/collections/list.rs | 7 +
rust/tvm-ffi/src/collections/map.rs | 7 +
rust/tvm-ffi/src/extra/unchanged.rs | 8 +
rust/tvm-ffi/src/function.rs | 59 +++++-
rust/tvm-ffi/src/function_internal.rs | 72 ++++++-
rust/tvm-ffi/src/macros.rs | 109 ++++++++++
rust/tvm-ffi/src/reflection/registry.rs | 154 +++++---------
rust/tvm-ffi/src/rvalue_ref.rs | 9 +
rust/tvm-ffi/src/type_traits.rs | 93 +++++++++
rust/tvm-ffi/tests/test_function_metadata.rs | 291 +++++++++++++++++++++++++++
17 files changed, 753 insertions(+), 121 deletions(-)
diff --git a/.github/workflows/ci_test.yml b/.github/workflows/ci_test.yml
index 7a1b735f..32d994eb 100644
--- a/.github/workflows/ci_test.yml
+++ b/.github/workflows/ci_test.yml
@@ -208,6 +208,7 @@ jobs:
working-directory: rust
run: |
cargo test
+ cargo test -p tvm-ffi --features export-metadata
# Build the Rust stubgen example against the installed package. Its
post-build
# step refreshes the bindings in place, so a diff means the checked-in
file was stale.
diff --git a/docs/guides/rust_lang_guide.md b/docs/guides/rust_lang_guide.md
index 32515785..8309bfd7 100644
--- a/docs/guides/rust_lang_guide.md
+++ b/docs/guides/rust_lang_guide.md
@@ -157,6 +157,47 @@ let my_func = Function::from_packed(|args: &[AnyView]| ->
Result<Any> {
Ok(Any::default())
});
Function::register_global("my_custom_func", my_func)?;
+
+// Register a typed function with a doc string
+Function::register_global_typed(
+ "my_add",
+ |x: i64, y: i64| -> Result<i64> { Ok(x + y) },
+ "Add two integers.",
+)?;
+```
+
+As C++ `refl::GlobalDef().def` does, a typed function's metadata records its
+type schema, so `tvm_ffi.get_global_func_metadata` and the stub generator see
+its parameter and return types. `Function::register_global` records the schema
+of an untyped function, as `def_packed` does.
+
+### Exporting Functions from a Library
+
+A Rust `cdylib` exports a typed function under the `__tvm_ffi_<name>` symbol
+that `Module::load_from_file` and `tvm_ffi.load_module` look up:
+
+```rust
+use tvm_ffi::*;
+
+fn add(a: i64, b: i64) -> Result<i64> {
+ Ok(a + b)
+}
+
+tvm_ffi_dll_export_typed_func!(add, add);
+tvm_ffi_dll_export_typed_func_doc!(add, "Add two integers and return the
sum.");
+```
+
+With the `export-metadata` feature of the `tvm-ffi` crate, the counterpart of
+C++ `TVM_FFI_DLL_EXPORT_INCLUDE_METADATA`, the export also writes the
+function's type schema as `__tvm_ffi__metadata_<name>`, and
+`tvm_ffi_dll_export_typed_func_doc!` writes its doc string as
+`__tvm_ffi__doc_<name>`, which `Module.get_function_metadata` and
+`Module.get_function_doc` read. As in C++, it is off by default, and without
+it `tvm_ffi_dll_export_typed_func_doc!` exports nothing:
+
+```toml
+[dependencies]
+tvm-ffi = { version = "...", features = ["export-metadata"] }
```
### Reflected Type Methods
diff --git a/rust/tvm-ffi-sys/src/c_api.rs b/rust/tvm-ffi-sys/src/c_api.rs
index 86f5e431..86ab0a51 100644
--- a/rust/tvm-ffi-sys/src/c_api.rs
+++ b/rust/tvm-ffi-sys/src/c_api.rs
@@ -494,6 +494,10 @@ unsafe extern "C" {
f: TVMFFIObjectHandle,
can_override: i32,
) -> i32;
+ pub fn TVMFFIFunctionSetGlobalFromMethodInfo(
+ method_info: *const TVMFFIMethodInfo,
+ allow_override: i32,
+ ) -> i32;
pub fn TVMFFIFunctionCreate(
self_ptr: *mut c_void,
safe_call: TVMFFISafeCallType,
diff --git a/rust/tvm-ffi/Cargo.toml b/rust/tvm-ffi/Cargo.toml
index 1fbd1c8a..7adf105f 100644
--- a/rust/tvm-ffi/Cargo.toml
+++ b/rust/tvm-ffi/Cargo.toml
@@ -39,8 +39,12 @@ tvm-ffi-macros = { version = "0.1.0-alpha.0", path =
"../tvm-ffi-macros" }
# The tests call the `testing.*` global functions.
tvm-ffi-sys = { version = "0.1.0-alpha.0", path = "../tvm-ffi-sys", features =
["testing"] }
-# only needed for example building
[features]
+# Export the metadata and doc string of each function that
+# `tvm_ffi_dll_export_typed_func!` and `tvm_ffi_dll_export_typed_func_doc!`
+# export, as C++ does with TVM_FFI_DLL_EXPORT_INCLUDE_METADATA.
+export-metadata = []
+# only needed for example building
example = []
[[example]]
diff --git a/rust/tvm-ffi/src/any.rs b/rust/tvm-ffi/src/any.rs
index 8b0d714c..1c5fc500 100644
--- a/rust/tvm-ffi/src/any.rs
+++ b/rust/tvm-ffi/src/any.rs
@@ -428,7 +428,7 @@ crate::impl_try_from_any_for_parametric!(Option<T>);
//------------------------------------------------------------
/// ArgTryFromAnyView: Helper for function argument passing
///-----------------------------------------------------------
-pub(crate) trait ArgTryFromAnyView: Sized {
+pub(crate) trait ArgTryFromAnyView: Sized + crate::type_traits::TypeSchema {
fn try_from_any_view(value: &AnyView, arg_index: usize) -> Result<Self,
Error>;
}
diff --git a/rust/tvm-ffi/src/big_int/mod.rs b/rust/tvm-ffi/src/big_int/mod.rs
index b0380d28..652dbd41 100644
--- a/rust/tvm-ffi/src/big_int/mod.rs
+++ b/rust/tvm-ffi/src/big_int/mod.rs
@@ -825,6 +825,10 @@ unsafe impl AnyCompatible for BigInt {
"BigInt".to_string()
}
+ fn type_schema() -> std::string::String {
+ crate::type_traits::type_schema(<BigIntObj as
crate::ObjectCore>::TYPE_KEY, &[])
+ }
+
unsafe fn copy_to_any_view(this: &Self, data: &mut TVMFFIAny) {
*data = this.data;
}
diff --git a/rust/tvm-ffi/src/collections/array.rs
b/rust/tvm-ffi/src/collections/array.rs
index 0a842bf8..5d6bb115 100644
--- a/rust/tvm-ffi/src/collections/array.rs
+++ b/rust/tvm-ffi/src/collections/array.rs
@@ -243,6 +243,13 @@ where
format!("Array<{}>", T::container_type_str())
}
+ fn type_schema() -> String {
+ crate::type_traits::type_schema(
+ <ArrayObj as crate::ObjectCore>::TYPE_KEY,
+ &[T::container_type_schema()],
+ )
+ }
+
unsafe fn check_any_strict(data: &TVMFFIAny) -> bool {
if data.type_index != TypeIndex::kTVMFFIArray as i32 {
return false;
diff --git a/rust/tvm-ffi/src/collections/list.rs
b/rust/tvm-ffi/src/collections/list.rs
index 68add977..2aabfacd 100644
--- a/rust/tvm-ffi/src/collections/list.rs
+++ b/rust/tvm-ffi/src/collections/list.rs
@@ -565,6 +565,13 @@ unsafe impl<T: ContainerElement> AnyCompatible for List<T>
{
format!("List<{}>", T::container_type_str())
}
+ fn type_schema() -> String {
+ crate::type_traits::type_schema(
+ <ListObj as crate::ObjectCore>::TYPE_KEY,
+ &[T::container_type_schema()],
+ )
+ }
+
unsafe fn check_any_strict(data: &TVMFFIAny) -> bool {
data.type_index == TypeIndex::kTVMFFIList as i32
&& seq_cells(data)
diff --git a/rust/tvm-ffi/src/collections/map.rs
b/rust/tvm-ffi/src/collections/map.rs
index 8a1a425d..c7da95eb 100644
--- a/rust/tvm-ffi/src/collections/map.rs
+++ b/rust/tvm-ffi/src/collections/map.rs
@@ -467,6 +467,13 @@ where
)
}
+ fn type_schema() -> String {
+ crate::type_traits::type_schema(
+ <MapObj as crate::ObjectCore>::TYPE_KEY,
+ &[K::container_type_schema(), V::container_type_schema()],
+ )
+ }
+
unsafe fn check_any_strict(data: &TVMFFIAny) -> bool {
// Mirrors C++ `CheckAnyStrict` (map_base.h): every entry must strictly
// match `K`/`V`. An FFI failure reads as "no match" — this fn has no
diff --git a/rust/tvm-ffi/src/extra/unchanged.rs
b/rust/tvm-ffi/src/extra/unchanged.rs
index cbf0f140..0ebdccad 100644
--- a/rust/tvm-ffi/src/extra/unchanged.rs
+++ b/rust/tvm-ffi/src/extra/unchanged.rs
@@ -63,6 +63,10 @@ unsafe impl AnyCompatible for Unchanged {
fn type_str() -> std::string::String {
"Unchanged".into()
}
+
+ fn type_schema() -> std::string::String {
+ crate::type_traits::type_schema("ffi.Unchanged", &[])
+ }
}
crate::impl_try_from_any!(Unchanged);
@@ -270,6 +274,10 @@ unsafe impl<T: ContainerElement> AnyCompatible for
UnchangedOr<T> {
fn type_str() -> std::string::String {
format!("UnchangedOr<{}>", T::container_type_str())
}
+
+ fn type_schema() -> std::string::String {
+ crate::type_traits::type_schema("UnchangedOr",
&[T::container_type_schema()])
+ }
}
impl<T: ContainerElement> TryFrom<Any> for UnchangedOr<T> {
diff --git a/rust/tvm-ffi/src/function.rs b/rust/tvm-ffi/src/function.rs
index 2c7b1827..f375ba7e 100644
--- a/rust/tvm-ffi/src/function.rs
+++ b/rust/tvm-ffi/src/function.rs
@@ -24,8 +24,8 @@ use crate::object::{Object, ObjectArc, ObjectCore};
use crate::type_traits::AnyCompatible;
use tvm_ffi_sys::{
TVMFFIAny, TVMFFIByteArray, TVMFFIFunctionCell, TVMFFIFunctionCreate,
TVMFFIFunctionGetGlobal,
- TVMFFIFunctionSetGlobal, TVMFFIGetTypeInfo, TVMFFIObjectHandle,
TVMFFISafeCallType,
- TVMFFITypeIndex, TVMFFITypeKeyToIndex,
+ TVMFFIFunctionSetGlobalFromMethodInfo, TVMFFIGetTypeInfo,
TVMFFIMethodInfo, TVMFFIObjectHandle,
+ TVMFFISafeCallType, TVMFFITypeIndex, TVMFFITypeKeyToIndex,
};
/// function object
@@ -298,6 +298,12 @@ impl Function {
}
/// Register a function as a global function
+ ///
+ /// The function's metadata records the type schema of an untyped
+ /// function, as C++ `refl::GlobalDef().def_packed` does. Use
+ /// [`Function::register_global_typed`] to record the parameter and return
+ /// types of a typed function, and a doc string.
+ ///
/// # Arguments
/// * `name` - The name of the function
/// * `func` - The function to register
@@ -305,15 +311,50 @@ impl Function {
/// # Returns
/// * `Result<()>` - The result of the registration
pub fn register_global(name: &str, func: Function) -> Result<()> {
+ let type_schema =
crate::type_traits::type_schema(FunctionObj::TYPE_KEY, &[]);
+ Self::register_global_with_info(name, &func, &type_schema, "")
+ }
+
+ /// Register a typed function as a global function, with a doc string
+ ///
+ /// As C++ `refl::GlobalDef().def` does, the function's metadata records
+ /// its type schema, which `tvm_ffi.get_global_func_metadata` and the stub
+ /// generator read. An empty `doc` records no doc string.
+ ///
+ /// # Arguments
+ /// * `name` - The name of the function
+ /// * `func` - The typed function to register, as for
[`Function::from_typed`]
+ /// * `doc` - The doc string of the function
+ ///
+ /// # Returns
+ /// * `Result<()>` - The result of the registration
+ pub fn register_global_typed<F, I, O>(name: &str, func: F, doc: &str) ->
Result<()>
+ where
+ F: AsPackedCallable<I, O> + 'static,
+ {
+ let type_schema = F::type_schema();
+ Self::register_global_with_info(name, &Self::from_typed(func),
&type_schema, doc)
+ }
+
+ fn register_global_with_info(
+ name: &str,
+ func: &Function,
+ type_schema: &str,
+ doc: &str,
+ ) -> Result<()> {
+ let metadata = crate::type_traits::type_schema_metadata(type_schema);
+ let mut method = TVMFFIAny::new();
unsafe {
- let name_arg = TVMFFIByteArray::from_str(name);
+ <Function as AnyCompatible>::copy_to_any_view(func, &mut method);
+ let info = TVMFFIMethodInfo {
+ name: TVMFFIByteArray::from_str(name),
+ doc: TVMFFIByteArray::from_str(doc),
+ metadata: TVMFFIByteArray::from_str(&metadata),
+ flags: 0,
+ method,
+ };
let can_override = 0;
- crate::check_safe_call!(TVMFFIFunctionSetGlobal(
- &name_arg,
- ObjectArc::as_raw(&func.data) as *mut FunctionObj as
TVMFFIObjectHandle,
- can_override
- ))?;
- Ok(())
+
crate::check_safe_call!(TVMFFIFunctionSetGlobalFromMethodInfo(&info,
can_override))
}
}
/// Construct a function from a packed function
diff --git a/rust/tvm-ffi/src/function_internal.rs
b/rust/tvm-ffi/src/function_internal.rs
index 33c9eb47..8f7e5c87 100644
--- a/rust/tvm-ffi/src/function_internal.rs
+++ b/rust/tvm-ffi/src/function_internal.rs
@@ -18,10 +18,12 @@
*/
use crate::any::{Any, AnyView, ArgTryFromAnyView};
use crate::error::{Error, Result, INTERNAL_ERROR};
-use crate::object::ObjectRefCore;
+use crate::function::FunctionObj;
+use crate::object::{ObjectCore, ObjectRefCore};
use crate::rvalue_ref::RValueRef;
use crate::string::{Bytes, String};
-use crate::type_traits::{AnyCompatible, ContainerElement};
+use crate::type_traits::{AnyCompatible, ContainerElement, TypeSchema};
+use tvm_ffi_sys::{TVMFFIAny, TVMFFIByteArray, TVMFFIStringFromByteArray};
//------------------------------------------------------------------------
// PackedCallable
@@ -45,6 +47,61 @@ pub fn panic_to_error(payload: Box<dyn std::any::Any +
Send>) -> Error {
pub trait AsPackedCallable<I, O> {
// Call the function in packed convention
fn call_packed(&self, packed_args: &[AnyView]) -> Result<Any>;
+
+ /// The type schema of the function, as C++ `FunctionInfo::TypeSchema`
+ /// writes it.
+ ///
+ /// A typed function lists its return and parameter types. The default,
+ /// for a callable whose signature is not known, is the schema of an
+ /// untyped function, as C++ `refl::GlobalDef().def_packed` records it.
+ fn type_schema() -> std::string::String
+ where
+ Self: Sized,
+ {
+ crate::type_traits::type_schema(FunctionObj::TYPE_KEY, &[])
+ }
+}
+
+/// The type schema of `func`, a value of a callable type.
+#[inline]
+pub fn type_schema_of<Fun, I, O>(_func: &Fun) -> std::string::String
+where
+ Fun: AsPackedCallable<I, O>,
+{
+ Fun::type_schema()
+}
+
+/// Writes `text` to `result` as a string allocated by the runtime, so that it
+/// outlives the library that wrote it, as the getters that
+/// `TVM_FFI_DLL_EXPORT_TYPED_FUNC` and `TVM_FFI_DLL_EXPORT_TYPED_FUNC_DOC`
+/// export do. Returns the safe call return code.
+///
+/// # Safety
+///
+/// `result` must be valid for writes of one `TVMFFIAny`.
+#[doc(hidden)]
+pub unsafe fn write_exported_str(text: Result<std::string::String>, result:
*mut TVMFFIAny) -> i32 {
+ match text {
+ Ok(text) => {
+ let text = TVMFFIByteArray::from_str(&text);
+ TVMFFIStringFromByteArray(&text, result)
+ }
+ Err(error) => {
+ Error::set_raised(&error);
+ -1
+ }
+ }
+}
+
+/// The metadata a library exports for a function `func` as
+/// `__tvm_ffi__metadata_<name>`, as `TVM_FFI_DLL_EXPORT_TYPED_FUNC` does.
+#[doc(hidden)]
+pub fn exported_metadata<Fun, I, O>(func: &Fun) -> Result<std::string::String>
+where
+ Fun: AsPackedCallable<I, O>,
+{
+ let type_schema = type_schema_of(func);
+ Ok(crate::type_traits::type_schema_metadata(&type_schema))
}
#[inline]
@@ -61,8 +118,19 @@ macro_rules! impl_as_packed_callable {
where
Fun: Fn($($t,)*) -> Result<Out> + 'static,
Any: From<Out>,
+ Out: TypeSchema,
$($t: ArgTryFromAnyView),*
{
+ fn type_schema() -> std::string::String {
+ let params: &[std::string::String] = &[$(<$t as
TypeSchema>::type_schema()),*];
+ format!(
+
r#"{{"type":"{}","named_args":{{"return":[{}],"params":[{}]}}}}"#,
+ FunctionObj::TYPE_KEY,
+ Out::type_schema(),
+ params.join(","),
+ )
+ }
+
fn call_packed(&self, packed_args: &[AnyView]) -> Result<Any>
{
crate::ensure!(
diff --git a/rust/tvm-ffi/src/macros.rs b/rust/tvm-ffi/src/macros.rs
index 5f32ecbf..8c88fba6 100644
--- a/rust/tvm-ffi/src/macros.rs
+++ b/rust/tvm-ffi/src/macros.rs
@@ -313,6 +313,17 @@ macro_rules! impl_arg_into_ref {
/// since unwinding into the caller would abort the process, but panicking
/// is discouraged.
///
+/// With the `export-metadata` feature, the counterpart of C++
+/// `TVM_FFI_DLL_EXPORT_INCLUDE_METADATA`, the macro also exports the
+/// function's metadata, `{"type_schema":...}`, as
`__tvm_ffi__metadata_<name>`,
+/// which `Module::GetFunctionMetadata` and the stub generator read. A
+/// function's schema lists its parameter and return types; a callable whose
+/// signature is not known, such as a custom [`AsPackedCallable`], has the
+/// schema of an untyped function. As in C++, the feature is off by default,
+/// so exports are minimal.
+///
+/// [`AsPackedCallable`]: crate::function_internal::AsPackedCallable
+///
/// # Arguments
/// * `$name` - The name of the function
/// * `$func` - The function to export
@@ -371,6 +382,104 @@ macro_rules! tvm_ffi_dll_export_typed_func {
}
}
}
+ $crate::__tvm_ffi_dll_export_metadata!($name, $func);
+ };
+}
+
+/// Macro to export the doc string of a function exported with
+/// [`tvm_ffi_dll_export_typed_func!`], as C++
`TVM_FFI_DLL_EXPORT_TYPED_FUNC_DOC`
+/// does
+///
+/// With the `export-metadata` feature, the doc string is exported as
+/// `__tvm_ffi__doc_<name>`, which `Module::GetFunctionDoc` and the stub
+/// generator read. Without it, as C++ without
+/// `TVM_FFI_DLL_EXPORT_INCLUDE_METADATA`, the macro exports nothing, but
+/// still checks that the doc string is a string constant. The macro never
+/// exports the function itself.
+///
+/// # Arguments
+/// * `$name` - The name the function is exported as
+/// * `$doc` - The doc string, a string literal
+///
+/// # Example
+/// ```rust
+/// use tvm_ffi::*;
+///
+/// fn add(a: i64, b: i64) -> Result<i64> { Ok(a + b) }
+///
+/// tvm_ffi_dll_export_typed_func!(add, add);
+/// tvm_ffi_dll_export_typed_func_doc!(add, "Add two integers and return the
sum.");
+/// ```
+#[macro_export]
+macro_rules! tvm_ffi_dll_export_typed_func_doc {
+ ($name:ident, $doc:expr) => {
+ $crate::__tvm_ffi_dll_export_doc!($name, $doc);
+ };
+}
+
+// The metadata and doc getters the export macros emit with the
+// `export-metadata` feature. A `#[cfg]` in a `macro_rules!` expansion would
+// test the features of the crate that invokes the macro, so the feature
+// selects which definition of these helpers this crate compiles instead.
+
+#[cfg(feature = "export-metadata")]
+#[doc(hidden)]
+#[macro_export]
+macro_rules! __tvm_ffi_dll_export_metadata {
+ ($name:ident, $func:expr) => {
+ $crate::macros::paste::paste! {
+ #[no_mangle]
+ pub unsafe extern "C" fn [<__tvm_ffi__metadata_ $name>](
+ _handle: *mut std::ffi::c_void,
+ _args: *const $crate::tvm_ffi_sys::TVMFFIAny,
+ _num_args: i32,
+ result: *mut $crate::tvm_ffi_sys::TVMFFIAny,
+ ) -> i32 {
+ let metadata = match
::std::panic::catch_unwind(::std::panic::AssertUnwindSafe(|| {
+ $crate::function_internal::exported_metadata(&$func)
+ })) {
+ Ok(metadata) => metadata,
+ Err(payload) =>
Err($crate::function_internal::panic_to_error(payload)),
+ };
+ $crate::function_internal::write_exported_str(metadata, result)
+ }
+ }
+ };
+}
+
+#[cfg(not(feature = "export-metadata"))]
+#[doc(hidden)]
+#[macro_export]
+macro_rules! __tvm_ffi_dll_export_metadata {
+ ($name:ident, $func:expr) => {};
+}
+
+#[cfg(feature = "export-metadata")]
+#[doc(hidden)]
+#[macro_export]
+macro_rules! __tvm_ffi_dll_export_doc {
+ ($name:ident, $doc:expr) => {
+ $crate::macros::paste::paste! {
+ #[no_mangle]
+ pub unsafe extern "C" fn [<__tvm_ffi__doc_ $name>](
+ _handle: *mut std::ffi::c_void,
+ _args: *const $crate::tvm_ffi_sys::TVMFFIAny,
+ _num_args: i32,
+ result: *mut $crate::tvm_ffi_sys::TVMFFIAny,
+ ) -> i32 {
+ const DOC: &str = $doc;
+
$crate::function_internal::write_exported_str(Ok(DOC.to_string()), result)
+ }
+ }
+ };
+}
+
+#[cfg(not(feature = "export-metadata"))]
+#[doc(hidden)]
+#[macro_export]
+macro_rules! __tvm_ffi_dll_export_doc {
+ ($name:ident, $doc:expr) => {
+ const _: &str = $doc;
};
}
diff --git a/rust/tvm-ffi/src/reflection/registry.rs
b/rust/tvm-ffi/src/reflection/registry.rs
index e7e59669..acbb0558 100644
--- a/rust/tvm-ffi/src/reflection/registry.rs
+++ b/rust/tvm-ffi/src/reflection/registry.rs
@@ -42,7 +42,7 @@ use crate::tvm_ffi_sys::{
TVMFFIGetTypeInfo, TVMFFIObject, TVMFFISEqHashKind, TVMFFITypeIndex,
TVMFFITypeMetadata,
TVMFFITypeRegisterAttr, TVMFFITypeRegisterField,
TVMFFITypeRegisterMetadata,
};
-use crate::type_traits::AnyCompatible;
+use crate::type_traits::{type_schema_metadata, AnyCompatible};
use crate::{Any, Error, Function, Result, INTERNAL_ERROR, TYPE_ERROR};
/// A copy of an object's own fields under a fresh object header (the C++ copy
@@ -207,8 +207,7 @@ impl<N: ObjectCore> ObjectDef<N> {
if !self.fresh {
return Ok(self);
}
- let (static_type_index, schema) = field_type_info(&F::type_str());
- let metadata = format!(r#"{{"type_schema":{}}}"#,
json_string(&schema));
+ let metadata = type_schema_metadata(&F::type_schema());
let info = TVMFFIFieldInfo {
name: unsafe { TVMFFIByteArray::from_str(name) },
doc: unsafe { TVMFFIByteArray::from_str("") },
@@ -220,7 +219,7 @@ impl<N: ObjectCore> ObjectDef<N> {
getter: Some(field_getter::<F>),
setter: field_setter::<F> as TVMFFIFieldSetter as *mut c_void,
default_value_or_factory: TVMFFIAny::new(),
- field_static_type_index: static_type_index,
+ field_static_type_index: field_static_type_index(&F::type_str()),
};
// The table copies the strings before returning.
crate::check_safe_call!(unsafe {
TVMFFITypeRegisterField(self.type_index, &info) })?;
@@ -364,135 +363,74 @@ unsafe extern "C" fn field_setter<F: AnyCompatible>(
-1
}
-/// The static type index and the JSON type schema of a field, from the Rust
type's `type_str()`
-/// (the names `AnyCompatible` reports, such as `int`, `ffi.String`,
`Array<my.Node>`), as C++
-/// `TypeTraits<T>::field_static_type_index` and `TypeTraits<T>::TypeSchema()`
state them.
-fn field_type_info(type_str: &str) -> (i32, String) {
+/// The static type index of a field, from the Rust type's `type_str()` (the
names
+/// `AnyCompatible` reports, such as `int`, `ffi.String`, `Array<my.Node>`),
as C++
+/// `TypeTraits<T>::field_static_type_index` states it.
+fn field_static_type_index(type_str: &str) -> i32 {
use TVMFFITypeIndex::*;
/// C++ `TypeIndex::kTVMFFIAny`, the static index of a field holding any
value.
const ANY: i32 = -1;
- let (name, args) = split_generic(type_str);
- let args: Vec<String> = args.iter().map(|arg|
field_type_info(arg).1).collect();
- let schema = |key: &str| match args.is_empty() {
- true => format!(r#"{{"type":"{key}"}}"#),
- false => format!(r#"{{"type":"{key}","args":[{}]}}"#, args.join(",")),
- };
- match name {
- "None" => (kTVMFFINone as i32, schema("None")),
- "int" => (kTVMFFIInt as i32, schema("int")),
- "bool" => (kTVMFFIBool as i32, schema("bool")),
- "float" => (kTVMFFIFloat as i32, schema("float")),
- "void*" => (kTVMFFIOpaquePtr as i32, schema("void*")),
- "DataType" => (kTVMFFIDataType as i32, schema("DataType")),
- "Device" => (kTVMFFIDevice as i32, schema("Device")),
- "DLTensor*" => (kTVMFFIDLTensorPtr as i32, schema("DLTensor*")),
- "Any" => (ANY, schema("Any")),
+ let origin = type_str.split('<').next().unwrap_or(type_str).trim();
+ match origin {
+ "None" => kTVMFFINone as i32,
+ "int" => kTVMFFIInt as i32,
+ "bool" => kTVMFFIBool as i32,
+ "float" => kTVMFFIFloat as i32,
+ "void*" => kTVMFFIOpaquePtr as i32,
+ "DataType" => kTVMFFIDataType as i32,
+ "Device" => kTVMFFIDevice as i32,
+ "DLTensor*" => kTVMFFIDLTensorPtr as i32,
// C++ `String`, `Bytes` and `BigInt` fields are typed `Any`
statically (they have
// inline forms), as is `Optional<T>`.
- "ffi.String" | "ffi.Bytes" => (ANY, schema(name)),
- "BigInt" => (ANY, schema("ffi.BigInt")),
- "Optional" => (ANY, schema("Optional")),
- "ffi.Function" => (kTVMFFIFunction as i32, schema(name)),
- "ffi.Shape" => (kTVMFFIShape as i32, schema(name)),
- "Array" => (kTVMFFIArray as i32, schema("ffi.Array")),
- "List" => (kTVMFFIList as i32, schema("ffi.List")),
- // Any other object reference, `Map` and `Dict` included, is typed
`Object` statically.
- "Map" => (kTVMFFIStaticObjectBegin as i32, schema("ffi.Map")),
- "Dict" => (kTVMFFIStaticObjectBegin as i32, schema("ffi.Dict")),
- _ => (kTVMFFIStaticObjectBegin as i32, schema(name)),
- }
-}
-
-/// `Name<A, B<C>>` as (`Name`, [`A`, `B<C>`]); a plain name has no arguments.
-fn split_generic(type_str: &str) -> (&str, Vec<&str>) {
- let Some(open) = type_str.find('<') else {
- return (type_str.trim(), Vec::new());
- };
- let Some(inner) = type_str[open + 1..].strip_suffix('>') else {
- return (type_str.trim(), Vec::new());
- };
- let mut args = Vec::new();
- let (mut depth, mut start) = (0usize, 0usize);
- for (index, byte) in inner.bytes().enumerate() {
- match byte {
- b'<' => depth += 1,
- b'>' => depth = depth.saturating_sub(1),
- b',' if depth == 0 => {
- args.push(inner[start..index].trim());
- start = index + 1;
- }
- _ => {}
- }
+ "ffi.String" | "ffi.Bytes" | "BigInt" | "Optional" => ANY,
+ "ffi.Function" => kTVMFFIFunction as i32,
+ "ffi.Shape" => kTVMFFIShape as i32,
+ "Array" => kTVMFFIArray as i32,
+ "List" => kTVMFFIList as i32,
+ // Any other object reference, `Map` included, is typed `Object`
statically.
+ _ => kTVMFFIStaticObjectBegin as i32,
}
- args.push(inner[start..].trim());
- (type_str[..open].trim(), args)
-}
-
-/// `text` as a JSON string literal.
-fn json_string(text: &str) -> String {
- let mut out = String::with_capacity(text.len() + 2);
- out.push('"');
- for ch in text.chars() {
- match ch {
- '"' => out.push_str("\\\""),
- '\\' => out.push_str("\\\\"),
- '\n' => out.push_str("\\n"),
- '\r' => out.push_str("\\r"),
- '\t' => out.push_str("\\t"),
- ch if (ch as u32) < 0x20 => out.push_str(&format!("\\u{:04x}", ch
as u32)),
- ch => out.push(ch),
- }
- }
- out.push('"');
- out
}
#[cfg(test)]
mod tests {
- use super::{field_type_info, json_string, split_generic};
+ use super::field_static_type_index;
+ use crate::type_traits::{type_schema_metadata, AnyCompatible};
+ use crate::{Array, List, Map};
#[test]
- fn generic_names_split_at_top_level_commas() {
- assert_eq!(split_generic("int"), ("int", vec![]));
- assert_eq!(split_generic("Array<my.Node>"), ("Array",
vec!["my.Node"]));
- assert_eq!(
- split_generic("Map<ffi.String, Array<my.Node>>"),
- ("Map", vec!["ffi.String", "Array<my.Node>"])
- );
+ fn field_static_type_indices_follow_the_cpp_type_traits() {
+ assert_eq!(field_static_type_index("int"), 1);
+ assert_eq!(field_static_type_index("ffi.String"), -1);
+ assert_eq!(field_static_type_index("Optional<my.Node>"), -1);
+ assert_eq!(field_static_type_index("my.Node"), 64);
+ assert_eq!(field_static_type_index("Array<my.Node>"), 71);
+ assert_eq!(field_static_type_index("List<int>"), 75);
+ assert_eq!(field_static_type_index("Map<ffi.String, int>"), 64);
}
#[test]
fn field_schemas_follow_the_cpp_type_traits() {
- assert_eq!(field_type_info("int"), (1, r#"{"type":"int"}"#.into()));
assert_eq!(
- field_type_info("ffi.String"),
- (-1, r#"{"type":"ffi.String"}"#.into())
+ <Map<crate::String, i64>>::type_schema(),
+
r#"{"type":"ffi.Map","args":[{"type":"ffi.String"},{"type":"int"}]}"#
);
assert_eq!(
- field_type_info("my.Node"),
- (64, r#"{"type":"my.Node"}"#.into())
+ <List<i64>>::type_schema(),
+ r#"{"type":"ffi.List","args":[{"type":"int"}]}"#
);
assert_eq!(
- field_type_info("Array<my.Node>"),
- (
- 71,
- r#"{"type":"ffi.Array","args":[{"type":"my.Node"}]}"#.into()
- )
+ <Option<Array<crate::String>>>::type_schema(),
+
r#"{"type":"Optional","args":[{"type":"ffi.Array","args":[{"type":"ffi.String"}]}]}"#
);
assert_eq!(
- field_type_info("Map<ffi.String, int>"),
- (
- 64,
-
r#"{"type":"ffi.Map","args":[{"type":"ffi.String"},{"type":"int"}]}"#.into()
- )
+ type_schema_metadata(r#"{"type":"int"}"#),
+ r#"{"type_schema":"{\"type\":\"int\"}"}"#
);
+ // As C++ `EscapeStringJSON`, `/`, `\b` and DEL are escaped too.
assert_eq!(
- field_type_info("Optional<my.Node>"),
- (
- -1,
- r#"{"type":"Optional","args":[{"type":"my.Node"}]}"#.into()
- )
+ type_schema_metadata("a/b\u{8}\u{7f}"),
+ r#"{"type_schema":"a\/b\b\u007f"}"#
);
- assert_eq!(json_string(r#"{"type":"int"}"#),
r#""{\"type\":\"int\"}""#);
}
}
diff --git a/rust/tvm-ffi/src/rvalue_ref.rs b/rust/tvm-ffi/src/rvalue_ref.rs
index f9d691b2..62e98887 100644
--- a/rust/tvm-ffi/src/rvalue_ref.rs
+++ b/rust/tvm-ffi/src/rvalue_ref.rs
@@ -126,6 +126,15 @@ where
}
}
+impl<T> crate::type_traits::TypeSchema for RValueRef<T>
+where
+ T: ObjectRefCore + AnyCompatible,
+{
+ fn type_schema() -> std::string::String {
+ crate::type_traits::type_schema("ObjectRValueRef", &[T::type_schema()])
+ }
+}
+
impl<T> ArgTryFromAnyView for RValueRef<T>
where
T: ObjectRefCore + AnyCompatible,
diff --git a/rust/tvm-ffi/src/type_traits.rs b/rust/tvm-ffi/src/type_traits.rs
index c298db9e..5f15a957 100644
--- a/rust/tvm-ffi/src/type_traits.rs
+++ b/rust/tvm-ffi/src/type_traits.rs
@@ -70,6 +70,16 @@ pub unsafe trait AnyCompatible: Sized {
/// the type string of the type
fn type_str() -> String;
+ /// The type schema of the type: the JSON object C++
+ /// `TypeTraits<T>::TypeSchema` writes, which the type schemas of functions
+ /// and fields embed.
+ ///
+ /// The default, `{"type":"<type_str>"}`, fits a type whose type string is
+ /// its type key, as for object references.
+ fn type_schema() -> String {
+ type_schema(&Self::type_str(), &[])
+ }
+
/// Borrow `self` into an owned [`Any`], increfing object-backed values.
///
/// The by-reference counterpart of `Any::from(value)`, for fields reached
@@ -83,6 +93,74 @@ pub unsafe trait AnyCompatible: Sized {
}
}
+/// Writes a type schema, `{"type":"<origin>"}` or, with type arguments,
+/// `{"type":"<origin>","args":[...]}`, as C++ `TypeTraits<T>::TypeSchema`
does.
+pub(crate) fn type_schema(origin: &str, args: &[String]) -> String {
+ if args.is_empty() {
+ format!(r#"{{"type":"{origin}"}}"#)
+ } else {
+ format!(r#"{{"type":"{origin}","args":[{}]}}"#, args.join(","))
+ }
+}
+
+/// The metadata C++ records for a function or field of type schema
+/// `type_schema`: `{"type_schema":<type_schema as a JSON string>}`.
+pub(crate) fn type_schema_metadata(type_schema: &str) -> String {
+ format!(r#"{{"type_schema":{}}}"#, escape_string_json(type_schema))
+}
+
+/// `text` as a JSON string literal, escaped as C++ `EscapeStringJSON`
+/// (`string.h`) escapes it, so that metadata matches the C++ one byte for
byte.
+fn escape_string_json(text: &str) -> String {
+ let mut out = String::with_capacity(text.len() + 2);
+ out.push('"');
+ for ch in text.chars() {
+ match ch {
+ '"' => out.push_str("\\\""),
+ '\\' => out.push_str("\\\\"),
+ '/' => out.push_str("\\/"),
+ '\u{8}' => out.push_str("\\b"),
+ '\u{c}' => out.push_str("\\f"),
+ '\n' => out.push_str("\\n"),
+ '\r' => out.push_str("\\r"),
+ '\t' => out.push_str("\\t"),
+ ch if (ch as u32) < 0x20 || ch == '\u{7f}' => {
+ out.push_str(&format!("\\u{:04x}", ch as u32))
+ }
+ ch => out.push(ch),
+ }
+ }
+ out.push('"');
+ out
+}
+
+/// The type schema of a function parameter or return type.
+///
+/// This is [`AnyCompatible::type_schema`] for a type that converts through
+/// [`Any`], and `{"type":"Any"}` for `Any` and `AnyView` themselves, as C++
+/// `details::TypeSchema<T>` writes them.
+pub(crate) trait TypeSchema {
+ fn type_schema() -> String;
+}
+
+impl<T: AnyCompatible> TypeSchema for T {
+ fn type_schema() -> String {
+ <T as AnyCompatible>::type_schema()
+ }
+}
+
+impl TypeSchema for Any {
+ fn type_schema() -> String {
+ type_schema("Any", &[])
+ }
+}
+
+impl TypeSchema for AnyView<'_> {
+ fn type_schema() -> String {
+ type_schema("Any", &[])
+ }
+}
+
/// Marker for a value that can be stored in an FFI container.
///
/// This is an implementation detail of [`crate::Array`] and [`crate::Map`].
@@ -112,6 +190,7 @@ mod container_element_ops {
unsafe fn container_try_cast_from_any_view(data: &TVMFFIAny) ->
Result<Self, ()>;
fn container_get_mismatch_type_info(data: &TVMFFIAny) -> String;
fn container_type_str() -> String;
+ fn container_type_schema() -> String;
}
impl<T: AnyCompatible> Ops for T {
@@ -154,6 +233,11 @@ mod container_element_ops {
fn container_type_str() -> String {
<T as AnyCompatible>::type_str()
}
+
+ #[inline]
+ fn container_type_schema() -> String {
+ <T as AnyCompatible>::type_schema()
+ }
}
impl Ops for Any {
@@ -198,6 +282,11 @@ mod container_element_ops {
fn container_type_str() -> String {
"Any".to_string()
}
+
+ #[inline]
+ fn container_type_schema() -> String {
+ <Any as super::TypeSchema>::type_schema()
+ }
}
}
@@ -297,6 +386,10 @@ unsafe impl<T: AnyCompatible> AnyCompatible for Option<T> {
"Optional<".to_string() + T::type_str().as_str() + ">"
}
+ fn type_schema() -> String {
+ type_schema("Optional", &[<T as AnyCompatible>::type_schema()])
+ }
+
unsafe fn copy_to_any_view(src: &Self, data: &mut TVMFFIAny) {
if let Some(ref value) = src {
T::copy_to_any_view(value, data);
diff --git a/rust/tvm-ffi/tests/test_function_metadata.rs
b/rust/tvm-ffi/tests/test_function_metadata.rs
new file mode 100644
index 00000000..e2370340
--- /dev/null
+++ b/rust/tvm-ffi/tests/test_function_metadata.rs
@@ -0,0 +1,291 @@
+/*
+ * 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 metadata and doc strings that Rust records for exported and global
+//! functions match the ones C++ records for the same signatures in
+//! `libtvm_ffi_testing`, byte for byte. Exported functions have them with the
+//! `export-metadata` feature, as C++ ones with
+//! `TVM_FFI_DLL_EXPORT_INCLUDE_METADATA`; global functions always do.
+
+use tvm_ffi::function_internal::AsPackedCallable;
+#[cfg(feature = "export-metadata")]
+use tvm_ffi::tvm_ffi_sys::TVMFFIAny;
+use tvm_ffi::*;
+
+// Exported by libtvm_ffi_testing with `TVM_FFI_DLL_EXPORT_TYPED_FUNC` and
+// `TVM_FFI_DLL_EXPORT_INCLUDE_METADATA`, for `int64_t(int64_t)`.
+#[cfg(feature = "export-metadata")]
+extern "C" {
+ fn __tvm_ffi__metadata_testing_dll_schema_id_int(
+ handle: *mut std::ffi::c_void,
+ args: *const TVMFFIAny,
+ num_args: i32,
+ result: *mut TVMFFIAny,
+ ) -> i32;
+}
+
+/// Calls an exported metadata or doc getter.
+#[cfg(feature = "export-metadata")]
+fn call_getter(getter: tvm_ffi::tvm_ffi_sys::TVMFFISafeCallType) -> String {
+ // SAFETY: the getters take no arguments and do not use the handle.
+ let getter = unsafe { Function::from_extern_c(std::ptr::null_mut(),
getter, None) };
+ getter.call_packed(&[]).unwrap().try_into().unwrap()
+}
+
+/// The dummy call keeps `libtvm_ffi_testing`, which registers the C++
+/// functions compared with, linked.
+fn global_metadata(name: &str) -> String {
+ assert_eq!(unsafe { tvm_ffi_sys::TVMFFITestingDummyTarget() }, 0);
+ Function::get_global("ffi.GetGlobalFuncMetadata")
+ .unwrap()
+ .call_tuple_with_len::<1, _>((String::from(name),))
+ .unwrap()
+ .try_into()
+ .unwrap()
+}
+
+fn rust_schema_id_int(x: i64) -> Result<i64> {
+ Ok(x)
+}
+tvm_ffi_dll_export_typed_func!(rust_schema_id_int, rust_schema_id_int);
+
+#[cfg(feature = "export-metadata")]
+#[test]
+fn test_exported_metadata_matches_cpp() {
+ let rust = call_getter(__tvm_ffi__metadata_rust_schema_id_int);
+ let cpp = call_getter(__tvm_ffi__metadata_testing_dll_schema_id_int);
+ assert_eq!(rust.as_str(), cpp.as_str());
+}
+
+fn rust_add(a: i64, b: i64) -> Result<i64> {
+ Ok(a + b)
+}
+tvm_ffi_dll_export_typed_func!(rust_add, rust_add);
+tvm_ffi_dll_export_typed_func_doc!(
+ rust_add,
+ "Add two integers and return the sum.\n\n\"a\" and \"b\" are integers."
+);
+
+#[cfg(feature = "export-metadata")]
+#[test]
+fn test_exported_doc_is_verbatim() {
+ let doc = call_getter(__tvm_ffi__doc_rust_add);
+ assert_eq!(
+ doc.as_str(),
+ "Add two integers and return the sum.\n\n\"a\" and \"b\" are integers."
+ );
+}
+
+/// A callable whose signature is not known, as a crate may define to export a
+/// packed function.
+struct Packed(fn(&[AnyView]) -> Result<Any>);
+struct PackedArgs;
+
+impl AsPackedCallable<PackedArgs, Any> for Packed {
+ fn call_packed(&self, packed_args: &[AnyView]) -> Result<Any> {
+ (self.0)(packed_args)
+ }
+}
+
+fn rust_packed(_args: &[AnyView]) -> Result<Any> {
+ Ok(Any::new())
+}
+tvm_ffi_dll_export_typed_func!(rust_packed, Packed(rust_packed));
+
+#[cfg(feature = "export-metadata")]
+#[test]
+fn test_exported_packed_metadata_is_untyped() {
+ register_globals();
+ let rust = call_getter(__tvm_ffi__metadata_rust_packed);
+ assert_eq!(
+ rust.as_str(),
+ global_metadata("testing.schema_packed").as_str()
+ );
+}
+
+// Without `export-metadata`, the exports define no metadata or doc getters:
+// these symbols of the same names would otherwise be defined twice.
+#[cfg(not(feature = "export-metadata"))]
+mod without_export_metadata {
+ #[no_mangle]
+ pub extern "C" fn __tvm_ffi__metadata_rust_schema_id_int() {}
+ #[no_mangle]
+ pub extern "C" fn __tvm_ffi__metadata_rust_packed() {}
+ #[no_mangle]
+ pub extern "C" fn __tvm_ffi__doc_rust_add() {}
+}
+
+#[cfg(not(feature = "export-metadata"))]
+#[test]
+fn test_exports_without_metadata() {
+ // The functions themselves are exported.
+ // SAFETY: rust_add takes two arguments and does not use the handle.
+ let add = unsafe { Function::from_extern_c(std::ptr::null_mut(),
__tvm_ffi_rust_add, None) };
+ let sum = add.call_tuple_with_len::<2, _>((1i64, 2i64)).unwrap();
+ assert_eq!(i64::try_from(sum).unwrap(), 3);
+ without_export_metadata::__tvm_ffi__metadata_rust_schema_id_int();
+ without_export_metadata::__tvm_ffi__metadata_rust_packed();
+ without_export_metadata::__tvm_ffi__doc_rust_add();
+}
+
+/// Registers `func` as `testing.rust.<name>` and records `name`.
+fn register_typed<F, I, O>(names: &mut Vec<&'static str>, name: &'static str,
func: F)
+where
+ F: AsPackedCallable<I, O> + 'static,
+{
+ Function::register_global_typed(&format!("testing.rust.{name}"), func,
"").unwrap();
+ names.push(name);
+}
+
+// The runtime does not guard its global function table against readers on
+// other threads, so each test registers every global here before it reads
+// the table. Returns the names of the typed globals that mirror a C++
+// `testing.<name>` global.
+fn register_globals() -> &'static [&'static str] {
+ static NAMES: std::sync::OnceLock<Vec<&'static str>> =
std::sync::OnceLock::new();
+ NAMES.get_or_init(|| {
+ let mut names = Vec::new();
+ let n = &mut names;
+ register_typed(n, "schema_id_float", |x: f64| -> Result<f64> { Ok(x)
});
+ register_typed(n, "schema_id_bool", |x: bool| -> Result<bool> { Ok(x)
});
+ register_typed(n, "schema_id_device", |x: DLDevice| ->
Result<DLDevice> {
+ Ok(x)
+ });
+ register_typed(
+ n,
+ "schema_id_dtype",
+ |x: DLDataType| -> Result<DLDataType> { Ok(x) },
+ );
+ register_typed(n, "schema_id_string", |x: String| -> Result<String> {
+ Ok(x)
+ });
+ register_typed(n, "schema_id_bytes", |x: Bytes| -> Result<Bytes> {
Ok(x) });
+ register_typed(n, "schema_id_func", |x: Function| -> Result<Function> {
+ Ok(x)
+ });
+ register_typed(
+ n,
+ "schema_id_object",
+ |x: object::ObjectRef| -> Result<object::ObjectRef> { Ok(x) },
+ );
+ register_typed(n, "schema_id_tensor", |x: Tensor| -> Result<Tensor> {
+ Ok(x)
+ });
+ register_typed(
+ n,
+ "schema_tensor_view_input",
+ |_: TensorView| -> Result<()> { Ok(()) },
+ );
+ register_typed(
+ n,
+ "schema_id_opt_int",
+ |x: Option<i64>| -> Result<Option<i64>> { Ok(x) },
+ );
+ register_typed(
+ n,
+ "schema_id_opt_str",
+ |x: Option<String>| -> Result<Option<String>> { Ok(x) },
+ );
+ register_typed(
+ n,
+ "schema_id_opt_obj",
+ |x: Option<object::ObjectRef>| ->
Result<Option<object::ObjectRef>> { Ok(x) },
+ );
+ register_typed(
+ n,
+ "schema_id_arr_int",
+ |x: Array<i64>| -> Result<Array<i64>> { Ok(x) },
+ );
+ register_typed(
+ n,
+ "schema_id_arr_str",
+ |x: Array<String>| -> Result<Array<String>> { Ok(x) },
+ );
+ register_typed(
+ n,
+ "schema_id_arr_obj",
+ |x: Array<object::ObjectRef>| -> Result<Array<object::ObjectRef>>
{ Ok(x) },
+ );
+ register_typed(
+ n,
+ "schema_id_map_str_int",
+ |x: Map<String, i64>| -> Result<Map<String, i64>> { Ok(x) },
+ );
+ register_typed(
+ n,
+ "schema_id_map_str_str",
+ |x: Map<String, String>| -> Result<Map<String, String>> { Ok(x) },
+ );
+ register_typed(
+ n,
+ "schema_id_map_str_obj",
+ |x: Map<String, object::ObjectRef>| -> Result<Map<String,
object::ObjectRef>> { Ok(x) },
+ );
+ register_typed(
+ n,
+ "schema_arr_map_opt",
+ |_: Array<Option<i64>>,
+ mp: Map<String, Array<i64>>,
+ _: Option<String>|
+ -> Result<Map<String, Array<i64>>> { Ok(mp) },
+ );
+ register_typed(n, "schema_no_args", || -> Result<i64> { Ok(1) });
+ register_typed(n, "schema_no_return", |_: i64| -> Result<()> { Ok(())
});
+ register_typed(n, "schema_no_args_no_return", || -> Result<()> {
Ok(()) });
+
+ let packed = Function::from_packed(|_: &[AnyView]| -> Result<Any> {
Ok(Any::new()) });
+ Function::register_global("testing.rust.schema_packed",
packed).unwrap();
+ Function::register_global_typed(
+ "testing.rust.add",
+ |a: i64, b: i64| -> Result<i64> { Ok(a + b) },
+ "Add two integers.",
+ )
+ .unwrap();
+ names
+ })
+}
+
+#[test]
+fn test_global_metadata_matches_cpp() {
+ let names = register_globals();
+ assert_eq!(names.len(), 23);
+ for name in names {
+ assert_eq!(
+ global_metadata(&format!("testing.rust.{name}")).as_str(),
+ global_metadata(&format!("testing.{name}")).as_str(),
+ "{name}"
+ );
+ }
+}
+
+#[test]
+fn test_global_untyped_metadata_matches_cpp() {
+ register_globals();
+ assert_eq!(
+ global_metadata("testing.rust.schema_packed").as_str(),
+ global_metadata("testing.schema_packed").as_str()
+ );
+}
+
+#[test]
+fn test_global_typed_registration_calls_the_function() {
+ register_globals();
+ let add = Function::get_global("testing.rust.add").unwrap();
+ let sum = add.call_tuple_with_len::<2, _>((1i64, 2i64)).unwrap();
+ assert_eq!(i64::try_from(sum).unwrap(), 3);
+}
