Efficiently split a sorted sequence with a list of sorted bounds

BoundsSegmentedSequence.swift

extension RandomAccessCollection {
    /// A sequence that partitions a sorted collection into non-overlapping segments based on a sorted sequence of boundaries.
    ///
    /// If N bounds are provided, the resulting sequence contains N + 1 segments, since the segment before the first boundary
    /// and the segment after the last boundary are both included.
    ///
    /// - Parameter bounds: A sequence of bounds used to partition the collection.
    /// - Parameter areInIncreasingOrder: A predicate that determines whether an element should be included
    ///   in the segment before the boundary.
    /// - Returns: A sequence of collection subseqences partitioned by the given bounds.
    /// - Warning: The base collection must be sorted according to the partitioning logic, or the behavior is undefined!
    func segmented<Bound, Bounds: Sequence>(
        by bounds: Bounds,
        using areInIncreasingOrder: @escaping (Element, Bound) -> Bool
    ) -> BoundsSegmentedSequence<Self, Bound, Bounds> where Bounds.Element == Bound {
        return BoundsSegmentedSequence(base: self, bounds: bounds, areInIncreasingOrder: areInIncreasingOrder)
    }

    /// A sequence that partitions a sorted collection into non-overlapping segments based on a sorted sequence of boundaries.
    ///
    /// If N bounds are provided, the resulting sequence contains N - 1 segments, since all elements must be contained between the first and last bound.
    ///
    /// - Parameter bounds: A sequence of bounds used to partition the collection.
    /// - Parameter areInIncreasingOrder: A predicate that determines whether an element should be included
    ///   in the segment before the boundary.
    /// - Returns: A sequence of collection subseqences partitioned by the given bounds.
    /// - Warning: The base collection must be sorted according to the partitioning logic, or the behavior is undefined.
    func segmented<Bound, Bounds: RandomAccessCollection>(
        between bounds: Bounds,
        using areInIncreasingOrder: @escaping (Element, Bound) -> Bool
    ) -> BoundsSegmentedSequence<Self, Bound, Bounds.SubSequence> where Bounds.Element == Bound {
        precondition(bounds.count >= 2, "Elements cannot exist between fewer than 2 bounds.")
        if let first {
            precondition(!areInIncreasingOrder(first, bounds.first!), "Element exists before the first bound.")
        }
        if let last {
            precondition(areInIncreasingOrder(last, bounds.last!), "Element exists after the last bound.")
        }
        return segmented(by: bounds.dropFirst().dropLast(), using: areInIncreasingOrder)
    }
}

struct BoundsSegmentedSequence<Base: RandomAccessCollection, Bound, Bounds: Sequence>: Sequence
where Bounds.Element == Bound {
    var base: Base
    var bounds: Bounds
    var areInIncreasingOrder: (Base.Element, Bound) -> Bool

    struct Iterator: IteratorProtocol {
        var base: Base.SubSequence
        var boundsIterator: Bounds.Iterator
        var areInIncreasingOrder: (Base.Element, Bound) -> Bool
        var isDone: Bool = false
        
        mutating func next() -> Base.SubSequence? {
            guard !isDone else { return nil }
            
            let index: Base.Index
            if let bound = boundsIterator.next() {
                index = base.firstIndex(where: { !areInIncreasingOrder($0, bound) }) ?? base.endIndex
            } else {
                isDone = true
                index = base.endIndex
            }
            
            defer { base = base[index...] }
            return base[..<index]
        }
    }
    
    func makeIterator() -> Iterator {
        .init(base: base[...], boundsIterator: bounds.makeIterator(), areInIncreasingOrder: areInIncreasingOrder)
    }
}