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 a0e43582 [FEAT][Rust] Register derived object types on first use (#824)
a0e43582 is described below
commit a0e43582f4a369078647cf510d3d665a4e91083b
Author: Gabriel Wu <[email protected]>
AuthorDate: Tue Oct 6 23:20:32 2026 +0800
[FEAT][Rust] Register derived object types on first use (#824)
## Why
`#[derive(Object)]` without `#[type_index]` resolves its type key with
`TVMFFITypeKeyToIndex` and panics when nothing has registered the key. A
Rust crate therefore cannot define an object type of its own: someone
else, usually C++, has to register the key first. The tests work around
this. `test_match_any_ast_lookup.rs` declares
`TVMFFITypeGetOrAllocIndex` itself and registers its types by hand
before they are used.
C++ does not have this gap. `TVM_FFI_DECLARE_OBJECT_INFO` calls
`TVMFFITypeGetOrAllocIndex` with the parent's index, the type depth and
the child slots. That returns the index the key already has, or
allocates one under the parent.
## What
- **Derive:** a type defined in Rust declares `#[type_register]`. Its
`type_index()` calls `TVMFFITypeGetOrAllocIndex` on first use, passing
the parent's `type_index()` (which registers the parent first if
needed), `TYPE_DEPTH` and the child slots. A key that is already
registered keeps its index.
- **Bindings are unchanged:** a type without `#[type_register]` binds a
key its defining library, usually C++, registers. As before, it looks
the key up with `TVMFFITypeKeyToIndex` and panics if that library is not
loaded yet. So a binding never registers a C++ type's key with the wrong
child slots, whatever the load order.
- **New attributes** `#[type_child_slots = N]` and
`#[type_child_slots_can_overflow = bool]` mirror C++ `_type_child_slots`
and `_type_child_slots_can_overflow`, with the same defaults of 0 and
true. They need `#[type_register]`, and a `#[type_final]` type cannot
reserve child slots.
- **Concurrency:** C++ registers types while a library loads, and the
type table relies on that (`src/ffi/object.cc`: updates happen during
loading "or explicitly locked from the caller"). Rust registers on first
use, possibly on any thread, so registrations from Rust hold one lock,
`tvm_ffi::object::TYPE_REGISTRATION`; the parent is registered before it
is taken. The docs state what stays the caller's duty, as for C++: not
registering while another thread loads a library that registers types.
Calling `type_index()` before starting such threads registers a type
eagerly.
- **tvm-ffi-sys** binds `TVMFFITypeGetOrAllocIndex`.
- **Test cleanup:** `test_match_any_ast_lookup.rs` drops its own `extern
"C"` declaration and registers its types with `#[type_register]`.
---
rust/tvm-ffi-macros/src/lib.rs | 33 +++-
rust/tvm-ffi-macros/src/object_macros.rs | 101 ++++++++++++-
rust/tvm-ffi-sys/src/c_api.rs | 8 +
rust/tvm-ffi/src/object.rs | 6 +
rust/tvm-ffi/tests/test_match_any_ast_lookup.rs | 35 +----
rust/tvm-ffi/tests/test_object_registration.rs | 190 ++++++++++++++++++++++++
6 files changed, 341 insertions(+), 32 deletions(-)
diff --git a/rust/tvm-ffi-macros/src/lib.rs b/rust/tvm-ffi-macros/src/lib.rs
index 054836b8..07026e02 100644
--- a/rust/tvm-ffi-macros/src/lib.rs
+++ b/rust/tvm-ffi-macros/src/lib.rs
@@ -49,8 +49,39 @@ pub fn match_any(input: TokenStream) -> TokenStream {
match_any::expand(input)
}
+/// Derive `ObjectCore` for an object struct whose first field is its parent.
+///
+/// `#[type_key = "..."]` names the type, and `#[type_index(...)]` gives it a
+/// static index. Otherwise a type defined in Rust declares `#[type_register]`:
+/// it is registered under its parent on first use of `type_index()`, or takes
+/// the index its key already has, as C++ `TVM_FFI_DECLARE_OBJECT_INFO` does.
+/// Without `#[type_register]`, the type binds a key its defining library,
+/// usually C++, has registered, and `type_index()` panics if that library is
+/// not loaded yet, so a binding never registers a key with the wrong child
+/// slots.
+///
+/// `#[type_final]`, and with `#[type_register]` `#[type_child_slots = N]` and
+/// `#[type_child_slots_can_overflow = bool]`, mirror C++ `_type_final`,
+/// `_type_child_slots` (default 0) and `_type_child_slots_can_overflow`
+/// (default true).
+///
+/// Registrations from Rust hold one lock, so they do not race each other. The
+/// runtime's type table is not locked otherwise: as for C++, a type must not
+/// be registered while another thread loads a library that registers types.
+/// Calling `type_index()` before starting such threads registers a type
+/// eagerly.
#[proc_macro_error]
-#[proc_macro_derive(Object, attributes(type_key, type_index, type_final))]
+#[proc_macro_derive(
+ Object,
+ attributes(
+ type_key,
+ type_index,
+ type_final,
+ type_register,
+ type_child_slots,
+ type_child_slots_can_overflow
+ )
+)]
pub fn derive_object(input: TokenStream) -> TokenStream {
object_macros::derive_object(input)
}
diff --git a/rust/tvm-ffi-macros/src/object_macros.rs
b/rust/tvm-ffi-macros/src/object_macros.rs
index dace824b..807cd476 100644
--- a/rust/tvm-ffi-macros/src/object_macros.rs
+++ b/rust/tvm-ffi-macros/src/object_macros.rs
@@ -37,8 +37,61 @@ pub fn derive_object(input: proc_macro::TokenStream) ->
TokenStream {
None => false,
};
- // type index can be optional
- // for now we make it required for static index
+ // Whether this type is defined here, so that it registers its key, as C++
+ // `TVM_FFI_DECLARE_OBJECT_INFO` does. Without it, the type binds a key
+ // that its defining library, usually C++, has registered.
+ let type_register = match get_attr(&derive_input, "type_register") {
+ Some(attr) if matches!(attr.parse_meta(), Ok(syn::Meta::Path(_))) =>
true,
+ Some(_) => panic!("Expect #[type_register] attribute"),
+ None => false,
+ };
+
+ // Reserved child slots and whether children may exceed them, as C++
+ // `_type_child_slots` and `_type_child_slots_can_overflow` declare them.
+ let child_slots = match get_attr(&derive_input,
"type_child_slots").map(|a| a.parse_meta()) {
+ Some(Ok(syn::Meta::NameValue(syn::MetaNameValue {
+ lit: syn::Lit::Int(n),
+ ..
+ }))) => n
+ .base10_parse::<i32>()
+ .expect("Expect #[type_child_slots = <non-negative integer>]"),
+ Some(_) => panic!("Expect #[type_child_slots = <non-negative
integer>]"),
+ None => 0,
+ };
+ let child_slots_can_overflow =
+ match get_attr(&derive_input, "type_child_slots_can_overflow").map(|a|
a.parse_meta()) {
+ Some(Ok(syn::Meta::NameValue(syn::MetaNameValue {
+ lit: syn::Lit::Bool(b),
+ ..
+ }))) => b.value,
+ Some(_) => panic!("Expect #[type_child_slots_can_overflow =
<bool>]"),
+ None => true,
+ };
+ assert!(
+ child_slots >= 0,
+ "Expect #[type_child_slots = <non-negative integer>]"
+ );
+ assert!(
+ !(type_final && child_slots > 0),
+ "a final object type cannot reserve child slots"
+ );
+ assert!(
+ type_register
+ || (get_attr(&derive_input, "type_child_slots").is_none()
+ && get_attr(&derive_input,
"type_child_slots_can_overflow").is_none()),
+ "#[type_child_slots] and #[type_child_slots_can_overflow] need
#[type_register]"
+ );
+ // we expect base always to be the first field
+ let base_ty = match &derive_input.data {
+ syn::Data::Struct(s) => s.fields.iter().next().map(|f| f.ty.clone()),
+ _ => None,
+ }
+ .expect("First field must be `<base_name>: <ObjectCoreType>`");
+
+ // A type with a static index has it; one defined here registers its key
+ // under its parent on first use, or takes the index the key already has,
+ // as C++ `TVM_FFI_DECLARE_OBJECT_INFO` does; any other type binds the
+ // index its defining library registered.
let type_index_tokens = match get_attr(&derive_input,
"type_index").map(attr_to_expr) {
Some(type_index) => {
let type_index_expr =
@@ -50,6 +103,42 @@ pub fn derive_object(input: proc_macro::TokenStream) ->
TokenStream {
}
}
}
+ None if type_register => {
+ quote! {
+ #[inline]
+ fn type_index() -> i32 {
+ static TYPE_INDEX: std::sync::LazyLock<i32> =
std::sync::LazyLock::new(||
+ unsafe {
+ // The parent first, since registering it takes the
+ // same lock.
+ let parent =
+ <#base_ty as
#tvm_ffi_crate::object::ObjectCore>::type_index();
+ let _registering =
#tvm_ffi_crate::object::TYPE_REGISTRATION
+ .lock()
+ .unwrap_or_else(|e| e.into_inner());
+ let type_key_arg =
+
#tvm_ffi_crate::tvm_ffi_sys::TVMFFIByteArray::from_str(#type_key);
+ let tindex =
#tvm_ffi_crate::tvm_ffi_sys::TVMFFITypeGetOrAllocIndex(
+ &type_key_arg,
+ -1,
+ <#struct_name as
#tvm_ffi_crate::object::ObjectCore>::TYPE_DEPTH,
+ #child_slots,
+ #child_slots_can_overflow as i32,
+ parent,
+ );
+ if tindex < 0 {
+ panic!(
+ "Failed to get or allocate type index for
type key: {}",
+ #type_key
+ );
+ }
+ tindex
+ }
+ );
+ *TYPE_INDEX
+ }
+ }
+ }
None => {
quote! {
#[inline]
@@ -59,11 +148,15 @@ pub fn derive_object(input: proc_macro::TokenStream) ->
TokenStream {
let type_key_arg =
#tvm_ffi_crate::tvm_ffi_sys::TVMFFIByteArray::from_str(#type_key);
let mut tindex = 0;
- let ret =
#tvm_ffi_crate::tvm_ffi_sys::TVMFFITypeKeyToIndex(
+ let ret =
#tvm_ffi_crate::tvm_ffi_sys::TVMFFITypeKeyToIndex(
&type_key_arg, &mut tindex
);
if ret != 0 {
- panic!("Failed to get type index for type key:
{}", #type_key);
+ panic!(
+ "Type key {} is not registered: load the
library that defines \
+ it first, or add #[type_register] if this
crate defines it",
+ #type_key
+ );
}
tindex
}
diff --git a/rust/tvm-ffi-sys/src/c_api.rs b/rust/tvm-ffi-sys/src/c_api.rs
index f7b7c334..7a0dc2f8 100644
--- a/rust/tvm-ffi-sys/src/c_api.rs
+++ b/rust/tvm-ffi-sys/src/c_api.rs
@@ -466,6 +466,14 @@ unsafe extern "C" {
pub fn TVMFFISetCustomAllocator(allocator: *mut TVMFFICustomAllocator) ->
i32;
pub fn TVMFFITypeKeyToIndex(type_key: *const TVMFFIByteArray, out_tindex:
*mut i32) -> i32;
+ pub fn TVMFFITypeGetOrAllocIndex(
+ type_key: *const TVMFFIByteArray,
+ static_type_index: i32,
+ type_depth: i32,
+ num_child_slots: i32,
+ child_slots_can_overflow: i32,
+ parent_type_index: i32,
+ ) -> i32;
pub fn TVMFFITypeRegisterAttr(
type_index: i32,
attr_name: *const TVMFFIByteArray,
diff --git a/rust/tvm-ffi/src/object.rs b/rust/tvm-ffi/src/object.rs
index f05d1a09..93306bdc 100644
--- a/rust/tvm-ffi/src/object.rs
+++ b/rust/tvm-ffi/src/object.rs
@@ -25,6 +25,12 @@ pub use tvm_ffi_sys::TVMFFITypeIndex as TypeIndex;
/// Object related ABI handling
use tvm_ffi_sys::{TVMFFIAny, TVMFFIGetTypeInfo, TVMFFIObject,
COMBINED_REF_COUNT_BOTH_ONE};
+/// Held while `#[derive(Object)]` with `#[type_register]` registers a type, so
+/// that registrations from Rust do not race each other: the runtime does not
+/// lock its type table and leaves that to callers.
+#[doc(hidden)]
+pub static TYPE_REGISTRATION: std::sync::Mutex<()> = std::sync::Mutex::new(());
+
/// Object type is by default the TVMFFIObject
#[repr(C)]
pub struct Object {
diff --git a/rust/tvm-ffi/tests/test_match_any_ast_lookup.rs
b/rust/tvm-ffi/tests/test_match_any_ast_lookup.rs
index 7d69189c..ad7ce02a 100644
--- a/rust/tvm-ffi/tests/test_match_any_ast_lookup.rs
+++ b/rust/tvm-ffi/tests/test_match_any_ast_lookup.rs
@@ -20,38 +20,19 @@
use tvm_ffi::derive::{Object, ObjectRef};
use tvm_ffi::object::{Object as ObjectBase, ObjectArc, ObjectCore};
use tvm_ffi::{match_any, Any, Array, Shape, TypeIndex};
-use tvm_ffi_sys::TVMFFIByteArray;
-
-unsafe extern "C" {
- fn TVMFFITypeGetOrAllocIndex(
- type_key: *const TVMFFIByteArray,
- static_type_index: i32,
- type_depth: i32,
- num_child_slots: i32,
- child_slots_can_overflow: i32,
- parent_type_index: i32,
- ) -> i32;
-}
#[repr(C)]
#[derive(Object)]
#[type_key = "testing.match_any.Expr"]
+#[type_register]
+#[type_child_slots = 20]
+#[type_child_slots_can_overflow = false]
struct ExprObj {
base: ObjectBase,
}
-fn register_type<T: ObjectCore>(num_child_slots: i32, parent_type_index: i32)
-> i32 {
- let type_key = unsafe { TVMFFIByteArray::from_str(T::TYPE_KEY) };
- let type_index = unsafe {
- TVMFFITypeGetOrAllocIndex(
- &type_key,
- -1,
- T::TYPE_DEPTH,
- num_child_slots,
- 0,
- parent_type_index,
- )
- };
+fn register_type<T: ObjectCore>() -> i32 {
+ let type_index = T::type_index();
assert!(type_index >= TypeIndex::kTVMFFIStaticObjectBegin as i32);
type_index
}
@@ -62,6 +43,7 @@ macro_rules! define_expr_leaves {
#[repr(C)]
#[derive(Object)]
#[type_key = $type_key]
+ #[type_register]
#[type_final]
struct $object {
base: ExprObj,
@@ -87,10 +69,9 @@ macro_rules! define_expr_leaves {
)+
fn register_expr_types() {
- let expr_type_index =
- register_type::<ExprObj>(20,
TypeIndex::kTVMFFIStaticObjectBegin as i32);
+ register_type::<ExprObj>();
$(
- register_type::<$object>(0, expr_type_index);
+ register_type::<$object>();
)+
}
};
diff --git a/rust/tvm-ffi/tests/test_object_registration.rs
b/rust/tvm-ffi/tests/test_object_registration.rs
new file mode 100644
index 00000000..2226c5ab
--- /dev/null
+++ b/rust/tvm-ffi/tests/test_object_registration.rs
@@ -0,0 +1,190 @@
+/*
+ * 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 tvm_ffi::derive::{Object, ObjectRef};
+use tvm_ffi::object::{is_instance_of, ObjectRef};
+use tvm_ffi::*;
+use tvm_ffi_sys::{TVMFFIByteArray, TVMFFIGetTypeInfo, TVMFFITypeKeyToIndex};
+
+// Type keys that nothing else registers: deriving `Object` with
+// `#[type_register]` registers them on first use, under their parents, as C++
+// `TVM_FFI_DECLARE_OBJECT_INFO` does.
+
+#[repr(C)]
+#[derive(Object)]
+#[type_key = "testing.rust.RegisteredBase"]
+#[type_register]
+#[type_child_slots = 2]
+#[type_child_slots_can_overflow = false]
+struct RegisteredBaseObj {
+ base: Object,
+ value: i64,
+}
+
+#[repr(C)]
+#[derive(ObjectRef, Clone)]
+struct RegisteredBase {
+ data: ObjectArc<RegisteredBaseObj>,
+}
+
+#[repr(C)]
+#[derive(Object)]
+#[type_key = "testing.rust.RegisteredLeaf"]
+#[type_register]
+#[type_final]
+struct RegisteredLeafObj {
+ base: RegisteredBaseObj,
+ extra: i64,
+}
+
+#[repr(C)]
+#[derive(ObjectRef, Clone)]
+struct RegisteredLeaf {
+ data: ObjectArc<RegisteredLeafObj>,
+}
+
+tvm_ffi::impl_object_upcast!(RegisteredLeaf => RegisteredBase);
+
+fn registered_index(type_key: &str) -> i32 {
+ let key = unsafe { TVMFFIByteArray::from_str(type_key) };
+ let mut index = -1;
+ assert_eq!(unsafe { TVMFFITypeKeyToIndex(&key, &mut index) }, 0);
+ index
+}
+
+#[test]
+fn test_derived_object_types_are_registered_on_first_use() {
+ // The leaf comes first: registering it registers its parent before it.
+ let leaf_index = RegisteredLeafObj::type_index();
+ let base_index = RegisteredBaseObj::type_index();
+ assert!(base_index >= TypeIndex::kTVMFFIDynObjectBegin as i32);
+ assert_eq!(registered_index("testing.rust.RegisteredBase"), base_index);
+ assert_eq!(registered_index("testing.rust.RegisteredLeaf"), leaf_index);
+ // The leaf takes the first of its parent's reserved child slots.
+ assert_eq!(leaf_index, base_index + 1);
+ // Registration happens once; the index is stable.
+ assert_eq!(RegisteredLeafObj::type_index(), leaf_index);
+
+ let info = unsafe { &*TVMFFIGetTypeInfo(leaf_index) };
+ assert_eq!(info.type_depth, RegisteredLeafObj::TYPE_DEPTH);
+ let parent = unsafe { &**info.type_acenstors.add(1) };
+ assert_eq!(parent.type_index, base_index);
+
+ assert!(is_instance_of::<RegisteredBaseObj>(leaf_index));
+ assert!(is_instance_of::<Object>(base_index));
+ assert!(!is_instance_of::<RegisteredLeafObj>(base_index));
+}
+
+#[test]
+fn test_registered_object_types_round_trip_through_any() {
+ let leaf = RegisteredLeaf {
+ data: ObjectArc::new(RegisteredLeafObj {
+ base: RegisteredBaseObj {
+ base: Object::new(),
+ value: 7,
+ },
+ extra: 8,
+ }),
+ };
+ let any = Any::from(leaf);
+ assert_eq!(any.type_index(), RegisteredLeafObj::type_index());
+ let base = RegisteredBase::try_from(any).unwrap();
+ assert_eq!(base.data.value, 7);
+ let obj: ObjectRef = base.try_cast().unwrap();
+ let leaf: RegisteredLeaf = obj.try_cast().unwrap();
+ assert_eq!(leaf.data.extra, 8);
+}
+
+// Types defined here, registered from several threads at once.
+macro_rules! concurrent_types {
+ ($($name:ident = $key:literal),*) => {
+ $(
+ #[repr(C)]
+ #[derive(Object)]
+ #[type_key = $key]
+ #[type_register]
+ struct $name {
+ base: Object,
+ }
+ )*
+ };
+}
+
+concurrent_types!(
+ Concurrent0 = "testing.rust.Concurrent0",
+ Concurrent1 = "testing.rust.Concurrent1",
+ Concurrent2 = "testing.rust.Concurrent2",
+ Concurrent3 = "testing.rust.Concurrent3",
+ Concurrent4 = "testing.rust.Concurrent4",
+ Concurrent5 = "testing.rust.Concurrent5",
+ Concurrent6 = "testing.rust.Concurrent6",
+ Concurrent7 = "testing.rust.Concurrent7"
+);
+
+#[test]
+fn test_registration_from_threads() {
+ let registrations: [fn() -> i32; 8] = [
+ <Concurrent0 as ObjectCore>::type_index,
+ <Concurrent1 as ObjectCore>::type_index,
+ <Concurrent2 as ObjectCore>::type_index,
+ <Concurrent3 as ObjectCore>::type_index,
+ <Concurrent4 as ObjectCore>::type_index,
+ <Concurrent5 as ObjectCore>::type_index,
+ <Concurrent6 as ObjectCore>::type_index,
+ <Concurrent7 as ObjectCore>::type_index,
+ ];
+ let barrier = std::sync::Barrier::new(registrations.len());
+ let indices: Vec<i32> = std::thread::scope(|scope| {
+ let handles: Vec<_> = registrations
+ .iter()
+ .map(|register| {
+ let barrier = &barrier;
+ scope.spawn(move || {
+ barrier.wait();
+ register()
+ })
+ })
+ .collect();
+ handles.into_iter().map(|h| h.join().unwrap()).collect()
+ });
+ for (i, &index) in indices.iter().enumerate() {
+ assert_eq!(
+ index,
+ registered_index(&format!("testing.rust.Concurrent{i}"))
+ );
+ assert!(
+ !indices[..i].contains(&index),
+ "type indices must be distinct"
+ );
+ }
+}
+
+// A type that binds a key its defining library registers, without
+// `#[type_register]`, does not register it.
+#[repr(C)]
+#[derive(Object)]
+#[type_key = "testing.rust.NeverRegistered"]
+struct UnregisteredBindingObj {
+ base: Object,
+}
+
+#[test]
+#[should_panic(expected = "is not registered: load the library that defines it
first")]
+fn test_binding_does_not_register() {
+ let _ = <UnregisteredBindingObj as ObjectCore>::type_index();
+}
