Fun With Streams

Or: a bit of Haskell in PLT Scheme.

In a one-off version of the n-body integrator I've been working on (seemingly in perpetuity, but probably finishing soon with a couple of papers), I had occasion to iterate a function to convergence. (i.e. to compute (f (f (f ... (f x)))) until the result converges.) Well, it's pretty easy to write:

(define (iter-to-convergence f)
 (lambda (x)
   (let loop ((last (f x)))
     (let ((next (f last)))
       (if (= last next)
           next
           (loop next))))))

but where's the fun in that? It's much more fun to use streams for this.

So, using the streams module appended below (shamelessly stolen from SICP), I can write

(define (iter-to-convergence f)
 (lambda (x)
   (letrec ((vals
             (stream-cons (f x) (stream-map f vals))))
     (stream-take-converged vals))))

This idiom is very common in Haskell (and other lazy languages)---the canonical example is, of course, (define ones (stream-cons 1 ones)). The fact that this is so easy to add to scheme is a real testament to the language's power. What other language lets you (lazily) have your cake and (eventually) eat it, too?

Here's the code for the stream module (feel free to steal it if you want):

(module streams mzscheme
 (require (only (lib "1.ss" "srfi") any))

 (provide stream? stream-car stream-cdr stream-ref stream-cons
          (rename my-stream stream) stream-map stream-filter stream-take
          stream-converged? stream-take-converged)
  
 (define-struct stream (a b) #f)

 (define stream-car stream-a)

 (define (stream-cdr s)
   (force (stream-b s)))

 (define (stream-ref s n)
   (if (= n 0)
       (stream-car s)
       (stream-ref (stream-cdr s) (- n 1))))

 (define-syntax stream-cons
   (syntax-rules ()
     ((cons-stream car cdr)
      (make-stream car (delay cdr)))))

 (define-syntax my-stream
   (syntax-rules ()
     ((stream a)
      (stream-cons a ()))
     ((stream a b)
      (stream-cons a b))
     ((stream a b ...)
      (stream-cons a (my-stream b ...)))))      

 (define (stream-map f . ss)
   (stream-cons (apply f (map stream-car ss)) (apply stream-map f (map stream-cdr ss))))

 (define (stream-for-each f . ss)
   (apply f (map stream-car ss))
   (stream-for-each f (map stream-cdr ss)))

 (define (stream-filter test? s)
   (let ((s1 (stream-car s)))
     (if (test? s1)
         (stream-cons s1 (stream-filter test? (stream-cdr s)))
         (stream-filter test? (stream-cdr s)))))

 (define (stream-take n s)
   (if (= n 0)
       '()
       (cons (stream-car s) (stream-take (- n 1) (stream-cdr s)))))

 (define (stream-converged? s)
   (= (stream-car s) (stream-ref s 1)))

 (define (stream-take-converged s)
   (if (stream-converged? s)
       (stream-car s)
       (stream-take-converged (stream-cdr s)))))

Here's a more involved example of what you can do with this: the sieve of Eratosthenes (SICP 3.5.2).

(define (first-n-primes n)
 (stream-take
  n
  (letrec ((integers-from (lambda (n) (stream-cons n (integers-from (+ n 1)))))
           (primes (stream-cons 2 (stream-filter prime? (integers-from 3))))
           (prime? (lambda (x)
                     (let iter ((ps primes))
                       (let ((p (stream-car ps)))
                         (cond
                           ((= (remainder x p) 0) #f)
                           ((> (* p p) x) #t)
                           (else (iter (stream-cdr ps)))))))))
    primes)))