tvm_ffi/
function_internal.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, ArgTryFromAnyView};
20use crate::error::{Error, Result, INTERNAL_ERROR};
21use crate::object::ObjectRefCore;
22use crate::rvalue_ref::RValueRef;
23use crate::string::{Bytes, String};
24use crate::type_traits::{AnyCompatible, ContainerElement};
25
26//------------------------------------------------------------------------
27// PackedCallable
28//------------------------------------------------------------------------
29/// The error to raise for a panic caught at an FFI boundary, where unwinding
30/// into the caller would abort the process.
31#[doc(hidden)]
32pub fn panic_to_error(payload: Box<dyn std::any::Any + Send>) -> Error {
33    let message = payload
34        .downcast_ref::<&str>()
35        .copied()
36        .or_else(|| {
37            payload
38                .downcast_ref::<std::string::String>()
39                .map(|s| s.as_str())
40        })
41        .unwrap_or("unknown payload");
42    Error::new(INTERNAL_ERROR, &format!("panicked: {message}"), "")
43}
44
45pub trait AsPackedCallable<I, O> {
46    // Call the function in packed convention
47    fn call_packed(&self, packed_args: &[AnyView]) -> Result<Any>;
48}
49
50#[inline]
51pub fn call_packed_callable<Fun, I, O>(func: Fun, packed_args: &[AnyView]) -> Result<Any>
52where
53    Fun: AsPackedCallable<I, O>,
54{
55    func.call_packed(packed_args)
56}
57
58macro_rules! impl_as_packed_callable {
59    ($len:literal; $($t:ident),*) => {
60        impl<Fun, $($t,)* Out> AsPackedCallable<($($t,)*), Out> for Fun
61        where
62            Fun: Fn($($t,)*) -> Result<Out> + 'static,
63            Any: From<Out>,
64            $($t: ArgTryFromAnyView),*
65        {
66            fn call_packed(&self, packed_args: &[AnyView]) -> Result<Any>
67            {
68                crate::ensure!(
69                    packed_args.len() == $len, crate::error::VALUE_ERROR,
70                    "Expected {} arguments, got {}", $len, packed_args.len()
71                );
72                // Expand the function call, consuming the iterator.
73                let mut _arg_iter = packed_args.iter().enumerate();
74                let ret_value = self(
75                    $({
76                        // unwrap is safe due to the length check above
77                        let (i, view) = _arg_iter.next().unwrap();
78                        $t::try_from_any_view(view, i)?
79                    }),*
80                )?;
81                Ok(Any::from(ret_value))
82            }
83        }
84    }
85}
86
87impl_as_packed_callable!(0;);
88impl_as_packed_callable!(1; T0);
89impl_as_packed_callable!(2; T0, T1);
90impl_as_packed_callable!(3; T0, T1, T2);
91impl_as_packed_callable!(4; T0, T1, T2, T3);
92impl_as_packed_callable!(5; T0, T1, T2, T3, T4);
93impl_as_packed_callable!(6; T0, T1, T2, T3, T4, T5);
94impl_as_packed_callable!(7; T0, T1, T2, T3, T4, T5, T6);
95impl_as_packed_callable!(8; T0, T1, T2, T3, T4, T5, T6, T7);
96
97//--------------------------------------------------------------
98// IntoArgHolder, helper to convert to canonical holding type
99//
100// This is needed sometimes for reference types that may need to
101// be converted to value types.
102//--------------------------------------------------------------
103pub trait IntoArgHolder {
104    type Target;
105    fn into_arg_holder(self) -> Self::Target;
106}
107
108crate::impl_into_arg_holder_default!(
109    (),
110    bool,
111    i8,
112    i16,
113    i32,
114    i64,
115    isize,
116    u8,
117    u16,
118    u32,
119    u64,
120    usize,
121    f32,
122    f64,
123    String,
124    Bytes,
125    Any,
126    crate::DLDataType,
127    crate::DLDevice
128);
129
130// string will be converted to String for argument passing
131impl IntoArgHolder for &str {
132    type Target = String;
133    fn into_arg_holder(self) -> Self::Target {
134        String::from(self)
135    }
136}
137
138// string will be converted to String for argument passing
139impl IntoArgHolder for &[u8] {
140    type Target = Bytes;
141    fn into_arg_holder(self) -> Self::Target {
142        Bytes::from(self)
143    }
144}
145
146// helper trait to implement IntoArgHolderTuple to apply into_arg_holder to each element
147pub trait IntoArgHolderTuple {
148    type Target;
149    fn into_arg_holder_tuple(self) -> Self::Target;
150}
151
152macro_rules! impl_into_arg_holder_tuple {
153    ( $($T:ident),* ; $($idx:tt),* ) => {
154        impl<$($T),*> $crate::function_internal::IntoArgHolderTuple for ($($T,)*)
155        where
156            $($T: IntoArgHolder),* {
157            type Target = ($($T::Target,)*);
158
159            fn into_arg_holder_tuple(self) -> Self::Target {
160                ($(self.$idx.into_arg_holder(),)*)
161            }
162        }
163    };
164}
165
166impl_into_arg_holder_tuple!(;);
167impl_into_arg_holder_tuple!(T0; 0);
168impl_into_arg_holder_tuple!(T0, T1; 0, 1);
169impl_into_arg_holder_tuple!(T0, T1, T2; 0, 1, 2);
170impl_into_arg_holder_tuple!(T0, T1, T2, T3; 0, 1, 2, 3);
171impl_into_arg_holder_tuple!(T0, T1, T2, T3, T4; 0, 1, 2, 3, 4);
172impl_into_arg_holder_tuple!(T0, T1, T2, T3, T4, T5; 0, 1, 2, 3, 4, 5);
173impl_into_arg_holder_tuple!(T0, T1, T2, T3, T4, T5, T6; 0, 1, 2, 3, 4, 5, 6);
174impl_into_arg_holder_tuple!(T0, T1, T2, T3, T4, T5, T6, T7; 0, 1, 2, 3, 4, 5, 6, 7);
175
176//------------------------------------------------------------
177// ArgIntoRef
178//
179// Helper to turn argument type to reference type
180// This is effectively AsRef<T> but removes the need of T
181//-----------------------------------------------------------
182pub trait ArgIntoRef {
183    type Target;
184    fn to_ref(&self) -> &Self::Target;
185}
186
187/// Convert a canonical argument holder into its packed ABI view.
188#[doc(hidden)]
189pub trait PackedArg {
190    fn as_packed_arg(&self) -> AnyView<'_>;
191}
192
193impl<T: AnyCompatible> PackedArg for T {
194    #[inline]
195    fn as_packed_arg(&self) -> AnyView<'_> {
196        AnyView::from(self)
197    }
198}
199
200impl PackedArg for Any {
201    #[inline]
202    fn as_packed_arg(&self) -> AnyView<'_> {
203        AnyView::from(self)
204    }
205}
206
207impl<T> PackedArg for RValueRef<T>
208where
209    T: ObjectRefCore + AnyCompatible,
210{
211    #[inline]
212    fn as_packed_arg(&self) -> AnyView<'_> {
213        AnyView::from(self)
214    }
215}
216
217crate::impl_arg_into_ref!(
218    (),
219    bool,
220    i8,
221    i16,
222    i32,
223    i64,
224    isize,
225    u8,
226    u16,
227    u32,
228    u64,
229    usize,
230    f32,
231    f64,
232    String,
233    Bytes,
234    Any,
235    crate::DLDataType,
236    crate::DLDevice
237);
238
239// Generic holders require explicit implementations rather than scalar macro entries.
240impl<T: AnyCompatible> IntoArgHolder for Option<T> {
241    type Target = Self;
242    fn into_arg_holder(self) -> Self::Target {
243        self
244    }
245}
246
247impl<'a, T: AnyCompatible> IntoArgHolder for &'a Option<T> {
248    type Target = &'a Option<T>;
249    fn into_arg_holder(self) -> Self::Target {
250        self
251    }
252}
253
254impl<T: AnyCompatible> ArgIntoRef for Option<T> {
255    type Target = Self;
256    fn to_ref(&self) -> &Self::Target {
257        self
258    }
259}
260
261impl<T: AnyCompatible> ArgIntoRef for &Option<T> {
262    type Target = Option<T>;
263    fn to_ref(&self) -> &Self::Target {
264        self
265    }
266}
267
268impl<T: ContainerElement + Clone> IntoArgHolder for crate::Array<T> {
269    type Target = crate::Array<T>;
270    fn into_arg_holder(self) -> Self::Target {
271        self
272    }
273}
274impl<'a, T: ContainerElement + Clone> IntoArgHolder for &'a crate::Array<T> {
275    type Target = &'a crate::Array<T>;
276    fn into_arg_holder(self) -> Self::Target {
277        self
278    }
279}
280impl<T: ContainerElement + Clone> ArgIntoRef for crate::Array<T> {
281    type Target = crate::Array<T>;
282    fn to_ref(&self) -> &Self::Target {
283        self
284    }
285}
286impl<T: ContainerElement + Clone> ArgIntoRef for &crate::Array<T> {
287    type Target = crate::Array<T>;
288    fn to_ref(&self) -> &Self::Target {
289        self
290    }
291}
292
293impl<T> IntoArgHolder for RValueRef<T>
294where
295    T: ObjectRefCore + AnyCompatible,
296{
297    type Target = Self;
298    fn into_arg_holder(self) -> Self::Target {
299        self
300    }
301}
302
303impl<T> ArgIntoRef for RValueRef<T>
304where
305    T: ObjectRefCore + AnyCompatible,
306{
307    type Target = Self;
308    fn to_ref(&self) -> &Self::Target {
309        self
310    }
311}
312
313impl<K: ContainerElement, V: ContainerElement> IntoArgHolder for crate::Map<K, V> {
314    type Target = crate::Map<K, V>;
315    fn into_arg_holder(self) -> Self::Target {
316        self
317    }
318}
319impl<'a, K: ContainerElement, V: ContainerElement> IntoArgHolder for &'a crate::Map<K, V> {
320    type Target = &'a crate::Map<K, V>;
321    fn into_arg_holder(self) -> Self::Target {
322        self
323    }
324}
325impl<K: ContainerElement, V: ContainerElement> ArgIntoRef for crate::Map<K, V> {
326    type Target = crate::Map<K, V>;
327    fn to_ref(&self) -> &Self::Target {
328        self
329    }
330}
331impl<K: ContainerElement, V: ContainerElement> ArgIntoRef for &crate::Map<K, V> {
332    type Target = crate::Map<K, V>;
333    fn to_ref(&self) -> &Self::Target {
334        self
335    }
336}
337
338//-----------------------------------------------------------
339// TupleAsPackedArgs
340//
341// Helper to turn tuple type to packed arguments
342//-----------------------------------------------------------
343pub trait TupleAsPackedArgs {
344    const LEN: usize;
345    fn fill_any_view<'a>(&'a self, any_view: &mut [AnyView<'a>]);
346}
347
348macro_rules! impl_tuple_as_packed_args {
349    ( $len:expr; $($T:ident),* ; $($idx:tt),* ) => {
350        impl<$($T),*> TupleAsPackedArgs for ($($T,)*)
351        where
352            $(
353                $T: ArgIntoRef,
354                $T::Target: PackedArg,
355            )*
356        {
357            const LEN: usize = $len;
358
359            fn fill_any_view<'a>(&'a self, _any_view: &mut [AnyView<'a>]) {
360                $(
361                    _any_view[$idx] = self.$idx.to_ref().as_packed_arg();
362                )*
363            }
364        }
365    };
366}
367
368impl_tuple_as_packed_args!(0;;);
369impl_tuple_as_packed_args!(1; T0; 0);
370impl_tuple_as_packed_args!(2; T0, T1; 0, 1);
371impl_tuple_as_packed_args!(3; T0, T1, T2; 0, 1, 2);
372impl_tuple_as_packed_args!(4; T0, T1, T2, T3; 0, 1, 2, 3);
373impl_tuple_as_packed_args!(5; T0, T1, T2, T3, T4; 0, 1, 2, 3, 4);
374impl_tuple_as_packed_args!(6; T0, T1, T2, T3, T4, T5; 0, 1, 2, 3, 4, 5);
375impl_tuple_as_packed_args!(7; T0, T1, T2, T3, T4, T5, T6; 0, 1, 2, 3, 4, 5, 6);
376impl_tuple_as_packed_args!(8; T0, T1, T2, T3, T4, T5, T6, T7; 0, 1, 2, 3, 4, 5, 6, 7);
tvm_ffi/function_internal.rs

tvm_ffi/
function_internal.rs
