ChunkedOn.swift

struct LazyChunkedOn<Base: Collection, Subject: Equatable> : Collection {
    typealias Element = Base.SubSequence
    var underlying: Base
    var projection: (Base.Element) -> Subject

    struct Index: Comparable {
        let start: Base.Index
        var subject: Subject?

        static func == (lhs: Index, rhs: Index) -> Bool { lhs.start == rhs.start }
        static func < (lhs: Index, rhs: Index) -> Bool { lhs.start < rhs.start }
    }

    var startIndex: Index { Index(start: underlying.startIndex, subject: underlying.first.map(projection)) }
    var endIndex: Index { Index(start: underlying.endIndex, subject: nil) }

    func index(after i: Index) -> Index {
        guard let subject = i.subject else { return i }
        var baseIndex = i.start
        while true {
            baseIndex = underlying.index(after: baseIndex)
            guard baseIndex != underlying.endIndex else { return self.endIndex }
            let newSubject = projection(underlying[baseIndex])
            if newSubject != subject {
                return Index(start: baseIndex, subject: newSubject)
            }
        }
    }

    subscript(position: Index) -> Base.SubSequence {
        return underlying[position.start ..< self.index(after: position).start]
    }

    struct Iterator : IteratorProtocol {
        var index: Index
        let chunkedOn: LazyChunkedOn
        mutating func next() -> Element? {
            guard index != chunkedOn.endIndex else { return nil }
            let next = chunkedOn.index(after: index)
            defer { index = next }
            return chunkedOn.underlying[index.start ..< next.start]
        }
    }

    func makeIterator() -> Iterator { Iterator(index: startIndex, chunkedOn: self) }
}

extension LazyCollectionProtocol {
    func chunked<Subject: Equatable>(on projection: @escaping (Element) -> Subject) -> LazyChunkedOn<Elements, Subject> {
        return LazyChunkedOn(underlying: self.elements, projection: projection)
    }
}