Add “iterate with separators” iterator function

benches_iters.rs

// Benchmarks for the Rust iterator extension discussion at
//   https://internals.rust-lang.org/t/add-iterate-with-separators-iterator-function/18781/13

// Place this page as /benches/iters.rs in a rust project created with `cargo new itertest --lib`

// Add to Cargo.toml:
//
// [dependencies]
// itertools = "0.10"
//
// [dev-dependencies]
// criterion = { version = "0.4", features = ["html_reports"] }
//
// [[bench]]
// name = "iters"
// harness = false

// To benchmark:   cargo +nightly bench

#![feature(iter_intersperse)]

use criterion::{criterion_group, criterion_main, Criterion};
use mydemo::IteratorExt;
use std::fmt::Write as _;

const ELEMENTS: usize = 10000;

//
// *************************************************************************************************
//    These benchmarks create a sum of the elements in a slice, interspersed with 1s.
// *************************************************************************************************
//

pub fn nones(elements: &[usize]) -> usize {
    let mut sum = 0;
    for v in elements.iter().nones() {
        if let Some(v) = v {
            sum += v;
        } else {
            sum += 1;
        }
    }
    sum
}

pub fn intersp(elements: &[usize]) -> usize {
    let mut sum = 0;
    for v in elements.iter().map(Some).intersperse(None) {
        if let Some(v) = v {
            sum += v;
        } else {
            sum += 1;
        }
    }
    sum
}

pub fn with_position(elements: &[usize]) -> usize {
    use itertools::{Itertools, Position};

    let mut sum = 0;
    for v in elements.iter().with_position() {
        if !matches!(v, Position::First(_) | Position::Only(_)) {
            sum += 1;
        }
        sum += v.into_inner();
    }
    sum
}

pub fn intersp_wth(elements: &[usize]) -> usize {
    let mut sum = 0;
    for v in elements.iter().map(Some).intersperse_with(|| None) {
        if let Some(v) = v {
            sum += v;
        } else {
            sum += 1;
        }
    }
    sum
}

pub fn iter_for(elements: &[usize]) -> usize {
    let mut sum = 0;
    let mut use_sep = false;
    for v in elements.iter() {
        if use_sep {
            sum += 1;
        } else {
            use_sep = true;
        }
        sum += v;
    }
    sum
}

pub fn foreach(elements: &[usize]) -> usize {
    let mut sum = 0;
    let mut use_sep = false;
    elements.iter().for_each(|v| {
        if use_sep {
            sum += 1;
        } else {
            use_sep = true;
        }
        sum += v;
    });
    sum
}

pub fn bench_iters(c: &mut Criterion) {
    let elements = (0..ELEMENTS).collect::<Vec<_>>();
    let mut group = c.benchmark_group("Sum");

    group.bench_function("nones", |b| b.iter(|| nones(&elements)));
    group.bench_function("foreach", |b| b.iter(|| foreach(&elements)));
    group.bench_function("iter_for", |b| b.iter(|| iter_for(&elements)));
    group.bench_function("intersp_wth", |b| b.iter(|| intersp_wth(&elements)));
    group.bench_function("intersp", |b| b.iter(|| intersp(&elements)));
    group.bench_function("with_position", |b| b.iter(|| with_position(&elements)));

    group.finish();
}

//
// *************************************************************************************************
//    These benchmarks are the same as the join() function, creating a CSV string with all values
// *************************************************************************************************
//

pub fn bench_iter_buffers(c: &mut Criterion) {
    // Pre-allocate the elements and the output buffer (it's cleared on each benchmark)
    let elements = (0..ELEMENTS).collect::<Vec<_>>();
    let mut buf = String::with_capacity(elements.len() * 10);

    let mut group = c.benchmark_group("StrJoin");

    group.bench_function("nones", |b| {
        b.iter(|| {
            buf.clear();
            for v in elements.iter().nones() {
                if let Some(v) = v {
                    let _ = write!(&mut buf, "{v}");
                } else {
                    buf.push(',')
                }
            }
            buf.len()
        })
    });

    group.bench_function("foreach", |b| {
        b.iter(|| {
            buf.clear();
            let mut use_sep = false;
            elements.iter().for_each(|v| {
                if use_sep {
                    buf.push(',');
                } else {
                    use_sep = true;
                }
                let _ = write!(buf, "{v}");
            });
            buf.len()
        })
    });

    group.bench_function("for", |b| {
        b.iter(|| {
            buf.clear();
            let mut use_sep = false;
            for v in elements.iter() {
                if use_sep {
                    buf.push(',');
                } else {
                    use_sep = true;
                }
                let _ = write!(buf, "{v}");
            }
            buf.len()
        })
    });

    group.bench_function("intersp_with", |b| {
        b.iter(|| {
            buf.clear();
            for v in elements.iter().map(Some).intersperse_with(|| None) {
                if let Some(v) = v {
                    let _ = write!(&mut buf, "{v}");
                } else {
                    buf.push(',')
                }
            }
            buf.len()
        })
    });

    group.bench_function("intersp", |b| {
        b.iter(|| {
            buf.clear();
            for v in elements.iter().map(Some).intersperse(None) {
                if let Some(v) = v {
                    let _ = write!(&mut buf, "{v}");
                } else {
                    buf.push(',')
                }
            }
            buf.len()
        })
    });

    group.bench_function("with_position", |b| {
        use itertools::{Itertools, Position};
        b.iter(|| {
            buf.clear();
            for v in elements.iter().with_position() {
                if !matches!(v, Position::First(_) | Position::Only(_)) {
                    buf.push(',')
                }
                let _ = write!(&mut buf, "{}", v.into_inner());
            }
            buf.len()
        })
    });

    group.finish();
}

criterion_group!(benches, bench_iters, bench_iter_buffers,);
criterion_main!(benches);

src_lib.rs

// Save this file as /src/lib.rs

pub struct IntersperseNones<I: Iterator> {
    next_item: Option<I::Item>,
    iter: I,
}

impl<I: Iterator> IntersperseNones<I> {
    pub fn new(mut iter: I) -> Self {
        Self {
            next_item: iter.next(),
            iter,
        }
    }
}

impl<I: Iterator> Iterator for IntersperseNones<I> {
    type Item = Option<I::Item>;

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        if let Some(v) = self.next_item.take() {
            Some(Some(v))
        } else {
            let next_item = self.iter.next();
            if next_item.is_some() {
                self.next_item = next_item;
                Some(None)
            } else {
                None
            }
        }
    }
}

pub trait IteratorExt: Iterator {
    #[inline]
    fn nones(self) -> IntersperseNones<Self>
    where
        Self: Sized,
    {
        IntersperseNones::new(self)
    }
}

impl<I: Iterator> IteratorExt for I {}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_nones() {
        let v: Vec<i32> = Vec::new();
        let mut it = v.iter().nones();
        assert_eq!(it.next(), None);
        assert_eq!(it.next(), None);

        let v = vec![1];
        let mut it = v.iter().nones();
        assert_eq!(it.next(), Some(Some(&1)));
        assert_eq!(it.next(), None);
        assert_eq!(it.next(), None);

        let v = vec![1, 2];
        let mut it = v.iter().nones();
        assert_eq!(it.next(), Some(Some(&1)));
        assert_eq!(it.next(), Some(None));
        assert_eq!(it.next(), Some(Some(&2)));
        assert_eq!(it.next(), None);
        assert_eq!(it.next(), None);

        let v = vec![1, 2, 3];
        let mut it = v.iter().nones();
        assert_eq!(it.next(), Some(Some(&1)));
        assert_eq!(it.next(), Some(None));
        assert_eq!(it.next(), Some(Some(&2)));
        assert_eq!(it.next(), Some(None));
        assert_eq!(it.next(), Some(Some(&3)));
        assert_eq!(it.next(), None);
        assert_eq!(it.next(), None);
    }
}