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.

Answer 2

Your explanations are mostly correct, you made only two mistakes:

takeWhile doesn't include the last checked element:

scala> List(1,2,3).takeWhile(_<2)
res1: List[Int] = List(1)

You assume that ps always contains only a two and a three but because Stream is lazy it is possible to add new elements to it. In fact each time a new prime is found it is added to ps and in the next step takeWhile will consider this new added element. Here, it is important to remember that the tail of a Stream is computed only when it is needed, thus takeWhile can't see it before forall is evaluated to true.

Keep these two things in mind and you should came up with this:

ps = [2]
i = 3
  takeWhile
    2*2 <= 3 -> false
  forall on []
    -> true
ps = [2,3]
i = 4
  takeWhile
    2*2 <= 4 -> true
    3*3 <= 4 -> false
  forall on [2]
    4%2 > 0 -> false
ps = [2,3]
i = 5
  takeWhile
    2*2 <= 5 -> true
    3*3 <= 5 -> false
  forall on [2]
    5%2 > 0 -> true
ps = [2,3,5]
i = 6
...

While these steps describe the behavior of the code, it is not fully correct because not only adding elements to the Stream is lazy but every operation on it. This means that when you call xs.takeWhile(f) not all values until the point when f is false are computed at once - they are computed when forall wants to see them (because it is the only function here that needs to look at all elements before it definitely can result to true, for false it can abort earlier). Here the computation order when laziness is considered everywhere (example only looking at 9):

ps = [2,3,5,7]
i = 9
  takeWhile on 2
    2*2 <= 9 -> true
  forall on 2
    9%2 > 0 -> true
  takeWhile on 3
    3*3 <= 9 -> true
  forall on 3
    9%3 > 0 -> false
ps = [2,3,5,7]
i = 10
...

Because forall is aborted when it evaluates to false, takeWhile doesn't calculate the remaining possible elements.