Skip to main content

tvm_ffi/
any.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::error::Error;
20use crate::object;
21use crate::type_traits::AnyCompatible;
22use tvm_ffi_sys::TVMFFITypeIndex as TypeIndex;
23use tvm_ffi_sys::{TVMFFIAny, TVMFFIAnyViewToOwnedAny};
24
25/// Unmanaged Any that can hold reference to values
26#[derive(Copy, Clone)]
27#[repr(C)]
28pub struct AnyView<'a> {
29    data: TVMFFIAny,
30    /// needs to explicit mark lifetime to avoid lifetime mismatch
31    _phantom: std::marker::PhantomData<&'a ()>,
32}
33
34/// Managed Any that can hold reference to values
35#[repr(C)]
36pub struct Any {
37    data: TVMFFIAny,
38}
39
40//---------------------
41// AnyView
42//---------------------
43impl<'a> AnyView<'a> {
44    pub fn new() -> Self {
45        Self {
46            data: TVMFFIAny::new(),
47            _phantom: std::marker::PhantomData,
48        }
49    }
50
51    #[inline]
52    pub fn type_index(&self) -> i32 {
53        self.data.type_index
54    }
55
56    /// More strict version than try_from/try_into
57    ///
58    /// This function will not try to cast the type
59    /// and ensures invariance that the return value is only Some(T)
60    /// Any::from(T) contains exactly the same value value
61    ///
62    /// will return Some(T) if the type is exactly compatible with T
63    /// will return None if the type is not compatible with T
64    #[inline]
65    pub fn try_as<T>(&self) -> Option<T>
66    where
67        T: AnyCompatible,
68    {
69        unsafe {
70            if T::check_any_strict(&self.data) {
71                Some(T::copy_from_any_view_after_check(&self.data))
72            } else {
73                None
74            }
75        }
76    }
77
78    /// Get the strong count of the underlying object for testing/debugging purposes
79    ///
80    /// If the underlying object is not ref counted, return None
81    pub fn debug_strong_count(&self) -> Option<usize> {
82        unsafe {
83            if self.data.type_index >= TypeIndex::kTVMFFIStaticObjectBegin as i32 {
84                Some(object::unsafe_::strong_count(self.data.data_union.v_obj))
85            } else {
86                None
87            }
88        }
89    }
90}
91
92impl<'a, T: AnyCompatible> From<&'a T> for AnyView<'a> {
93    #[inline]
94    fn from(value: &'a T) -> Self {
95        unsafe {
96            let mut data = TVMFFIAny::new();
97            T::copy_to_any_view(&value, &mut data);
98            Self {
99                data: data,
100                _phantom: std::marker::PhantomData,
101            }
102        }
103    }
104}
105
106impl<'a> From<&AnyView<'a>> for AnyView<'a> {
107    #[inline]
108    fn from(value: &AnyView<'a>) -> Self {
109        *value
110    }
111}
112
113impl Default for AnyView<'_> {
114    fn default() -> Self {
115        Self::new()
116    }
117}
118
119/// Holder for Any value
120///
121/// This is used to define try_from rule while conforming to orphan rule
122/// Users should not use this directly
123pub struct TryFromTemp<T> {
124    value: T,
125}
126
127impl<T> TryFromTemp<T> {
128    /// Create a new holder for the value
129    #[inline(always)]
130    pub fn new(value: T) -> Self {
131        Self { value }
132    }
133
134    /// Move the value out of the holder
135    #[inline(always)]
136    pub fn into_value(this: Self) -> T {
137        this.value
138    }
139}
140
141//---------------------
142// Any
143//---------------------
144impl Any {
145    pub fn new() -> Self {
146        Self {
147            data: TVMFFIAny::new(),
148        }
149    }
150    #[inline]
151    pub fn type_index(&self) -> i32 {
152        self.data.type_index
153    }
154    /// Try to query if stored typed in Any exactly matches the type T
155    ///
156    /// This function is fast in the case of failure and can be used to check
157    /// if the type is compatible with T
158    ///
159    /// This function will not try to cast the type
160    /// and ensures invariance that the return value is only Some(T)
161    /// `Any::from(T)` contains exactly the same value value
162    /// `Any::try_as<T>()` contains exactly the same value value
163    ///
164    /// will return Some(T) if the type is exactly compatible with T
165    /// will return None if the type is not compatible with T
166    #[inline]
167    pub fn try_as<T>(&self) -> Option<T>
168    where
169        T: AnyCompatible,
170    {
171        unsafe {
172            if T::check_any_strict(&self.data) {
173                Some(T::copy_from_any_view_after_check(&self.data))
174            } else {
175                None
176            }
177        }
178    }
179
180    #[inline]
181    pub unsafe fn as_data_ptr(&mut self) -> *mut TVMFFIAny {
182        &mut self.data
183    }
184
185    #[inline]
186    pub unsafe fn into_raw_ffi_any(this: Self) -> TVMFFIAny {
187        let this = std::mem::ManuallyDrop::new(this);
188        this.data
189    }
190
191    #[inline]
192    pub unsafe fn from_raw_ffi_any(data: TVMFFIAny) -> Self {
193        Self { data }
194    }
195
196    /// Get the strong count of the underlying object for testing/debugging purposes
197    ///
198    /// If the underlying object is not ref counted, return None
199    pub fn debug_strong_count(&self) -> Option<usize> {
200        unsafe {
201            if self.data.type_index >= TypeIndex::kTVMFFIStaticObjectBegin as i32 {
202                Some(object::unsafe_::strong_count(self.data.data_union.v_obj))
203            } else {
204                None
205            }
206        }
207    }
208}
209
210impl Default for Any {
211    fn default() -> Self {
212        Self::new()
213    }
214}
215
216impl Clone for Any {
217    #[inline]
218    fn clone(&self) -> Self {
219        if self.data.type_index >= TypeIndex::kTVMFFIStaticObjectBegin as i32 {
220            unsafe { object::unsafe_::inc_ref(self.data.data_union.v_obj) }
221        }
222        Self { data: self.data }
223    }
224}
225
226impl Drop for Any {
227    #[inline]
228    fn drop(&mut self) {
229        if self.data.type_index >= TypeIndex::kTVMFFIStaticObjectBegin as i32 {
230            unsafe { object::unsafe_::dec_ref(self.data.data_union.v_obj) }
231        }
232    }
233}
234
235// convert Any ref to AnyView
236impl<'a> From<&'a Any> for AnyView<'a> {
237    #[inline]
238    fn from(value: &'a Any) -> Self {
239        Self {
240            data: value.data,
241            _phantom: std::marker::PhantomData,
242        }
243    }
244}
245
246// convert AnyView to Any
247impl From<AnyView<'_>> for Any {
248    #[inline]
249    fn from(value: AnyView<'_>) -> Self {
250        unsafe {
251            let mut data = TVMFFIAny::new();
252            crate::check_safe_call!(TVMFFIAnyViewToOwnedAny(&value.data, &mut data)).unwrap();
253            Self { data }
254        }
255    }
256}
257
258impl<T: AnyCompatible> From<T> for Any {
259    #[inline]
260    fn from(value: T) -> Self {
261        unsafe {
262            let mut data = TVMFFIAny::new();
263            T::move_to_any(value, &mut data);
264            Self { data }
265        }
266    }
267}
268
269impl<'a, T: AnyCompatible> TryFrom<AnyView<'a>> for TryFromTemp<T> {
270    type Error = crate::error::Error;
271    #[inline]
272    fn try_from(value: AnyView<'a>) -> Result<Self, Self::Error> {
273        unsafe {
274            if T::check_any_strict(&value.data) {
275                Ok(TryFromTemp::new(T::copy_from_any_view_after_check(
276                    &value.data,
277                )))
278            } else {
279                T::try_cast_from_any_view(&value.data)
280                    .map_err(|_| {
281                        let msg = format!(
282                            "Cannot convert from type `{}` to `{}`",
283                            T::get_mismatch_type_info(&value.data),
284                            T::type_str()
285                        );
286                        crate::error::Error::new(crate::error::TYPE_ERROR, &msg, "")
287                    })
288                    .map(TryFromTemp::new)
289            }
290        }
291    }
292}
293
294impl<T: AnyCompatible> TryFrom<Any> for TryFromTemp<T> {
295    type Error = crate::error::Error;
296    #[inline]
297    fn try_from(value: Any) -> Result<Self, Self::Error> {
298        unsafe {
299            if T::check_any_strict(&value.data) {
300                let mut value = std::mem::ManuallyDrop::new(value);
301                Ok(TryFromTemp::new(T::move_from_any_after_check(
302                    &mut value.data,
303                )))
304            } else {
305                T::try_cast_from_any_view(&value.data)
306                    .map_err(|_| {
307                        let msg = format!(
308                            "Cannot convert from type `{}` to `{}`",
309                            T::get_mismatch_type_info(&value.data),
310                            T::type_str()
311                        );
312                        crate::error::Error::new(crate::error::TYPE_ERROR, &msg, "")
313                    })
314                    .map(TryFromTemp::new)
315            }
316        }
317    }
318}
319
320crate::impl_try_from_any!(
321    bool,
322    i8,
323    i16,
324    i32,
325    i64,
326    isize,
327    u8,
328    u16,
329    u32,
330    u64,
331    usize,
332    f32,
333    f64,
334    (),
335    *mut core::ffi::c_void,
336    crate::string::String,
337    crate::string::Bytes,
338    crate::object::ObjectRef,
339    tvm_ffi_sys::dlpack::DLDataType,
340    tvm_ffi_sys::dlpack::DLDevice,
341);
342
343crate::impl_try_from_any_for_parametric!(Option<T>);
344
345//------------------------------------------------------------
346/// ArgTryFromAnyView: Helper for function argument passing
347///-----------------------------------------------------------
348pub(crate) trait ArgTryFromAnyView: Sized {
349    fn try_from_any_view(value: &AnyView, arg_index: usize) -> Result<Self, Error>;
350}
351
352impl<T: AnyCompatible> ArgTryFromAnyView for T {
353    fn try_from_any_view(value: &AnyView, arg_index: usize) -> Result<Self, Error> {
354        unsafe {
355            if T::check_any_strict(&value.data) {
356                Ok(T::copy_from_any_view_after_check(&value.data))
357            } else {
358                T::try_cast_from_any_view(&value.data).map_err(|_| {
359                    let msg = format!(
360                        "Argument #{}: Cannot convert from type `{}` to `{}`",
361                        arg_index,
362                        T::get_mismatch_type_info(&value.data),
363                        T::type_str()
364                    );
365                    crate::error::Error::new(crate::error::TYPE_ERROR, &msg, "")
366                })
367            }
368        }
369    }
370}