Home
🔢

Prime streams

This is a continuation of my post Operating on infinite lists which walks you through writing a primitive “Stream” datatype to represent infinite lists of values. Consider scanning through it, but I’ll do a quick refresher.
In this post I’ll use them for one of my favorite recursion tricks – generating lists of prime numbers (SICP 3.5.2).

🔗 Recap

A stream is a function which returns a value and a stream. 
type Stream<T> = () => {
  value: T;
  next: Stream<T>;
};
We can build a stream of the number “1” 
const ones: Stream<number> = () => ({
  value: 1,
  next: ones,
});
For testing purposes, we can define a function which takes the first n elements of a stream. 
function take<T>(
  n: number,
  stream: Stream<T>
): T[] {
  if (n === 0) {
    return [];
  }

  const { value, next } = stream();
  return [value].concat(take(n - 1, next));
}
And we can test things as we go. 
import { expect, test } from "bun:test";

test("take", () => {
  expect(take(5, ones)).toEqual([
    1, 1, 1, 1, 1,
  ]);
});

/*
  $ bun test ./index.ts
  bun test v1.0.22 (b400b36c)

  index.ts:
  ✓ take [0.10ms]

  1 pass
  0 fail
  1 expect() calls
  Ran 1 tests across 1 files. [12.00ms]
*/
We generate a stream of the natural numbers. 
function nats(start = 1): Stream<number> {
  return () => ({
    value: start,
    next: nats(start + 1),
  });
}

test("nats", () => {
  expect(take(5, nats())).toEqual([
    1, 2, 3, 4, 5,
  ]);
});
We can start this stream from wherever we’d like. 
test("nats(7)", () => {
  expect(take(5, nats(7))).toEqual([
    7, 8, 9, 10, 11,
  ]);
});
We can map a function over a stream. 
function map<T, U>(
  fn: (x: T) => U,
  stream: Stream<T>
): Stream<U> {
  const { value, next } = stream();
  return () => ({
    value: fn(value),
    next: map(fn, next),
  });
}

test("map", () => {
  expect(
    take(
      5,
      map((x) => x * x, nats(7))
    )
  ).toEqual([49, 64, 81, 100, 121]);
});
We can filter a stream. 
function filter<T>(
  fn: (x: T) => boolean,
  stream: Stream<T>
): Stream<T> {
  const { value, next } = stream();
  if (fn(value)) {
    return () => ({
      value,
      next: filter(fn, next),
    });
  }
  return filter(fn, next);
}

test("filter", () => {
  expect(
    take(
      5,
      filter((x) => x % 2 === 0, nats())
    )
  ).toEqual([2, 4, 6, 8, 10]);
});

🔗 Using our new tools

We can generate a stream of the natural numbers starting from 2. 
test("nats(2)", () => {
  expect(take(5, nats(2))).toEqual([
    2, 3, 4, 5, 6,
  ]);
});
We can filter out numbers divisible by 2 in this stream. 
const twoOnwardNoEvens = filter(
  (x) => x % 2 !== 0,
  nats(2)
);

test("twoOnwardNoEvens", () => {
  expect(take(5, twoOnwardNoEvens)).toEqual([
    3, 5, 7, 9, 11,
  ]);
});
We can keep the 2 around if we want, though. 
const twoButNoOtherEvens = () => ({
  value: 2,
  next: filter((x) => x % 2 !== 0, nats(2)),
});

test("twoButNoOtherEvens", () => {
  expect(
    take(5, twoButNoOtherEvens)
  ).toEqual([2, 3, 5, 7, 9]);
});
We can abstract this away to a function which keeps the first value, but removes any multiples of it in the future. 
function sieveOutMultiples(
  stream: Stream<number>
): Stream<number> {
  const { value, next } = stream();
  return () => ({
    value,
    next: filter(
      (x) => x % value !== 0,
      next
    ),
  });
}
We can rewrite twoButNoOtherEvens to use this abstraction. 
const twoButNoOtherEvensAgain =
  sieveOutMultiples(nats(2));

test("twoButNoOtherEvensAgain", () => {
  expect(
    take(5, twoButNoOtherEvensAgain)
  ).toEqual([2, 3, 5, 7, 9]);
});
Why limit ourselves to filtering the rest of the stream? We can also sieve it. 
	function sieveOutMultiples(
	  stream: Stream<number>
	): Stream<number> {
	  const { value, next } = stream();
	  return () => ({
	    value,
-     next: filter(
-       (x) => x % value !== 0,
-       next
+     next: sieveOutMultiples(
+       filter((x) => x % value !== 0, next)
	    ),
	  });
	}
The name’s getting pretty long, let’s just call it sieve. 
function sieve(
  stream: Stream<number>
): Stream<number> {
  const { value, next } = stream();
  return () => ({
    value,
    next: sieve(
      filter((x) => x % value !== 0, next)
    ),
  });
}

test("sieve", () => {
  expect(take(10, sieve(nats(2)))).toEqual([
    2, 3, 5, 7, 11, 13, 17, 19, 23, 29,
  ]);
});
We can now generate an infinite list of primes. 
const primes = sieve(nats(2));
// [2, 3, 5, 7, 11, 13, 17, 19, 23, 29, ...]

🔗 Exercises