Apple BTree Collection extended

_BTree++.swift

//
//  customization.swift
//  BTree
//
//  Created by 박병관 on 3/9/25.
//

import Foundation


extension _BTree {
    
    @inlinable
    internal func lookupGreaterOrEqual(forKey key: Key) -> Element? {
        var node: Node? = self.root
        var result: Element? = nil
        while let currentNode = node {
            currentNode.read { handle in
                let slot = handle.startSlot(forKey: key)

                // If exact match found, return immediately
                if slot < handle.elementCount, handle[keyAt: slot] == key {
                    result = (handle[keyAt: slot], handle[valueAt: slot])
                    node = nil
                    return
                }
                
                // If slot is within bounds, check for the next greater key
                if slot < handle.elementCount {
                    result = (handle[keyAt: slot], handle[valueAt: slot])
                }

                // Continue searching in the child node
                if handle.isLeaf {
                    node = nil
                } else {
                    node = handle[childAt: slot]
                }
            }
        }
        
        return result
    }
    
    @inlinable
    internal func lookupGreater(forKey key: Key) -> Element? {
        var node: Node? = self.root
        var result: Element? = nil
        while let currentNode = node {
            currentNode.read { handle in
                let slot = handle.endSlot(forKey: key)

                // If exact match found, return immediately
                if slot < handle.elementCount, handle[keyAt: slot] == key {
                    result = (handle[keyAt: slot], handle[valueAt: slot])
                    node = nil
                    return
                }
                
                // If slot is within bounds, check for the next greater key
                if slot < handle.elementCount {
                    result = (handle[keyAt: slot], handle[valueAt: slot])
                }

                // Continue searching in the child node
                if handle.isLeaf {
                    node = nil
                } else {
                    node = handle[childAt: slot]
                }
            }
        }
        
        return result
    }
    
    @inlinable
    internal func lookupLessThan(forKey key: Key) -> Element? {
        var node: Node? = self.root
        var result: Element? = nil
        
        while let currentNode = node {
            currentNode.read { handle in
                let slot = handle.startSlot(forKey: key) - 1
                
                // If there is a key < `key`, update result
                if slot >= 0 {
                    result = (handle[keyAt: slot], handle[valueAt: slot])
                }
                
                // If `key` is smaller than the smallest key in this node, move left
                // // Continue searching in the appropriate child
                if handle.isLeaf {
                    node = nil
                } else {
                    node = slot < 0 ? handle[childAt: 0] : handle[childAt: slot + 1]
                }
            }
        }
        
        return result
    }
    
    @inlinable
    internal func lookupLessThanOrEqual(forKey key: Key) -> Element? {
        var node: Node? = self.root
        var result: Element? = nil
        
        while let currentNode = node {
            currentNode.read { handle in
                let slot = handle.endSlot(forKey: key) - 1
                
                // If there is a key <= `key`, update result
                if slot >= 0 {
                    result = (handle[keyAt: slot], handle[valueAt: slot])
                }

                // If `key` is smaller than the smallest key in this node, move left
                // Continue searching in the appropriate child
                if handle.isLeaf {
                    node = nil
                } else {
                    node = slot < 0 ? handle[childAt: 0] : handle[childAt: slot + 1]
                }
            }
        }
        
        return result
    }
    
    
    
    @inlinable
    internal func lookup(forKey key: Key, search: BTreeSearch) -> Element? {
//        self.findAnyValue(forKey: key)
        var node: Node? = self.root
        var result: Element? = nil
        while let currentNode = node {
            currentNode.read { handle in
                let slot:Int
                do {
                    var mutable:Int
                    switch search {
                    case .lessThan, .greaterThanOrEqual, .equal:
                        mutable = handle.startSlot(forKey: key)
                        mutable -= search == .lessThan ? 1 : 0
                    case .greaterThan, .lessThanOrEqual:
                        mutable = handle.endSlot(forKey: key)
                        mutable -= search == .lessThanOrEqual ? 1 : 0
                    }
                    slot = mutable
                }
                
                switch search {
                case .lessThan, .lessThanOrEqual:
                    // If there is a key <= `key`, update result
                    if slot >= 0 {
                        result = (handle[keyAt: slot], handle[valueAt: slot])
                    }

                    // If `key` is smaller than the smallest key in this node, move left
                    // Continue searching in the appropriate child
                    if handle.isLeaf {
                        node = nil
                    } else {
                        node = slot < 0 ? handle[childAt: 0] : handle[childAt: slot + 1]
                    }
                case .greaterThan, .greaterThanOrEqual, .equal:
                    // If slot is within bounds, check for the next greater key
                    if slot < handle.elementCount {
                        result = (handle[keyAt: slot], handle[valueAt: slot])
                        if search == .equal, result?.key != key {
                            result = nil
                        }
                    }

                    // Continue searching in the child node
                    if handle.isLeaf {
                        node = nil
                    } else {
                        node = handle[childAt: slot]
                    }
                }
            }
        }
        return result
    }
    
}


public enum BTreeSearch: Hashable, BitwiseCopyable {
    
    case lessThan
    case greaterThan
    case lessThanOrEqual
    case greaterThanOrEqual
    case equal

}