stream - Calculating prime numbers in Scala: how does this code work?

Question

So I've spent hours trying to work out exactly how this code produces prime numbers.

lazy val ps: Stream[Int] = 2 #:: Stream.from(3).filter(i =>
   ps.takeWhile{j => j * j <= i}.forall{ k => i % k > 0});

I've used a number of printlns etc, but nothings making it clearer.

This is what I think the code does:

/**
 * [2,3] 
 * 
 * takeWhile 2*2 <= 3 
 * takeWhile 2*2 <= 4 found match
 *      (4 % [2,3] > 1) return false.
 * takeWhile 2*2 <= 5 found match
 *      (5 % [2,3] > 1) return true 
 *          Add 5 to the list
 * takeWhile 2*2 <= 6 found match
 *      (6 % [2,3,5] > 1) return false
 * takeWhile 2*2 <= 7
 *      (7 % [2,3,5] > 1) return true
 *          Add 7 to the list
 */

But If I change j*j in the list to be 2*2 which I assumed would work exactly the same, it causes a stackoverflow error.

I'm obviously missing something fundamental here, and could really use someone explaining this to me like I was a five year old.

Any help would be greatly appreciated.

Answer 1

I'm not sure that seeking a procedural/imperative explanation is the best way to gain understanding here. Streams come from functional programming and they're best understood from that perspective. The key aspects of the definition you've given are:

1. It's lazy. Other than the first element in the stream, nothing is computed until you ask for it. If you never ask for the 5th prime, it will never be computed.

2. It's recursive. The list of prime numbers is defined in terms of itself.

3. It's infinite. Streams have the interesting property (because they're lazy) that they can represent a sequence with an infinite number of elements. Stream.from(3) is an example of this: it represents the list [3, 4, 5, ...].

Let's see if we can understand why your definition computes the sequence of prime numbers.

The definition starts out with 2 #:: .... This just says that the first number in the sequence is 2 - simple enough so far.

The next part defines the rest of the prime numbers. We can start with all the counting numbers starting at 3 (Stream.from(3)), but we obviously need to filter a bunch of these numbers out (i.e., all the composites). So let's consider each number i. If i is not a multiple of a lesser prime number, then i is prime. That is, i is prime if, for all primes k less than i, i % k > 0. In Scala, we could express this as

nums.filter(i => ps.takeWhile(k => k < i).forall(k => i % k > 0))

However, it isn't actually necessary to check all lesser prime numbers -- we really only need to check the prime numbers whose square is less than or equal to i (this is a fact from number theory*). So we could instead write

nums.filter(i => ps.takeWhile(k => k * k <= i).forall(k => i % k > 0))

So we've derived your definition.

Now, if you happened to try the first definition (with k < i), you would have found that it didn't work. Why not? It has to do with the fact that this is a recursive definition.

Suppose we're trying to decide what comes after 2 in the sequence. The definition tells us to first determine whether 3 belongs. To do so, we consider the list of primes up to the first one greater than or equal to 3 (takeWhile(k => k < i)). The first prime is 2, which is less than 3 -- so far so good. But we don't yet know the second prime, so we need to compute it. Fine, so we need to first see whether 3 belongs ... BOOM!

* It's pretty easy to see that if a number n is composite then the square of one of its factors must be less than or equal to n. If n is composite, then by definition n == a * b, where 1 < a <= b < n (we can guarantee a <= b just by labeling the two factors appropriately). From a <= b it follows that a^2 <= a * b, so it follows that a^2 <= n.