diff --git a/src/lib.rs b/src/lib.rs index 745e82048..bdc104241 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -3701,6 +3701,77 @@ pub trait Itertools: Iterator { (left, right) } + /// Collect all iterator elements into one of two partitions. + /// Unlike [`Iterator::partition`], when predicate returns true for the first time, + /// it collects this element and all rest elements into `B`. + /// + /// ``` + /// use itertools::Itertools; + /// + /// let nums = vec![0, 1, 2, 3, 4, 5]; + /// + /// let (a, b): (Vec<_>, Vec<_>) = nums.into_iter().split(|n| *n == 3); + /// + /// assert_eq!(a, [0, 1, 2]); + /// assert_eq!(b, [3, 4, 5]); + /// ``` + fn split(mut self, mut predicate: P) -> (A, B) + where + Self: Sized, + P: FnMut(&Self::Item) -> bool, + A: Default + Extend, + B: Default + Extend, + { + let mut collect_a = A::default(); + let mut collect_b = B::default(); + + while let Some(elem) = self.next() { + if predicate(&elem) { + collect_b.extend(once(elem)); + break; + } + collect_a.extend(once(elem)); + } + collect_b.extend(self); + + (collect_a, collect_b) + } + + /// Collect all iterator elements into one of two partitions. + /// Unlike [`Itertools::split`], when predicate returns true, the element is collected into `A`. + /// + /// ``` + /// use itertools::Itertools; + /// + /// let nums = vec![0, 1, 2, 3, 4, 5]; + /// + /// let (a, b): (Vec<_>, Vec<_>) = nums.into_iter().split_inclusive(|n| *n == 3); + /// + /// assert_eq!(a, [0, 1, 2, 3]); + /// assert_eq!(b, [4, 5]); + /// ``` + fn split_inclusive(mut self, mut predicate: P) -> (A, B) + where + Self: Sized, + P: FnMut(&Self::Item) -> bool, + A: Default + Extend, + B: Default + Extend, + { + let mut collect_a = A::default(); + let mut collect_b = B::default(); + + while let Some(elem) = self.next() { + let is_split_point = predicate(&elem); + collect_a.extend(once(elem)); + if is_split_point { + break; + } + } + collect_b.extend(self); + + (collect_a, collect_b) + } + /// Partition a sequence of `Result`s into one list of all the `Ok` elements /// and another list of all the `Err` elements. ///