Python+cython prime generators

_primegen.pyx

# distutils: language=c++

from cython.operator cimport dereference
from libcpp.unordered_map cimport unordered_map


def primes():
    # Yield first two primes
    yield 2
    yield 3

    # Initialize composite lookup
    cdef unordered_map[int, int] composites

    # Initialize the inner generator
    prime_gen = primes()
    next(prime_gen)
    next(prime_gen)
    cdef int next_prime = 3
    cdef int next_square = 9

    # Loop over candidates
    cdef int candidate = 3
    cdef int step, next_composite
    while True:
        candidate += 2

        # Reached next prime square
        if candidate == next_square:
            step = 2 * next_prime
            composites[candidate + step] = step
            next_prime = next(prime_gen)
            next_square = next_prime * next_prime
            continue

        # Reached a prime
        iterator = composites.find(candidate)
        if iterator == composites.end():
            yield candidate
            continue

        # Reached a composite
        step = dereference(iterator).second
        composites.erase(iterator)
        next_composite = candidate + step
        while composites.count(next_composite):
            next_composite += step
        composites[next_composite] = step

primegen.py

import time
import argparse
import itertools


def primes():
    # Yield first two primes
    yield 2
    yield 3

    # Initialize composite lookup
    composites = {}

    # Initialize the inner generator
    prime_gen = primes()
    assert next(prime_gen) == 2
    assert next(prime_gen) == 3
    next_prime = 3
    next_square = 9

    # Loop over candidates
    for candidate in itertools.count(5, 2):

        # Reached next prime square
        if candidate == next_square:
            step = 2 * next_prime
            composites[candidate + step] = step
            next_prime = next(prime_gen)
            next_square = next_prime * next_prime
            continue

        # Reached a prime
        if candidate not in composites:
            yield candidate
            continue

        # Reached a composite
        step = composites.pop(candidate)
        next_composite = candidate + step
        while next_composite in composites:
            next_composite += step
        composites[next_composite] = step


def test():
    for _, prime in zip(range(10**6), primes()):
        pass
    assert prime == 15485863


def main(args=None):
    parser = argparse.ArgumentParser()
    parser.add_argument("N", type=lambda x: int(eval(x)))
    parser.add_argument("-c", "--cython", action="store_true")
    namespace = parser.parse_args(args)

    if namespace.cython:
        global primes
        from _primegen import primes

    start = time.time()
    for _, prime in zip(range(namespace.N), primes()):
        pass
    delta = time.time() - start

    print(f"The {namespace.N}th prime is {prime} (in {delta:.3f} seconds)")


if __name__ == "__main__":
    main()

primegen.rs

#![feature(generators, generator_trait)]

use std::env;
use std::time::Instant;
use std::collections::HashMap;
use std::ops::{Generator, GeneratorState};
use std::pin::Pin;

// Generator to Iterator

struct GeneratorToIterator<G>(G);

impl<G> Iterator for GeneratorToIterator<G>
where
    G: Generator<Return = ()>,
{
    type Item = G::Yield;

    fn next(&mut self) -> Option<Self::Item> {
        let me = unsafe { Pin::new_unchecked(&mut self.0) };
        match me.resume() {
            GeneratorState::Yielded(x) => Some(x),
            GeneratorState::Complete(_) => None,
        }
    }
}

// Prime generator

fn primes() -> impl Iterator<Item = u64> {
    GeneratorToIterator(move || {
        // Yield first two primes
        yield 2;
        yield 3;

        // Initialize the inner prime generator
        let mut next_prime = 3;
        let mut next_square = 9;
        let mut prime_gen: Box<dyn Iterator<Item = u64>> = Box::new(primes());
        assert_eq!(prime_gen.next().unwrap(), 2);
        assert_eq!(prime_gen.next().unwrap(), next_prime);

        // Initialize composite hash map
        let mut composites = HashMap::new();

        // Initialize loop variables
        let mut step;
        let mut next_composite;

        // Loop over candidates
        for candidate in (5..).step_by(2) {

            // Reached next prime square
            if candidate == next_square {
                step = 2 * next_prime;
                composites.insert(candidate + step, step);
                next_prime = prime_gen.next().unwrap();
                next_square = next_prime * next_prime;
                continue;
            }

            // Reached a prime
            match composites.remove(&candidate){
                Some(x) => step = x,
                None => {yield candidate; continue},
            }

            // Reached a composite
            next_composite = candidate + step;
            while composites.contains_key(&next_composite) {
                next_composite += step;
            }
            composites.insert(next_composite, step);

        }

    })
}

// Main

fn main() {
    let arg = env::args().nth(1).unwrap();
    let param: usize = arg.parse().unwrap();
    let start = Instant::now();
    let x = primes().nth(param - 1);
    let delta = start.elapsed().as_secs_f32();
    println!("The {}th prime is {} (in {:.3} seconds)", param, x.unwrap(), delta);

}

setup.py

from setuptools import setup
from Cython.Build import cythonize

import Cython.Compiler.Options
Cython.Compiler.Options.annotate = True

setup(
    ext_modules=cythonize("_primegen.pyx", annotate=True),
    extra_compile_args=["-O3"],
)