Red Black Tree with indirect and pattern matching for Swift 2.0b4

rbt.swift

enum Color { case R, B }

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 {
        switch self {
        case .Empty: return false
        case let .Node(_, left, y, right):
            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> {
    switch into {
    case .Empty: return Tree(x, color: .R)
    case let .Node(c, l, y, 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 {
        switch ins(self, x) {
        case let .Node(_,l,y,r):
            return .Node(.B,l,y,r)
        default:
            fatalError("ins should never return an empty tree")
        }
    }
}


extension Tree: SequenceType {
    func generate() -> AnyGenerator<Element> {
        // stack-based iterative inorder traversal to
        // make it easy to use with anyGenerator
        var stack: [Tree] = []
        var current: Tree = self
        return anyGenerator { _ -> Element? in
            while true {
                switch current {
                case let .Node(_,l,_,_):
                    stack.append(current)
                    current = l
                case .Empty where !stack.isEmpty:
                    switch stack.removeLast() {
                    case let .Node(_,_,x,r):
                        current = r
                        return x
                    default:
                        break
                    }
                case .Empty:
                    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
            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(),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(",".join(t))
}