This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm-ffi.git
The following commit(s) were added to refs/heads/main by this push:
new c7a8cead [RUST] Add borrowed object node cast (#734)
c7a8cead is described below
commit c7a8cead2523c8a26bdd99f3a2aa8db7ca98919e
Author: Shushi Hong <[email protected]>
AuthorDate: Thu Sep 3 16:16:14 2026 -0400
[RUST] Add borrowed object node cast (#734)
This PR adds ObjectRefCore::as_node<N>() for zero-copy borrowed
object-node casts.
Checks such as is_zero currently clone an owning object handle before
calling try_cast. This performs an atomic reference-count increment and
decrement even though the caller only needs to inspect the node
temporarily.
The new API checks the runtime type directly from the existing object
header and returns a reference tied to the source handle. It does not
clone the handle, construct an owning cast result, or update the
reference count. Final node types, such as IntImmObj, use a fast exact
type-index comparison.
---
rust/tvm-ffi/src/object.rs | 20 ++++++++++++++++++++
rust/tvm-ffi/tests/test_cast.rs | 21 +++++++++++++++++++++
2 files changed, 41 insertions(+)
diff --git a/rust/tvm-ffi/src/object.rs b/rust/tvm-ffi/src/object.rs
index ddfbcf55..d690a64b 100644
--- a/rust/tvm-ffi/src/object.rs
+++ b/rust/tvm-ffi/src/object.rs
@@ -144,6 +144,21 @@ pub unsafe trait ObjectRefCore: Sized + Clone {
== ObjectArc::as_raw(Other::data(other)).cast::<()>()
}
}
+
+ /// Borrow the underlying object as node type `N` when its runtime type
matches.
+ ///
+ /// Unlike [`ObjectRefCast::try_cast`], this method neither consumes the
+ /// reference nor changes the object's reference count. The returned node
+ /// cannot outlive `self`.
+ #[inline(always)]
+ fn as_node<N: ObjectCore>(&self) -> Option<&N> {
+ let object = unsafe { ObjectArc::as_raw(Self::data(self)) };
+ let type_index = unsafe { (*object.cast::<TVMFFIObject>()).type_index
};
+ if !is_instance_of::<N>(type_index) {
+ return None;
+ }
+ Some(unsafe { &*object.cast::<N>() })
+ }
}
/// An owning, hashable identity key for an FFI object.
@@ -217,6 +232,11 @@ pub fn is_instance_of<Target:
ObjectCore>(object_type_index: i32) -> bool {
if object_type_index == target_type_index {
return true;
}
+ // A final type cannot have a separately registered subtype. Keep common
+ // borrowed checks, such as `IntImmObj`, to one integer comparison.
+ if Target::TYPE_FINAL {
+ return false;
+ }
let object_begin = TypeIndex::kTVMFFIStaticObjectBegin as i32;
// Only object types participate in the type hierarchy.
if object_type_index < object_begin || target_type_index < object_begin {
diff --git a/rust/tvm-ffi/tests/test_cast.rs b/rust/tvm-ffi/tests/test_cast.rs
index 29689808..4cd1ce90 100644
--- a/rust/tvm-ffi/tests/test_cast.rs
+++ b/rust/tvm-ffi/tests/test_cast.rs
@@ -54,6 +54,7 @@ struct TestBase {
#[repr(C)]
#[derive(Object)]
#[type_key = "testing.TestObjectDerived"]
+#[type_final]
struct TestDerivedObj {
base: TestBaseObj,
extra: i64,
@@ -142,6 +143,26 @@ fn test_upcast_downcast_roundtrip() {
assert_eq!(delete_counter.load(Ordering::Relaxed), 1);
}
+#[test]
+fn test_borrowed_node_cast_preserves_reference_count() {
+ let delete_counter = Arc::new(AtomicU32::new(0));
+ let base: TestBase = new_derived(7, 8, delete_counter.clone())
+ .try_cast()
+ .unwrap();
+ let strong_count = ObjectArc::strong_count(TestBase::data(&base));
+
+ let derived = base.as_node::<TestDerivedObj>().unwrap();
+ assert_eq!(derived.base.value, 7);
+ assert_eq!(derived.extra, 8);
+ let base_node = base.as_node::<TestBaseObj>().unwrap();
+ assert_eq!(base_node.value, 7);
+ assert_eq!(ObjectArc::strong_count(TestBase::data(&base)), strong_count);
+
+ let base_only = new_base(1, delete_counter.clone());
+ assert!(base_only.as_node::<TestDerivedObj>().is_none());
+ assert_eq!(ObjectArc::strong_count(TestBase::data(&base_only)), 1);
+}
+
#[test]
fn test_generated_borrow_and_upcast_conversions() {
let delete_counter = Arc::new(AtomicU32::new(0));
