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