Union-Find disjoint-set data structure in Swift

UnificationSystem.swift

public protocol UnificationSystemProtocol {
    associatedtype Variable
    mutating func makeFreshVariable() -> Variable
    mutating func canonicalVariable(for variable: Variable) -> Variable
    @discardableResult 
    mutating func unify(_ first: Variable, _ second: Variable) -> Bool
}

public struct UnificationSystem: UnificationSystemProtocol {
    public typealias Variable = Int
    
    private var parent: [Variable] = []
    private var rank: [Int] = []
    
    public init() {}
    
    public mutating func makeFreshVariable() -> Variable {
        let variable = parent.count
        parent.append(variable)
        rank.append(0)
        return variable
    }
    
    public mutating func canonicalVariable(for variable: Variable) -> Variable {
        guard parent[variable] != variable else { return variable }
        parent[variable] = canonicalVariable(for: parent[variable])
        return parent[variable]
    }
    
    @discardableResult
    public mutating func unify(_ lhs: Variable, _ rhs: Variable) -> Bool {
        let lhs = canonicalVariable(for: lhs)
        let rhs = canonicalVariable(for: rhs)
        guard lhs != rhs else { return false }
        
        if rank[lhs] < rank[rhs] {
            parent[lhs] = rhs
        } else if rank[lhs] > rank[rhs] {
            parent[rhs] = lhs
        } else {
            parent[rhs] = lhs
            rank[lhs] += 1
        }
        return true
    }
}