diff --git a/docs/guides/rust_lang_guide.md b/docs/guides/rust_lang_guide.md index f2ae5eff5..226b5e310 100644 --- a/docs/guides/rust_lang_guide.md +++ b/docs/guides/rust_lang_guide.md @@ -137,6 +137,30 @@ let my_func = Function::from_packed(|args: &[AnyView]| -> Result { Function::register_global("my_custom_func", my_func)?; ``` +### Reflected Type Methods + +Libraries that register their API through the C++ reflection registry +(`refl::ObjectDef().def(...)`) store methods in a per-type method table +rather than the global function table. Resolve them by type key (or type +index) and method name; constructors registered via `refl::init` are +reachable under the reserved name `__ffi_init__`: + +```rust +use tvm_ffi::{AnyView, Function}; + +// Resolve the reflected constructor and construct an instance +let ctor = Function::from_type_key_method("testing.TestIntPair", "__ffi_init__")?; +let pair = ctor.call_tuple((1i64, 2i64))?; + +// Resolve an instance method; the first packed argument is the object itself +let sum = Function::from_type_key_method("testing.TestIntPair", "sum")?; +let result = sum.call_packed(&[AnyView::from(&pair)])?; +assert_eq!(i64::try_from(result)?, 3); +``` + +`Function::from_type_method(type_index, name)` performs the same lookup when +the type index is already known (e.g. from `Any::type_index`). + ### Type-Erased Functions Create functions from Rust closures: diff --git a/rust/tvm-ffi/src/function.rs b/rust/tvm-ffi/src/function.rs index 4af971bdf..2d37a2bdb 100644 --- a/rust/tvm-ffi/src/function.rs +++ b/rust/tvm-ffi/src/function.rs @@ -21,9 +21,11 @@ use crate::derive::{Object, ObjectRef}; use crate::error::{Error, Result}; use crate::function_internal::{AsPackedCallable, TupleAsPackedArgs}; use crate::object::{Object, ObjectArc, ObjectCore}; +use crate::type_traits::AnyCompatible; use tvm_ffi_sys::{ TVMFFIAny, TVMFFIByteArray, TVMFFIFunctionCell, TVMFFIFunctionCreate, TVMFFIFunctionGetGlobal, - TVMFFIFunctionSetGlobal, TVMFFIObjectHandle, TVMFFISafeCallType, TVMFFITypeIndex, + TVMFFIFunctionSetGlobal, TVMFFIGetTypeInfo, TVMFFIObjectHandle, TVMFFISafeCallType, + TVMFFITypeIndex, TVMFFITypeKeyToIndex, }; /// function object @@ -196,6 +198,83 @@ impl Function { } } + /// Look up a reflected method of a type by type index and method name + /// + /// Methods registered through the C++ reflection registry + /// (`refl::ObjectDef().def(...)`) live in the per-type method table + /// rather than the global function table. Constructors registered via + /// `refl::init` are reachable under the reserved name `__ffi_init__`. + /// For instance methods, the first packed argument is the object itself. + /// + /// `type_index` must be a registered type index (e.g. obtained from a + /// live object via `Any::type_index` or from a type key); the underlying + /// C API treats an unregistered index as a fatal error. + /// + /// # Arguments + /// * `type_index` - The type index of the type that owns the method + /// * `method_name` - The name of the method + /// + /// # Returns + /// * `Function` - The reflected method + pub fn from_type_method(type_index: i32, method_name: &str) -> Result { + unsafe { + let type_info = TVMFFIGetTypeInfo(type_index); + if type_info.is_null() { + crate::bail!( + crate::error::TYPE_ERROR, + "Cannot find type info for type_index={}", + type_index + ); + } + let type_info = &*type_info; + for i in 0..type_info.num_methods as usize { + let method_info = &*type_info.methods.add(i); + if method_info.name.as_str() != method_name { + continue; + } + if !::check_any_strict(&method_info.method) { + crate::bail!( + crate::error::TYPE_ERROR, + "Method `{}` of type `{}` is not a Function", + method_name, + type_info.type_key.as_str() + ); + } + // the table entry stores the method as a non-owning AnyView; + // copy out a strong reference + return Ok(::copy_from_any_view_after_check( + &method_info.method, + )); + } + crate::bail!( + crate::error::TYPE_ERROR, + "Cannot find method `{}` of type `{}`", + method_name, + type_info.type_key.as_str() + ); + } + } + + /// Look up a reflected method of a type by type key and method name + /// + /// Same as [`Function::from_type_method`], but resolves `type_key` to a + /// type index first. + /// + /// # Arguments + /// * `type_key` - The type key of the type that owns the method + /// * `method_name` - The name of the method + /// + /// # Returns + /// * `Function` - The reflected method + pub fn from_type_key_method(type_key: &str, method_name: &str) -> Result { + unsafe { + let type_key_arg = TVMFFIByteArray::from_str(type_key); + let mut type_index: i32 = 0; + crate::check_safe_call!(TVMFFITypeKeyToIndex(&type_key_arg, &mut type_index))?; + Self::from_type_method(type_index, method_name) + } + } + /// Register a function as a global function /// # Arguments /// * `name` - The name of the function diff --git a/rust/tvm-ffi/tests/test_function.rs b/rust/tvm-ffi/tests/test_function.rs index 42205eeba..9ed531727 100644 --- a/rust/tvm-ffi/tests/test_function.rs +++ b/rust/tvm-ffi/tests/test_function.rs @@ -131,6 +131,44 @@ fn test_function_echo_tensor_typed() { assert_eq!(result_data[3], 4.0); } +#[test] +fn test_function_from_type_key_method_ctor_and_method() { + // constructors registered via refl::init are reachable as `__ffi_init__` + let ctor = Function::from_type_key_method("testing.TestIntPair", "__ffi_init__").unwrap(); + let pair = ctor.call_tuple((1i64, 2i64)).unwrap(); + // instance method: the first packed argument is the object itself + let sum = Function::from_type_key_method("testing.TestIntPair", "sum").unwrap(); + let result = sum.call_packed(&[AnyView::from(&pair)]).unwrap(); + assert_eq!(i64::try_from(result).unwrap(), 3); +} + +#[test] +fn test_function_from_type_method_by_index() { + let ctor = Function::from_type_key_method("testing.TestIntPair", "__ffi_init__").unwrap(); + let pair = ctor.call_tuple((5i64, 7i64)).unwrap(); + let sum = Function::from_type_method(pair.type_index(), "sum").unwrap(); + let result = sum.call_packed(&[AnyView::from(&pair)]).unwrap(); + assert_eq!(i64::try_from(result).unwrap(), 12); +} + +#[test] +fn test_function_from_type_method_unknown_method() { + let error = Function::from_type_key_method("testing.TestIntPair", "nonexistent_method") + .err() + .unwrap(); + assert_eq!(error.kind(), TYPE_ERROR); + assert!(error.message().contains("nonexistent_method")); + assert!(error.message().contains("testing.TestIntPair")); +} + +#[test] +fn test_function_from_type_key_method_unknown_type_key() { + let error = Function::from_type_key_method("testing.NonExistentType", "sum") + .err() + .unwrap(); + assert!(error.message().contains("testing.NonExistentType")); +} + fn testing_add_one(x: i32) -> Result { Ok(x + 1) }