Unbounded Sieve of Eratosthenes in Swift

Question

I've spent a while teaching myself Swift, and decided to take on the challenge of writing an unbounded Sieve of Eratosthenes to challenge myself. This is actually the first time I've written an unbounded sieve, as well as the first large Swift program I've written, so all comments are appreciated!

This was written using Swift 3 in Xcode 8 beta 3, so it should work on Swift 3.0 Preview 2 from swift.org, but I have not tested it.

public struct PrimeSequence: Sequence {
    private let stop: StoppingPoint

    /// Construct a sequence of primes up to a max.
    ///
    /// - note: This constructor is eager.
    ///
    /// - complexity: `O(log n)`
    ///
    public init(max: Int) {
        self.stop = .precompute([1] + sieveOfEratosthenes(n: max))
    }

    public init(terms: Int? = nil) {
        self.stop = .terms(terms)
    }

    public func makeIterator() -> AnyIterator<Int> {
        switch (stop) {
        case .precompute(let primes):
            return AnyIterator(primes.makeIterator())
        case .terms(let terms):
            var sieve = UnboundedSieve()
            var term = 0
            return AnyIterator {
                term += 1
                if let terms = terms, term > terms {
                    return nil
                }
                if term == 1 {
                    return 1
                } else {
                    return sieve.next()
                }
            }
        }
    }

    private enum StoppingPoint {
        case precompute([Int])
        case terms(Int?)
    }
}

///
/// Lazy, Unbounded, Sieve of Erathosthenes.
///
/// See https://www.cs.hmc.edu/~oneill/papers/Sieve-JFP.pdf
///
private struct UnboundedSieve: IteratorProtocol {

    var sieve: PriorityQueue<PrimeCounter>
    var counter = 2

    init() {
        sieve = PriorityQueue { $0.composite < $1.composite }
    }

    private struct PrimeCounter {
        let prime: Int
        var composite: Int
        func next() -> PrimeCounter {
            return PrimeCounter(prime: prime, composite: composite + prime)
        }
    }

    mutating func next() -> Int? {
        while let catcher = sieve.peek(), catcher.composite == counter {
            while let catcher = sieve.peek(), catcher.composite == counter {
                sieve.replaceTop(with: catcher.next())
            }
            counter += 1
        }
        defer {
            sieve.enqueue(PrimeCounter(prime: counter, composite: counter + counter))
            counter += 1
        }
        return counter
    }
}

///
/// Mutating Sieve of Erathosthenes.
///
/// Source: https://developer.apple.com/videos/play/wwdc2015/414/ @ 8:15
///
/// I would use the more "swifty" version @ 10:00, but this is the "true" sieve.
///
/// Complexity: `O(n log log n)`
///
private func sieveOfEratosthenes(n: Int) -> [Int] {
    var numbers = [Int](2..<n)
    for i in 0..<n-2 {
        let prime = numbers[i]
        guard prime > 0 else { continue }
        for multiple in stride(from: 2 * prime-2, to: n-2, by: prime) {
            numbers[multiple] = 0
        }
    }
    return numbers.filter { $0 > 0 }
}

I'm using a heap-based PriorityQueue, whose implementation doesn't matter, but here's the protocol pseudocode:

protocol PriorityQueue<T> {
    init(sort: (T, T) -> Bool)
    enqueue(_ element: T)
    dequeue() -> T?
    peek() -> T?
    replaceTop(with value: T)
}

For the purposes of this review, you can assume that it is accurate (but you can see it here if you want).

Answer 1

1 is not a prime number!

By definition, a prime number is an integer greater than one which has no (positive) divisor other than 1 and itself.

Both the UnboundedSieve iterator and the sieveOfEratosthenes function – correctly – produce 2 as the first prime. But then in PrimeSequence, the number 1 is inserted and I see no reason to do so.

---

From the description and naming of the max parameter in the

/// Construct a sequence of primes up to a max.
/// ...
public init(max: Int)

constructor I would assume that it produces a sequence of all primes up to and including max, but that is not true: It produces all primes below max because

private func sieveOfEratosthenes(n: Int) -> [Int] 

returns all primes below n. That might be unexpected. I would suggest to change the function so that the upper limit is included, and give the parameters a more descriptive name:

/// Creates a sequence of primes up to (and including) `limit`.
public init(upTo limit: Int) {

---

I don't see the benefit of the terms parameter in

public init(terms: Int? = nil)

because PrimeSequence(terms: numberOfTerms) can be replaced by PrimeSequence().prefix(numberOfTerms), using the prefix() method from the Swift standard library. I would suggest to remove that parameter.

That leaves us with only two cases: creating a finite sequence (using the sieve of Eratosthenes), or an infinite sequence (using a priority queue), and the PrimeSequence definition can be simplified to

public struct PrimeSequence: Sequence {
    private let iterator: AnyIterator<Int>

    /// Creates a sequence of primes up to (and including) `limit`.
    public init(upTo limit: Int) {
        self.iterator = AnyIterator(sieveOfEratosthenes(upTo: limit).makeIterator())
   }

    /// Creates an "infinite" sequence of prime numbers.
    public init() {
        self.iterator = AnyIterator(UnboundedSieve())
    }

    public func makeIterator() -> AnyIterator<Int> {
        return iterator
    }
}

---

struct UnboundedSieve looks like a clean implementation of the priority queue-based algorithm from https://www.cs.hmc.edu/~oneill/papers/Sieve-JFP.pdf to me. However,

sieve.enqueue(PrimeCounter(prime: counter, composite: counter + counter))

can be improved to

sieve.enqueue(PrimeCounter(prime: counter, composite: counter * counter))

since all primes less than counter have already been found. This reduced the execution time of my test (code below) from 2.26 sec to 1.17 sec.

A possible further improvement is to treat the only even prime 2 separately and then handle only odd numbers:

private struct UnboundedSieve: IteratorProtocol {

    var sieve: PriorityQueue<PrimeCounter>
    var counter = 2

    init() {
        sieve = PriorityQueue { $0.composite < $1.composite }
    }

    private struct PrimeCounter {
        let prime: Int
        var composite: Int
        func next() -> PrimeCounter {
            return PrimeCounter(prime: prime, composite: composite + 2 * prime)
        }
    }

    mutating func next() -> Int? {
        if counter == 2 {
            counter = 3
            return 2
        }

        while let catcher = sieve.peek(), catcher.composite == counter {
            while let catcher = sieve.peek(), catcher.composite == counter {
                sieve.replaceTop(with: catcher.next())
            }
            counter += 2
        }
        defer {
            sieve.enqueue(PrimeCounter(prime: counter, composite: counter * counter))
            counter += 2
        }
        return counter
    }
}

The reduced the execution time of the test code to 0.58 sec.

---

In sieveOfEratosthenes() the same improvement as mentioned above applies: If a prime number is found then "crossing out multiples" can start at the square of that prime:

private func sieveOfEratosthenes(upTo n: Int) -> [Int] {
    var numbers = [Int](2...n)
    for i in 0...n-2 {
        let prime = numbers[i]
        guard prime > 0 else { continue }
        for multiple in stride(from: prime * prime-2, through: n-2, by: prime) {
            numbers[multiple] = 0
        }
    }
    return numbers.filter { $0 > 0 }
}

The reduced the execution time for my test code from 0.37 sec to 0.28 sec.

The numbers array in this function holds redundant information, memory can be saved by replacing it a by an array of booleans which marks all numbers which have been found to be composite:

private func sieveOfEratosthenes(upTo n: Int) -> [Int] {
    var result = [Int]()
    var composite = [Bool](repeating: false, count: n + 1)
    for i in 2...n {
        if !composite[i] {
            result.append(i)
            for multiple in stride(from: i * i, through: n, by: i) {
                composite[multiple] = true
            }
        }
    }
    return result
}

In my test this also turned out to be faster: 0.13 sec.

---

You could also implement the Eratosthenes sieve as an iterator, which makes it non-eager and removes the results array:

private struct EratosthenesIterator: IteratorProtocol {
    let n: Int
    var composite: [Bool]
    var current = 2

    init(upTo n: Int) {
        self.n = n
        self.composite = [Bool](repeating: false, count: n + 1)
    }

    mutating func next() -> Int? {
        while current <= self.n {
            if !composite[current] {
                let prime = current
                for multiple in stride(from: current * current, through: self.n, by: current) {
                    composite[multiple] = true
                }
                current += 1
                return prime
            }
            current += 1
        }
        return nil
    }
}

The corresponding init method in struct PrimeSequence then becomes

    /// Creates a sequence of primes up to (and including) `limit`.
    public init(upTo limit: Int) {
        self.iterator = AnyIterator(EratosthenesIterator(upTo: limit))
   }

The execution time for the test code is now 0.12 sec.

---

Benchmarking code for the infinite generator:

Calculate the sum of the first 10,000 primes:

let numPrimes = 10_000
do {
    let startTime = Date()
    let sum = PrimeSequence().prefix(numPrimes).reduce(0, combine: +)
    let endTime = Date()
    print(sum, endTime.timeIntervalSince(startTime), "sec")
}

Benchmarking code for the Eratosthenes generator:

Calculates the sum of all primes below 10,000,000:

let N = 10_000_000

do {
    let startTime = Date()
    let sum = PrimeSequence(upTo: N).reduce(0, combine: +)
    let endTime = Date()
    print(sum, endTime.timeIntervalSince(startTime))
}

The tests were done on a MacBook, with the program compiled in Release mode.