Compiler Explorer

Source code

#![feature(try_trait, core_intrinsics)]
use std::ops::Try;
use std::intrinsics::{unlikely, likely};

#[derive(Clone, Debug)]
#[must_use = "iterators are lazy and do nothing unless consumed"]
pub struct Skip {
 iter: I,
 n: usize,
}
impl Skip {
 pub fn new(iter: I, n: usize) -> Skip {
 Skip { iter, n }
 }
}

impl Iterator for Skip
where
 I: Iterator,
{
 type Item = ::Item;

#[inline]
 fn next(&mut self) -> Option<I::Item> {
 if self.n == 0 {
 self.iter.next()
 } else {
 let old_n = self.n;
 self.n = 0;
 self.iter.nth(old_n)
 }
 }

#[inline]
 fn nth(&mut self, n: usize) -> Option<I::Item> {
 // Can't just add n + self.n due to overflow.
 if unlikely(self.n > 0) {
 let to_skip = self.n;
 self.n = 0;
 // nth(n) skips n+1
 self.iter.nth(to_skip - 1)?;
 }
 self.iter.nth(n)
 }

#[inline]
    fn count(mut self) -> usize {
        if self.n > 0 {
            // nth(n) skips n+1
            if self.iter.nth(self.n - 1).is_none() {
                return 0;
            }
        }
        self.iter.count()
    }

#[inline]
 fn last(mut self) -> Option<I::Item> {
 if self.n > 0 {
 // nth(n) skips n+1
 self.iter.nth(self.n - 1)?;
 }
 self.iter.last()
 }

#[inline]
 fn size_hint(&self) -> (usize, Option<usize>) {
 let (lower, upper) = self.iter.size_hint();

let lower = lower.saturating_sub(self.n);
        let upper = match upper {
            Some(x) => Some(x.saturating_sub(self.n)),
            None => None,
        };

(lower, upper)
    }

#[inline]
 fn fold<Acc, Fold>(mut self, init: Acc, fold: Fold) -> Acc
 where
 Fold: FnMut(Acc, Self::Item) -> Acc,
 {
 if self.n > 0 {
 // nth(n) skips n+1
 if self.iter.nth(self.n - 1).is_none() {
 return init;
 }
 }
 self.iter.fold(init, fold)
 }
}

pub struct CustomSkip {
 iter: I,
 n: Option<usize>,
}

impl CustomSkip {
 pub fn new(iter: I, n: usize) -> CustomSkip {
 CustomSkip { iter, n: Some(n) }
 }
}

impl<I: Iterator> Iterator for CustomSkip {

type Item = I::Item;

#[inline]
 fn next(&mut self) -> Option<I::Item> {
 if let Some(n) = self.n.take() {
 self.iter.nth(n)
 } else {
 self.iter.next()
 }
 }
}

trait SkipExt: Iterator + Sized {
 fn custom_skip(self, n: usize) -> CustomSkip<Self> {
 CustomSkip::new(self, n)
 }
}

impl<I: Iterator> SkipExt for I {}

pub fn test() -> usize {
    let iter = 0..100000;
    let mut sum = 0;
    for i in Skip::new(iter, 10) {
        sum += i;
    }
    sum
}

pub fn test2() -> usize {
    let iter = 0..100000;
    let mut sum = 0;
    for i in CustomSkip::new(iter, 10) {
        sum += i;
    }
    sum
}