Skip to main content

hapi_rs/attribute/
array.rs

1use crate::errors::Result;
2use crate::stringhandle::{StringArray, StringHandle};
3use std::borrow::Cow;
4
5/// Groups _data_ and _sizes_ arrays for working with geometry attributes.
6/// See [`crate::attribute::NumericArrayAttr::get`].
7pub struct DataArray<'a, T>
8where
9    [T]: ToOwned<Owned = Vec<T>>,
10{
11    data: Cow<'a, [T]>,
12    sizes: Cow<'a, [i32]>,
13}
14impl<'a, T> DataArray<'a, T>
15where
16    [T]: ToOwned<Owned = Vec<T>>,
17{
18    /// Create a new data array
19    pub fn new(data: &'a [T], sizes: &'a [i32]) -> DataArray<'a, T> {
20        debug_assert_eq!(sizes.iter().sum::<i32>() as usize, data.len());
21        DataArray {
22            data: Cow::Borrowed(data),
23            sizes: Cow::Borrowed(sizes),
24        }
25    }
26
27    // Owned variant returned by APIs
28    pub(crate) fn new_owned(data: Vec<T>, sizes: Vec<i32>) -> DataArray<'static, T> {
29        debug_assert_eq!(sizes.iter().sum::<i32>() as usize, data.len());
30        DataArray {
31            data: Cow::Owned(data),
32            sizes: Cow::Owned(sizes),
33        }
34    }
35
36    /// Get reference to the data buffer.
37    #[must_use]
38    pub fn data(&self) -> &[T] {
39        self.data.as_ref()
40    }
41    /// Get reference to the sizes array.
42    #[must_use]
43    pub fn sizes(&self) -> &[i32] {
44        self.sizes.as_ref()
45    }
46
47    pub fn data_mut(&mut self) -> &mut [T] {
48        self.data.to_mut().as_mut()
49    }
50    pub fn sizes_mut(&mut self) -> &mut [i32] {
51        self.sizes.to_mut().as_mut()
52    }
53
54    /// Create an iterator over the data .
55    #[must_use]
56    pub fn iter(&'a self) -> ArrayIter<'a, T> {
57        ArrayIter {
58            sizes: self.sizes.iter(),
59            data: self.data.as_ref(),
60            cursor: 0,
61        }
62    }
63    /// Create an mutable iterator over the data .
64    pub fn iter_mut(&'a mut self) -> ArrayIterMut<'a, T> {
65        ArrayIterMut {
66            sizes: self.sizes.to_mut().iter_mut(),
67            data: self.data.to_mut().as_mut(),
68            cursor: 0,
69        }
70    }
71}
72
73impl<'a, T> IntoIterator for &'a DataArray<'a, T>
74where
75    [T]: ToOwned<Owned = Vec<T>>,
76{
77    type Item = &'a [T];
78    type IntoIter = ArrayIter<'a, T>;
79
80    fn into_iter(self) -> Self::IntoIter {
81        self.iter()
82    }
83}
84
85impl<'a, T> IntoIterator for &'a mut DataArray<'a, T>
86where
87    [T]: ToOwned<Owned = Vec<T>>,
88{
89    type Item = &'a mut [T];
90    type IntoIter = ArrayIterMut<'a, T>;
91
92    fn into_iter(self) -> Self::IntoIter {
93        self.iter_mut()
94    }
95}
96
97/// Represents multi-array string data. Used as storage for string and dictionary array attributes.
98/// Each element of this array is itself a [`StringArray`]
99#[derive(Debug, Clone)]
100pub struct StringMultiArray {
101    pub(crate) handles: Vec<StringHandle>,
102    pub(crate) sizes: Vec<i32>,
103    pub(crate) session: debug_ignore::DebugIgnore<crate::session::Session>,
104}
105
106/// Returned by [`DataArray::iter`]
107pub struct ArrayIter<'a, T> {
108    data: &'a [T],
109    sizes: std::slice::Iter<'a, i32>,
110    cursor: usize,
111}
112
113/// Returned by [`DataArray::iter_mut`]
114pub struct ArrayIterMut<'a, T> {
115    data: &'a mut [T],
116    sizes: std::slice::IterMut<'a, i32>,
117    cursor: usize,
118}
119
120pub struct MultiArrayIter<'a> {
121    handles: std::slice::Iter<'a, StringHandle>,
122    sizes: std::slice::Iter<'a, i32>,
123    session: &'a crate::session::Session,
124    cursor: usize,
125}
126
127impl<'a, T> Iterator for ArrayIter<'a, T> {
128    type Item = &'a [T];
129
130    fn next(&mut self) -> Option<Self::Item> {
131        match self.sizes.next() {
132            None => None,
133            Some(size) => {
134                let start = self.cursor;
135                let end = self.cursor + (*size as usize);
136                self.cursor = end;
137                // SAFETY: The data and the sizes arrays are both provided by HAPI
138                // are expected to match. Also bounds are checked in debug build.
139                Some(unsafe { self.data.get_unchecked(start..end) })
140            }
141        }
142    }
143}
144
145impl<'a, T> Iterator for ArrayIterMut<'a, T> {
146    type Item = &'a mut [T];
147
148    fn next(&mut self) -> Option<Self::Item> {
149        match self.sizes.next() {
150            None => None,
151            Some(size) => {
152                let start = self.cursor;
153                let end = self.cursor + (*size as usize);
154                self.cursor = end;
155                // SAFETY: Compiler can't know that we're never return overlapping references
156                // so we "erase" the lifetime by casting to pointer and back.
157                Some(unsafe {
158                    &mut *std::ptr::from_mut::<[T]>(self.data.get_unchecked_mut(start..end))
159                })
160            }
161        }
162    }
163}
164
165impl StringMultiArray {
166    #[must_use]
167    pub fn iter(&self) -> MultiArrayIter<'_> {
168        MultiArrayIter {
169            handles: self.handles.iter(),
170            sizes: self.sizes.iter(),
171            session: &self.session,
172            cursor: 0,
173        }
174    }
175    /// Convenient method to flatten the data and the sizes multidimensional arrays into individual vectors
176    pub fn flatten(self) -> Result<(Vec<String>, Vec<usize>)> {
177        let mut flat_array = Vec::with_capacity(self.sizes.iter().sum::<i32>() as usize);
178        let mut iter = self.iter();
179        while let Some(Ok(string_array)) = iter.next() {
180            flat_array.extend(string_array);
181        }
182        Ok((flat_array, self.sizes.iter().map(|v| *v as usize).collect()))
183    }
184}
185
186impl<'a> IntoIterator for &'a StringMultiArray {
187    type Item = Result<StringArray>;
188    type IntoIter = MultiArrayIter<'a>;
189
190    fn into_iter(self) -> Self::IntoIter {
191        self.iter()
192    }
193}
194
195impl Iterator for MultiArrayIter<'_> {
196    type Item = Result<StringArray>;
197
198    fn next(&mut self) -> Option<Self::Item> {
199        match self.sizes.next() {
200            None => None,
201            Some(size) => {
202                let start = self.cursor;
203                let end = self.cursor + (*size as usize);
204                self.cursor = end;
205                let handles = &self.handles.as_slice()[start..end];
206                Some(crate::stringhandle::get_string_array(handles, self.session))
207            }
208        }
209    }
210}
211
212#[cfg(test)]
213mod tests {
214    use super::*;
215
216    #[test]
217    fn data_array_accessors() {
218        let data = [1, 2, 3, 4, 5];
219        let sizes = [2, 3];
220        let ar = DataArray::new(&data, &sizes);
221
222        assert_eq!(ar.data(), &data);
223        assert_eq!(ar.sizes(), &sizes);
224    }
225
226    #[test]
227    fn data_array_mut_accessors_are_copy_on_write() {
228        let data = [1, 2, 3, 4, 5];
229        let sizes = [2, 3];
230        let mut ar = DataArray::new(&data, &sizes);
231
232        ar.data_mut()[1] = 20;
233        ar.sizes_mut()[0] = 1;
234        ar.sizes_mut()[1] = 4;
235
236        assert_eq!(ar.data(), &[1, 20, 3, 4, 5]);
237        assert_eq!(ar.sizes(), &[1, 4]);
238        assert_eq!(data, [1, 2, 3, 4, 5]);
239        assert_eq!(sizes, [2, 3]);
240    }
241
242    #[test]
243    fn data_array_iter() {
244        let ar = DataArray::new_owned(vec![1, 2, 3, 4, 5, 6], vec![2, 1, 3]);
245        let mut iter = ar.iter();
246        assert_eq!(iter.next(), Some([1, 2].as_slice()));
247        assert_eq!(iter.next(), Some([3].as_slice()));
248        assert_eq!(iter.next(), Some([4, 5, 6].as_slice()));
249    }
250
251    #[test]
252    fn data_array_into_iter() {
253        let ar = DataArray::new(&[1, 2, 3, 4, 5, 6], &[2, 1, 3]);
254        let mut iter = (&ar).into_iter();
255
256        assert_eq!(iter.next(), Some([1, 2].as_slice()));
257        assert_eq!(iter.next(), Some([3].as_slice()));
258        assert_eq!(iter.next(), Some([4, 5, 6].as_slice()));
259        assert_eq!(iter.next(), None);
260    }
261
262    #[test]
263    fn data_array_mutate() {
264        let mut ar = DataArray::new(&[1, 2, 3, 4, 5, 6], &[2, 1, 3]);
265        let mut iter = ar.iter_mut().map(|array| {
266            array.iter_mut().for_each(|v| *v *= 2);
267            array
268        });
269        assert_eq!(iter.next(), Some([2, 4].as_mut_slice()));
270        assert_eq!(iter.next(), Some([6].as_mut_slice()));
271        assert_eq!(iter.next(), Some([8, 10, 12].as_mut_slice()));
272    }
273
274    #[test]
275    fn data_array_into_iter_mut() {
276        let mut ar = DataArray::new(&[1, 2, 3, 4, 5, 6], &[2, 1, 3]);
277        let mut iter = (&mut ar).into_iter();
278
279        assert_eq!(iter.next(), Some([1, 2].as_mut_slice()));
280        assert_eq!(iter.next(), Some([3].as_mut_slice()));
281        assert_eq!(iter.next(), Some([4, 5, 6].as_mut_slice()));
282        assert_eq!(iter.next(), None);
283    }
284}