1use crate::errors::Result;
2use crate::stringhandle::{StringArray, StringHandle};
3use std::borrow::Cow;
4
5pub 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 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 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 #[must_use]
38 pub fn data(&self) -> &[T] {
39 self.data.as_ref()
40 }
41 #[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 #[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 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#[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
106pub struct ArrayIter<'a, T> {
108 data: &'a [T],
109 sizes: std::slice::Iter<'a, i32>,
110 cursor: usize,
111}
112
113pub 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 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 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 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}