red black tree updated for 2.1b2

RBTree.swift

private enum ListNode<Element> {
    case End
    indirect case Node(Element, next: ListNode<Element>)
    
    func cons(x: Element) -> ListNode<Element> {
        return .Node(x, next: self)
    }
}

public struct ListIndex<Element> {
    private let node: ListNode<Element>
    private let tag: Int
}

extension ListIndex: ForwardIndexType {
    public func successor() -> ListIndex<Element> {
        guard case let .Node(_, next: next) = node
            else { fatalError("cannot increment endIndex") }
        
        return ListIndex(node: next, tag: tag.predecessor())
    }
}

public func == <T>(lhs: ListIndex<T>, rhs: ListIndex<T>) -> Bool {
    return lhs.tag == rhs.tag
}

public struct List<Element> {
    public typealias Index = ListIndex<Element>
    public var startIndex: Index
    public var endIndex: Index { return Index(node: .End, tag: 0) }
    
    public init() {
        startIndex = Index(node: .End, tag: 0)
    }
}

extension List {
    public mutating func push(x: Element) {
        startIndex = Index(node: startIndex.node.cons(x),
            tag:  startIndex.tag.successor())
    }
    
    public mutating func pop() -> Element? {
        guard case let .Node(x, next: _) = startIndex.node
            else { return nil }
        
        ++startIndex
        return x
    }
}

extension List: ArrayLiteralConvertible {
    public init(arrayLiteral elements: Element...) {
        self = List()
        for x in elements.reverse() { push(x) }
    }
}

private enum Color { case R, B }

private indirect enum Tree<Element: Comparable> {
    case Empty
    case Node(Color,Tree<Element>,Element,Tree<Element>)
    
    init() { self = .Empty }
    
    init(_ x: Element, color: Color = .B,
        left: Tree<Element> = .Empty, right: Tree<Element> = .Empty)
    {
        self = .Node(color, left, x, right)
    }
}


extension Tree {
    func contains(x: Element) -> Bool {
        guard case let .Node(_,left,y,right) = self
            else { return false }
        
        if x < y { return left.contains(x) }
        if y < x { return right.contains(x) }
        return true
    }
}

private func balance<T>(tree: Tree<T>) -> Tree<T> {
    switch tree {
    case let .Node(.B, .Node(.R, .Node(.R, a, x, b), y, c), z, d):
        return .Node(.R, .Node(.B,a,x,b),y,.Node(.B,c,z,d))
    case let .Node(.B, .Node(.R, a, x, .Node(.R, b, y, c)), z, d):
        return .Node(.R, .Node(.B,a,x,b),y,.Node(.B,c,z,d))
    case let .Node(.B, a, x, .Node(.R, .Node(.R, b, y, c), z, d)):
        return .Node(.R, .Node(.B,a,x,b),y,.Node(.B,c,z,d))
    case let .Node(.B, a, x, .Node(.R, b, y, .Node(.R, c, z, d))):
        return .Node(.R, .Node(.B,a,x,b),y,.Node(.B,c,z,d))
    default:
        return tree
    }
}


private func ins<T>(into: Tree<T>, _ x: T) -> Tree<T> {
    guard case let .Node(c, l, y, r) = into
        else { return Tree(x, color: .R) }

    if x < y { return balance(Tree(y, color: c, left: ins(l,x), right: r)) }
    if y < x { return balance(Tree(y, color: c, left: l, right: ins(r, x))) }
    return into
}


extension Tree {
    func insert(x: Element) -> Tree {
        guard case let .Node(_,l,y,r) = ins(self, x)
            else { fatalError("ins should never return an empty tree") }

        return .Node(.B,l,y,r)
    }
}



extension Tree: SequenceType {
    func generate() -> AnyGenerator<Element> {
        var stack: List<Tree> = []
        var current: Tree = self

        return anyGenerator { _ -> Element? in
            while true {
                // if there's a left-hand node, head down it
                if case let .Node(_,l,_,_) = current {
                    stack.push(current)
                    current = l
                }
                // if there isn’t, head back up, going right as
                // soon as you can:
                else if case let .Node(_,_,x,r)? = stack.pop() {
                    current = r
                    return x
                }
                else {
                    // otherwise, we’re done
                    return nil
                }
            }
        }
    }
}

extension Tree: ArrayLiteralConvertible {
    init <S: SequenceType where S.Generator.Element == Element>(_ source: S) {
        self = source.reduce(Tree()) { $0.insert($1) }
    }
    
    init(arrayLiteral elements: Element...) {
        self = Tree(elements)
    }
}

import Darwin

extension Array {
    func shuffle() -> [Element] {
        var list = self
        for i in 0..<(list.count - 1) {
            let j = Int(arc4random_uniform(UInt32(list.count - i))) + i
            guard i != j else { continue }
            swap(&list[i], &list[j])
        }
        return list
    }
}


let engines = [
    "Daisy", "Salty", "Harold", "Cranky",
    "Thomas", "Henry", "James", "Toby",
    "Belle", "Diesel", "Stepney", "Gordon",
    "Captain", "Percy", "Arry", "Bert",
    "Spencer",
]

// test various inserting engines in various different permutations
for permutation in [engines, engines.sort(), engines.sort(>),engines.shuffle(),engines.shuffle()] {
    let t = Tree(permutation)
    assert(!t.contains("Fred"))
    assert(t.elementsEqual(t.insert("Thomas")))
    assert(!engines.contains { !t.contains($0) })
    assert(t.elementsEqual(engines.sort()))
    print(t.joinWithSeparator(","))
}