An Implementation of Streams in Swift

Streams.swift

// We're going to show off some examples first--implementation later!

// Generates a stream of integers starting at n
func IntegersFrom<I : IntegerType>(n: I) -> Stream<I> {
    return Stream(head: n, tail: IntegersFrom(n + 1))
}

// One potential implementation of a fibonacci stream; initial input (1,1)
func MakeFibonacci(a: Int, b: Int) -> Stream<Int> {
    return Stream<Int>(head: a, tail: MakeFibonacci(b, a + b))
}

// An even more mindblowing implementation of a fibonacci stream
func MakeFibonacci2() -> Stream<Int> {
    return Stream(head: 0, tail: Stream(head: 1, tail: MakeFibonacci2().tail.zip(MakeFibonacci2(), combine: +)))
}

// Filters the primes out of a stream
func Primes() -> Stream<Int> {
    return PrimeSieve(IntegersFrom(2), { a, b in b < 2 || a % b != 0 })
}
// Helper for filter primes
func PrimeSieve<T>(stream: Stream<T>, accepts: (T, T) -> Bool) -> Stream<T> {
    return Stream(head: stream.head, tail: stream.tail.filter({ x in accepts(x, stream.head)}))
}

for i in finite(100)(Primes()) {
    // Iterate over the first 100 prime numbers
}

// An array of the 1st 12 fibonacci numbers greater than 50
let arr = MakeFibonacci(1, 1).filter({ x in x > 50 }).take(12)

// The first prime integer that is greater than 100
let num = Primes().dropWhile({ x <= 100 }).head

// The product of the first 5 fibonacci numbers
let product = MakeFibonacci2().reduceWhile(1, combine: *, count: 5)

// Result is a tuple of a value from squares and a vlue from cubes
// such that these are the first values for an integer x such that
// cubes is greater than squares
let cubes = IntegersFrom(1).map({x in x * x * x - 10 })
let squares = IntegersFrom(1).map({x in 4 * x * x + 3})
let result = squares.zip(cubes).dropWhile({ (square, cube) in square > cube}).head

// Another cool example would be approximating infinite series
// As an exercise, perhaps try to approximate the value of Pi uisng a Stream

// Note that a lot of these recursive functions, such as filter, must be rewritten as
// iterative functions in Swift 1.1 to work reliably because Swift does not guarentee tail
// call optimization. In hope that Swift fixes this issue in the future, we have left these
// definitions in their recursive forms---it's much prettier.

/*
 * Used to represent a lazy (often recursively computed) list
 */
class Stream<T> : SequenceType {

    let head: T
    private let unevaluatedTail: () -> Stream<T>
    
    var tail: Stream<T> {
        get {
            return unevaluatedTail()
        }
    }
    
    init(head: T, tail: @autoclosure () -> Stream<T>){
        self.head = head
        self.unevaluatedTail = tail
    }
    
    func dropWhile(test: T -> Bool) -> Stream<T> {
        return test(head) ? tail.dropWhile(test) : self
    }
    
    func takeWhile(test: T-> Bool) -> [T] {
        return test(head) ? [head] + tail.takeWhile(test) : []
    }
    
    func reduceWhile<T, U>(initial: U, combine: (U, T) -> U, test: T -> Bool) -> U {
        return test(head as T) ? combine(tail.reduceWhile(initial, combine: combine, test: test), head as T) : initial
    }
    
    func drop(var count: Int) -> Stream<T> {
        assert(count >= 0, "Drop count must be positive")
        
        return dropWhile({ _ in count-- > 0 })
    }
    
    func take(var count: Int) -> [T] {
        assert(count >= 0, "Take count must be positive")

        return takeWhile({ _ in count-- > 0 })
    }
    
    func reduce<T, U>(initial: U, combine: (U, T) -> U, var count: Int) -> U {
        assert(count >= 0, "Reduce count must be positive")

        return reduceWhile(initial, combine: combine, test: { _ in count-- > 0 })
    }
    
    subscript(index: Int) -> Stream<T> {
        return drop(index)
    }
    
    func zip<U>(stream: Stream<U>) -> Stream<(T,U)> {
        return Stream<(T,U)>(head: (head, stream.head), tail: tail.zip(stream.tail))
    }
    
    func zip<U,V>(stream: Stream<U>, combine: (T, U) -> V) -> Stream<V> {
        return Stream<V>(head: combine(head, stream.head), tail: tail.zip(stream.tail, combine))
    }
    
    func map<U>(transform: T -> U) -> Stream<U> {
        return Stream<U>(head: transform(head), tail: tail.map(transform))
    }
    
    func filter(includeElement: T -> Bool) -> Stream<T> {
        return includeElement(head) ? Stream(head: head, tail: tail.filter(includeElement)) : tail.filter(includeElement)
    }
    
    func generate() -> StreamGenerator<T> {
        return StreamGenerator(stream: self)
    }
}

struct StreamGenerator<T> : GeneratorType {
    
    private var stream: Stream<T>
    
    init(stream: Stream<T>){
        self.stream = stream
    }
    
    mutating func next() -> T? {
        let head = stream.head
        stream = stream.tail
        return head
    }
}

struct FiniteSequence<T : SequenceType> : SequenceType {
    
    let sequence: T
    let count: Int
    
    func generate() -> FiniteGenerator<T.Generator> {
        return FiniteGenerator(generator: sequence.generate(), count: count)
    }
}

struct FiniteGenerator<T : GeneratorType> : GeneratorType {
    
    private var generator: T
    private var remainingCount: Int
    
    init(generator: T, count: Int){
        self.generator = generator
        self.remainingCount = count
    }
    
    mutating func next() -> T.Element? {
        return remainingCount-- > 0 ? generator.next() : nil
    }
}

func finite<T : SequenceType>(count: Int) -> T -> FiniteSequence<T> {
    return { sequence in
        return FiniteSequence(sequence: sequence, count: count)
    }
}