A simple stream implementation in sml

stream.sml

(* a stream is a thunk that when called produces a pair
   #1 of the pair will be of type 'a and #2 of the pair will be another stream *)
(* can only have recursive type using datatype binding *)
datatype 'a StreamPair = Stream of 'a * (unit -> 'a StreamPair)
			    
(* a ones stream *)
fun ones () = Stream (1, ones)
		   
(* a fibonacci number stream *)
fun fibs () =
    let fun fibs prev curr =
	    Stream (prev, fn () => fibs curr (prev+curr))
    in
	fibs 0 1
    end

(* a natural number stream *)
val nats =
    let fun nats n =
	    Stream (n, fn () => nats (n+1))
    in
	fn () => nats 0
    end

(* funny number stream:
   a stream like the natural number stream except numbers divisble by 5 are
   negated *)
val funny_number_stream =
    let fun funny n =
	    Stream (if n mod 5 = 0 then ~n else n, fn () => funny (n+1))
    in
	fn () => funny 0
    end

(* a stream where the elements alternate between the strings "dan.jpg"
   and "dog.jpg" (starting with "dan.jpg") *)
val dan_then_dog =
    let fun dan () = Stream ("dan.jpg", dog)
	and dog () = Stream ("dog.jpg", dan)
    in
	dan
    end

(* takes a stream and an non-negative integer n, 
   produce a list of n elements from the stream *)	
fun stream_for_n_steps s n =
    if n = 0
    then []
    else
	let val Stream (curr, next) = s ()
	in
	    curr :: (stream_for_n_steps next (n-1))
	end

(* takes a function f and a stream s, and returns a new stream whose values are the result
   of applying f to the values produced by s *)
fun stream_map f s =
    let val Stream (curr, next) = s ()
    in
	fn () => Stream (f curr, stream_map f next)
    end

(* takes in two stream s1 and s2 and returns a stream that produces the pairs that result from
   the other two streams (so the first value for the result stream will be the pair of the first
   value of s1 and first value of s2 *)
fun stream_zip s1 s2 =
    let val Stream (curr1, next1) = s1 ()
	val Stream (curr2, next2) = s2 ()
    in
	fn () => Stream ((curr1, curr2), stream_zip next1 next2)
    end

(* takes a stream s and returns another stream where the ith element is a pair (0, v), 
   v is the ith element of s *)
fun stream_add_zero s =
    stream_map (fn x => (0, x)) s

(* takes two lists xs and ys and returns a stream
   the lists may or may not to be the same length, but assume they are both non-empty
   the elements produced by the stream are pairs where the first part is from xs and second from ys
   the stream cycles foreer through the lists
   e.g. xs = [1,2,3], ys = ["a","b"], stream produces (1,"a"),(2,"b"),(3,"a"),(1,"b"),(2,"a"),etc. *)
fun cycle_lists xs ys =
    let fun cycle l1 l2 =
	    case (l1, l2) of
		([], l2') => cycle xs l2'
	      | (l1', []) => cycle l1' ys
	      | (h1::l1', h2::l2') =>
		fn () => Stream ((h1, h2), cycle l1' l2')
    in
	cycle xs ys
    end

(* takes a list of streams and produces a new stream that takes one element from each stream
   in sequence
   so it will first produce the first value of the first stream
   then the first value of the second stream and so on
   and it will go back to the first stream when it reaches the end of the list *)
fun interleave ss =
    let fun stream_nth s n =
	    let val Stream (curr, next) = s ()
	    in
		if n = 0
		then curr
		else stream_nth next (n-1)
	    end
	fun loop i xs =
	    case xs of
		[] => loop (i+1) ss 
	      | x::xs' => fn () => Stream (stream_nth x i, loop i xs')
    in
	loop 0 ss
    end

(* takes an integer n and a stream s, and returns a stream that produces the same values as s
   but packed in lists of n elements
   So the first value of the new stream will be the list consisting of the first n values of s
   the second value of the new stream will contain the next n values, and so on *)
fun pack n =
    let fun take n s =
	    if n = 0
	    then ([], s)
	    else let val Stream (curr, next) = s ()
		     val (taken_so_far, next') = take (n-1) next
		 in (curr::taken_so_far, next')
		 end
	fun pack s =
	    let val (taken, next) = take n s
	    in fn () => Stream (taken, pack next)
	    end
    in
	pack
    end

(* takes a real and produce a stream of approximating square root of n using Newton's Method *)	
fun sqrt_stream n =
    let fun sqrt x =
	    fn () => Stream (x, sqrt (0.5 * (x + n/x)))
    in
	sqrt n
    end

(* takes two real n and e and returns a real x such that x*x is within e of n *)	
fun approx_sqrt n e =
    let fun approx s =
	    let val Stream (curr, next) = s()
	    in
		if abs (curr * curr - n) < e
		then curr
		else approx next
	    end
    in
	approx (sqrt_stream n)
    end