Skip to content

Commit

Permalink
New Itertools::tail
Browse files Browse the repository at this point in the history
Note that all cases consume the iterator (even when n=0).

It could alternatively return a vector and not `VecIntoIter` but most of our methods do this (and immediately collect to a vector won't reallocate).

I don't think we should write a non-lazy `head` method (`.take(n)` is a lazy one) but it would be easy: `.take(n).collect_vec().into_iter()`.
  • Loading branch information
Philippe-Cholet committed Mar 8, 2024
1 parent 8ed734b commit 42cac22
Show file tree
Hide file tree
Showing 2 changed files with 52 additions and 0 deletions.
47 changes: 47 additions & 0 deletions src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3133,6 +3133,53 @@ pub trait Itertools: Iterator {
self.k_largest_by(k, k_smallest::key_to_cmp(key))
}

/// Consumes the iterator and return an iterator of the last `n` elements.
///
/// It allocates up to `n` elements.
/// The iterator, if directly collected to a `Vec`, is converted
/// without any extra copying or allocation cost.
///
/// ```
/// use itertools::{assert_equal, Itertools};
///
/// let v = vec![5, 9, 8, 4, 2, 12, 0];
/// assert_equal(v.iter().tail(3), &[2, 12, 0]);
/// assert_equal(v.iter().tail(10), &v);
///
/// assert_equal((0..100).tail(10), 90..100);
/// ```
///
/// For double ended iterators without side-effects, you might prefer
/// `.rev().take(n).collect_vec().into_iter().rev()`
/// to have the same result without consuming the entire iterator.
#[cfg(feature = "use_alloc")]
fn tail(mut self, n: usize) -> VecIntoIter<Self::Item>
where
Self: Sized,
{
match n {
0 => {
self.last();
Vec::new()
}
1 => self.last().into_iter().collect(),
_ => {
let mut data: Vec<_> = self.by_ref().take(n).collect();
// If the iterator is not exhausted yet, update the data cyclically and
// finally rotate the data so that the order of the tail is respected.
if data.len() == n {
let idx = self.fold(0, |i, val| {
data[i] = val;
(i + 1) % n
});
data.rotate_left(idx);
}
data
}
}
.into_iter()
}

/// Collect all iterator elements into one of two
/// partitions. Unlike [`Iterator::partition`], each partition may
/// have a distinct type.
Expand Down
5 changes: 5 additions & 0 deletions tests/quick.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1949,4 +1949,9 @@ quickcheck! {
result_set.is_empty()
}
}

fn tail(v: Vec<i32>, n: u8) -> bool {
let n = n as usize;
itertools::equal(v.iter().tail(n), &v[v.len().saturating_sub(n)..])
}
}

0 comments on commit 42cac22

Please sign in to comment.