tvm_ffi/collections/
array.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 std::fmt::Debug;
20use std::marker::PhantomData;
21use std::ops::Deref;
22
23use crate::any::TryFromTemp;
24use crate::derive::Object;
25use crate::object::{Object, ObjectArc};
26use crate::type_traits::ContainerElement;
27use crate::{Any, AnyCompatible, AnyView, ObjectCoreWithExtraItems, ObjectRefCore};
28use tvm_ffi_sys::TVMFFITypeIndex as TypeIndex;
29use tvm_ffi_sys::{TVMFFIAny, TVMFFIObject};
30
31#[repr(C)]
32#[derive(Object)]
33#[type_key = "ffi.Array"]
34#[type_index(TypeIndex::kTVMFFIArray)]
35pub struct ArrayObj {
36    pub object: Object,
37    /// Pointer to the start of the element buffer (AddressOf(0)).
38    pub data: *mut core::ffi::c_void,
39    pub size: i64,
40    pub capacity: i64,
41    /// Optional custom deleter for the data pointer.
42    pub data_deleter: Option<unsafe extern "C" fn(*mut core::ffi::c_void)>,
43}
44
45unsafe impl ObjectCoreWithExtraItems for ArrayObj {
46    type ExtraItem = TVMFFIAny;
47    fn extra_items_count(this: &Self) -> usize {
48        this.size as usize
49    }
50}
51
52#[repr(C)]
53#[derive(Clone)]
54pub struct Array<T: ContainerElement + Clone> {
55    data: ObjectArc<ArrayObj>,
56    _marker: PhantomData<T>,
57}
58
59impl<T: ContainerElement + Clone> Debug for Array<T> {
60    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
61        let full_name = std::any::type_name::<T>();
62        let short_name = full_name.split("::").last().unwrap_or(full_name);
63        write!(f, "Array<{}>[{}]", short_name, self.len())
64    }
65}
66
67impl<T: ContainerElement + Clone> Default for Array<T> {
68    fn default() -> Self {
69        Self::new(vec![])
70    }
71}
72
73unsafe impl<T: ContainerElement + Clone> ObjectRefCore for Array<T> {
74    type ContainerType = ArrayObj;
75
76    fn data(this: &Self) -> &ObjectArc<Self::ContainerType> {
77        &this.data
78    }
79
80    fn into_data(this: Self) -> ObjectArc<Self::ContainerType> {
81        this.data
82    }
83
84    unsafe fn from_data(data: ObjectArc<Self::ContainerType>) -> Self {
85        Self {
86            data,
87            _marker: PhantomData,
88        }
89    }
90}
91
92impl<T: ContainerElement + Clone> Array<T> {
93    /// Creates a new Array from a vector of items.
94    pub fn new(items: Vec<T>) -> Self {
95        let capacity = items.len();
96        Self::new_with_capacity(items, capacity)
97    }
98
99    /// Internal helper to allocate an ArrayObj with specific headroom.
100    fn new_with_capacity(items: Vec<T>, capacity: usize) -> Self {
101        let size = items.len();
102
103        // Allocate with capacity
104        let arc = ObjectArc::<ArrayObj>::new_with_extra_items(ArrayObj {
105            object: Object::new(),
106            data: core::ptr::null_mut(),
107            size: size as i64,
108            capacity: capacity as i64,
109            data_deleter: None,
110        });
111
112        unsafe {
113            let raw_ptr = ObjectArc::as_raw(&arc) as *mut ArrayObj;
114            let container = &mut *raw_ptr;
115
116            let base_ptr = ArrayObj::extra_items_mut(container).as_ptr() as *mut TVMFFIAny;
117            container.data = base_ptr as *mut _;
118
119            for (i, item) in items.into_iter().enumerate() {
120                let mut raw = TVMFFIAny::new();
121                T::container_move_to_any(item, &mut raw);
122                core::ptr::write(base_ptr.add(i), raw);
123            }
124        }
125        // SAFETY: `arc` was allocated and initialized above as an `Array<T>`.
126        unsafe { Self::from_data(arc) }
127    }
128
129    pub fn len(&self) -> usize {
130        self.data.size as usize
131    }
132
133    pub fn is_empty(&self) -> bool {
134        self.len() == 0
135    }
136
137    /// Retrieves an item at the given index.
138    pub fn get(&self, index: usize) -> Result<T, crate::Error> {
139        if index >= self.len() {
140            crate::bail!(crate::error::INDEX_ERROR, "Array get index out of bound");
141        }
142        unsafe {
143            let container = self.data.deref();
144            let base_ptr = container.data as *const TVMFFIAny;
145            let raw_any_ref = &*base_ptr.add(index);
146
147            match T::container_try_cast_from_any_view(raw_any_ref) {
148                Ok(val) => Ok(val),
149                Err(_) => crate::bail!(
150                    crate::error::TYPE_ERROR,
151                    "Failed to cast element at {} to {}",
152                    index,
153                    T::container_type_str()
154                ),
155            }
156        }
157    }
158
159    pub fn iter(&'_ self) -> ArrayIterator<'_, T> {
160        ArrayIterator {
161            array: self,
162            index: 0,
163            len: self.len(),
164        }
165    }
166
167    #[inline]
168    fn as_container(&self) -> &ArrayObj {
169        unsafe {
170            let ptr = ObjectArc::as_raw(&self.data) as *const ArrayObj;
171            &*ptr
172        }
173    }
174}
175
176// --- Index Implementation ---
177
178impl<T: ContainerElement + Clone> std::ops::Index<usize> for Array<T> {
179    type Output = AnyView<'static>;
180
181    fn index(&self, index: usize) -> &Self::Output {
182        let container = self.as_container();
183        let len = container.size as usize;
184        if index >= len {
185            panic!(
186                "Index out of bounds: the len is {} but the index is {}",
187                len, index
188            );
189        }
190        unsafe {
191            let ptr = (container.data as *const AnyView<'static>).add(index);
192            &*ptr
193        }
194    }
195}
196
197// --- Iterator Implementations ---
198
199pub struct ArrayIterator<'a, T: ContainerElement + Clone> {
200    array: &'a Array<T>,
201    index: usize,
202    len: usize,
203}
204
205impl<'a, T: ContainerElement + Clone> Iterator for ArrayIterator<'a, T> {
206    type Item = T;
207
208    fn next(&mut self) -> Option<Self::Item> {
209        if self.index < self.len {
210            let item = self.array.get(self.index).ok();
211            self.index += 1;
212            item
213        } else {
214            None
215        }
216    }
217}
218
219impl<'a, T: ContainerElement + Clone> IntoIterator for &'a Array<T> {
220    type Item = T;
221    type IntoIter = ArrayIterator<'a, T>;
222
223    fn into_iter(self) -> Self::IntoIter {
224        self.iter()
225    }
226}
227
228impl<T: ContainerElement + Clone> FromIterator<T> for Array<T> {
229    fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
230        let items: Vec<T> = iter.into_iter().collect();
231        Self::new(items)
232    }
233}
234
235// --- Any Type System Conversions ---
236
237unsafe impl<T> AnyCompatible for Array<T>
238where
239    T: ContainerElement + Clone,
240{
241    fn type_str() -> String {
242        format!("Array<{}>", T::container_type_str())
243    }
244
245    unsafe fn check_any_strict(data: &TVMFFIAny) -> bool {
246        if data.type_index != TypeIndex::kTVMFFIArray as i32 {
247            return false;
248        }
249
250        let container = &*(data.data_union.v_obj as *const ArrayObj);
251        let base_ptr = container.data as *const TVMFFIAny;
252        for i in 0..container.size {
253            let elem_any = &*base_ptr.add(i as usize);
254            if !T::container_check_any_strict(elem_any) {
255                return false;
256            }
257        }
258        true
259    }
260
261    unsafe fn copy_to_any_view(src: &Self, data: &mut TVMFFIAny) {
262        data.type_index = TypeIndex::kTVMFFIArray as i32;
263        data.data_union.v_obj = ObjectArc::as_raw(Self::data(src)) as *mut TVMFFIObject;
264        data.small_str_len = 0;
265    }
266
267    unsafe fn move_to_any(src: Self, data: &mut TVMFFIAny) {
268        data.type_index = TypeIndex::kTVMFFIArray as i32;
269        data.data_union.v_obj = ObjectArc::into_raw(Self::into_data(src)) as *mut TVMFFIObject;
270        data.small_str_len = 0;
271    }
272
273    unsafe fn copy_from_any_view_after_check(data: &TVMFFIAny) -> Self {
274        let ptr = data.data_union.v_obj as *const ArrayObj;
275        crate::object::unsafe_::inc_ref(ptr as *mut TVMFFIObject);
276        Self::from_data(ObjectArc::from_raw(ptr))
277    }
278
279    unsafe fn move_from_any_after_check(data: &mut TVMFFIAny) -> Self {
280        let ptr = data.data_union.v_obj as *const ArrayObj;
281        let obj = Self::from_data(ObjectArc::from_raw(ptr));
282
283        data.type_index = TypeIndex::kTVMFFINone as i32;
284        data.data_union.v_int64 = 0;
285
286        obj
287    }
288
289    unsafe fn try_cast_from_any_view(data: &TVMFFIAny) -> Result<Self, ()> {
290        if data.type_index != TypeIndex::kTVMFFIArray as i32 {
291            return Err(());
292        }
293
294        // Fast path: if types match exactly, we can just copy the reference.
295        if <Self as AnyCompatible>::check_any_strict(data) {
296            return Ok(<Self as AnyCompatible>::copy_from_any_view_after_check(
297                data,
298            ));
299        }
300
301        // Slow path: try to convert element by element.
302        let container = &*(data.data_union.v_obj as *const ArrayObj);
303        let base_ptr = container.data as *const TVMFFIAny;
304        let mut items = Vec::with_capacity(container.size as usize);
305
306        for i in 0..container.size {
307            let any_v = &*base_ptr.add(i as usize);
308            if let Ok(item) = T::container_try_cast_from_any_view(any_v) {
309                items.push(item);
310            } else {
311                return Err(());
312            }
313        }
314
315        Ok(Array::new(items))
316    }
317}
318
319impl<T> TryFrom<Any> for Array<T>
320where
321    T: ContainerElement + Clone,
322{
323    type Error = crate::error::Error;
324
325    fn try_from(value: Any) -> Result<Self, Self::Error> {
326        let temp: TryFromTemp<Self> = TryFromTemp::try_from(value)?;
327        Ok(TryFromTemp::into_value(temp))
328    }
329}
330
331impl<'a, T> TryFrom<AnyView<'a>> for Array<T>
332where
333    T: ContainerElement + Clone,
334{
335    type Error = crate::error::Error;
336
337    fn try_from(value: AnyView<'a>) -> Result<Self, Self::Error> {
338        let temp: TryFromTemp<Self> = TryFromTemp::try_from(value)?;
339        Ok(TryFromTemp::into_value(temp))
340    }
341}
tvm_ffi/collections/array.rs

tvm_ffi/collections/
array.rs
