Lazy infinite streams in Haskell, Swift, and Python

Stream.hs

module Main where

import Prelude hiding (
    iterate, repeat, map, filter, zipWith, zip,
    take, takeWhile, drop, dropWhile)
import qualified Prelude as P
import Text.Printf (printf)
import Data.List (intercalate)

infixr 5 :>
data Stream a = a :> (Stream a)

first :: Stream a -> a
first (x :> _) = x

rest :: Stream a -> Stream a
rest (_ :> xs) = xs

iterate :: (a -> a) -> a -> Stream a
iterate f a = a :> iterate f (f a)

count :: Integer -> Integer -> Stream Integer
count start step = iterate (+step) start

repeat :: a -> Stream a
repeat a = a :> repeat a

map :: (a -> b) -> Stream a -> Stream b
map f (x :> xs) = f x :> map f xs

filter :: (a -> Bool) -> Stream a -> Stream a
filter p (x :> xs)
    | p x       = x :> filter p xs
    | otherwise = filter p xs

zipWith :: (a -> b -> c) -> Stream a -> Stream b -> Stream c
zipWith f (x :> xs) (y :> ys) = f x y :> zipWith f xs ys

zip :: Stream a -> Stream b -> Stream (a, b)
zip = zipWith (,)

scan :: (b -> a -> b) -> b -> Stream a -> Stream b
scan f z (x :> xs) = z :> scan f (f z x) xs

scan1 :: (a -> a -> a) -> Stream a -> Stream a
scan1 f (x :> xs) = scan f x xs

take :: Integer -> Stream a -> [a]
take n (x :> xs)
    | n <= 0    = []
    | otherwise = x : take (n - 1) xs

drop :: Integer -> Stream a -> Stream a
drop n xs
    | n <= 0    = xs
    | otherwise = drop (n - 1) (rest xs)

takeWhile :: (a -> Bool) -> Stream a -> [a]
takeWhile p (x :> xs)
    | p x       = x : takeWhile p xs
    | otherwise = []

dropWhile :: (a -> Bool) -> Stream a -> Stream a
dropWhile p (x :> xs)
    | p x       = dropWhile p xs
    | otherwise = x :> xs

annotate :: Show a => String -> [a] -> IO ()
annotate description body = putStrLn message
  where
    message   = printf "%s:\n  %s\n" description formatted
    formatted = "[" ++ intercalate ", " (P.map show body) ++ "]"

nats    = iterate (+1) 0
ones    = repeat 1
strings = map (flip replicate '*') nats
squares = map (^2) nats
evens   = filter even nats
odds    = filter odd nats
sums    = zipWith (+) nats nats
pairs   = zip evens odds
facts   = scan1 (*) (count 1 1)
fibs    = 0 :> 1 :> zipWith (+) fibs (drop 1 fibs)

main = do
    annotate
        "First ten nats (iterate)"
        (take 10 nats)

    annotate
        "First ten nats after dropping ten nats (drop)"
         (take 10 (drop 10 nats))

    annotate
        "First ten ones (repeat)"
        (take 10 ones)

    annotate
        "Type-changing map (map int -> string)"
        (take 10 strings)

    annotate
        "First ten squares (map int -> int)"
        (take 10 squares)

    annotate
        "First ten evens (filter int -> bool)"
        (take 10 evens)

    annotate
        "First ten odds (filter int -> bool)"
        (take 10 odds)

    annotate
        "First ten sums ((int, int) -> int)"
        (take 10 sums)

    annotate
        "Even/odd pairs (zip)"
        (take 10 pairs)

    annotate
        "First ten factorial (scan)"
        (take 10 facts)

    annotate
        "First ten fibs (forward/tie/zipwith)"
        (take 10 fibs)

    annotate
        "Take while <10"
        (takeWhile (<10) nats)

    annotate
        "First ten after drop while <10"
        (take 10 (dropWhile (<10) nats))

Stream.py

from __future__ import unicode_literals
import operator

class Stream(object):
    sentinel = object()
    __slots__ = ('_first', '_rest', '_step')

    def __init__(self, first, rest):
        self._first = first
        if callable(rest):
            self._step = rest
            self._rest = self.sentinel
        elif isinstance(rest, Stream):
            self._step = lambda: rest
            self._rest = rest
        else:
            raise TypeError('Second argument to Stream must be a callable or a Stream')

    @property
    def first(self):
        return self._first

    @property
    def rest(self):
        if self._rest is self.sentinel:
            self._rest = self._step()
        return self._rest

    @classmethod
    def cons(cls, head, rest):
        return cls(head, rest)

    @classmethod
    def iterate(cls, f, x):
        return cls(x, lambda: cls.iterate(f, f(x)))

    @classmethod
    def count(cls, start=0, step=1):
        return cls.iterate(lambda i: i + step, start)

    @classmethod
    def repeat(cls, x):
        return cls(x, lambda: cls.repeat(x))

    def map(self, f):
        return Stream(f(self.first), lambda: self.rest.map(f))

    def filter(self, p):
        if p(self.first):
            return Stream(self.first, lambda: self.rest.filter(p))
        return self.rest.filter(p)

    def zipWith(self, f, stream):
        return Stream(
            f(self.first, stream.first),
            lambda: self.rest.zipWith(f, stream.rest)
        )

    def zip(self, stream):
        return self.zipWith(lambda x, y: (x, y), stream)

    def scan(self, f, first):
        return Stream(first, lambda: self.rest.scan(f, f(first, self.first)))

    def scan1(self, f):
        return self.rest.scan(f, self.first)

    def take(self, n):
        out = []
        stream = self
        i = 0
        while i < n:
            out.append(stream.first)
            stream = stream.rest
            i += 1
        return out

    def drop(self, n):
        stream = self
        i = 0
        while i < n:
            stream = stream.rest
            i += 1
        return stream

    def takeWhile(self, p):
        out = []
        stream = self
        while 1:
            if p(stream.first):
                out.append(stream.first)
                stream = stream.rest
            else:
                break
        return out

    def dropWhile(self, p):
        stream = self
        while 1:
            if p(stream.first):
                stream = stream.rest
            else:
                break
        return stream

def annotate(description, body):
    formatted = ', '.join(map(str, body)).join('[]')
    print("{}:\n  {}\n".format(description, formatted))

nats    = Stream.iterate(lambda x: x + 1, 0)
ones    = Stream.repeat(1)
strings = nats.map(lambda x: '*' * x)
squares = nats.map(lambda x: x * x)
evens   = nats.filter(lambda x: x % 2 == 0)
odds    = nats.filter(lambda x: x % 2 == 1)
sums    = nats.zipWith(operator.add, nats)
pairs   = evens.zip(odds)
facts   = Stream.count(1).scan1(operator.mul)
fibs    = Stream.cons(0, Stream.cons(1, lambda: fibs.zipWith(operator.add, fibs.drop(1))))

if __name__ == '__main__':
    annotate(
        "First ten nats (iterate)",
        nats.take(10),
    )

    annotate(
        "First ten nats after dropping ten nats (drop)",
        nats.drop(10).take(10),
    )

    annotate(
        "First ten ones (repeat)",
        ones.take(10),
    )

    annotate(
        "Type-changing map (map int -> string)",
        strings.take(10),
    )

    annotate(
        "First ten squares (map int -> int)",
        squares.take(10),
    )

    annotate(
        "First ten evens (filter int -> bool)",
        evens.take(10),
    )

    annotate(
        "First ten odds (filter int -> bool)",
        odds.take(10),
    )

    annotate(
        "First ten sums ((int, int) -> int)",
        sums.take(10),
    )

    annotate(
        "Even/odd pairs (zip)",
        pairs.take(10),
    )

    annotate(
        "First ten factorial (scan)",
        facts.take(10),
    )

    annotate(
        "First ten fibs (forward/tie/zipwith)",
        fibs.take(10),
    )

    annotate(
        "Take while <10",
        nats.takeWhile(lambda x: x < 10),
    )

    annotate(
        "First ten after drop while <10",
        nats.dropWhile(lambda x: x < 10).take(10),
    )

Stream.swift

public class Stream<A> {
    private var _first : A
    private var _rest : Stream<A>?
    private var _step : () -> Stream<A>

    init(_ first : A, _ step : @autoclosure () -> Stream<A>) {
        self._first = first
        self._rest = nil
        self._step = step
    }

    public func first() -> A {
        return _first
    }

    public func rest() -> Stream<A> {
        if let stream = _rest {
            return stream
        } else {
            let result = _step()
            _rest = result
            return result
        }
    }

    public class func cons(head : A, _ tail: @autoclosure () -> Stream<A>) -> Stream<A> {
        return Stream(head, tail)
    }

    public class func iterate(x : A, _ f : A -> A) -> Stream<A> {
        return Stream(x, Stream.iterate(f(x), f))
    }

    public class func count(start : Int = 0, step : Int = 1) -> Stream<Int> {
        return Stream<Int>.iterate(start) { $0 + step }
    }

    public class func repeat(x : A) -> Stream<A> {
        return Stream(x, Stream.repeat(x))
    }

    public func map<B>(f : A -> B) -> Stream<B> {
        return Stream<B>(f(self.first()), self.rest().map(f))
    }

    public func filter(predicate : A -> Bool) -> Stream<A> {
        let x = self.first()
        if predicate(x) {
            return Stream(x, self.rest().filter(predicate))
        } else {
            return self.rest().filter(predicate)
        }
    }

    public func zipWith<B, C>(other : Stream<B>, _ f : (A, B) -> C) -> Stream<C> {
        return Stream<C>(
            f(self.first(), other.first()),
            self.rest().zipWith(other.rest(), f)
        )
    }

    public func zip<B>(other : Stream<B>) -> Stream<(A, B)> {
        return self.zipWith(other) { ($0, $1) }
    }

    public func scan<B>(first : B, _ f : (B, A) -> B) -> Stream<B> {
        return Stream<B>(first, self.rest().scan(f(first, self.first()), f))
    }

    public func scan1(f : (A, A) -> A) -> Stream<A> {
        return self.rest().scan(self.first(), f)
    }

    public func take(n : Int) -> [A] {
        var i = 0
        var result : [A] = []
        var stream = self
        while i < n {
            result.append(stream.first())
            stream = stream.rest()
            i += 1
        }
        return result
    }

    public func drop(n : Int) -> Stream<A> {
        var i = 0
        var stream = self
        while i < n {
            stream = stream.rest()
            i += 1
        }
        return stream
    }

    public func takeWhile(predicate : A -> Bool) -> [A] {
        var result : [A] = []
        var stream = self
        while predicate(stream.first()) {
            result.append(stream.first())
            stream = stream.rest() 
        }
        return result
    }

    public func dropWhile(predicate : A -> Bool) -> Stream<A> {
        var stream = self
        while predicate(stream.first()) {
            stream = stream.rest()
        }
        return stream
    }
}

public class RecursiveStream<A> : Stream<A> {
    override init(_ first : A, _ step : @autoclosure () -> Stream<A>) {
        super.init(first, step)
    }

    public class func forward(head : A) -> RecursiveStream<A> {
        return RecursiveStream(head, Stream.repeat(head))
    }

    public func tie(step : @autoclosure () -> Stream<A>) {
        self._step = step
    }
}

func stars(n : Int) -> String {
    let star : Character = "*"
    return String(count: n, repeatedValue: star)
}

let nats    = Stream.iterate(0) { $0 + 1 }
let ones    = Stream.repeat(1)
let strings = nats.map { stars($0) }
let squares = nats.map { $0 * $0 }
let evens   = nats.filter { $0 % 2 == 0 }
let odds    = nats.filter { $0 % 2 == 1 }
let sums    = nats.zipWith(nats) { $0 + $1 }
let pairs   = evens.zip(odds)

/* Type inference for facts blows up if we don't specify <Int>
     stream.swift:152:37: error: cannot invoke '*' with an argument of type
        (($T8, ($T8, $T9) -> ($T8, $T9) -> $T7) -> ($T8, ($T8, $T9) -> $T7) ...
*/
let facts   = Stream<Int>.count(start:1).scan1 { $0 * $1 }

/* Cannot use fibs inside of its own definition:
     stream.swift:166:45: error: variable used within its own initial value
   Fake with with
     stream = RecursiveStream.forward()
     stream.tie(...)
*/

let fibs    = RecursiveStream.forward(0)
fibs.tie(Stream.cons(1, fibs.zipWith(fibs.drop(1)) { $0 + $1 }))

func annotate<A : Printable>(description : String, body : () -> A) {
    println("\(description):\n  \(body())\n")
}

annotate("First ten nats (iterate)") {
    nats.take(10)
}

annotate("First ten nats after dropping ten nats (drop)") {
    nats.drop(10).take(10)
}

annotate("First ten ones (repeat)") {
    ones.take(10)
}

annotate("Type-changing map (map int -> string)") {
    strings.take(10)
}

annotate("First ten squares (map int -> int)") {
    squares.take(10)
}

annotate("First ten evens (filter int -> bool)") {
    evens.take(10)
}

annotate("First ten odds (filter int -> bool)") {
    odds.take(10)
}

annotate("First ten sums ((int, int) -> int)") {
    sums.take(10)
}

annotate("Even/odd pairs (zip)") {
    pairs.take(10)
}

annotate("First ten factorial (scan)") {
    facts.take(10)
}

annotate("First ten fibs (forward/tie/zipwith)") {
    fibs.take(10)
}

annotate("Take while <10") {
    nats.takeWhile { $0 < 10 }
}

annotate("First ten after drop while <10") {
    (nats.dropWhile { $0 < 10 }).take(10)
}