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    #[inline]
57    pub(crate) fn as_raw_ffi_any(&self) -> &TVMFFIAny {
58        &self.data
59    }
60
61    /// Construct a borrowed view from its ABI representation.
62    ///
63    /// # Safety
64    ///
65    /// `data.type_index` must describe the payload, and the caller must keep
66    /// every resource it references alive for the view's complete lifetime.
67    ///
68    /// `kTVMFFIObjectRValueRef` carries a further obligation: `v_ptr` must
69    /// point to a writable slot the caller uniquely owns, holding one strong
70    /// reference. Owning the view takes that reference and writes null back
71    /// through the pointer, so the view must be converted at most once and no
72    /// other access to the slot may overlap the conversion.
73    #[inline]
74    pub(crate) unsafe fn from_raw_ffi_any(data: TVMFFIAny) -> Self {
75        Self {
76            data,
77            _phantom: std::marker::PhantomData,
78        }
79    }
80
81    /// More strict version than try_from/try_into
82    ///
83    /// This function will not try to cast the type
84    /// and ensures invariance that the return value is only Some(T)
85    /// Any::from(T) contains exactly the same value value
86    ///
87    /// will return Some(T) if the type is exactly compatible with T
88    /// will return None if the type is not compatible with T
89    #[inline]
90    pub fn try_as<T>(&self) -> Option<T>
91    where
92        T: AnyCompatible,
93    {
94        unsafe {
95            if T::check_any_strict(&self.data) {
96                Some(T::copy_from_any_view_after_check(&self.data))
97            } else {
98                None
99            }
100        }
101    }
102
103    /// Get the strong count of the underlying object for testing/debugging purposes
104    ///
105    /// If the underlying object is not ref counted, return None
106    pub fn debug_strong_count(&self) -> Option<usize> {
107        unsafe {
108            if self.data.type_index >= TypeIndex::kTVMFFIStaticObjectBegin as i32 {
109                Some(object::unsafe_::strong_count(self.data.data_union.v_obj))
110            } else {
111                None
112            }
113        }
114    }
115}
116
117impl<'a, T: AnyCompatible> From<&'a T> for AnyView<'a> {
118    #[inline]
119    fn from(value: &'a T) -> Self {
120        unsafe {
121            let mut data = TVMFFIAny::new();
122            T::copy_to_any_view(&value, &mut data);
123            Self {
124                data: data,
125                _phantom: std::marker::PhantomData,
126            }
127        }
128    }
129}
130
131impl<'a> From<&AnyView<'a>> for AnyView<'a> {
132    #[inline]
133    fn from(value: &AnyView<'a>) -> Self {
134        *value
135    }
136}
137
138impl Default for AnyView<'_> {
139    fn default() -> Self {
140        Self::new()
141    }
142}
143
144/// Holder for Any value
145///
146/// This is used to define try_from rule while conforming to orphan rule
147/// Users should not use this directly
148pub struct TryFromTemp<T> {
149    value: T,
150}
151
152impl<T> TryFromTemp<T> {
153    /// Create a new holder for the value
154    #[inline(always)]
155    pub fn new(value: T) -> Self {
156        Self { value }
157    }
158
159    /// Move the value out of the holder
160    #[inline(always)]
161    pub fn into_value(this: Self) -> T {
162        this.value
163    }
164}
165
166//---------------------
167// Any
168//---------------------
169impl Any {
170    pub fn new() -> Self {
171        Self {
172            data: TVMFFIAny::new(),
173        }
174    }
175    #[inline]
176    pub fn type_index(&self) -> i32 {
177        self.data.type_index
178    }
179    #[inline]
180    pub(crate) fn as_raw_ffi_any(&self) -> &TVMFFIAny {
181        &self.data
182    }
183    /// Try to query if stored typed in Any exactly matches the type T
184    ///
185    /// This function is fast in the case of failure and can be used to check
186    /// if the type is compatible with T
187    ///
188    /// This function will not try to cast the type
189    /// and ensures invariance that the return value is only Some(T)
190    /// `Any::from(T)` contains exactly the same value value
191    /// `Any::try_as<T>()` contains exactly the same value value
192    ///
193    /// will return Some(T) if the type is exactly compatible with T
194    /// will return None if the type is not compatible with T
195    #[inline]
196    pub fn try_as<T>(&self) -> Option<T>
197    where
198        T: AnyCompatible,
199    {
200        unsafe {
201            if T::check_any_strict(&self.data) {
202                Some(T::copy_from_any_view_after_check(&self.data))
203            } else {
204                None
205            }
206        }
207    }
208
209    #[inline]
210    pub unsafe fn as_data_ptr(&mut self) -> *mut TVMFFIAny {
211        &mut self.data
212    }
213
214    #[inline]
215    pub unsafe fn into_raw_ffi_any(this: Self) -> TVMFFIAny {
216        let this = std::mem::ManuallyDrop::new(this);
217        this.data
218    }
219
220    #[inline]
221    pub unsafe fn from_raw_ffi_any(data: TVMFFIAny) -> Self {
222        Self { data }
223    }
224
225    /// Get the strong count of the underlying object for testing/debugging purposes
226    ///
227    /// If the underlying object is not ref counted, return None
228    pub fn debug_strong_count(&self) -> Option<usize> {
229        unsafe {
230            if self.data.type_index >= TypeIndex::kTVMFFIStaticObjectBegin as i32 {
231                Some(object::unsafe_::strong_count(self.data.data_union.v_obj))
232            } else {
233                None
234            }
235        }
236    }
237}
238
239impl Default for Any {
240    fn default() -> Self {
241        Self::new()
242    }
243}
244
245impl Clone for Any {
246    #[inline]
247    fn clone(&self) -> Self {
248        if self.data.type_index >= TypeIndex::kTVMFFIStaticObjectBegin as i32 {
249            unsafe { object::unsafe_::inc_ref(self.data.data_union.v_obj) }
250        }
251        Self { data: self.data }
252    }
253}
254
255impl Drop for Any {
256    #[inline]
257    fn drop(&mut self) {
258        if self.data.type_index >= TypeIndex::kTVMFFIStaticObjectBegin as i32 {
259            unsafe { object::unsafe_::dec_ref(self.data.data_union.v_obj) }
260        }
261    }
262}
263
264/// Convert an [`AnyView`] without constructing a diagnostic error on mismatch.
265#[inline]
266pub(crate) fn try_cast_from_any_view<T>(value: &AnyView<'_>) -> Result<T, ()>
267where
268    T: AnyCompatible,
269{
270    unsafe { T::try_cast_from_any_view(&value.data) }
271}
272
273/// Copy a value after exact-leaf lookup has established compatibility.
274///
275/// # Safety
276///
277/// `T::MATCH_ANY_EXACT` must be true and `value.type_index()` must equal
278/// `T::match_any_exact_type_index()`.
279#[inline(always)]
280pub(crate) unsafe fn copy_from_any_view_after_check<T>(value: &AnyView<'_>) -> T
281where
282    T: AnyCompatible,
283{
284    debug_assert!(T::MATCH_ANY_EXACT);
285    debug_assert!(T::check_any_strict(&value.data));
286    T::copy_from_any_view_after_check(&value.data)
287}
288
289// convert Any ref to AnyView
290impl<'a> From<&'a Any> for AnyView<'a> {
291    #[inline]
292    fn from(value: &'a Any) -> Self {
293        Self {
294            data: value.data,
295            _phantom: std::marker::PhantomData,
296        }
297    }
298}
299
300/// Whether a `TVMFFIAny` cell owns everything it holds, so its owning form is a
301/// bitwise copy. Mirrors C++ `details::InplaceConvertAnyViewToAny`.
302#[inline]
303pub(crate) fn is_plain_inline(type_index: i32) -> bool {
304    type_index < TypeIndex::kTVMFFIRawStr as i32
305        || type_index == TypeIndex::kTVMFFISmallStr as i32
306        || type_index == TypeIndex::kTVMFFISmallBytes as i32
307        || type_index == TypeIndex::kTVMFFIUnchanged as i32
308}
309
310// convert AnyView to Any
311impl From<AnyView<'_>> for Any {
312    #[inline]
313    fn from(value: AnyView<'_>) -> Self {
314        let data = value.data;
315        // Owning a borrowed object is the same incref `Any::clone` does below.
316        if data.type_index >= TypeIndex::kTVMFFIStaticObjectBegin as i32 {
317            unsafe { object::unsafe_::inc_ref(data.data_union.v_obj) };
318            return Self { data };
319        }
320        if is_plain_inline(data.type_index) {
321            return Self { data };
322        }
323        // What is left borrows foreign storage and needs the runtime.
324        any_view_to_owned_via_runtime(data)
325    }
326}
327
328/// Out of line so the inlined conversions above stay leaf code, without the
329/// stack frame and unwind path this call needs.
330#[cold]
331#[inline(never)]
332fn any_view_to_owned_via_runtime(view: TVMFFIAny) -> Any {
333    unsafe {
334        let mut data = TVMFFIAny::new();
335        crate::check_safe_call!(TVMFFIAnyViewToOwnedAny(&view, &mut data)).unwrap();
336        Any { data }
337    }
338}
339
340impl<T: AnyCompatible> From<T> for Any {
341    #[inline]
342    fn from(value: T) -> Self {
343        unsafe {
344            let mut data = TVMFFIAny::new();
345            T::move_to_any(value, &mut data);
346            Self { data }
347        }
348    }
349}
350
351impl<'a, T: AnyCompatible> TryFrom<AnyView<'a>> for TryFromTemp<T> {
352    type Error = crate::error::Error;
353    #[inline]
354    fn try_from(value: AnyView<'a>) -> Result<Self, Self::Error> {
355        unsafe {
356            if T::check_any_strict(&value.data) {
357                Ok(TryFromTemp::new(T::copy_from_any_view_after_check(
358                    &value.data,
359                )))
360            } else {
361                T::try_cast_from_any_view(&value.data)
362                    .map_err(|_| {
363                        let msg = format!(
364                            "Cannot convert from type `{}` to `{}`",
365                            T::get_mismatch_type_info(&value.data),
366                            T::type_str()
367                        );
368                        crate::error::Error::new(crate::error::TYPE_ERROR, &msg, "")
369                    })
370                    .map(TryFromTemp::new)
371            }
372        }
373    }
374}
375
376impl<T: AnyCompatible> TryFrom<Any> for TryFromTemp<T> {
377    type Error = crate::error::Error;
378    #[inline]
379    fn try_from(value: Any) -> Result<Self, Self::Error> {
380        unsafe {
381            if T::check_any_strict(&value.data) {
382                let mut value = std::mem::ManuallyDrop::new(value);
383                Ok(TryFromTemp::new(T::move_from_any_after_check(
384                    &mut value.data,
385                )))
386            } else {
387                T::try_cast_from_any_view(&value.data)
388                    .map_err(|_| {
389                        let msg = format!(
390                            "Cannot convert from type `{}` to `{}`",
391                            T::get_mismatch_type_info(&value.data),
392                            T::type_str()
393                        );
394                        crate::error::Error::new(crate::error::TYPE_ERROR, &msg, "")
395                    })
396                    .map(TryFromTemp::new)
397            }
398        }
399    }
400}
401
402crate::impl_try_from_any!(
403    bool,
404    i8,
405    i16,
406    i32,
407    i64,
408    isize,
409    u8,
410    u16,
411    u32,
412    u64,
413    usize,
414    f32,
415    f64,
416    (),
417    *mut core::ffi::c_void,
418    crate::string::String,
419    crate::string::Bytes,
420    crate::object::ObjectRef,
421    tvm_ffi_sys::dlpack::DLDataType,
422    tvm_ffi_sys::dlpack::DLDevice,
423);
424
425crate::impl_try_from_any_for_parametric!(Option<T>);
426
427//------------------------------------------------------------
428/// ArgTryFromAnyView: Helper for function argument passing
429///-----------------------------------------------------------
430pub(crate) trait ArgTryFromAnyView: Sized {
431    fn try_from_any_view(value: &AnyView, arg_index: usize) -> Result<Self, Error>;
432}
433
434impl<T: AnyCompatible> ArgTryFromAnyView for T {
435    fn try_from_any_view(value: &AnyView, arg_index: usize) -> Result<Self, Error> {
436        unsafe {
437            if T::check_any_strict(&value.data) {
438                Ok(T::copy_from_any_view_after_check(&value.data))
439            } else {
440                T::try_cast_from_any_view(&value.data).map_err(|_| {
441                    let msg = format!(
442                        "Argument #{}: Cannot convert from type `{}` to `{}`",
443                        arg_index,
444                        T::get_mismatch_type_info(&value.data),
445                        T::type_str()
446                    );
447                    crate::error::Error::new(crate::error::TYPE_ERROR, &msg, "")
448                })
449            }
450        }
451    }
452}
tvm_ffi/any.rs

tvm_ffi/
any.rs
