BatchedSequence

BatchedSequence.swift

import Foundation

/// A sequence whose elements are emitted in batches of the specified size
struct BatchedSequence<Base: Sequence>: Sequence {
    private let size: Int
    private let base: Base
    
    init(base: Base, size: Int) {
        self.size = size
        self.base = base
    }
    func makeIterator() -> AnyIterator<[Base.Element]> {
        var iterator = base.makeIterator()
        var current: [Base.Element] = []
        
        return .init {
            while current.count < self.size, let value = iterator.next() {
                current.append(value)
            }
            
            defer { current = [] }
            return current.isEmpty ? nil : current
        }
    }
}

extension Sequence {
    /// Return `self` as a `BatchedSequence`
    ///
    /// - Parameter size: The maximum size of batches to create
    /// - Returns: A `Sequence` whose elements are batches of elements from `self`
    func batched(by size: Int) -> BatchedSequence<Self> {
        return BatchedSequence(base: self, size: size)
    }
}