use std::collections::HashMap;

struct Sieve {
    composites: HashMap<u64, Vec<u64>>,
    current: u64,
}

impl Default for Sieve {
    fn default() -> Self {
        Sieve {
            composites: HashMap::new(),
            current: 2,
        }
    }
}

impl Sieve {
    pub fn new() -> Sieve {
        Default::default()
    }
}

impl Iterator for Sieve {
    type Item = u64;

    fn next(&mut self) -> Option<u64> {
        fn next_prime(composites: &mut HashMap<u64, Vec<u64>>, x: u64) -> u64 {
            match composites.get(&x) {
                Some(numbers) => {
                    for num in numbers.to_owned() {
                        composites
                            .entry(x + num)
                            .and_modify(|v| v.push(num))
                            .or_insert_with(|| vec![num]);
                    }
                    composites.remove(&x);
                    next_prime(composites, x + 1)
                }
                None => {
                    composites.insert(x * x, vec![x]);
                    x
                }
            }
        }

        let prime = next_prime(&mut self.composites, self.current);
        self.current = prime + 1; // This number will be the next to be tested

        Some(prime)
    }
}

fn main() {
    let mut sieve = Sieve::new();

    assert_eq!(sieve.next(), Some(2));
    assert_eq!(sieve.next(), Some(3));
    assert_eq!(sieve.next(), Some(5));
    assert_eq!(sieve.next(), Some(7));
}