prime_data/data/iterators/
prime.rs

1use std::ops::RangeInclusive;
2use crate::{PrimeData, PrimeByte, data::{error::*, utils::{Divisible, ContainsRange}}};
3
4/// Struct that iterates over prime numbers from some data.
5/// 
6/// To learn more, read the [guide](crate::guide::iterators::_2_prime).
7/// 
8/// Iterator takes a reference to [PrimeData](crate::PrimeData), therefore, it cannot
9/// outlive the given data.
10/// 
11/// # Examples
12/// 
13/// ```
14/// use prime_data::{PrimeData, PrimeIter};
15/// let data = PrimeData::new();
16/// let mut iter = PrimeIter::new(&data, 3..=27).unwrap();
17/// 
18/// assert_eq!(iter.next(), Some(3));
19/// assert_eq!(iter.next(), Some(5));
20/// assert_eq!(iter.next(), Some(7));
21/// assert_eq!(iter.next(), Some(11));
22/// assert_eq!(iter.next(), Some(13));
23/// assert_eq!(iter.next(), Some(17));
24/// assert_eq!(iter.next(), Some(19));
25/// assert_eq!(iter.next(), Some(23));
26/// assert_eq!(iter.next(), None);
27/// ```
28pub struct PrimeIter<'a> {
29    data: &'a [PrimeByte],
30    primes: Option<Vec<u64>>,
31    current: (u64, usize),
32    data_offset: u64,
33    stop_at: u64,
34}
35
36impl<'a> PrimeIter<'a> {
37    /// Creates an iterator over some [PrimeData](crate::PrimeData) within a given range.
38    /// 
39    /// Returns a [NotEnoughData](crate::error::ErrorType::NotEnoughData) error if the given range
40    /// is not contained in the PrimeData's range.
41    pub fn new(prime_data: &'a PrimeData, range: RangeInclusive<u64>) -> PrimeResult<Self> {
42        if let Err(out_of_bounds) = prime_data.range.contains_range(&range) {
43
44            let error = PrimeError {
45                context: ErrorContext { action: ErrorAction::Reading, source: ErrorSource::PrimeData },
46                error: ErrorType::NotEnoughData(out_of_bounds)
47            };
48
49            return Err(error)
50        }
51
52        let (range_start, stop_at) = range.into_inner();
53
54        let original_offset = prime_data.offset();
55        let data_start = range_start.div_floor(30) as usize - original_offset;
56        let data_end = stop_at.div_ceil(30) as usize - original_offset;
57
58        let data = &prime_data.data[data_start..data_end];
59        let data_offset = (data_start + original_offset) as u64;
60        let mut current = (0u64, 0usize);
61
62        let primes = loop {
63
64            if let Some(byte) = data.get(current.0 as usize) {
65                let byte_primes = if current.0 == 0 {
66                    let range_offset = data_offset + current.0;
67                    let start = (range_start - 30 * range_offset) as u8;
68                    byte.as_primes_in_range(range_offset, start..=(start + 30))
69                } else {
70                    byte.as_primes(data_offset + current.0)
71                };
72                if byte_primes.len() > 0 {
73                    // PrimeData does not store the primes {2, 3, 5}, so if the range includes any
74                    // of those, we need to manually add them to the first vector
75                    let byte_primes = vec![2u64, 3u64, 5u64].into_iter()
76                    .filter(|&x| x >= range_start)
77                    .chain(byte_primes.into_iter())
78                    .collect();
79
80                    break Some(byte_primes);
81                }
82
83                current.0 += 1;
84            } else {
85                break None;
86            }
87
88        };
89
90        Ok(Self { data, primes, current, data_offset, stop_at })
91    }
92}
93
94impl<'a> Iterator for PrimeIter<'a> {
95    type Item = u64;
96
97    fn next(&mut self) -> Option<Self::Item> {
98        if let Some(vector) = &self.primes {
99
100            let current_prime = vector[self.current.1];
101            if current_prime > self.stop_at { return None }
102
103            if self.current.1 + 1 < vector.len() {
104                self.current.1 += 1;
105            } else {
106                self.current.1 = 0;
107                self.current.0 += 1;
108
109                self.primes = loop {
110
111                    if let Some(byte) = self.data.get(self.current.0 as usize) {
112                        let byte_primes = byte.as_primes(self.data_offset + self.current.0);
113        
114                        if byte_primes.len() > 0 {
115                            break Some(byte_primes);
116                        }
117        
118                        self.current.0 += 1;
119                    } else {
120                        break None;
121                    }
122        
123                };
124            }
125
126            Some(current_prime)
127
128        } else {
129            None
130        }
131    }
132}