<commit_before>use primal_bit::{BitVec};
use std::{cmp};

use wheel;
use primal_smallsieve::Primes;

/// A segmented sieve that yields only a small run of primes at a
/// time.
///
/// This is heavily inspired by this [segmented
/// sieve](http://primesieve.org/segmented_sieve.html) code.
#[derive(Debug)]
pub struct StreamingSieve {
    small: Primes,
    // stores which numbers *aren't* prime, i.e. true == composite.
    sieve: BitVec,
    primes: Vec<wheel::WheelInfo>,

    low: usize,
    current: usize,
    limit: usize,
}

const CACHE: usize = (32 << 10) - 2;
// 8 for the bit vector, 2 for storing odd numbers only
const SEG_ELEMS: usize = 8 * CACHE;
const SEG_LEN: usize = SEG_ELEMS * wheel::MODULO / wheel::SIZE;

impl StreamingSieve {
    /// Create a new instance of the streaming sieve that will
    /// correctly progressively filter primes up to `limit`.
    pub fn new(limit: usize) -> StreamingSieve {
        let small = Primes::sieve((limit as f64).sqrt() as usize + 1);
        let current = match wheel::MODULO {
            6 => 5,
            30 => 7,
            210 => 11,
            _ => unimplemented!(),
        };
        let low = 0;

        let elems = cmp::min((limit * wheel::SIZE + wheel::MODULO - 1) / wheel::MODULO, SEG_ELEMS);
        StreamingSieve {
            small: small,
            sieve: BitVec::from_elem(elems, false),
            primes: vec![],

            low: low,
            current: current,
            limit: limit
        }
    }

    fn direct_sieve(&mut self) {
        let bytes = self.sieve.as_bytes_mut();
        let top = bytes.len();

        for wi in self.primes.iter_mut() {
            let mut si = wi.sieve_index;
            let mut wi_ = wi.wheel_index;
            let p = wi.prime;
            while si < top {
                wheel::set_bit(bytes, &mut si, &mut wi_, p);
            }
            // if this wraps, we've hit the limit, and so won't be
            // continuing, so whatever, it can be junk.
            wi.sieve_index = si.wrapping_sub(top);
            wi.wheel_index = wi_;
        }
    }

    /// Extract the next chunk of filtered primes, the return value is
    /// `Some((low, v))` or `None` if the sieve has reached the limit.
    ///
    /// The vector stores bits for each odd number starting at `low`.
    /// Bit `n` of `v` is set if and only if `low + 2 * n + 1` is
    /// prime.
    ///
    /// NB. the prime 2 is not included in any of these sieves and so
    /// needs special handling.
    pub fn next(&mut self) -> Option<(usize, &BitVec)> {
        if self.low >= self.limit {
            return None
        }

        let low = self.low;
        self.low += SEG_LEN;
        let high = cmp::min(low + SEG_LEN - 1, self.limit);
        self.sieve.clear();

        let mut s = self.current;

        while s * s <= high {
            if self.small.is_prime(s) {
                self.primes.push(wheel::compute_wheel_elem(s, low));
            }
            s += 1
        }

        self.current = s;
        self.direct_sieve();

        if low == 0 {
            // 1 is not prime.
            self.sieve.set(0, true);
        }

        Some((low, &self.sieve))
    }
}

#[cfg(test)]
mod tests {
    use wheel;
    use super::StreamingSieve;
    fn gcd(x: usize, y: usize) -> usize {
        if y == 0 { x }
        else { gcd(y, x % y) }
    }
    fn coprime_to(x: usize) -> Vec<usize> {
        (1..x).filter(|&n| gcd(n, x) == 1).collect()
    }
    #[test]
    fn test() {
        let coprime = coprime_to(wheel::MODULO);
        const LIMIT: usize = 2_000_000;
        let mut sieve = StreamingSieve::new(LIMIT);
        let primes = ::primal_smallsieve::Primes::sieve(LIMIT);

        let mut base = 0;
        let mut index = 0;

        while let Some((_low, next)) = sieve.next() {
            for val in next {
                let i = wheel::MODULO * base + coprime[index];
                if i >= LIMIT { break }
                assert!(primes.is_prime(i) == !val,
                        "failed for {} (is prime = {})", i, primes.is_prime(i));

                index += 1;
                if index == wheel::SIZE {
                    index = 0;
                    base += 1
                }
            }
        }
    }
}

#[cfg(all(test, feature = "unstable"))]
mod benches {
    use test::Bencher;
    use super::StreamingSieve;

    fn run(b: &mut Bencher, n: usize) {
        b.iter(|| {
            let mut sieve = StreamingSieve::new(n);
            while sieve.next().is_some() {}
        })
    }

    #[bench]
    fn sieve_small(b: &mut Bencher) {
        run(b, 100)
    }
    #[bench]
    fn sieve_medium(b: &mut Bencher) {
        run(b, 10_000)
    }
    #[bench]
    fn sieve_large(b: &mut Bencher) {
        run(b, 100_000)
    }
    #[bench]
    fn sieve_larger(b: &mut Bencher) {
        run(b, 1_000_000)
    }
    #[bench]
    fn sieve_huge(b: &mut Bencher) {
        run(b, 10_000_000)
    }
}
<commit_msg>sieve: implement presieving for wheel30.<commit_after>use primal_bit::{BitVec};
use std::{cmp};

use wheel;
use primal_smallsieve::Primes;

/// A segmented sieve that yields only a small run of primes at a
/// time.
///
/// This is heavily inspired by this [segmented
/// sieve](http://primesieve.org/segmented_sieve.html) code.
#[derive(Debug)]
pub struct StreamingSieve {
    small: Primes,
    // stores which numbers *aren't* prime, i.e. true == composite.
    sieve: BitVec,
    primes: Vec<wheel::WheelInfo>,
    presieve: BitVec,

    low: usize,
    current: usize,
    limit: usize,
}

const CACHE: usize = (32 << 10) - 2 * (wheel::MODULO == 210) as usize;
// 8 for the bit vector, 2 for storing odd numbers only
const SEG_ELEMS: usize = 8 * CACHE;
const SEG_LEN: usize = SEG_ELEMS * wheel::MODULO / wheel::SIZE;

const PRESIEVE_PROD: usize = 2 * 3 * 5 * 7 * 11;
const PRESIEVE_PRIMES: &'static [usize] = &[2, 3, 5, 7, 11];
const PRESIEVE_NEXT: usize = 13;
const PRESIEVE_ACTIVE: bool = wheel::MODULO == 30;

#[inline(never)]
fn compute_presieve() -> BitVec {
    assert!(PRESIEVE_ACTIVE);
    let len = (PRESIEVE_PROD * wheel::SIZE + wheel::MODULO - 1) / wheel::MODULO;
    let mut bitv = BitVec::from_elem(len, false);

    // this is silly and should be done with a sieve that only uses
    // the primes in PRESIEVE_PRIMES.
    for i in 0..len {
        let true_ = wheel::from_bit_index(i);
        fn gcd(x: usize, y: usize) -> usize {
            if y == 0 { x }
            else { gcd(y, x % y) }
        }
        if gcd(true_, PRESIEVE_PROD) != 1 {
            bitv.set(i, true)
        }
    }
    bitv
}

impl StreamingSieve {
    /// Create a new instance of the streaming sieve that will
    /// correctly progressively filter primes up to `limit`.
    pub fn new(limit: usize) -> StreamingSieve {
        let small = Primes::sieve((limit as f64).sqrt() as usize + 1);
        let mut current = match wheel::MODULO {
            6 => 5,
            30 => 7,
            210 => 11,
            _ => unimplemented!(),
        };
        let low = 0;

        let presieve = if PRESIEVE_ACTIVE && limit > PRESIEVE_PROD * 2 {
            current = PRESIEVE_NEXT;
            compute_presieve()
        } else {
            BitVec::new()
        };

        let elems = cmp::min((limit * wheel::SIZE + wheel::MODULO - 1) / wheel::MODULO, SEG_ELEMS);
        StreamingSieve {
            small: small,
            sieve: BitVec::from_elem(elems, false),
            primes: vec![],
            presieve: presieve,

            low: low,
            current: current,
            limit: limit
        }
    }

    pub fn very_small_sieve(&mut self, low: usize) {
        let offset = (low % PRESIEVE_PROD) * wheel::SIZE / wheel::MODULO / 8;

        copy_all(self.sieve.as_bytes_mut(),
                 self.presieve.as_bytes(),
                 offset);
        fn copy_all(mut dst: &mut [u8], src: &[u8], init_offset: usize) {
            let mut pos = 0;
            // pre-fill data at the start, as a rotation of `src`.
            pos += memcpy(&mut dst[pos..], &src[init_offset..]);
            if pos >= dst.len() { return }

            pos += memcpy(&mut dst[pos..], &src[..init_offset]);
            if pos >= dst.len() { return }

            // progressively copy the prefix to the rest of the
            // vector, doubling each time.
            while pos < dst.len() {
                let (filled, unfilled) = dst.split_at_mut(pos);
                pos += memcpy(unfilled, filled);
            }
        }
        fn memcpy<'d>(dst: &'d mut [u8], src: &[u8]) -> usize {
            use std::ptr;
            let l = cmp::min(dst.len(), src.len());
            unsafe {
                ptr::copy(src.as_ptr(), dst.as_mut_ptr(), l);
            }
            l
        }
    }

    fn direct_sieve(&mut self) {
        let bytes = self.sieve.as_bytes_mut();
        let top = bytes.len();

        for wi in self.primes.iter_mut() {
            let mut si = wi.sieve_index;
            let mut wi_ = wi.wheel_index;
            let p = wi.prime;
            while si < top {
                wheel::set_bit(bytes, &mut si, &mut wi_, p);
            }
            // if this wraps, we've hit the limit, and so won't be
            // continuing, so whatever, it can be junk.
            wi.sieve_index = si.wrapping_sub(top);
            wi.wheel_index = wi_;
        }
    }

    /// Extract the next chunk of filtered primes, the return value is
    /// `Some((low, v))` or `None` if the sieve has reached the limit.
    ///
    /// The vector stores bits for each odd number starting at `low`.
    /// Bit `n` of `v` is set if and only if `low + 2 * n + 1` is
    /// prime.
    ///
    /// NB. the prime 2 is not included in any of these sieves and so
    /// needs special handling.
    pub fn next(&mut self) -> Option<(usize, &BitVec)> {
        if self.low >= self.limit {
            return None
        }

        let low = self.low;
        self.low += SEG_LEN;
        let high = cmp::min(low + SEG_LEN - 1, self.limit);
        self.sieve.clear();

        let mut s = self.current;

        while s * s <= high {
            if self.small.is_prime(s) {
                self.primes.push(wheel::compute_wheel_elem(s, low));
            }
            s += 1
        }

        self.current = s;
        if PRESIEVE_ACTIVE && self.presieve.len() > 0 {
            self.very_small_sieve(low);
        }
        self.direct_sieve();

        if low == 0 {
            // 1 is not prime.
            self.sieve.set(0, true);
            if PRESIEVE_ACTIVE {
                // but these are
                for &x in PRESIEVE_PRIMES {
                    let (use_, idx) = wheel::bit_index(x);
                    if use_ {
                        self.sieve.set(idx, false)
                    }
                }
            }
        }

        Some((low, &self.sieve))
    }
}

#[cfg(test)]
mod tests {
    use wheel;
    use super::StreamingSieve;
    fn gcd(x: usize, y: usize) -> usize {
        if y == 0 { x }
        else { gcd(y, x % y) }
    }
    fn coprime_to(x: usize) -> Vec<usize> {
        (1..x).filter(|&n| gcd(n, x) == 1).collect()
    }
    #[test]
    fn test() {
        let coprime = coprime_to(wheel::MODULO);
        const LIMIT: usize = 2_000_000;
        let mut sieve = StreamingSieve::new(LIMIT);
        let primes = ::primal_smallsieve::Primes::sieve(LIMIT);

        let mut base = 0;
        let mut index = 0;

        while let Some((_low, next)) = sieve.next() {
            for val in next {
                let i = wheel::MODULO * base + coprime[index];
                if i >= LIMIT { break }
                assert!(primes.is_prime(i) == !val,
                        "failed for {} (is prime = {})", i, primes.is_prime(i));

                index += 1;
                if index == wheel::SIZE {
                    index = 0;
                    base += 1
                }
            }
        }
    }
}

#[cfg(all(test, feature = "unstable"))]
mod benches {
    use test::Bencher;
    use super::StreamingSieve;

    fn run(b: &mut Bencher, n: usize) {
        b.iter(|| {
            let mut sieve = StreamingSieve::new(n);
            while sieve.next().is_some() {}
        })
    }

    #[bench]
    fn sieve_small(b: &mut Bencher) {
        run(b, 100)
    }
    #[bench]
    fn sieve_medium(b: &mut Bencher) {
        run(b, 10_000)
    }
    #[bench]
    fn sieve_large(b: &mut Bencher) {
        run(b, 100_000)
    }
    #[bench]
    fn sieve_larger(b: &mut Bencher) {
        run(b, 1_000_000)
    }
    #[bench]
    fn sieve_huge(b: &mut Bencher) {
        run(b, 10_000_000)
    }
}

{'fraction_non_alphanumeric': 0.13422447203986396, 'fraction_numerical': 0.014474412718500355, 'mean_word_length': 1.633617996250781, 'pattern_counts': {'":': 0, '<': 19, '<?xml version=': 0, '>': 37, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3283715', 'n_tokens_mistral': 4387, 'n_tokens_neox': 4193, 'n_words': 1170}