Fast eratosthenes prime generator with a cached wheel

primes.py

"""Fast eratosthenes prime generator with a cached wheel"""

import sys
import time
import operator
from itertools import cycle, chain, accumulate, islice, count
from functools import lru_cache, reduce, partial
from contextlib import contextmanager

CACHE_LEVEL = 5


# Generating the wheel

@lru_cache(maxsize=None)
def cached_wheel(n):
    if n <= 0:
        return (2,), (1,), {0:0, 1:1}, 1
    primes, wheel, accumulated, total = cached_wheel(n-1)
    primes, wheel = next_wheel(primes, wheel)
    accumulated = {x: i for i, x in enumerate(accumulate(wheel), 1)}
    accumulated[0] = 0
    return primes, wheel, accumulated, sum(wheel)


def next_wheel(primes, wheel):
    wheel = cycle(wheel)
    prime = primes[-1]
    new_prime = prime + next(wheel)
    product = reduce(operator.mul, primes, 1)
    new_wheel = cancel_wheel(product, prime, new_prime, wheel)
    return primes + (new_prime,), tuple(new_wheel)


def cancel_wheel(product, prime, current, wheel, candidate=0):
    while product > 0:
        step = next(wheel)
        candidate += step
        current += step
        if current % prime:
            yield candidate
            product -= candidate
            candidate = 0


# Spining the wheel 

def spin(n, start=None, factor=1):
    primes, raw_wheel, accumulated, total = cached_wheel(n)
    # Initialize
    first = primes[-1]
    current = first if start is None else start
    index = accumulated[(current-first) % total]
    # Make generator
    half_wheel = (x * factor for x in islice(raw_wheel, index, None))
    infinite_wheel = cycle(x * factor for x in raw_wheel)
    wheel_chain = chain([factor*current], half_wheel, infinite_wheel)
    generator = accumulate(wheel_chain)
    # Discard first element
    next(generator)
    return generator


# Generating the primes 

def primes(level=CACHE_LEVEL):
    sieve = {}
    first_primes, *_ = cached_wheel(level)
    # First primes
    for prime in first_primes:
        yield prime
    # Initialize prime generator         
    prime_generator = primes()
    for _ in first_primes:
        next_prime = next(prime_generator)
    next_square = next_prime * next_prime
    # Loop over candidates
    for candidate in spin(level):
        # The candidate is a composite
        if candidate in sieve:
            composites = sieve.pop(candidate)
        # The candidate is a prime
        elif candidate != next_square:  
             yield candidate  
             continue
        # The candidate is a squared prime
        else:
            composites = spin(level, next_prime, next_prime)
            next_prime = next(prime_generator)             
            next_square = next_prime * next_prime
        # Update sieve
        for composite in composites:   
            if composite not in sieve:
                break
        sieve[composite] = composites


# Main execution

def main(args=None):
    if args is None:
        args = sys.argv[1:]
    n = int(args[0]) if args else 10**6
    level = int(args[1] if args[1:] else CACHE_LEVEL)
    prime_gen = primes()
    start = time.time()
    result = next(islice(prime_gen, n-1, None))
    stop = time.time()
    print('{}th prime = {} (level {}, {:.3f}s)'
          .format(n, result, level, stop-start))


if __name__ == '__main__':
    main()