Skip to content

Commit

Permalink
Itertools::tail: return VecDeque rather than Vec
Browse files Browse the repository at this point in the history
`rotate_left` is more efficient on `VecDeque` than on a slice.
`VecDeque::from(vec)` is O(1) on recent rust.

Co-Authored-By: scottmcm <scottmcm@users.noreply.github.com>
  • Loading branch information
Philippe-Cholet and scottmcm committed Mar 16, 2024
1 parent cd6f902 commit e9d8e90
Showing 1 changed file with 12 additions and 5 deletions.
17 changes: 12 additions & 5 deletions src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -55,7 +55,7 @@ extern crate core as std;
extern crate alloc;

#[cfg(feature = "use_alloc")]
use alloc::{string::String, vec::Vec};
use alloc::{collections::VecDeque, string::String, vec::Vec};

pub use either::Either;

Expand All @@ -72,6 +72,8 @@ use std::fmt::Write;
use std::hash::Hash;
use std::iter::{once, IntoIterator};
#[cfg(feature = "use_alloc")]
type VecDequeIntoIter<T> = alloc::collections::vec_deque::IntoIter<T>;
#[cfg(feature = "use_alloc")]
type VecIntoIter<T> = alloc::vec::IntoIter<T>;
use std::iter::FromIterator;

Expand Down Expand Up @@ -3136,8 +3138,10 @@ pub trait Itertools: Iterator {
/// 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
/// The iterator, if directly collected to a `VecDeque`, is converted
/// without any extra copying or allocation cost.
/// If directly collected to a `Vec`, it may need some data movement
/// but no re-allocation.
///
/// ```
/// use itertools::{assert_equal, Itertools};
Expand All @@ -3155,20 +3159,22 @@ pub trait Itertools: Iterator {
/// `.rev().take(n).rev()` to have a similar result (lazy and non-allocating)
/// without consuming the entire iterator.
#[cfg(feature = "use_alloc")]
fn tail(self, n: usize) -> VecIntoIter<Self::Item>
fn tail(self, n: usize) -> VecDequeIntoIter<Self::Item>
where
Self: Sized,
{
match n {
0 => {
self.last();
Vec::new()
VecDeque::new()
}
1 => self.last().into_iter().collect(),
_ => {
// Skip the starting part of the iterator if possible.
let (low, _) = self.size_hint();
let mut iter = self.fuse().skip(low.saturating_sub(n));
// TODO: If VecDeque has a more efficient method than
// `.pop_front();.push_back(val)` in the future then maybe revisit this.
let mut data: Vec<_> = iter.by_ref().take(n).collect();
// Update `data` cyclically.
let idx = iter.fold(0, |i, val| {
Expand All @@ -3179,7 +3185,8 @@ pub trait Itertools: Iterator {
i + 1
}
});
// Respect the insertion order.
// Respect the insertion order, efficiently.
let mut data = VecDeque::from(data);
data.rotate_left(idx);
data
}
Expand Down

0 comments on commit e9d8e90

Please sign in to comment.