diff --git a/shared/comm-lib/src/tools.rs b/shared/comm-lib/src/tools.rs --- a/shared/comm-lib/src/tools.rs +++ b/shared/comm-lib/src/tools.rs @@ -74,6 +74,67 @@ } } +pub trait IntoChunks { + /// Splits the vec into `num_chunks` chunks and returns an iterator + /// over these chunks. The chunks do not overlap. + /// + /// Chunks size is given by `ceil(vector_length / num_chunks)`. + /// If vector length is not divisible by `num_chunks`, + /// the last chunk will have less elements. + /// + /// If you're looking for chunks of given size, use [`chunks`] instead. + /// + /// # Panics + /// + /// Panics if `num_chunks` is 0. + /// + /// # Examples + /// + /// ``` + /// use comm_lib::tools::IntoChunks; + /// + /// let items = vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]; + /// let mut iter = items.into_n_chunks(3); + /// assert_eq!(&iter.next().unwrap(), &[1, 2, 3, 4]); + /// assert_eq!(&iter.next().unwrap(), &[5, 6, 7, 8]); + /// assert_eq!(&iter.next().unwrap(), &[9, 10]); + /// assert!(iter.next().is_none()); + /// ``` + /// + /// [`chunks`]: slice::chunks + fn into_n_chunks(self, num_chunks: usize) -> impl Iterator>; +} + +impl IntoChunks for Vec { + fn into_n_chunks(self, num_chunks: usize) -> impl Iterator> { + struct ChunksIterator { + pub slice: Vec, + pub chunk_size: usize, + } + impl Iterator for ChunksIterator { + type Item = Vec; + fn next(&mut self) -> Option> { + let next_size = std::cmp::min(self.slice.len(), self.chunk_size); + if next_size == 0 { + None + } else { + let next_chunk = self.slice.drain(0..next_size).collect(); + Some(next_chunk) + } + } + } + + assert!(num_chunks > 0, "Number of chunks cannot be 0"); + let len = self.len(); + let rem = len % num_chunks; + let chunk_size = len / num_chunks + if rem > 0 { 1 } else { 0 }; + ChunksIterator { + slice: self, + chunk_size, + } + } +} + pub fn generate_random_string( length: usize, rng: &mut (impl Rng + CryptoRng), @@ -128,6 +189,11 @@ fn empty_is_invalid() { assert!(!is_valid_identifier("")); } +} + +#[cfg(test)] +mod defer_tests { + use super::*; #[test] fn defer_runs() { @@ -164,3 +230,50 @@ assert!(v) } } + +#[cfg(test)] +mod vec_utils_tests { + use super::*; + + #[test] + fn test_chunks_without_remainder() { + let items = vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]; + let mut iter = items.into_n_chunks(3); + assert_eq!(&iter.next().unwrap(), &[1, 2, 3, 4]); + assert_eq!(&iter.next().unwrap(), &[5, 6, 7, 8]); + assert_eq!(&iter.next().unwrap(), &[9, 10, 11, 12]); + assert!(iter.next().is_none()); + } + + #[test] + fn test_chunks_with_remainder() { + let items = vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]; + let mut iter = items.into_n_chunks(3); + assert_eq!(&iter.next().unwrap(), &[1, 2, 3, 4]); + assert_eq!(&iter.next().unwrap(), &[5, 6, 7, 8]); + assert_eq!(&iter.next().unwrap(), &[9, 10]); + assert!(iter.next().is_none()); + } + + #[test] + fn test_one_chunk() { + let items: Vec = vec![1, 2, 3]; + let mut iter = items.into_n_chunks(1); + assert_eq!(&iter.next().unwrap(), &[1, 2, 3]); + assert!(iter.next().is_none()); + } + + #[test] + fn test_empty_vec() { + let items: Vec = vec![]; + let mut iter = items.into_n_chunks(2); + assert!(iter.next().is_none()); + } + + #[test] + #[should_panic] + fn into_n_chunks_panics_with_0_chunks() { + let items = vec![1, 2, 3]; + let _ = items.into_n_chunks(0); + } +}