aho corasick in swift

main.swift
class TrieNode {
    var value: Character?
    var children = [Character: TrieNode]()
    var isEndOfWord = false
    var failure: TrieNode?
    var output = [String]()
    
    init(_ value: Character? = nil) {
        self.value = value
    }
}

class AhoCorasick {
    private var root = TrieNode()
    
    func addPattern(_ pattern: String) {
        var node = root
        for char in pattern {
            if let child = node.children[char] {
                node = child
            } else {
                let newNode = TrieNode(char)
                node.children[char] = newNode
                node = newNode
            }
        }
        node.isEndOfWord = true
    }
    
    private func buildFailureFunction() {
        let queue = Queue<TrieNode>()
        
        for child in root.children.values {
            child.failure = root
            queue.enqueue(child)
        }
        
        while let current = queue.dequeue() {
            for child in current.children.values {
                queue.enqueue(child)
                var failure = current.failure
                while failure != nil, failure!.children[child.value!] == nil {
                    failure = failure?.failure
                }
                child.failure = failure?.children[child.value!] ?? root
                child.output = child.failure!.output
                if child.isEndOfWord {
                    child.output.append(String(child.value!))
                }
            }
        }
    }
    
    func search(in text: String) -> [String: [Int]] {
        buildFailureFunction()
        
        var result = [String: [Int]]()
        var node = root
        
        for (index, char) in text.enumerated() {
            while node.children[char] == nil, node != root {
                node = node.failure!
            }
            if node.children[char] != nil {
                node = node.children[char]!
            }
            for pattern in node.output {
                let endIndex = index - pattern.count + 1
                if result[pattern] == nil {
                    result[pattern] = [endIndex]
                } else {
                    result[pattern]?.append(endIndex)
                }
            }
        }
        
        return result
    }
}

struct Queue<T> {
    private var elements = [T]()
    
    mutating func dequeue() -> T? {
        return elements.isEmpty ? nil : elements.removeFirst()
    }
    
    mutating func enqueue(_ element: T) {
        elements.append(element)
    }
    
    func isEmpty() -> Bool {
        return elements.isEmpty
    }
}
2613 chars
96 lines

main.swift
let ac = AhoCorasick()
ac.addPattern("he")
ac.addPattern("she")
ac.addPattern("his")
ac.addPattern("hers")
let results = ac.search(in: "ushers")
print(results) // Output: ["he": [1], "she": [2], "hers": [3]]
208 chars
8 lines