A too complicated implementation of NestedDiff (diffing collections containing collections (i.e. arrays containing arrays).

NestedDiff.swift

public struct NestedDiff: DiffProtocol {

    public enum Element {
        case deleteSection(Int)
        case insertSection(Int)
        case deleteRow(Int, section: Int)
        case insert(row: Int, section: Int)
    }

    /// Returns the position immediately after the given index.
    ///
    /// - Parameter i: A valid index of the collection. `i` must be less than
    ///   `endIndex`.
    /// - Returns: The index value immediately after `i`.
    public func index(after i: Int) -> Int {
        return i + 1
    }

    public let elements: [Element]
}

public protocol Keyed {
    associatedtype KeyType: Equatable
    var key: KeyType { get }
}

public extension Collection
    where Iterator.Element: Collection,
    Iterator.Element: Keyed,
    Iterator.Element.Iterator.Element: Equatable {
    func nestedDiff(to: Self) -> NestedDiff {

        let fromKeys = flatMap { $0.key }
        let toKeys = to.flatMap { $0.key }
        let diffTraces = fromKeys.outputDiffPathTraces(to: toKeys)
        
        // Diff sections
        var elements = Diff(traces: diffTraces).map { element -> NestedDiff.Element in
            switch(element) {
                case let .delete(at):
                    return .deleteSection(at)
                case let .insert(at):
                    return .insertSection(at)
            }
        }
        
        // Diff matching sections (moves, deletions, insertions)
        let filterMatchPoints = { (trace: Trace) -> Bool in
            if case .matchPoint = trace.type() {
                return true
            }
            return false
        }
        
        // offset & section
        
        let fromMatchingSectionIndices = diffTraces
            .filter(filterMatchPoints)
            .map { $0.from.x }
        let toMatchingSectionIndices = diffTraces
            .filter(filterMatchPoints)
            .map { $0.from.y }
        
        let fromElements = fromMatchingSectionIndices
            .map { self[index(startIndex, offsetBy: IndexDistance(IntMax($0)))] }
        let toElements = toMatchingSectionIndices
            .map { to[to.index(to.startIndex, offsetBy: IndexDistance(IntMax($0)))] }

        let fromSectionCounts = fromElements.map { Int($0.count.toIntMax()) }
        let toSectionCounts = toElements.map { Int($0.count.toIntMax()) }
        let fromSectionOffsets = sectionOffsets(in: fromElements)
        let toSectionOffsets = sectionOffsets(in: toElements)
        
        let diff =
            fromElements
            .flatMap { $0 }
            .diff(
                toElements
                .flatMap { $0 }
            )
        
        var sectionDifferences: [Int] = []
        var deletedIndicesInSection: [Set<Int>] = []
        
        if var section = diff.indices.first {
            var currentSectionDiff = fromSectionCounts[0] - toSectionCounts[0]
            var deletedIndices = Set<Int>()
            
            for element in diff {
                
                switch element {
                case let .delete(at):
                    while section < fromSectionOffsets.count - 1
                        && fromSectionOffsets[section + 1] <= at {
                            sectionDifferences.append(currentSectionDiff)
                            section += 1
                            currentSectionDiff = toSectionCounts[section] - fromSectionCounts[section]
                            deletedIndicesInSection.append(deletedIndices)
                            deletedIndices = Set<Int>()
                    }
                    
                    let offset = fromSectionOffsets[section]
                    elements.append(.deleteRow(at - offset, section: fromMatchingSectionIndices[section]))
                    deletedIndices.insert(at - offset)
                    currentSectionDiff -= 1
                case let .insert(at):
                    while section < toSectionOffsets.count - 1
                        && toSectionOffsets[section + 1] <= at {
                            sectionDifferences.append(currentSectionDiff)
                            section += 1
                            currentSectionDiff = toSectionCounts[section] - fromSectionCounts[section]
                            deletedIndicesInSection.append(deletedIndices)
                            deletedIndices = Set<Int>()
                    }
                    
                    let offset = toSectionOffsets[section]
                    elements.append(.insert(row: at - offset, section: toMatchingSectionIndices[section]))
                    
                    currentSectionDiff += 1
                }
            }
            
            deletedIndicesInSection.append(deletedIndices)
            sectionDifferences.append(currentSectionDiff)
        }
        
        for section in sectionDifferences.count ..< fromSectionOffsets.count {
            deletedIndicesInSection.append(Set())
            sectionDifferences.append(toSectionCounts[section] - fromSectionCounts[section])
        }
        
        
        if var section = diffTraces.indices.first {
            
            var prevTrace: Trace?
            for trace in diffTraces {
                
                guard let prev = prevTrace,
                    case .matchPoint = prev.type(),
                    case .matchPoint = trace.type()
                    else {
                        if case .matchPoint = trace.type() {
                            prevTrace = trace
                        }
                        continue
                }
                
                let lastIndex = fromSectionCounts[section] - 1
                let numberOfElementsToShift = sectionDifferences[section]
                
                if numberOfElementsToShift > 0 {
                    // More insertions than deletions
                    
                    for index in lastIndex + 1 ... (lastIndex + numberOfElementsToShift) {
                        elements.append(.insert(row: index, section: toMatchingSectionIndices[section]))
                    }
                    
                    var elementShift = 0
                    var sectionShift = 1
                    for index in 0 ..< abs(numberOfElementsToShift) {
                        while (deletedIndicesInSection[section + sectionShift].contains(index)) {
                            elementShift += 1
                        }
                        while fromSectionOffsets.count - 1 > section + sectionShift && index + elementShift >= fromSectionOffsets[section + sectionShift + 1] {
                            sectionShift += 1
                            elementShift = 0
                        }
                        
                        elements.append(.deleteRow(index + elementShift, section: fromMatchingSectionIndices[section + sectionShift]))
                    }

                } else if numberOfElementsToShift < 0 {
                    // More deletions than insertions - take last `numberOfElementsToShift` of items and move to the next section
                    
//                    for index in 0 ..< abs(numberOfElementsToShift) {
//                        elements.append(.deleteRow(index, section: toMatchingSectionIndices[section]))
//                    }
//                    
//                    
//                    for index in (lastIndex + numberOfElementsToShift + 1) ... lastIndex  {
//                        elements.append(.insert(row: index, section: fromMatchingSectionIndices[section + 1]))
//                    }
                }
                
                section += 1
                prevTrace = trace
            }
        }
        
        return NestedDiff(elements: elements)
    }
}

func sectionOffsets<T: Collection>(in array: Array<T>) -> [Int] {
    return array
        .flatMap { $0.count }
        .dropLast()
        .reduce([0]) { prev, item in
            let prevCount = prev.last ?? 0
            return prev + [(prevCount + Int(item.toIntMax()))]
    }
}

extension NestedDiff.Element: CustomDebugStringConvertible {
    public var debugDescription: String {
        switch self {
        case let .deleteRow(row, section):
            return "DR(\(row),\(section))"
        case let .deleteSection(section):
            return "DS(\(section))"
        case let .insert(row, section):
            return "I(\(row),\(section))"
        case let .insertSection(section):
            return "IS(\(section))"        }
    }
}