Producing Combinations in Swift based on Eric Lippert's post

combinations.swift

// Producing combinations
// Porting Eric Lippert's code to Swift
// http://ericlippert.com/2014/10/13/producing-combinations-part-one/

import Cocoa

public class ImmutableStack<T> {

    public func push(t:T) -> ImmutableStack<T> {
        return NonEmptyStack(top: t, tail: self)
    }

    public func top() -> T? {
        return nil
    }

    public func pop() -> ImmutableStack<T>? {
        return nil
    }

    public func isEmpty() -> Bool {
        return true
    }

    public class func emptyStack() -> ImmutableStack<T> {
        return ImmutableStack()
    }

}

private class NonEmptyStack<T> : ImmutableStack<T> {
    let _top:T
    let _tail:ImmutableStack<T>

    init(top:T, tail:ImmutableStack<T>) {
        _top = top
        _tail = tail
    }

    override func top() -> T {
        return _top
    }

    override func pop() -> ImmutableStack<T> {
        return _tail
    }

    override func isEmpty() -> Bool {
        return false
    }

}

extension ImmutableStack : SequenceType {
    public func generate() -> GeneratorOf<T> {
        var stack = Optional.Some(self)

        return GeneratorOf {
            let result = stack?.top()
            stack = stack?.pop()
            return result
        }
    }
}



extension SequenceOf {

    static func empty() -> SequenceOf<T> {
        return SequenceOf([])
    }

    static func singleton(element:T) -> SequenceOf<T> {
        return SequenceOf([element])
    }

    func map<U>(transform:T -> U) -> SequenceOf<U> {

        return SequenceOf<U> { () -> GeneratorOf<U> in
            var g = self.generate()
            return GeneratorOf {
                return g.next().map(transform)
            }
        }
    }

    func rightNotSet<L>(t:(L,Bool)?) -> Bool? {
        return t == nil ? nil : !t!.1;
    }

    func zipWhere<S:SequenceType where S.Generator.Element == Bool>(bools:S) -> SequenceOf<T> {

        return SequenceOf { () -> GeneratorOf<T> in
            var generator = Zip2(self, bools).generate()
            return GeneratorOf<T> {
                var next = generator.next()

                while let (value, isSet) = next {
                    if (isSet) {
                        return value
                    } else {
                        next = generator.next()
                    }
                }

                return nil
            }
        }
    }

    func extend<S1:SequenceType where S1.Generator.Element == T>(s1:S1) -> SequenceOf<T> {

        return SequenceOf{ () -> GeneratorOf<T> in
            var g0 = self.generate()
            var g1 = s1.generate()
            return GeneratorOf {
                g0.next() ?? g1.next()
            }
        }

    }

    func count() -> UInt {
        var i:UInt = 0
        for _ in self { i++ }
        return i
    }
}

typealias BoolStack = ImmutableStack<Bool>

let singletonEmptyBoolStack = SequenceOf.singleton(BoolStack.emptyStack())
let emptySequenceOfBoolStack:SequenceOf<BoolStack> = SequenceOf.empty()

func combinations(n:UInt, k:UInt) -> SequenceOf<BoolStack> {
    if (k == 0 && n == 0) {
        return singletonEmptyBoolStack
    }

    if (n < k) {
        return emptySequenceOfBoolStack
    }

    let seq1 = (k>0) ? combinations(n-1,k-1).map({$0.push(true)}) : emptySequenceOfBoolStack

    return seq1.extend(combinations(n-1, k).map({$0.push(false)}))
}

func combinations<S:SequenceType>(s:S,k:UInt) -> SequenceOf<SequenceOf<S.Generator.Element>> {

    let sseq = SequenceOf(s)
    return combinations(sseq.count(), k).map { sseq.zipWhere($0) }
}




var combs = combinations([50,60,70,80,90],3)

for comb in combs {
    println(Array(comb))
}
println("-----")

for comb in combs {
    println(Array(comb))
}
println("=====")