Skip to content

Instantly share code, notes, and snippets.

@edwardw
Created December 20, 2019 14:42
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save edwardw/e42c50b6fa1c9112b6fff282835341ba to your computer and use it in GitHub Desktop.
Save edwardw/e42c50b6fa1c9112b6fff282835341ba to your computer and use it in GitHub Desktop.
Cycle a Rust iterator a given number of times
[package]
name = "cycle_n"
version = "0.1.0"
authors = ["whoever"]
edition = "2018"
[dependencies]
[dev-dependencies]
criterion = "0.3"
[[bench]]
name = "cycle_n_benchmark"
harness = false
use ::cycle_n::{cycle_n, cycle_n_manual, cycle_n_trait};
use criterion::{black_box, criterion_group, criterion_main, BenchmarkId, Criterion, Throughput};
pub fn bench_cycle_n(c: &mut Criterion) {
let mut group = c.benchmark_group("Bench CycleN");
for size in [16, 32, 64, 128, 256, 512, 1024].iter() {
let input = (0..*size).collect::<Vec<_>>();
group.throughput(Throughput::Elements(*size as u64));
group.bench_with_input(
BenchmarkId::new("repeat_take_flatten", *size),
&input,
|b, v| b.iter(|| cycle_n(black_box(v.clone()), 4)),
);
group.bench_with_input(BenchmarkId::new("manual_loop", *size), &input, |b, v| {
b.iter(|| cycle_n_manual(black_box(v.clone()), 4))
});
group.bench_with_input(BenchmarkId::new("cycle_n_trait", *size), &input, |b, v| {
b.iter(|| cycle_n_trait(black_box(v.clone()), 4))
});
}
group.finish();
}
criterion_group!(benches, bench_cycle_n,);
criterion_main!(benches);
pub fn cycle_n(v: Vec<i32>, n: usize) -> Vec<i32> {
std::iter::repeat(v.into_iter()).take(n).flatten().collect()
}
pub fn cycle_n_manual(v: Vec<i32>, n: usize) -> Vec<i32> {
let it = v.into_iter();
let mut res = Vec::new();
for _ in 0..n {
for x in it.clone() {
res.push(x);
}
}
res
}
pub fn cycle_n_trait(v: Vec<i32>, n: usize) -> Vec<i32> {
v.into_iter().cycle_n(n).collect()
}
pub trait Itermisc: std::iter::Iterator {
fn cycle_n(self, n: usize) -> CycleN<Self>
where
Self: Clone,
{
CycleN::new(self, n)
}
}
impl<T: ?Sized> Itermisc for T where T: std::iter::Iterator {}
pub struct CycleN<I> {
orig: I,
iter: I,
tick: usize,
}
impl<I: Clone> CycleN<I> {
pub fn new(iter: I, n: usize) -> CycleN<I> {
CycleN {
orig: iter.clone(),
iter,
tick: n,
}
}
}
impl<I> Iterator for CycleN<I>
where
I: Clone + Iterator,
{
type Item = <I as Iterator>::Item;
#[inline]
fn next(&mut self) -> Option<<I as Iterator>::Item> {
match self.iter.next() {
None if self.tick > 0 => {
self.tick -= 1;
self.iter = self.orig.clone();
self.iter.next()
}
y => y,
}
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment