Collection+Split.swift

extension Collection {
    /// Returns an array of subsequences of maximum size `chunkSize`.
    ///
    /// For example:
    ///
    ///     // ["ABCD", "EFGH", "IJ"]
    ///     "ABCDEFGHIJ".split(chunkSize: 4)
    ///
    /// - parameter chunkSize: the maximum size of the returned subsequences.
    /// - precondition: chunkSize > 0
    func split(chunkSize: Int) -> [SubSequence] {
        precondition(chunkSize > 0, "chunkSize must be greater than zero")
        var result: [SubSequence] = []
        result.reserveCapacity((count + chunkSize - 1) / chunkSize)
        var chunkStart = self.startIndex
        while chunkStart != endIndex {
            let chunkEnd = index(chunkStart, offsetBy: chunkSize, limitedBy: endIndex) ?? endIndex
            result.append(self[chunkStart..<chunkEnd])
            chunkStart = chunkEnd
        }
        return result
    }
}

// Custom implementation for RandomAccessCollection, because it has O(1)
// `index(offsetBy:limitedBy:)`.
extension RandomAccessCollection where Index: Strideable, Index.Stride == Int {
    /// Returns an array of subsequences of maximum size `chunkSize`.
    ///
    /// For example:
    ///
    ///     // ["ABCD", "EFGH", "IJ"]
    ///     "ABCDEFGHIJ".split(chunkSize: 4)
    ///
    /// - parameter chunkSize: the maximum size of the returned subsequences.
    /// - precondition: chunkSize > 0
    func split(chunkSize: Int) -> [SubSequence] {
        precondition(chunkSize > 0, "chunkSize must be greater than zero")
        var result: [SubSequence] = []
        result.reserveCapacity((count + chunkSize - 1) / chunkSize)
        var chunkStart = self.startIndex
        while chunkStart != endIndex {
            let chunkEnd = index(chunkStart, offsetBy: chunkSize, limitedBy: endIndex) ?? endIndex
            result.append(self[chunkStart..<chunkEnd])
            chunkStart = chunkEnd
        }
        return result
    }
}

CollectionSplitTests.swift

import XCTest

class CollectionSplitTests: XCTestCase {
    func testCollectionSplit() {
        XCTAssertEqual((0..<6).split(chunkSize: 1), [0..<1, 1..<2, 2..<3, 3..<4, 4..<5, 5..<6])
        XCTAssertEqual((0..<6).split(chunkSize: 2), [0..<2, 2..<4, 4..<6])
        XCTAssertEqual((0..<6).split(chunkSize: 3), [0..<3, 3..<6])
        XCTAssertEqual((0..<6).split(chunkSize: 4), [0..<4, 4..<6])
        XCTAssertEqual((0..<6).split(chunkSize: 5), [0..<5, 5..<6])
        XCTAssertEqual((0..<6).split(chunkSize: 6), [0..<6])
        XCTAssertEqual((0..<6).split(chunkSize: 7), [0..<6])
        
        XCTAssertEqual(Array(0..<6).split(chunkSize: 1), [[0], [1], [2], [3], [4], [5]])
        XCTAssertEqual(Array(0..<6).split(chunkSize: 2), [[0, 1], [2, 3], [4, 5]])
        XCTAssertEqual(Array(0..<6).split(chunkSize: 3), [[0, 1, 2], [3, 4, 5]])
        XCTAssertEqual(Array(0..<6).split(chunkSize: 4), [[0, 1, 2, 3], [4, 5]])
        XCTAssertEqual(Array(0..<6).split(chunkSize: 5), [[0, 1, 2, 3, 4], [5]])
        XCTAssertEqual(Array(0..<6).split(chunkSize: 6), [[0, 1, 2, 3, 4, 5]])
        XCTAssertEqual(Array(0..<6).split(chunkSize: 7), [[0, 1, 2, 3, 4, 5]])
        
        XCTAssertEqual("Hello!".split(chunkSize: 1), ["H", "e", "l", "l", "o", "!"])
        XCTAssertEqual("Hello!".split(chunkSize: 2), ["He", "ll", "o!"])
        XCTAssertEqual("Hello!".split(chunkSize: 3), ["Hel", "lo!"])
        XCTAssertEqual("Hello!".split(chunkSize: 4), ["Hell", "o!"])
        XCTAssertEqual("Hello!".split(chunkSize: 5), ["Hello", "!"])
        XCTAssertEqual("Hello!".split(chunkSize: 6), ["Hello!"])
        XCTAssertEqual("Hello!".split(chunkSize: 7), ["Hello!"])
        
        XCTAssertEqual([Int]().split(chunkSize: 1), [])
        XCTAssertEqual([Int]().split(chunkSize: 2), [])
        
        XCTAssertEqual("Hello!".dropFirst(1).split(chunkSize: 2), ["el", "lo", "!"])
        XCTAssertEqual("Hello!".reversed().split(chunkSize: 4), [["!", "o", "l", "l"], ["e", "H"]])
    }
}