Maintaining visible scroll position while inserting items in a UICollectionView (Swift 4.0 playground)

Playground.swift

import Foundation
import UIKit
import XCPlayground
import PlaygroundSupport

PlaygroundPage.current.needsIndefiniteExecution = true

class Layout: UICollectionViewLayout {
    private var attributes: [[UICollectionViewLayoutAttributes]] = []
    
    private var topmostIndexPathBeforeUpdates: IndexPath? = nil
    private var originOfTopmostIndexPath: CGFloat = 0.0
    
    
    // This is the important part: the collection view will always let the layout know about
    // upcoming changes. You can use this method to take any notes about the current state
    // of things.
    override func prepare(forCollectionViewUpdates updateItems: [UICollectionViewUpdateItem]) {
        super.prepare(forCollectionViewUpdates: updateItems)
        
        guard let collectionView = collectionView else { fatalError() }

        // Get the layout attributes of the item closest to the top of the collection view
        let topmostLayoutAttributes = attributes.flatMap { (attributes) -> [UICollectionViewLayoutAttributes] in
            return attributes
        }.sorted{ (a, b) -> Bool in
            return abs(a.center.y - collectionView.contentOffset.y) < abs(b.center.y - collectionView.contentOffset.y)
        }.first
        
        
        // Run through the updateItems to see if the indexPath will change. This is not comprehensive,
        // you'll need to handle all of the other potential actions.
        var indexPath = topmostLayoutAttributes?.indexPath
        for item in updateItems {
            guard indexPath != nil else { break }
            switch item.updateAction {
            case .insert where item.indexPathAfterUpdate!.item <= indexPath!.item:
                indexPath = IndexPath(item: indexPath!.item+1, section: indexPath!.section)
            default:
                // Handle the rest of the cases here
                break
            }
        }
        
        // Remember the position
        topmostIndexPathBeforeUpdates = indexPath
        originOfTopmostIndexPath = topmostLayoutAttributes?.frame.origin.y ?? 0.0
    }
    
    override func prepare() {
        super.prepare()
        guard let collectionView = collectionView else { return }
        var globalIndex = 0
        attributes = (0..<collectionView.numberOfSections).map({ (section) -> [UICollectionViewLayoutAttributes] in
            return (0..<collectionView.numberOfItems(inSection: section)).map({ (item) -> UICollectionViewLayoutAttributes in
                let l = UICollectionViewLayoutAttributes(forCellWith: IndexPath(item: item, section: section))
                l.frame = CGRect(x: 0.0, y: CGFloat(globalIndex) * 60.0, width: collectionView.bounds.width, height: 44.0)
                globalIndex += 1
                return l
            })
        })
    }
    
    override func layoutAttributesForItem(at indexPath: IndexPath) -> UICollectionViewLayoutAttributes? {
        return attributes[indexPath.section][indexPath.item]
    }
    
    override func layoutAttributesForElements(in rect: CGRect) -> [UICollectionViewLayoutAttributes]? {
        return attributes.flatMap({ (attributes) -> [UICollectionViewLayoutAttributes] in
            return attributes
        })
    }
    
    override func targetContentOffset(forProposedContentOffset proposedContentOffset: CGPoint) -> CGPoint {
        if let topmost = topmostIndexPathBeforeUpdates {
            let top = attributes[topmost.section][topmost.item].frame.origin.y
            return CGPoint(x: proposedContentOffset.x, y: top)
        }
        return proposedContentOffset
    }
}


class TestCell: UICollectionViewCell {
    let label = UILabel(frame: CGRect.zero)
    
    override init(frame: CGRect) {
        super.init(frame: frame)
        contentView.addSubview(label)
    }

    required init?(coder aDecoder: NSCoder) {
        fatalError("init(coder:) has not been implemented")
    }
    
    override func layoutSubviews() {
        super.layoutSubviews()
        label.frame = contentView.bounds
    }
}


class DataSource: NSObject, UICollectionViewDataSource {
    var labels = "ABCDEFGHIJKLMNOPQRSTUVWXYZ123456789".map({ (c) -> String in
        return "\(c)"
    })
    
    func numberOfSectionsInCollectionView(collectionView: UICollectionView) -> Int {
        return 1
    }
    
    func collectionView(_ collectionView: UICollectionView, numberOfItemsInSection section: Int) -> Int {
        return labels.count
    }
    
    func collectionView(_ collectionView: UICollectionView, cellForItemAt indexPath: IndexPath) -> UICollectionViewCell {
        let cell = collectionView.dequeueReusableCell(withReuseIdentifier: "cell", for: indexPath) as! TestCell
        cell.backgroundColor = UIColor.green
        cell.label.text = labels[indexPath.item]
        return cell
    }
}



let cv = UICollectionView(frame: CGRect(x: 0.0, y: 0.0, width: 400.0, height: 400.0), collectionViewLayout: Layout())
let dataSource = DataSource()
cv.register(TestCell.self, forCellWithReuseIdentifier: "cell")
cv.dataSource = dataSource


PlaygroundPage.current.liveView = cv

// Wait a second after the initial load
DispatchQueue.main.asyncAfter(deadline: .now() + 1) {
    // Make all of the currently visible cells yellow
    for c in cv.visibleCells {
        c.backgroundColor = UIColor.yellow
    }
    
    // Then insert a cell at the top.
    cv.performBatchUpdates({ () -> Void in
        dataSource.labels.insert("New element!", at: 0)
        cv.insertItems(at: [IndexPath(item: 0, section: 0)])
        }, completion: { (animated) in
            cv.scrollToItem(at: IndexPath(item: 0, section: 0), at: .top, animated: true)
            PlaygroundPage.current.finishExecution()
        }
    )
}