Persistent Red/Black Tree in Swift using 2.0 Pattern Matching

rbt.swift

final class Box<T> {
    let unbox: T
    init(_ x: T) { self.unbox = x }
}

enum Color { case R, B }

enum Tree { //<T: Comparable> { // something not working when I make this generic
    case Empty
    case Node(Color,Box<Tree>,String,Box<Tree>)
}

private func balance(tree: Tree) -> Tree {
    // hoping indirect will mean these can be done in a single let with nested patterns
    
    /**
    Rebalance logic based on the ML version from Chris Okasaki's Purely Functional Data Structures book
    
    Transformations of the 4 possible cases of a red child and grandchild under a black node,
    each one looking like this (right-hand is same for all 4):
    
          z: black        becomes        y: red
         /       \                     /        \
        x: red    d                   x: black   z: black
       / \                           / \        / \
      a   y: red                    a   b      c   d
     / \
    b   c
    
    */
    
    // (B,T(R,T(R,a,x,b),y,c),z,d)
    if case let .Node(.B,l,z,d) = tree, case let .Node(.R,ll,y,c) = l.unbox, case let .Node(.R,a,x,b) = ll.unbox {
        print("a,", appendNewline: false)
        return .Node(.R, Box(.Node(.B,a,x,b)),y,Box(.Node(.B,c,z,d)))
    }
    
    // (B,T(R,a,x,T(R,b,y,c)),z,d)
    if case let .Node(.B,l,z,d) = tree, case let .Node(.R,a,x,ll) = l.unbox, case let .Node(.R,b,y,c) = ll.unbox {
        print("b,", appendNewline: false)
        return .Node(.R,Box(.Node(.B,a,x,b)),y,Box(.Node(.B,c,z,d)))
    }
    
    // (B,a,x,T(R,T(R,b,y,c),z,d))
    if case let .Node(.B,a,x,r) = tree, case let .Node(.R,rr,z,d) = r.unbox, case let .Node(.R,b,y,c) = rr.unbox {
        print("c,", appendNewline: false)
        return .Node(.R,Box(.Node(.B,a,x,b)),y,Box(.Node(.B,c,z,d)))
    }
    
    // (B,a,x,T(R,b,y,T(R,c,z,d)))
    if case let .Node(.B,a,x,r) = tree, case let .Node(.R,b,y,rr) = r.unbox, case let .Node(.R,c,z,d) = rr.unbox {
        print("d,", appendNewline: false)
        return .Node(.R,Box(.Node(.B,a,x,b)),y,Box(.Node(.B,c,z,d)))
    }
    
    return tree
}

private func ins(into: Tree, _ x: String) -> Tree {
    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.unbox,x), right: r.unbox)) }
    if y < x { return balance(Tree(y, color: c, left: l.unbox, right: ins(r.unbox, x))) }
    return into
}


extension Tree {
    init() { self = .Empty }
    init(_ x: String, color: Color = .B, left: Tree = .Empty, right: Tree = .Empty) {
        self = Tree.Node(color, Box(left), x, Box(right))
    }
    
    func contains(x: String) -> Bool {
        guard case let .Node(_,l,y,r) = self else { return false }
        
        if x < y { return l.unbox.contains(x) }
        if y < x { return r.unbox.contains(x) }
        return true
    }
    
    func insert(x: String) -> Tree {
        guard case let .Node(_,l,y,r) = ins(self, x) else { fatalError() }
        return .Node(.B,l,y,r)
    }
}

extension Tree: SequenceType {
    func generate() -> AnyGenerator<String> {
        // stack-based iterative inorder traversal to
        // make it easy to use with anyGenerator
        var stack: [Tree] = []
        var current: Tree = self
        return anyGenerator { _ -> String? in
            while true {
                if case let .Node(_,l,_,_) = current {
                    stack.append(current)
                    current = l.unbox
                }
                else if !stack.isEmpty, case let .Node(_,_,x,r) = stack.removeLast() {
                    current = r.unbox
                    return x
                }
                else {
                    return nil
                }
            }
        }
    }
}

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

import Darwin

extension Array {
    func shuffle() -> [T] {
        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))
}