Merge #616

616: Add `take_while_inclusive` method r=jswrenn a=e-rhodes

Adds `take_while_inclusive` method that returns elements as long as a predicate is satisfied, but including the first element that doesn't satisfy the predicate.

First discussed addition to std in <rust-lang/rust#62208>.

Closes #597.

Co-authored-by: Erik Rhodes <erik@space-nav.com>

src/lib.rs

@@ -146,6 +146,7 @@ pub mod structs {
     pub use crate::repeatn::RepeatN;
     #[allow(deprecated)]
     pub use crate::sources::{RepeatCall, Unfold, Iterate};
+    pub use crate::take_while_inclusive::TakeWhileInclusive;
     #[cfg(feature = "use_alloc")]
     pub use crate::tee::Tee;
     pub use crate::tuple_impl::{TupleBuffer, TupleWindows, CircularTupleWindows, Tuples};
@@ -233,6 +234,7 @@ mod rciter_impl;
 mod repeatn;
 mod size_hint;
 mod sources;
+mod take_while_inclusive;
 #[cfg(feature = "use_alloc")]
 mod tee;
 mod tuple_impl;
@@ -1457,6 +1459,74 @@ pub trait Itertools : Iterator {
         adaptors::take_while_ref(self, accept)
     }
 
+    /// Returns an iterator adaptor that consumes elements while the given
+    /// predicate is `true`, *including* the element for which the predicate
+    /// first returned `false`.
+    ///
+    /// The [`.take_while()`][std::iter::Iterator::take_while] adaptor is useful
+    /// when you want items satisfying a predicate, but to know when to stop
+    /// taking elements, we have to consume that first element that doesn't
+    /// satisfy the predicate. This adaptor includes that element where
+    /// [`.take_while()`][std::iter::Iterator::take_while] would drop it.
+    ///
+    /// The [`.take_while_ref()`][crate::Itertools::take_while_ref] adaptor
+    /// serves a similar purpose, but this adaptor doesn't require [`Clone`]ing
+    /// the underlying elements.
+    ///
+    /// ```rust
+    /// # use itertools::Itertools;
+    /// let items = vec![1, 2, 3, 4, 5];
+    /// let filtered: Vec<_> = items
+    ///     .into_iter()
+    ///     .take_while_inclusive(|&n| n % 3 != 0)
+    ///     .collect();
+    ///
+    /// assert_eq!(filtered, vec![1, 2, 3]);
+    /// ```
+    ///
+    /// ```rust
+    /// # use itertools::Itertools;
+    /// let items = vec![1, 2, 3, 4, 5];
+    ///
+    /// let take_while_inclusive_result: Vec<_> = items
+    ///     .iter()
+    ///     .copied()
+    ///     .take_while_inclusive(|&n| n % 3 != 0)
+    ///     .collect();
+    /// let take_while_result: Vec<_> = items
+    ///     .into_iter()
+    ///     .take_while(|&n| n % 3 != 0)
+    ///     .collect();
+    ///
+    /// assert_eq!(take_while_inclusive_result, vec![1, 2, 3]);
+    /// assert_eq!(take_while_result, vec![1, 2]);
+    /// // both iterators have the same items remaining at this point---the 3
+    /// // is lost from the `take_while` vec
+    /// ```
+    ///
+    /// ```rust
+    /// # use itertools::Itertools;
+    /// #[derive(Debug, PartialEq)]
+    /// struct NoCloneImpl(i32);
+    ///
+    /// let non_clonable_items: Vec<_> = vec![1, 2, 3, 4, 5]
+    ///     .into_iter()
+    ///     .map(NoCloneImpl)
+    ///     .collect();
+    /// let filtered: Vec<_> = non_clonable_items
+    ///     .into_iter()
+    ///     .take_while_inclusive(|n| n.0 % 3 != 0)
+    ///     .collect();
+    /// let expected: Vec<_> = vec![1, 2, 3].into_iter().map(NoCloneImpl).collect();
+    /// assert_eq!(filtered, expected);
+    fn take_while_inclusive<F>(&mut self, accept: F) -> TakeWhileInclusive<Self, F>
+    where
+        Self: Sized,
+        F: FnMut(&Self::Item) -> bool,
+    {
+        take_while_inclusive::TakeWhileInclusive::new(self, accept)
+    }
+
     /// Return an iterator adaptor that filters `Option<A>` iterator elements
     /// and produces `A`. Stops on the first `None` encountered.
     ///

src/take_while_inclusive.rs

@@ -0,0 +1,68 @@
+use core::iter::FusedIterator;
+use std::fmt;
+
+/// An iterator adaptor that consumes elements while the given predicate is
+/// `true`, including the element for which the predicate first returned
+/// `false`.
+///
+/// See [`.take_while_inclusive()`](crate::Itertools::take_while_inclusive)
+/// for more information.
+#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
+pub struct TakeWhileInclusive<'a, I: 'a, F> {
+    iter: &'a mut I,
+    predicate: F,
+    done: bool,
+}
+
+impl<'a, I, F> TakeWhileInclusive<'a, I, F>
+where
+    I: Iterator,
+    F: FnMut(&I::Item) -> bool,
+{
+    /// Create a new [`TakeWhileInclusive`] from an iterator and a predicate.
+    pub fn new(iter: &'a mut I, predicate: F) -> Self {
+        Self { iter, predicate, done: false}
+    }
+}
+
+impl<'a, I, F> fmt::Debug for TakeWhileInclusive<'a, I, F>
+    where I: Iterator + fmt::Debug,
+{
+    debug_fmt_fields!(TakeWhileInclusive, iter);
+}
+
+impl<'a, I, F> Iterator for TakeWhileInclusive<'a, I, F>
+where
+    I: Iterator,
+    F: FnMut(&I::Item) -> bool
+{
+    type Item = I::Item;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        if self.done {
+            None
+        } else {
+            self.iter.next().map(|item| {
+                if !(self.predicate)(&item) {
+                    self.done = true;
+                }
+                item
+            })
+        }
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        if self.done {
+            (0, Some(0))
+        } else {
+            (0, self.iter.size_hint().1)
+        }
+    }
+}
+
+impl<I, F> FusedIterator for TakeWhileInclusive<'_, I, F>
+where
+    I: Iterator,
+    F: FnMut(&I::Item) -> bool
+{
+}