Implement AVLTree with conditional insert where element is comparable

AVLTree.swift

import Foundation

final class MutableAVLTree<T> {
    typealias Process = ((MutableAVLTree<T>) -> ())
    var value: T?
    var height: Int = 0
    var left: MutableAVLTree?
    var right: MutableAVLTree?
    
    init() {
        self.value = nil
    }
    
    init(_ value: T) {
        self.value = value
    }
}

// MARK: - Computed Properties

extension MutableAVLTree {
    
    var isEmpty: Bool {
        value == nil
    }
    
    var leftHeight: Int {
        left?.height ?? -1
    }
    var rightHeight: Int {
        right?.height ?? -1
    }
    
    var balance: Int {
        leftHeight - rightHeight
    }
}

// MARK: - Rotation
// - Note: https://www.tutorialspoint.com/data_structures_algorithms/avl_tree_algorithm.htm
extension MutableAVLTree {
    
    // - returns: new root
    func rebalance() -> MutableAVLTree {
        switch balance {
        case 2:
            if let leftTree = left, leftTree.balance == -1 {
                return rotateLeftRight()
            } else {
                return rotateRight()
            }
        case -2:
            if let rightTree = right, rightTree.balance == 1 {
                return rotateRightLeft()
            } else {
                return rotateLeft()
            }
        default:
            return self
        }
    }
    
    /// - returns: new root
    func rotateLeft() -> MutableAVLTree {
        guard let newRoot = right else { return self }

        self.right = newRoot.left
        newRoot.left = self
        
        self.updateHeight()
        newRoot.updateHeight()
        
        return newRoot
    }
    
    /// - returns: new root
    func rotateRight() -> MutableAVLTree {
        guard let newRoot = left else { return self }

        self.left = newRoot.right
        newRoot.right = self
        
        self.updateHeight()
        newRoot.updateHeight()
        
        return newRoot
    }
    
    /// - returns: new root
    func rotateRightLeft() -> MutableAVLTree {
        guard let rightTree = right else { return self }
        
        let newRoot = rightTree.rotateRight()
        self.right = newRoot
        return self.rotateLeft()
    }
    
    /// - returns: new root
    func rotateLeftRight() -> MutableAVLTree {
        guard let leftTree = left else { return self }
        
        let newRoot = leftTree.rotateLeft()
        self.left = newRoot
        return self.rotateRight()
    }
}

// MARK: - Internal Update

private extension MutableAVLTree {
    func updateHeight() {
        height = max(leftHeight, rightHeight) + 1
    }
}

// MARK: - Traverse

extension MutableAVLTree {
    func traverse(process: Process) {
        guard !isEmpty else { return }
        left?.traverse(process: process)
        process(self)
        right?.traverse(process: process)
    }
}

extension MutableAVLTree where T: Comparable {
    func insert(_ newValue: T, except condition: ((T, T) -> Bool)) -> Bool {
        var result: Bool
        if let value = value {
            if condition(value, newValue) {
                result = false
            } else if newValue < value {
                if let left = left {
                    result = left.insert(newValue, except: condition)
                } else {
                    self.left = MutableAVLTree(newValue)
                    result = true
                }
            } else {
                if let right = right {
                    result = right.insert(newValue, except: condition)
                } else {
                    self.right = MutableAVLTree(newValue)
                    result = true
                }
            }
        } else {
            self.value = newValue
            result = true
        }
        if result {
            _ = rebalance()
        }
        return result
    }
}

main.swift

import Foundation

var tree = MutableAVLTree<Subscription>()

let formatter = DateFormatter()
formatter.dateFormat = "yyyy-MM-dd HH:mm"
let now = formatter.date(from: "2021-01-01 00:00")!
for i in (1..<100) {
    let start: Double = 3600 * Double(i)
    let end: Double = 3600 * Double(i + 1)
    let s = Subscription(name: "\(i)",
                         start: now.advanced(by: start),
                         end: now.advanced(by: end))
    _ = tree.insert(s) { $0.overlap(with: $1) }
}
for i in (0..<100) {
    let start: Double = -3600 * Double(i + 1)
    let end: Double = -3600 * Double(i)
    let s = Subscription(name: "\(-i)",
                         start: now.advanced(by: start),
                         end: now.advanced(by: end))
    _ = tree.insert(s) { $0.overlap(with: $1) }
}

let test1 = Subscription(name: "test1", start: now.advanced(by: -1800), end: now.advanced(by: 1800))
let test2 = Subscription(name: "test2", start: now, end: now.advanced(by: 1800))
let test3 = Subscription(name: "test3", start: test2.end, end: test2.end.advanced(by: 1800))
let test4 = Subscription(name: "test4", start: now, end: test3.end)

print(tree.insert(test1) { $0.overlap(with: $1) } == false)
print(tree.insert(test2) { $0.overlap(with: $1) } == true)
print(tree.insert(test3) { $0.overlap(with: $1) } == true)
print(tree.insert(test4) { $0.overlap(with: $1) } == false)

print(tree.balance)

Subscription.swift

struct Subscription: Comparable, CustomStringConvertible {
    let name: String
    let start: Date
    let end: Date
    
    static func < (lhs: Subscription, rhs: Subscription) -> Bool {
        lhs.start < rhs.start
    }
    
    var description: String {
        "name: \(name), start: \(start), end: \(end)"
    }
    
    func overlap(with otherSubscription: Subscription) -> Bool {
        otherSubscription.start <= self.start  && self.start < otherSubscription.end ||
            self.start <= otherSubscription.start && otherSubscription.start < self.end
    }
}