List and OptimizedList, a sectioned Collection.

List.swift

struct List<Element>: Sequence, Collection where Element: Hashable {
    private var contents: [[Element]]
    
    init(_ sections: [[Element]] = [[]]) {
        contents = sections
    }
    
    
    // MARK: - Sequence
    
    typealias Iterator = AnyIterator<Element>
    
    func makeIterator() -> AnyIterator<Element> {
        let end = endIndex
        var idx = startIndex
        
        return AnyIterator<Element> {
            var element: Element?
            
            if idx < end {
                element = self[idx]
                idx = self.index(after: idx)
            }
            
            return element
        }
    }
    
    
    // MARK: - Collection
    
    typealias Index = ListIndex
    
    var startIndex: ListIndex {
        return ListIndex(section: contents.startIndex, row: contents[contents.startIndex].startIndex)
    }
    
    var endIndex: ListIndex {
        return ListIndex(section: contents.endIndex, row: 0)
    }
    
    func index(after i: ListIndex) -> ListIndex {
        //                          |*********|--> This is my concern.
        guard (i.row + 1) < contents[i.section].endIndex else {
            return ListIndex(section: i.section + 1, row: 0)
        }
        
        return ListIndex(section: i.section, row: i.row + 1)
    }
    
    subscript(position: ListIndex) -> Element {
        return contents[position.section][position.row]
    }
}

ListIndex.swift

struct ListIndex: Comparable, Hashable {
    var section: Int
    var row: Int
    
    init(section: Int, row: Int) {
        self.section = section
        self.row = row
    }
    
    // MARK: Hashable
    
    var hashValue: Int {
        return section.hashValue ^ row.hashValue
    }
    
    
    // MARK: Equatable
    
    static func ==(lhs: ListIndex, rhs: ListIndex) -> Bool {
        return lhs.section == rhs.section && lhs.row == rhs.row
    }
    
    
    // MARK: Comparable
    
    static func <(lhs: ListIndex, rhs: ListIndex) -> Bool {
        if lhs.section == rhs.section {
            return lhs.row < rhs.row
        }
        
        return lhs.section < rhs.section
    }
}

OptimizedList.swift

struct OptimizedList<Element>: Sequence, BidirectionalCollection where Element: Hashable {
    
    init() {
        self.init([])
    }
    
    init<S: Sequence>(_ section: S) where S.Element == Element {
        self.init([Array(section)])
    }
    
    init(_ sections: [[Element]]) {
        contents = sections
        endIndices = sections.endListIndices()
    }
    
    private var contents: [[Element]] {
        didSet {
            endIndices = contents.endListIndices()
        }
    }
    
    private var endIndices: Set<ListIndex>

    
    // MARK: - Inserting
    
    mutating func prepend<S: Sequence>(section newItems: S) where S.Element == Element {
        insert(section: newItems, at: .start)
    }
    
    mutating func append<S: Sequence>(section newItems: S) where S.Element == Element {
        insert(section: newItems, at: .end)
    }
    
    mutating func insert<S: Sequence>(section newItems: S, at section: Int) where S.Element == Element {
        insert(section: newItems, at: .index(section))
    }
    
    mutating func prepend<S: Sequence>(items newItems: S, at section: Int) where S.Element == Element {
        insert(items: newItems, into: section, at: .start)
    }
    
    mutating func append<S: Sequence>(items newItems: S, at section: Int) where S.Element == Element {
        insert(items: newItems, into: section, at: .end)
    }
    
    mutating func insert<S: Sequence>(items newItems: S, at index: ListIndex) where  S.Element == Element {
        insert(items: newItems, into: index.section, at: .index(index.row))
    }
    
    
    // MARK: - Removing
    
    mutating func remove(at section: Int) -> [Element] {
        return contents.remove(at: section)
    }
    
    mutating func removeFirst(_ n: Int, from section: Int) {
        contents[section].removeFirst(n)
    }
    
    mutating func removeFirst(from section: Int) -> Element {
        return contents[section].removeFirst()
    }
    
    mutating func removeLast(_ n: Int, from section: Int) {
        contents[section].removeLast(n)
    }
    
    mutating func removeLast(from section: Int) -> Element {
        return contents[section].removeLast()
    }
    
    
    // MARK: Sequence
    
    typealias Iterator = AnyIterator<Element>
    
    func makeIterator() -> AnyIterator<Element> {
        let end = endIndex
        var index = startIndex
        
        return AnyIterator {
            var element: Element?
            
            if index < end {
                element = self[index]
                index = self.index(after: index)
            }
            
            return element
        }
    }
    
    // MARK: Collection
    
    typealias Index = ListIndex
    
    var startIndex: ListIndex {
        return ListIndex(section: contents.startIndex, row: contents[contents.startIndex].startIndex)
    }
    
    var endIndex: ListIndex {
        return ListIndex(section: contents.endIndex, row: 0)
    }
    
    func index(after i: ListIndex) -> ListIndex {
        var next = ListIndex(section: i.section, row: i.row + 1)
        
        if endIndices.contains(next) {
            next.section += 1
            next.row = 0
        }
        
        return next
    }
    
    func index(before i: ListIndex) -> ListIndex {
        var prev = ListIndex(section: i.section, row: i.row - 1)
        
        if prev.row < 0 {
            guard let endIndex = endIndices.first(where: { $0.section == (prev.section - 1) }) else {
                fatalError("section (\(prev.section - 1)) out of bounds")
            }
            
            prev = endIndex
            prev.row -= 1
        }
        
        return prev
    }
    
    subscript(position: ListIndex) -> Element {
        return contents[position.section][position.row]
    }
    
    
    // MARK: - Private
    
    /// Index location for new items
    private enum Insert {
        case start
        case index(Int)
        case end
    }
    
    private mutating func insert<S: Sequence>(section newItems: S, at location: Insert) where S.Element == Element {
        let index: Int
        
        switch location {
        case .start:
            index = contents.startIndex
        case .index(let i):
            index = i
        case .end:
            index = contents.endIndex
        }
        
        contents.insert(Array(newItems), at: index)
    }
    
    private mutating func insert<S: Sequence>(items: S, into section: Int, at location: Insert) where S.Element == Element {
        let index: Int
        
        switch location {
        case .start:
            index = contents[section].startIndex
        case .index(let i):
            index = i
        case .end:
            index = contents[section].endIndex
        }
        
        contents[section].insert(contentsOf: Array(items), at: index)
    }
}


private extension Collection where Index == Int, Element: Collection, Element.Index == Index {
    func endListIndices() -> Set<ListIndex> {
        var result = Set<ListIndex>()
        var index = startIndex
        
        for element in self {
            result.insert(ListIndex(section: index, row: element.endIndex))
            index += 1
        }
        
        return result
    }
}