trk_io/
array_sequence.rs

1use std::{
2    mem,
3    ops::{Index, Range},
4    slice,
5    vec::Vec,
6};
7
8#[derive(Clone, PartialEq)]
9pub struct ArraySequence<T> {
10    pub offsets: Vec<usize>,
11    pub data: Vec<T>,
12}
13
14impl<'a, T> IntoIterator for &'a ArraySequence<T> {
15    type Item = &'a [T];
16    type IntoIter = ArraySequenceIterator<'a, T>;
17
18    fn into_iter(self) -> Self::IntoIter {
19        ArraySequenceIterator { arr: self, index: 0..self.len() }
20    }
21}
22
23pub struct ArraySequenceIterator<'a, T: 'a> {
24    arr: &'a ArraySequence<T>,
25    index: Range<usize>,
26}
27
28impl<'a, T> Iterator for ArraySequenceIterator<'a, T> {
29    type Item = &'a [T];
30
31    fn next(&mut self) -> Option<Self::Item> {
32        let idx = self.index.next()?;
33        Some(&self.arr[idx])
34    }
35
36    fn size_hint(&self) -> (usize, Option<usize>) {
37        (0, Some(self.arr.len()))
38    }
39}
40
41impl<'a, T> IntoIterator for &'a mut ArraySequence<T> {
42    type Item = &'a mut [T];
43    type IntoIter = ArraySequenceIteratorMut<'a, T>;
44
45    fn into_iter(self) -> Self::IntoIter {
46        let mut offsets = self.offsets.iter();
47        let last_offset = *offsets.next().unwrap();
48        ArraySequenceIteratorMut { data: &mut self.data, offsets, last_offset }
49    }
50}
51
52pub struct ArraySequenceIteratorMut<'a, T: 'a> {
53    data: &'a mut [T],
54    offsets: slice::Iter<'a, usize>,
55    last_offset: usize,
56}
57
58impl<'a, T> Iterator for ArraySequenceIteratorMut<'a, T> {
59    type Item = &'a mut [T];
60
61    fn next(&mut self) -> Option<Self::Item> {
62        let current_offset = *self.offsets.next()?;
63        let nb_elements = current_offset - self.last_offset;
64        self.last_offset = current_offset;
65
66        let data = mem::replace(&mut self.data, &mut []);
67        let (slice, remaining_data) = data.split_at_mut(nb_elements);
68        self.data = remaining_data;
69        Some(slice)
70    }
71}
72
73impl<'data, T> ExactSizeIterator for ArraySequenceIterator<'data, T> {}
74
75impl<'data, T> DoubleEndedIterator for ArraySequenceIterator<'data, T> {
76    fn next_back(&mut self) -> Option<Self::Item> {
77        let idx = self.index.next_back()?;
78        Some(&self.arr[idx])
79    }
80}
81
82impl<T> Index<usize> for ArraySequence<T> {
83    type Output = [T];
84
85    fn index<'a>(&'a self, i: usize) -> &'a Self::Output {
86        let start = unsafe { *self.offsets.get_unchecked(i) };
87        let end = unsafe { *self.offsets.get_unchecked(i + 1) };
88        &self.data[start..end]
89    }
90}
91
92impl<T> Default for ArraySequence<T> {
93    fn default() -> Self {
94        ArraySequence::empty()
95    }
96}
97
98impl<T> ArraySequence<T> {
99    pub fn empty() -> ArraySequence<T> {
100        ArraySequence { offsets: vec![0], data: vec![] }
101    }
102
103    pub fn with_capacity(n: usize) -> ArraySequence<T> {
104        ArraySequence { offsets: vec![0], data: Vec::with_capacity(n) }
105    }
106
107    pub fn new(lengths: Vec<usize>, data: Vec<T>) -> ArraySequence<T> {
108        // CumSum over lengths. [0, ..., ..., data.len()]
109        // There's an additional offset at the end because we want a
110        // branchless `Index<usize>` function.
111        let offsets = [0]
112            .iter()
113            .chain(&lengths)
114            .scan(0, |state, x| {
115                *state += *x;
116                Some(*state)
117            })
118            .collect::<Vec<usize>>();
119
120        // Check if `offsets` fits with the numbers of points in `data`
121        let expected_points = *offsets.last().unwrap();
122        if expected_points != data.len() {
123            panic!(
124                "`offsets` declares {} points but `data` contains {} points.",
125                expected_points,
126                data.len()
127            );
128        }
129
130        ArraySequence { offsets, data: data }
131    }
132
133    pub fn push(&mut self, val: T) {
134        self.data.push(val);
135    }
136
137    pub fn nb_push_done(&self) -> usize {
138        self.data.len() - self.offsets.last().unwrap()
139    }
140
141    pub fn end_push(&mut self) {
142        let nb = self.nb_push_done();
143        if nb > 0 {
144            self.offsets.push(self.data.len());
145        }
146    }
147
148    /// Returns `true` if the array contains no elements.
149    ///
150    /// The array will be considered non empty if there was one or more
151    /// `push()`, even without an `end_push()`. Use `len()` instead to ignore
152    /// all pushed elements.
153    pub fn is_empty(&self) -> bool {
154        self.data.is_empty()
155    }
156
157    pub fn len(&self) -> usize {
158        self.offsets.len() - 1
159    }
160
161    /// Same as obj[i].len(), without building a slice
162    pub fn length_of_array(&self, i: usize) -> usize {
163        let current = unsafe { *self.offsets.get_unchecked(i) };
164        let next = unsafe { *self.offsets.get_unchecked(i + 1) };
165        next - current
166    }
167
168    pub fn filter<P>(&self, predicate: &mut P) -> ArraySequence<T>
169    where
170        P: FnMut(&[T]) -> bool,
171        T: Clone,
172    {
173        let mut new = ArraySequence::<T>::empty();
174        for array in self {
175            if predicate(array) {
176                new.extend(array.iter().cloned());
177            }
178        }
179        new
180    }
181
182    pub fn iter(&self) -> ArraySequenceIterator<T> {
183        self.into_iter()
184    }
185
186    pub fn iter_mut(&mut self) -> ArraySequenceIteratorMut<T> {
187        self.into_iter()
188    }
189}
190
191impl<T> Extend<T> for ArraySequence<T> {
192    fn extend<I: IntoIterator<Item = T>>(&mut self, iter: I) {
193        self.data.extend(iter);
194        self.end_push();
195    }
196}
197
198impl<T: Clone> ArraySequence<T> {
199    pub fn extend_from_slice(&mut self, other: &[T]) {
200        self.data.extend_from_slice(other);
201        self.end_push();
202    }
203}
trk_io/array_sequence.rs

trk_io/
array_sequence.rs
