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

Playground.swift

import Foundation
import UIKit
import XCPlayground

XCPlaygroundPage.currentPage.needsIndefiniteExecution = true



class Layout: UICollectionViewLayout {
    private var attributes: [[UICollectionViewLayoutAttributes]] = []
    
    private var topmostIndexPathBeforeUpdates: NSIndexPath? = 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 prepareForCollectionViewUpdates(updateItems: [UICollectionViewUpdateItem]) {
        super.prepareForCollectionViewUpdates(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 { (section) -> [UICollectionViewLayoutAttributes] in
            return section
        }.sort { (a, b) -> Bool in
            return fabs(a.center.y - collectionView.contentOffset.y) < fabs(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 = NSIndexPath(forItem: indexPath!.item+1, inSection: 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 prepareLayout() {
        super.prepareLayout()
        guard let collectionView = collectionView else { return }
        var globalIndex = 0
        attributes = (0..<collectionView.numberOfSections()).map({ (section) -> [UICollectionViewLayoutAttributes] in
            return (0..<collectionView.numberOfItemsInSection(section)).map({ (item) -> UICollectionViewLayoutAttributes in
                let l = UICollectionViewLayoutAttributes(forCellWithIndexPath: NSIndexPath(forItem: item, inSection: 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 layoutAttributesForItemAtIndexPath(indexPath: NSIndexPath) -> UICollectionViewLayoutAttributes? {
        return attributes[indexPath.section][indexPath.item]
    }
    
    override func layoutAttributesForElementsInRect(rect: CGRect) -> [UICollectionViewLayoutAttributes]? {
        return attributes.flatMap({ (attributes) -> [UICollectionViewLayoutAttributes] in
            return attributes
        })
    }
    
    override func targetContentOffsetForProposedContentOffset(proposedContentOffset: CGPoint) -> CGPoint {
        // If we have a cached topmost index path, use it.
        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: CGRectZero)
    
    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".characters.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, cellForItemAtIndexPath indexPath: NSIndexPath) -> UICollectionViewCell {
        let cell = collectionView.dequeueReusableCellWithReuseIdentifier("cell", forIndexPath: indexPath) as! TestCell
        cell.backgroundColor = UIColor.greenColor()
        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.registerClass(TestCell.self, forCellWithReuseIdentifier: "cell")
cv.dataSource = dataSource


XCPlaygroundPage.currentPage.liveView = cv

// Wait a second after the initial load
let t = dispatch_time(DISPATCH_TIME_NOW, Int64(NSEC_PER_SEC) * 1)
dispatch_after(t, dispatch_get_main_queue()) { () -> Void in
    
    // Make all of the currently visible cells yellow
    for c in cv.visibleCells() {
        c.backgroundColor = UIColor.yellowColor()
    }
    
    // Then insert a cell at the top.
    cv.performBatchUpdates({ () -> Void in
        dataSource.labels.insert("New element!", atIndex: 0)
        cv.insertItemsAtIndexPaths([NSIndexPath(forItem: 0, inSection: 0)])
        }, completion: nil)
}

// Manually scroll to the top to make sure the inserted item is really there
let t2 = dispatch_time(DISPATCH_TIME_NOW, Int64(NSEC_PER_SEC) * 2)
dispatch_after(t2, dispatch_get_main_queue()) { () -> Void in
    cv.scrollToItemAtIndexPath(NSIndexPath(forItem: 0, inSection: 0), atScrollPosition: .Top, animated: true)
}