tvm_ffi/
function.rs

1/*
2 * Licensed to the Apache Software Foundation (ASF) under one
3 * or more contributor license agreements.  See the NOTICE file
4 * distributed with this work for additional information
5 * regarding copyright ownership.  The ASF licenses this file
6 * to you under the Apache License, Version 2.0 (the
7 * "License"); you may not use this file except in compliance
8 * with the License.  You may obtain a copy of the License at
9 *
10 *   http://www.apache.org/licenses/LICENSE-2.0
11 *
12 * Unless required by applicable law or agreed to in writing,
13 * software distributed under the License is distributed on an
14 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15 * KIND, either express or implied.  See the License for the
16 * specific language governing permissions and limitations
17 * under the License.
18 */
19use crate::any::{Any, AnyView};
20use crate::derive::{Object, ObjectRef};
21use crate::error::{Error, Result};
22use crate::function_internal::{AsPackedCallable, TupleAsPackedArgs};
23use crate::object::{Object, ObjectArc, ObjectCore};
24use tvm_ffi_sys::{
25    TVMFFIAny, TVMFFIByteArray, TVMFFIFunctionCell, TVMFFIFunctionCreate, TVMFFIFunctionGetGlobal,
26    TVMFFIFunctionSetGlobal, TVMFFIObjectHandle, TVMFFISafeCallType, TVMFFITypeIndex,
27};
28
29/// function object
30#[repr(C)]
31#[derive(Object)]
32#[type_key = "ffi.Function"]
33#[type_index(TVMFFITypeIndex::kTVMFFIFunction)]
34pub struct FunctionObj {
35    object: Object,
36    cell: TVMFFIFunctionCell,
37}
38
39/// Error reference class
40#[derive(Clone, ObjectRef)]
41pub struct Function {
42    data: ObjectArc<FunctionObj>,
43}
44
45//------------------------------------------------------------------------
46// CallbackFunctionObjImpl
47//------------------------------------------------------------------------
48/// Special helper class to hold a generic callback state as Object
49/// Logically this Impl can be viewed as a FunctionObj
50/// We can create an ObjectArc<CallbackFunctionObjImpl<F>> so the deleter
51/// can correctly delete the entire object including callback part
52/// then we will convert to ObjectArc<FunctionObj> to be used as function
53#[repr(C)]
54struct CallbackFunctionObjImpl<F: Fn(&[AnyView]) -> Result<Any> + 'static> {
55    function: FunctionObj,
56    callback: F,
57}
58
59impl<F: Fn(&[AnyView]) -> Result<Any> + 'static> CallbackFunctionObjImpl<F> {
60    pub fn from_callback(callback: F) -> Self {
61        Self {
62            function: FunctionObj {
63                object: Object::new(),
64                cell: TVMFFIFunctionCell {
65                    // specfic callback for F
66                    safe_call: Self::invoke_callback,
67                    cxx_call: std::ptr::null_mut(),
68                },
69            },
70            callback,
71        }
72    }
73
74    unsafe extern "C" fn invoke_callback(
75        handle: *mut std::ffi::c_void,
76        args: *const TVMFFIAny,
77        num_args: i32,
78        result: *mut TVMFFIAny,
79    ) -> i32 {
80        let this = &*(handle as *mut Self);
81        let packed_args = std::slice::from_raw_parts(args as *const AnyView, num_args as usize);
82        let ret_value = (this.callback)(packed_args);
83        match ret_value {
84            Ok(value) => {
85                *result = Any::into_raw_ffi_any(value);
86                0
87            }
88            Err(error) => {
89                Error::set_raised(&error);
90                -1
91            }
92        }
93    }
94}
95
96unsafe impl<F: Fn(&[AnyView]) -> Result<Any> + 'static> ObjectCore for CallbackFunctionObjImpl<F> {
97    const TYPE_KEY: &'static str = FunctionObj::TYPE_KEY;
98    const TYPE_DEPTH: i32 = FunctionObj::TYPE_DEPTH;
99    fn type_index() -> i32 {
100        FunctionObj::type_index()
101    }
102    unsafe fn object_header_mut(this: &mut Self) -> &mut tvm_ffi_sys::TVMFFIObject {
103        FunctionObj::object_header_mut(&mut this.function)
104    }
105}
106
107impl Function {
108    /// Call the function in packed format.
109    pub fn call_packed(&self, packed_args: &[AnyView]) -> Result<Any> {
110        unsafe {
111            let packed_args_ptr = packed_args.as_ptr() as *const TVMFFIAny;
112            let mut result = Any::new();
113            let ret_code = (self.data.cell.safe_call)(
114                ObjectArc::as_raw(&self.data) as *mut FunctionObj as *mut std::ffi::c_void,
115                packed_args_ptr,
116                packed_args.len() as i32,
117                Any::as_data_ptr(&mut result),
118            );
119            if ret_code == 0 {
120                Ok(result)
121            } else {
122                Err(Error::from_raised())
123            }
124        }
125    }
126
127    pub fn call_tuple<TupleType>(&self, tuple_args: TupleType) -> Result<Any>
128    where
129        TupleType: TupleAsPackedArgs,
130    {
131        // This is a workaround for Rust's requirement that stack allocation size
132        // must be known at compile time for generic types.
133        // While we know args_len is a constant, Rust doesn't allow us to directly
134        // declare [AnyView::new(); args_len] in generic contexts.
135        //
136        // We use a small vector optimization pattern:
137        // 1. First allocate a small stack buffer (stack_args)
138        // 2. If args_len exceeds STACK_LEN, allocate a heap buffer (heap_args)
139        // 3. Use the appropriate buffer based on size
140        //
141        // Since args_len is a compile-time constant, the compiler should optimize
142        // away the unused branch, making this approach efficient.
143        const STACK_LEN: usize = 4;
144        let mut stack_args = [AnyView::new(); STACK_LEN];
145        let mut heap_args = Vec::<AnyView>::new();
146        let args_len = <TupleType as TupleAsPackedArgs>::LEN;
147        // get packed arguments
148        let packed_args: &mut [AnyView] = if args_len <= STACK_LEN {
149            &mut stack_args[..args_len]
150        } else {
151            heap_args.resize(args_len, AnyView::new());
152            &mut heap_args[..args_len]
153        };
154        (&tuple_args).fill_any_view(packed_args);
155        self.call_packed(packed_args)
156    }
157    /// Call function with compile-time known argument count
158    /// This is an optimized version of call_tuple for when the argument count
159    /// is known at compile time, avoiding the small vector optimization overhead.
160    ///
161    /// # Arguments
162    /// * `tuple_args` - The tuple arguments
163    ///
164    /// # Returns
165    /// * `Any` - The result
166    pub fn call_tuple_with_len<const LEN: usize, TupleType>(
167        &self,
168        tuple_args: TupleType,
169    ) -> Result<Any>
170    where
171        TupleType: TupleAsPackedArgs,
172    {
173        let mut packed_args = [AnyView::new(); LEN];
174        (&tuple_args).fill_any_view(&mut packed_args);
175        self.call_packed(&packed_args)
176    }
177    /// Get global function by name
178    /// This function will throw an error if the function is not found.
179    ///
180    /// # Arguments
181    /// * `name` - The name of the function
182    ///
183    /// # Returns
184    /// * `Function` - The global function
185    pub fn get_global(name: &str) -> Result<Function> {
186        unsafe {
187            let name_arg = TVMFFIByteArray::from_str(name);
188            let mut result: TVMFFIObjectHandle = ::std::ptr::null_mut();
189            crate::check_safe_call!(TVMFFIFunctionGetGlobal(&name_arg, &mut result))?;
190            if result.is_null() {
191                crate::bail!(crate::error::RUNTIME_ERROR, "Function {} not found", name);
192            }
193            Ok(Self {
194                data: ObjectArc::<FunctionObj>::from_raw(result as *mut FunctionObj),
195            })
196        }
197    }
198
199    /// Register a function as a global function
200    /// # Arguments
201    /// * `name` - The name of the function
202    /// * `func` - The function to register
203    ///
204    /// # Returns
205    /// * `Result<()>` - The result of the registration
206    pub fn register_global(name: &str, func: Function) -> Result<()> {
207        unsafe {
208            let name_arg = TVMFFIByteArray::from_str(name);
209            let can_override = 0;
210            crate::check_safe_call!(TVMFFIFunctionSetGlobal(
211                &name_arg,
212                ObjectArc::as_raw(&func.data) as *mut FunctionObj as TVMFFIObjectHandle,
213                can_override
214            ))?;
215            Ok(())
216        }
217    }
218    /// Construct a function from a packed function
219    /// # Arguments
220    /// * `func` - The packed function in signature of `Fn(&[AnyView]) -> Result<Any>`
221    ///
222    /// # Returns
223    /// * `Function` - The function
224    pub fn from_packed<F>(func: F) -> Self
225    where
226        F: Fn(&[AnyView]) -> Result<Any> + 'static,
227    {
228        unsafe {
229            let callback_arc = ObjectArc::new(CallbackFunctionObjImpl::from_callback(func));
230            let func_arc = ObjectArc::<FunctionObj>::from_raw(
231                ObjectArc::into_raw(callback_arc) as *mut FunctionObj
232            );
233            Self { data: func_arc }
234        }
235    }
236
237    /// Construct a function from a typed function
238    /// # Arguments
239    /// * `func` - The typed function with function signature of `F(T0, T1, ...) -> Result<O>`
240    ///
241    /// # Returns
242    /// * `Function` - The function
243    pub fn from_typed<F, I, O>(func: F) -> Self
244    where
245        F: AsPackedCallable<I, O> + 'static,
246    {
247        let closure = move |packed_args: &[AnyView]| -> Result<Any> {
248            let ret_value = func.call_packed(packed_args)?;
249            Ok(ret_value)
250        };
251        Self::from_packed(closure)
252    }
253
254    /// # Safety
255    ///
256    /// `handle` must be a valid pointer (or null) that is compatible with
257    /// `safe_call` and `deleter`. The caller must ensure the handle outlives
258    /// the returned `Function` (or that `deleter` properly frees it).
259    pub unsafe fn from_extern_c(
260        handle: *mut std::ffi::c_void,
261        safe_call: TVMFFISafeCallType,
262        deleter: Option<unsafe extern "C" fn(*mut std::ffi::c_void)>,
263    ) -> Self {
264        unsafe {
265            let mut out_handle: TVMFFIObjectHandle = std::ptr::null_mut();
266            crate::check_safe_call!(TVMFFIFunctionCreate(
267                handle,
268                safe_call,
269                deleter,
270                &mut out_handle
271            ))
272            .unwrap();
273            Self {
274                data: ObjectArc::<FunctionObj>::from_raw(out_handle as *mut FunctionObj),
275            }
276        }
277    }
278}
tvm_ffi/function.rs

tvm_ffi/
function.rs
