implement a aho corasick graph in swift

main.swift
class TrieNode {
    var value: Character?
    weak var parent: TrieNode?
    var children: [Character: TrieNode] = [:]
    var isTerminating = false
    var pattern: String?
    
    init(value: Character?, parent: TrieNode?) {
        self.value = value
        self.parent = parent
    }
    
    func add(child: Character) {
        guard children[child] == nil else { return }
        children[child] = TrieNode(value: child, parent: self)
    }
}
453 chars
18 lines

main.swift
class AhoCorasick {
    private let root = TrieNode(value: nil, parent: nil) // create root node
    private var failureLinks = [TrieNode]()
    
    func add(pattern: String) {
        guard !pattern.isEmpty else { return }
        
        var currentNode = root
        
        // Add nodes to the Trie
        for character in pattern {
            if let childNode = currentNode.children[character] {
                currentNode = childNode
            } else {
                currentNode.add(child: character)
                currentNode = currentNode.children[character]!
            }
        }
        
        currentNode.isTerminating = true
        currentNode.pattern = pattern
    }
    
    func buildFailures() {
        failureLinks = [TrieNode?](repeating: nil, count: root.children.count)
        var queue = [TrieNode]()
        
        // Initialize queue with children of the root
        for child in root.children.values {
            queue.append(child)
            failureLinks.append(root)
        }
        
        // Compute failure links
        while !queue.isEmpty {
            let node = queue.removeFirst()
            
            for child in node.children.values {
                queue.append(child)
                
                var failure = node.parent
                while failure?.children[node.value!] == nil, failure !== root {
                    failure = failure?.failure
                }
                
                if let f = failure?.children[node.value!], f !== node {
                    child.failure = f
                } else {
                    child.failure = root
                }
                
                failureLinks.append(child.failure)
            }
        }
    }
    
    func search(text: String) -> Set<String> {
        var currentNode = root
        var setOfMatches = Set<String>()
        var i = 0
        let characters = Array(text)
        
        while i < characters.count {
            let character = characters[i]
            
            while currentNode.children[character] == nil, currentNode !== root {
                currentNode = currentNode.failure!
            }
            
            if let child = currentNode.children[character] {
                currentNode = child
            }
            
            var node = currentNode
            while node !== root {
                if node.isTerminating {
                    if let pattern = node.pattern {
                        setOfMatches.insert(pattern)
                    }
                }
                node = node.failure!
            }
            
            i += 1
        }
        
        return setOfMatches
    }
}
2711 chars
90 lines

main.swift
let ahoCorasick = AhoCorasick()

ahoCorasick.add(pattern: "he")
ahoCorasick.add(pattern: "she")
ahoCorasick.add(pattern: "his")
ahoCorasick.add(pattern: "hers")

ahoCorasick.buildFailures()

let text = "ushers"
let matches = ahoCorasick.search(text: text)
print(matches) // Output: ["hers", "he"]
297 chars
13 lines