Trie.swift

//
//  Trie.swift
//  SceneTest
//
//  Created by Greg Titus on 6/7/14.
//  Copyright (c) 2014 The Omni Group. All rights reserved.
//

// This is essentially the same data structure as OFTrie in OmniFoundation, except that the OmniFoundation version is mutable and ephemeral,
// and this version is immutable and persistent. Every time you add a string to the trie, it replaces all the nodes along the string's path
// to the root, and the only mutable change is reassigning the root. Old copies of the trie are unaffected. 
// 100 or so lines of code here, 600 or so in 3 classes (6 .h or .m files) in the Objective-C version.

// some conveniences, to treat the Sliceable UnicodeScalarView of a String as if it's a sort of cons'd list type...
extension String.UnicodeScalarView {
    var isEmpty : Bool { return self.startIndex == self.endIndex }
    var head : UnicodeScalar { return self[self.startIndex] }
    var tail : String.UnicodeScalarView { return self[self.startIndex.succ() .. self.endIndex] }
}

struct TrieNode {
    enum TrieChildren {
        // either a single leaf value, with a trailing suffix
        case Leaf(String.UnicodeScalarView)
        // or a branch, mapping first character to child node
        case Branch(Dictionary<UnicodeScalar, TrieNode>)
    }
    
    var value : AnyObject?
    var children : TrieChildren
    
    func insertValue(newValue: AnyObject, forChars: String.UnicodeScalarView) -> TrieNode {
        switch children {
        case .Leaf(let suffix):
            if (suffix == forChars) {
                return TrieNode(value: newValue, children: .Leaf(forChars))
            } else if (forChars.isEmpty) {
                let oldHead = suffix.head
                let oldLeaf = TrieNode(value: value, children: .Leaf(suffix.tail))
                return TrieNode(value: newValue, children: .Branch([oldHead: oldLeaf]))
            } else if (suffix.isEmpty) {
                let newHead = forChars.head
                let newLeaf = TrieNode(value: newValue, children: .Leaf(forChars.tail))
                return TrieNode(value: value, children: .Branch([newHead: newLeaf]))
            } else {
                let oldHead = suffix.head
                let newHead = forChars.head
                let oldLeaf = TrieNode(value: value, children: .Leaf(suffix.tail))
                
                if (oldHead == newHead) {
                    let newChild = oldLeaf.insertValue(newValue, forChars: forChars.tail)
                    return TrieNode(value: nil, children: .Branch([newHead : newChild]))
                } else {
                    let newLeaf = TrieNode(value: newValue, children: .Leaf(forChars.tail))
                    return TrieNode(value: nil, children: .Branch([oldHead: oldLeaf, newHead: newLeaf]))
                }
            }
            
        case .Branch(let subnodes):
            if (forChars.isEmpty) {
                return TrieNode(value: newValue, children: children)
            }
            
            var newChild : TrieNode
            let charHead = forChars.head
            if let child = subnodes[charHead] {
                newChild = child.insertValue(newValue, forChars: forChars.tail)
            } else {
                newChild = TrieNode(value: newValue, children: .Leaf(forChars.tail))
            }
            
            var newSubnodes = subnodes
            newSubnodes[charHead] = newChild
            return TrieNode(value: value, children: .Branch(newSubnodes))
        }
    }
    
    func lookup (str : String.UnicodeScalarView) -> AnyObject? {
        if (str.isEmpty) {
            return value
        }
        
        switch children {
        case .Leaf(let suffix):
            return suffix == str ? value : nil
        case .Branch(let subnodes):
            return subnodes[str.head]?.lookup(str.tail)
        }
    }
}

class Trie {
    var root : TrieNode? = nil
    
    func insert(value: AnyObject, forString: String) {
        if root {
            root = root!.insertValue(value, forChars: forString.unicodeScalars)
        } else {
            root = TrieNode(value: value, children: .Leaf(forString.unicodeScalars))
        }
    }
    
    func lookup(str : String) -> AnyObject? {
        return root?.lookup(str.unicodeScalars)
    }
}