arbitrary_chunks/
lib.rs

1//! # arbitrary-chunks
2//! An iterator that allows specifying an input array of arbitrary chunk-sizes with which to split a vector or array. As with the standard `.chunks()`, this iterator also includes `_mut()`, `_exact()` and `_exact_mut()` variants.
3//!
4//! ## Usage
5//!
6//! By default, this iterator is implemented for `[T]`, meaning it works for both arrays and Vec of any type.
7//!
8//! If there is not enough data to satisfy the provided chunk length, you will get all the remaining data for that chunk, so it will be shorter than expected. To instead stop early if there is not enough data, see the `.arbitrary_chunks_exact()` variant.
9//!
10//! ```rust
11//! use arbitrary_chunks::ArbitraryChunks;
12//!
13//! let chunks: Vec<usize> = vec![1, 3, 1];
14//! let data: Vec<i32> = vec![0, 1, 2, 3, 4];
15//!
16//! let chunked_data: Vec<Vec<i32>> = data
17//!   .arbitrary_chunks(&chunks)
18//!   .map(|chunk| chunk.to_vec())
19//!   .collect();
20//!
21//! assert_eq!(vec![0], chunked_data[0]);
22//! assert_eq!(vec![1, 2, 3], chunked_data[1]);
23//! assert_eq!(vec![4], chunked_data[2]);
24//! ```
25//!
26//! ## Exact Variant
27//!
28//! Unlike the regular variant, the exact variant's iterator will end early if there is not enough data to satisfy the chunk. Instead, you will be able to access the remainder of the data with `.remainder()`.
29//!
30//! ```rust
31//! use arbitrary_chunks::ArbitraryChunks;
32//!
33//! let chunks: Vec<usize> = vec![1, 3];
34//! let data: Vec<i32> = vec![0, 1, 2];
35//! let mut iter = data.arbitrary_chunks_exact(&chunks);
36//!
37//! assert_eq!(vec![0], iter.next().unwrap());
38//! assert_eq!(None, iter.next());
39//! assert_eq!(vec![1, 2], iter.remainder());
40//! ```
41//!
42//! ### Mutable Variants
43//!
44//! Each of the regular and exact variants also have their own mutable variants. These allow you to mutably modify slices and vectors in arbitrarily-sized chunks.
45//!
46//! ```ignore
47//! use arbitrary_chunks::ArbitraryChunks;
48//!
49//! let chunks: Vec<usize> = vec![1, 3, 1];
50//! let data: Vec<i32> = vec![0, 1, 2, 3, 4];
51//!
52//! let iter_1 = data.arbitrary_chunks_mut(chunks.clone());
53//! let iter_2 = data.arbitrary_chunks_exact_mut(chunks);
54//! ```
55//!
56//! ### Parallel Iterator
57//!
58//! This can be used in parallel with rayon using `.par_bridge()`, for example:
59//!
60//! ```rust
61//! use arbitrary_chunks::ArbitraryChunks;
62//! use rayon::prelude::*;
63//!
64//! let chunks: Vec<usize> = vec![1, 3, 1];
65//! let data: Vec<i32> = vec![0, 1, 2, 3, 4];
66//!
67//! data
68//!   .arbitrary_chunks(&chunks)
69//!   .par_bridge()
70//!   .for_each(|chunk| {
71//!     assert!(chunk.len() >= 1 && chunk.len() <= 3);
72//!     println!("{:?}", chunk);
73//!   });
74//!
75//! // Prints (in pseudo-random order):
76//! // [1, 2, 3]
77//! // [0]
78//! // [4]
79//! ```
80//!
81//! ## Motivation
82//!
83//! This library was inspired by the need to mutably modify many sections of the same slice concurrently. With this library plus Rayon, you are able to mutably borrow and modify slices safely from many threads at the same time, without upsetting the borrow checker.
84//!
85//! ## License
86//!
87//! Licensed under either of
88//!
89//! * Apache License, Version 2.0, ([LICENSE-APACHE](LICENSE-APACHE) or <http://www.apache.org/licenses/LICENSE-2.0>)
90//! * MIT license ([LICENSE-MIT](LICENSE-MIT) or <http://opensource.org/licenses/MIT>)
91//!
92//! at your option.
93//!
94//! ### Contribution
95//!
96//! Unless you explicitly state otherwise, any contribution intentionally submitted for inclusion in the work by you, as defined in the Apache-2.0 license, shall be dual licensed as above, without any additional terms or conditions.
97//!
98
99use std::cmp::min;
100use std::mem;
101
102pub struct ArbitraryChunkMut<'a, 'b, T: 'a> {
103    data: &'a mut [T],
104    counts: &'b [usize],
105    cursor: usize,
106}
107
108impl<'a, 'b, T> Iterator for ArbitraryChunkMut<'a, 'b, T> {
109    type Item = &'a mut [T];
110
111    fn next(&mut self) -> Option<Self::Item> {
112        if self.cursor > self.counts.len() - 1 {
113            return None;
114        }
115
116        let c = self.counts[self.cursor];
117        self.cursor += 1;
118
119        if c == 0 {
120            return Some(&mut []);
121        }
122
123        if self.data.is_empty() {
124            return None;
125        }
126
127        let point = min(c, self.data.len());
128        let slice = mem::take(&mut self.data);
129        let (l, r) = slice.split_at_mut(point);
130        self.data = r;
131
132        Some(l)
133    }
134
135    fn size_hint(&self) -> (usize, Option<usize>) {
136        let remaining = self.counts.len() - self.cursor;
137        (remaining, Some(remaining))
138    }
139}
140
141pub struct ArbitraryChunk<'a, 'b, T: 'a> {
142    data: &'a [T],
143    counts: &'b [usize],
144    cursor: usize,
145}
146
147impl<'a, 'b, T> Iterator for ArbitraryChunk<'a, 'b, T> {
148    type Item = &'a [T];
149
150    fn next(&mut self) -> Option<Self::Item> {
151        if self.cursor > self.counts.len() - 1 {
152            return None;
153        }
154
155        let c = self.counts[self.cursor];
156        self.cursor += 1;
157
158        if c == 0 {
159            return Some(&[]);
160        }
161
162        if self.data.is_empty() {
163            return None;
164        }
165
166        let point = min(c, self.data.len());
167        let slice = mem::take(&mut self.data);
168        let (l, r) = slice.split_at(point);
169        self.data = r;
170
171        Some(l)
172    }
173
174    fn size_hint(&self) -> (usize, Option<usize>) {
175        let remaining = self.counts.len() - self.cursor;
176        (remaining, Some(remaining))
177    }
178}
179
180pub struct ArbitraryChunkExactMut<'a, 'b, T: 'a> {
181    data: &'a mut [T],
182    counts: &'b [usize],
183    cursor: usize,
184}
185
186impl<'a, 'b, T> Iterator for ArbitraryChunkExactMut<'a, 'b, T> {
187    type Item = &'a mut [T];
188
189    fn next(&mut self) -> Option<Self::Item> {
190        if self.cursor > self.counts.len() - 1 {
191            return None;
192        }
193
194        let c = self.counts[self.cursor];
195        self.cursor += 1;
196
197        if c == 0 {
198            return Some(&mut []);
199        }
200
201        if self.data.is_empty() || c > self.data.len() {
202            return None;
203        }
204
205        let slice = mem::take(&mut self.data);
206        let (l, r) = slice.split_at_mut(c);
207        self.data = r;
208
209        Some(l)
210    }
211
212    fn size_hint(&self) -> (usize, Option<usize>) {
213        let remaining = self.counts.len() - self.cursor;
214        (remaining, Some(remaining))
215    }
216}
217
218impl<'a, 'b, T> ArbitraryChunkExactMut<'a, 'b, T> {
219    pub fn remainder(&'a mut self) -> &'a mut [T] {
220        self.data
221    }
222}
223
224pub struct ArbitraryChunkExact<'a, 'b, T: 'a> {
225    data: &'a [T],
226    counts: &'b [usize],
227    cursor: usize,
228}
229
230impl<'a, 'b, T> Iterator for ArbitraryChunkExact<'a, 'b, T> {
231    type Item = &'a [T];
232
233    fn next(&mut self) -> Option<Self::Item> {
234        if self.cursor > self.counts.len() - 1 {
235            return None;
236        }
237
238        let c = self.counts[self.cursor];
239        self.cursor += 1;
240
241        if c == 0 {
242            return Some(&[]);
243        }
244
245        if self.data.is_empty() || c > self.data.len() {
246            return None;
247        }
248
249        let slice = mem::take(&mut self.data);
250        let (l, r) = slice.split_at(c);
251        self.data = r;
252
253        Some(l)
254    }
255
256    fn size_hint(&self) -> (usize, Option<usize>) {
257        let remaining = self.counts.len() - self.cursor;
258        (remaining, Some(remaining))
259    }
260}
261
262impl<'a, 'b, T> ArbitraryChunkExact<'a, 'b, T> {
263    pub fn remainder(&'a self) -> &'a [T] {
264        self.data
265    }
266}
267
268pub trait ArbitraryChunks<'a, 'b, T> {
269    /// `arbitrary_chunks` returns an iterator over chunks of sizes defined in `counts`.
270    ///
271    /// ```rust
272    /// use arbitrary_chunks::ArbitraryChunks;
273    ///
274    /// let chunks: Vec<usize> = vec![1, 3, 1];
275    /// let data: Vec<i32> = vec![0, 1, 2, 3, 4];
276    ///
277    /// let chunked_data: Vec<Vec<i32>> = data
278    ///     .arbitrary_chunks(&chunks)
279    ///     .map(|chunk| chunk.to_vec())
280    ///     .collect();
281    ///
282    /// assert_eq!(vec![0], chunked_data[0]);
283    /// assert_eq!(vec![1, 2, 3], chunked_data[1]);
284    /// assert_eq!(vec![4], chunked_data[2]);
285    /// ```
286    fn arbitrary_chunks(&'a self, counts: &'b [usize]) -> ArbitraryChunk<'a, 'b, T>;
287
288    /// `arbitrary_chunks_mut` returns an iterator over mutable chunks of sizes defined in `counts`.
289    ///
290    /// ```rust
291    /// use arbitrary_chunks::ArbitraryChunks;
292    ///
293    /// let chunks: Vec<usize> = vec![1, 3, 1];
294    /// let mut data: Vec<i32> = vec![0, 1, 2, 3, 4];
295    ///
296    /// data
297    ///     .arbitrary_chunks_mut(&chunks)
298    ///     .for_each(|chunk| {
299    ///         chunk[0] = chunk[0] * 2;
300    ///     });
301    ///
302    /// assert_eq!(vec![0, 2, 2, 3, 8], data);
303    /// ```
304    fn arbitrary_chunks_mut(&'a mut self, counts: &'b [usize]) -> ArbitraryChunkMut<'a, 'b, T>;
305
306    /// `arbitrary_chunks_exact` returns chunks sized exactly as requested, or not at all.
307    /// If there is not enough data to satisfy the chunk, the iterator will end. You will then be
308    /// able to get the remainder of the data using `.remainder()` on the iterator.
309    ///
310    /// ```rust
311    /// use arbitrary_chunks::ArbitraryChunks;
312    ///
313    /// let chunks: Vec<usize> = vec![1, 3];
314    /// let data: Vec<i32> = vec![0, 1, 2];
315    /// let mut iter = data.arbitrary_chunks_exact(&chunks);
316    ///
317    /// assert_eq!(vec![0], iter.next().unwrap());
318    /// assert_eq!(None, iter.next());
319    /// assert_eq!(vec![1, 2], iter.remainder());
320    /// ```
321    fn arbitrary_chunks_exact(&'a self, counts: &'b [usize]) -> ArbitraryChunkExact<'a, 'b, T>;
322
323    /// `arbitrary_chunks_exact_mut` returns chunks sized exactly as requested, or not at all.
324    /// If there is not enough data to satisfy the chunk, the iterator will end. You will then be
325    /// able to get the remainder of the data using `.remainder()` on the iterator.
326    ///
327    /// ```rust
328    /// use arbitrary_chunks::ArbitraryChunks;
329    ///
330    /// let chunks: Vec<usize> = vec![1, 3];
331    /// let mut data: Vec<i32> = vec![0, 1, 2];
332    /// let mut iter = data.arbitrary_chunks_exact_mut(&chunks);
333    ///
334    /// assert_eq!(vec![0], iter.next().unwrap());
335    /// assert_eq!(None, iter.next());
336    /// assert_eq!(vec![1, 2], iter.remainder());
337    /// ```
338    fn arbitrary_chunks_exact_mut(&'a mut self, counts: &'b [usize]) -> ArbitraryChunkExactMut<'a, 'b, T>;
339}
340
341impl<'a, 'b, T> ArbitraryChunks<'a, 'b, T> for [T] {
342    fn arbitrary_chunks(&'a self, counts: &'b [usize]) -> ArbitraryChunk<'a, 'b, T> {
343        ArbitraryChunk {
344            data: self,
345            counts,
346            cursor: 0,
347        }
348    }
349
350    fn arbitrary_chunks_mut(&'a mut self, counts: &'b [usize]) -> ArbitraryChunkMut<'a, 'b, T> {
351        ArbitraryChunkMut {
352            data: self,
353            counts,
354            cursor: 0,
355        }
356    }
357
358    fn arbitrary_chunks_exact(&'a self, counts: &'b [usize]) -> ArbitraryChunkExact<'a, 'b, T> {
359        ArbitraryChunkExact {
360            data: self,
361            counts,
362            cursor: 0,
363        }
364    }
365
366    fn arbitrary_chunks_exact_mut(&'a mut self, counts: &'b [usize]) -> ArbitraryChunkExactMut<'a, 'b, T> {
367        ArbitraryChunkExactMut {
368            data: self,
369            counts,
370            cursor: 0,
371        }
372    }
373}
374
375#[cfg(test)]
376mod tests {
377    use crate::ArbitraryChunks;
378
379    #[test]
380    fn it_stops_when_chunks_run_out() {
381        let chunks: Vec<usize> = vec![0, 1, 2, 3];
382        let data = vec![8, 7, 6, 5, 4, 3, 2, 1];
383        let chunk_data: Vec<Vec<i32>> = data
384            .arbitrary_chunks(&chunks)
385            .map(|chunk| chunk.to_vec())
386            .collect();
387
388        assert_eq!(Vec::<i32>::new(), chunk_data[0]);
389        assert_eq!(vec![8], chunk_data[1]);
390        assert_eq!(vec![7, 6], chunk_data[2]);
391        assert_eq!(vec![5, 4, 3], chunk_data[3]);
392        assert_eq!(None, chunk_data.get(4));
393    }
394
395    #[test]
396    fn mut_stops_when_chunks_run_out() {
397        let chunks: Vec<usize> = vec![0, 1, 2, 3];
398        let mut data = vec![8, 7, 6, 5, 4, 3, 2, 1];
399        let chunk_data: Vec<Vec<i32>> = data
400            .arbitrary_chunks_mut(&chunks)
401            .map(|chunk| chunk.to_vec())
402            .collect();
403
404        assert_eq!(Vec::<i32>::new(), chunk_data[0]);
405        assert_eq!(vec![8], chunk_data[1]);
406        assert_eq!(vec![7, 6], chunk_data[2]);
407        assert_eq!(vec![5, 4, 3], chunk_data[3]);
408        assert_eq!(None, chunk_data.get(4));
409    }
410
411    #[test]
412    fn exact_stops_when_chunks_run_out() {
413        let chunks: Vec<usize> = vec![0, 1, 2, 3];
414        let data = vec![8, 7, 6, 5, 4, 3, 2, 1];
415        let mut iter = data.arbitrary_chunks_exact(&chunks);
416
417        assert_eq!(&[0i32; 0], iter.next().unwrap());
418        assert_eq!(&[8i32], iter.next().unwrap());
419        assert_eq!(&[7i32, 6], iter.next().unwrap());
420        assert_eq!(&[5i32, 4, 3], iter.next().unwrap());
421        assert_eq!(None, iter.next());
422        assert_eq!(&[2i32, 1], iter.remainder());
423    }
424
425    #[test]
426    fn exact_mut_stops_when_chunks_run_out() {
427        let chunks: Vec<usize> = vec![0, 1, 2, 3];
428        let mut data = vec![8, 7, 6, 5, 4, 3, 2, 1];
429        let mut iter = data.arbitrary_chunks_exact_mut(&chunks);
430
431        assert_eq!(&mut [0i32; 0], iter.next().unwrap());
432        assert_eq!(&mut [8i32], iter.next().unwrap());
433        assert_eq!(&mut [7i32, 6], iter.next().unwrap());
434        assert_eq!(&mut [5i32, 4, 3], iter.next().unwrap());
435        assert_eq!(None, iter.next());
436        assert_eq!(&mut [2i32, 1], iter.remainder());
437    }
438
439    #[test]
440    fn it_accounts_for_trailing_zeros() {
441        let chunks: Vec<usize> = vec![0, 1, 2, 3, 0, 0];
442        let data = vec![8, 7, 6, 5, 4, 3, 2, 1];
443        let chunk_data: Vec<Vec<i32>> = data
444            .arbitrary_chunks(&chunks)
445            .map(|chunk| chunk.to_vec())
446            .collect();
447
448        assert_eq!(Vec::<i32>::new(), chunk_data[0]);
449        assert_eq!(vec![8], chunk_data[1]);
450        assert_eq!(vec![7, 6], chunk_data[2]);
451        assert_eq!(vec![5, 4, 3], chunk_data[3]);
452        assert_eq!(Vec::<i32>::new(), chunk_data[4]);
453        assert_eq!(Vec::<i32>::new(), chunk_data[5]);
454        assert_eq!(None, chunk_data.get(6));
455    }
456
457    #[test]
458    fn mut_accounts_for_trailing_zeros() {
459        let chunks: Vec<usize> = vec![0, 1, 2, 3, 0, 0];
460        let mut data = vec![8, 7, 6, 5, 4, 3, 2, 1];
461        let chunk_data: Vec<Vec<i32>> = data
462            .arbitrary_chunks_mut(&chunks)
463            .map(|chunk| chunk.to_vec())
464            .collect();
465
466        assert_eq!(Vec::<i32>::new(), chunk_data[0]);
467        assert_eq!(vec![8], chunk_data[1]);
468        assert_eq!(vec![7, 6], chunk_data[2]);
469        assert_eq!(vec![5, 4, 3], chunk_data[3]);
470        assert_eq!(Vec::<i32>::new(), chunk_data[4]);
471        assert_eq!(Vec::<i32>::new(), chunk_data[5]);
472        assert_eq!(None, chunk_data.get(6));
473    }
474
475    #[test]
476    fn exact_accounts_for_trailing_zeros() {
477        let chunks: Vec<usize> = vec![0, 1, 2, 3, 0, 0];
478        let data = vec![8, 7, 6, 5, 4, 3, 2, 1];
479        let mut iter = data.arbitrary_chunks_exact(&chunks);
480
481        assert_eq!(&[0i32; 0], iter.next().unwrap());
482        assert_eq!(&[8i32], iter.next().unwrap());
483        assert_eq!(&[7i32, 6], iter.next().unwrap());
484        assert_eq!(&[5i32, 4, 3], iter.next().unwrap());
485        assert_eq!(&[0i32; 0], iter.next().unwrap());
486        assert_eq!(&[0i32; 0], iter.next().unwrap());
487        assert_eq!(None, iter.next());
488        assert_eq!(&[2i32, 1], iter.remainder());
489    }
490
491    #[test]
492    fn exact_mut_accounts_for_trailing_zeros() {
493        let chunks: Vec<usize> = vec![0, 1, 2, 3, 0, 0];
494        let mut data = vec![8, 7, 6, 5, 4, 3, 2, 1];
495        let mut iter = data.arbitrary_chunks_exact_mut(&chunks);
496
497        assert_eq!(&mut [0i32; 0], iter.next().unwrap());
498        assert_eq!(&mut [8i32], iter.next().unwrap());
499        assert_eq!(&mut [7i32, 6], iter.next().unwrap());
500        assert_eq!(&mut [5i32, 4, 3], iter.next().unwrap());
501        assert_eq!(&[0i32; 0], iter.next().unwrap());
502        assert_eq!(&[0i32; 0], iter.next().unwrap());
503        assert_eq!(None, iter.next());
504        assert_eq!(&mut [2i32, 1], iter.remainder());
505    }
506
507    #[test]
508    fn it_stops_when_data_runs_out() {
509        let chunks: Vec<usize> = vec![0, 1, 2, 3];
510        let data = vec![8, 7, 6, 5, 4];
511
512        let chunk_data: Vec<Vec<i32>> = data
513            .arbitrary_chunks(&chunks)
514            .map(|chunk| chunk.to_vec())
515            .collect();
516
517        assert_eq!(Vec::<i32>::new(), chunk_data[0]);
518        assert_eq!(vec![8], chunk_data[1]);
519        assert_eq!(vec![7, 6], chunk_data[2]);
520        assert_eq!(vec![5, 4], chunk_data[3]);
521        assert_eq!(None, chunk_data.get(4));
522    }
523
524    #[test]
525    fn mut_stops_when_data_runs_out() {
526        let chunks: Vec<usize> = vec![0, 1, 2, 3];
527        let mut data = vec![8, 7, 6, 5, 4];
528        let chunk_data: Vec<Vec<i32>> = data
529            .arbitrary_chunks_mut(&chunks)
530            .map(|chunk| chunk.to_vec())
531            .collect();
532
533        assert_eq!(Vec::<i32>::new(), chunk_data[0]);
534        assert_eq!(vec![8], chunk_data[1]);
535        assert_eq!(vec![7, 6], chunk_data[2]);
536        assert_eq!(vec![5, 4], chunk_data[3]);
537        assert_eq!(None, chunk_data.get(4));
538    }
539
540    #[test]
541    fn exact_stops_when_data_runs_out() {
542        let chunks: Vec<usize> = vec![0, 1, 2, 3];
543        let data = vec![8, 7, 6, 5, 4];
544        let mut iter = data.arbitrary_chunks_exact(&chunks);
545
546        assert_eq!(&[0i32; 0], iter.next().unwrap());
547        assert_eq!(&[8i32], iter.next().unwrap());
548        assert_eq!(&[7i32, 6], iter.next().unwrap());
549        assert_eq!(None, iter.next());
550        assert_eq!(&[5i32, 4], iter.remainder());
551    }
552
553    #[test]
554    fn exact_mut_stops_when_data_runs_out() {
555        let chunks: Vec<usize> = vec![0, 1, 2, 3];
556        let mut data = vec![8, 7, 6, 5, 4];
557        let mut iter = data.arbitrary_chunks_exact_mut(&chunks);
558
559        assert_eq!(&mut [0i32; 0], iter.next().unwrap());
560        assert_eq!(&mut [8i32], iter.next().unwrap());
561        assert_eq!(&mut [7i32, 6], iter.next().unwrap());
562        assert_eq!(None, iter.next());
563        assert_eq!(&mut [5i32, 4], iter.remainder());
564    }
565}
arbitrary_chunks/lib.rs

arbitrary_chunks/
lib.rs
