A prefix-matching trie in Swift

prefix-trie.swift

//  prefix-trie.swift  A generic prefix-matching trie in Swift
//  By Andrew Bulhak http://dev.null.org/acb/

/** The prefix trie has two types: the type of input element, C, and the type of final result for a successful prefix match, T. */
enum Trie<C:Hashable,T> {
    typealias IndexType = C
    typealias ResultType = T
    case Success(value:T)
    case Node([C:Trie<C,T>])
    
    init() {
        self = .Node([:])
    }
}

struct AlreadyMatches: ErrorType {}

extension Trie {
    func match<G: GeneratorType where G.Element == IndexType>(inout g:  G) -> ResultType? {
        switch (self) {
        case .Success(let v) : return v
        case .Node(let n): 
            return g.next().flatMap { n[$0]?.match(&g) }
        }       
    }
    
    func match<S: SequenceType where S.Generator.Element == IndexType>(s: S) -> ResultType? {
        var g = s.generate()
        return self.match(&g)
    }
    
    func inserting<G: GeneratorType where G.Element == IndexType>(inout g: G, value: ResultType) throws -> Trie {
        
        func sequenceForGenerator(inout g: G, value: ResultType) -> Trie {
            if let v = g.next() {
                return Trie.Node([v:sequenceForGenerator(&g, value: value)])
            } else {
                return Trie.Success(value: value)
            }
        }
        
        if let c = g.next() {
            switch(self) {
            case .Success(_): throw AlreadyMatches()
            case .Node(var d):
                if var n = d[c] {
                    try n.insert(&g, value:value)
                    d[c] = n
                } else {
                    d[c] = sequenceForGenerator(&g, value:value)
                }
                return .Node(d)
            }
        } else {
            return .Success(value: value)
        }
    }

    func inserting<S: SequenceType where S.Generator.Element == IndexType>(key: S, value: ResultType) throws -> Trie {
        var g = key.generate()
        return try self.inserting(&g, value: value)
    }
    
    mutating func insert<G: GeneratorType where G.Element == IndexType>(inout g: G, value: ResultType) throws {
        self = try self.inserting(&g, value: value)
    }
    
    mutating func insert<S: SequenceType where S.Generator.Element == IndexType>(key: S, value: ResultType) throws {
        var g = key.generate()
        try self.insert(&g, value: value)
    }
}

extension Trie : CustomDebugStringConvertible {
    var debugDescription: String {
        switch(self) {
        case .Success(let v): return "\(v)"
        case .Node(let d): 
            return "{" + (d.map { (k,v) in "\(k):\(v)" }).joinWithSeparator(",") + "}"
        }
    }
}

// Now use the trie for a sample application: matching stop words at the start of a string

var t = Trie<Character,Int>()

do {
    for word in ["the", "a", "an"] {
        for space in [" ", "_"] {
            try t.insert((word+space).characters, value:word.characters.count+1)
        }
    }
} catch {
    print(error)
}

t.match("a something".characters) // 2
t.match("the_word".characters)    // 4
t.match("this is not a match".characters) // nil