BlogCustomLayout.swift

// MARK: CustomLayoutInvalidationContext
class CustomLayoutInvalidationContext: UICollectionViewLayoutInvalidationContext {
    var invalidatedBecauseOfBoundsChange: Bool = false
}

// MARK: CustomCollectionViewLayout
class CustomCollectionViewLayout: UICollectionViewLayout {
    
    private enum ContentUpdateValue{
        case fixed(value: CGFloat)
        case offset(value: CGFloat)
    }
    
    var estimatedRowHeight: CGFloat = 80
    
    var estimatedInlineItemWidth: CGFloat = {
        
        switch UIScreen.main.bounds.width{
        case ...400:
            return UIScreen.main.bounds.width/2
        default:
            return 200
        }
    }()
    
    init(useLargeLayout: Bool){
        numberOfColumns = useLargeLayout ? 2 : 1
        super.init()
    }
    
    required init?(coder: NSCoder) {
        fatalError("init(coder:) has not been implemented")
    }
    
    private var layoutAttributesForItems: [IndexPath:CustomCollectionViewLayoutAttribute] = [:]
    
    private var contentHeight: CGFloat = 0
    private var contentWidth: CGFloat {
        guard let collectionView = collectionView else {
            return 0
        }
        
        let insets = collectionView.contentInset
        return collectionView.bounds.width - (insets.left + insets.right)
    }
    
    private var numberOfColumns: Int
    private var inlineColumns: Int{
        return Int(floor((contentWidth/estimatedInlineItemWidth)))
    }
    
    private var cellWidth: CGFloat {
        guard let collectionView = collectionView else {
            return 0
        }
        return floor((collectionView.frame.width - CGFloat(numberOfColumns - 1)) / CGFloat(numberOfColumns) - 1)
    }
    
    override var collectionViewContentSize: CGSize {
        return CGSize(width: contentWidth, height: contentHeight)
    }
    
    override class var invalidationContextClass: AnyClass {
        return CustomLayoutInvalidationContext.self
    }
    
//     MARK: Methods
    override func invalidateLayout(with context: UICollectionViewLayoutInvalidationContext) {
        let invalidationContext = context as! CustomLayoutInvalidationContext
        
        if invalidationContext.invalidatedBecauseOfBoundsChange || invalidationContext.invalidateEverything {
            layoutAttributesForItems.removeAll()
        }
        super.invalidateLayout(with: invalidationContext)
        
    }
    
    override func prepare() {
        guard layoutAttributesForItems.isEmpty else {
            return
        }
        initialLayout()
    }
    
    override func layoutAttributesForElements(in rect: CGRect) -> [UICollectionViewLayoutAttributes]? {
        var visibleLayoutAttributes: [UICollectionViewLayoutAttributes] = []
        
        for attributes in layoutAttributesForItems.values {
            if attributes.frame.intersects(rect) {
                visibleLayoutAttributes.append(attributes)
            }
        }
        return visibleLayoutAttributes
    }
    
    override func shouldInvalidateLayout(forBoundsChange newBounds: CGRect) -> Bool {
        guard let collectionView = collectionView else {
            return false
        }
        
        guard newBounds.size != collectionView.frame.size else {
            return false
        }
        return true
    }
    
    override func invalidationContext(forBoundsChange newBounds: CGRect) -> UICollectionViewLayoutInvalidationContext {
        let context = super.invalidationContext(forBoundsChange: newBounds) as! CustomLayoutInvalidationContext
        context.invalidatedBecauseOfBoundsChange = true
        return context
    }
    
    override func shouldInvalidateLayout(forPreferredLayoutAttributes preferredAttributes: UICollectionViewLayoutAttributes, withOriginalAttributes originalAttributes: UICollectionViewLayoutAttributes) -> Bool {
        
        guard let originalProductLayoutAttribute = originalAttributes as? CustomCollectionViewLayoutAttribute else {
            return false
        }
        
        switch originalProductLayoutAttribute.collectionPosition{
        case .inline:
            if (preferredAttributes.frame.size.height == originalAttributes.frame.size.height && originalAttributes.frame.size.width == estimatedInlineItemWidth ){
                return false
            }
        default:
            
            if (preferredAttributes.frame.size.height == originalAttributes.frame.size.height && originalAttributes.frame.size.width == cellWidth){
                return false
            }
        }
        return true
    }
    
    
    override func invalidationContext(forPreferredLayoutAttributes preferredAttributes: UICollectionViewLayoutAttributes, withOriginalAttributes originalAttributes: UICollectionViewLayoutAttributes) -> UICollectionViewLayoutInvalidationContext {
        
        let context = super.invalidationContext(forPreferredLayoutAttributes: preferredAttributes, withOriginalAttributes: originalAttributes) as! CustomLayoutInvalidationContext
        
        let contentHeightAdjustment: CGFloat = preferredAttributes.frame.size.height - originalAttributes.frame.size.height
        
        let attributes = layoutAttributesForItems[originalAttributes.indexPath]!
        
        if attributes.collectionPosition == .inline{
            attributes.frame.size.width = estimatedInlineItemWidth
        }else{
            attributes.frame.size.width = cellWidth
        }
        attributes.frame.size.height += contentHeightAdjustment
        
        layoutAttributesForItems[originalAttributes.indexPath] = attributes
        
        invalidateContext(byAttributeAdjustment: originalAttributes, in: attributes, contentHeightAdjustment, context)
        collectionViewHeight(from: getLastRowAttributes(attributes: layoutAttributesForItems), context: context)
        
        return context
    }
    
    func getLastRowAttributes(attributes: [IndexPath : CustomCollectionViewLayoutAttribute]) -> [CustomCollectionViewLayoutAttribute]{
        
        guard let maxIndexPath = attributes.keys.max() else {
            return []
        }
        var lastIndexes: [IndexPath] = []
        
        if maxIndexPath.row > 0 {
            lastIndexes.append(IndexPath(row: maxIndexPath.row - 1, section: maxIndexPath.section))
        }
        lastIndexes.append(maxIndexPath)
        
        return lastIndexes.compactMap({ (indexes) -> CustomCollectionViewLayoutAttribute? in
            return attributes[indexes]
        })
    }
    
    private func collectionViewHeight(from lastRowCellAttributes: [CustomCollectionViewLayoutAttribute], context: CustomLayoutInvalidationContext) {
        let maxColumn = lastRowCellAttributes.max { (attribute1, attribute2) -> Bool in
            attribute1.frame.maxY < attribute2.frame.maxY
        }
        
        let diff =  maxColumn!.frame.maxY - contentHeight
        
        guard diff != 0 else {
            return
        }
        
        contentHeight = maxColumn!.frame.maxY
        
        context.contentSizeAdjustment = CGSize(width: 0, height: diff)
    }
    
    override func layoutAttributesForItem(at indexPath: IndexPath) -> UICollectionViewLayoutAttributes? {
        return layoutAttributesForItems[indexPath]
    }
    
    func initialLayout() {
        layoutAttributesForItems = [:]
        
        guard let collectionView = collectionView,
            let dataSource = collectionView.dataSource else {
                return
        }
        
        contentHeight = collectionView.contentInset.top
        
        let numberOfSections = dataSource.numberOfSections?(in: collectionView) ?? 1
        let columnWidth = contentWidth / CGFloat(numberOfColumns)
        
        
        var yOffset = [CGFloat](repeating: 0, count: numberOfColumns)
        var xColumnOffset = [CGFloat]()
        for column in 0 ..< numberOfColumns {
            xColumnOffset.append(CGFloat(column) * columnWidth)
        }
        
        var xInlineOffset = [CGFloat]()
        var yInlineOffset = [CGFloat](repeating: 0, count: inlineColumns)
        for column in 0 ..< inlineColumns {
            xInlineOffset.append(CGFloat(column) * estimatedInlineItemWidth)
        }
        
        for section in (0..<numberOfSections) {
            
            //set inline offset to last y offset value to continue the array.
            if  getCollectionPosition(for: IndexPath(row:0,section: section)) == .inline{
                yInlineOffset.removeAll()
                let firstColumnYValue: CGFloat = yOffset[0]
                for _ in 0 ..< inlineColumns{
                    yInlineOffset.append(CGFloat(firstColumnYValue))
                }
            }
            
            for row in (0..<dataSource.collectionView(collectionView, numberOfItemsInSection: section)) {
                
                let indexPath = IndexPath(row: row, section: section)
                let itemPosition: CollectionPosition = getCollectionPosition(for: indexPath)
                let estimatedHeightForItem: CGFloat
                estimatedHeightForItem = estimatedRowHeight
                
                let xValue: CGFloat
                let yValue: CGFloat
                let cellWidth: CGFloat
                var column = 0
                
                switch itemPosition{
                case .inline:
                    let inlinePosition = row % inlineColumns
                    
                    xValue = xInlineOffset[inlinePosition]
                    yValue = yInlineOffset[inlinePosition]
                    cellWidth = estimatedInlineItemWidth
                case .start,.end:
                    column = (numberOfColumns > 1 && itemPosition == .end) ? 1 : 0
                    
                    xValue = xColumnOffset[column]
                    yValue = yOffset[column]
                    cellWidth = self.cellWidth
                }
                let frame = CGRect(x: xValue, y: yValue, width: cellWidth, height: estimatedHeightForItem)
                let insetFrame = frame.insetBy(dx: collectionView.contentInset.left, dy: collectionView.contentInset.right)
                
                let attributes = CustomCollectionViewLayoutAttribute(forCellWith: indexPath)
                attributes.frame = insetFrame
                attributes.collectionPosition = itemPosition
                
                layoutAttributesForItems[IndexPath(row:row, section: section)] = attributes
                contentHeight = max(contentHeight, frame.maxY)
                
                if itemPosition == .inline{
                    let inlinePosition = row % inlineColumns
                    yInlineOffset[inlinePosition] = yInlineOffset[inlinePosition] + estimatedHeightForItem
                    
                    for columnIndex in  0 ..< yOffset.count{
                        yOffset[columnIndex] = yInlineOffset[inlinePosition]
                    }
                    
                    attributes.column = inlinePosition
                    column = 0
                }else{
                    attributes.column = column
                    yOffset[column] = yOffset[column] + estimatedHeightForItem
                }
            }
        }
        contentHeight += collectionView.contentInset.bottom
    }
    
    fileprivate func getCollectionPosition(for indexPath: IndexPath) -> CollectionPosition{
        let collectionPosition: CollectionPosition
        
        if let collectionView = collectionView,
            let delegate = collectionView.delegate as? CustomCollectionViewDelegate {
            collectionPosition = delegate.collectionView(sectionPosition: collectionView, for: indexPath)
        }else {
            collectionPosition = .start
        }
        return collectionPosition
    }
    
    fileprivate func invalidateContext(byAttributeAdjustment originalAttributes: UICollectionViewLayoutAttributes, in attributes: CustomCollectionViewLayoutAttribute, _ contentHeightAdjustment: CGFloat, _ context: CustomLayoutInvalidationContext) {
        var invalidationIndexes: [IndexPath] = []
        invalidationIndexes.append(attributes.indexPath)
        
        for (indexPath, layoutAttributesForItem) in layoutAttributesForItems {
            
            if indexPath <= originalAttributes.indexPath{
                continue
            }
            
            if layoutAttributesForItem.collectionPosition == .inline{
                
                if layoutAttributesForItem.column == 0,
                    layoutAttributesForItem.column == attributes.column{
                    updateAttributes(forIndexPath: [indexPath], byContentValue: .offset(value: contentHeightAdjustment))
                    invalidationIndexes.append(indexPath)
                    
                    let rowsForUpdate = getInlineColumsIndexes(forRow: indexPath.row/inlineColumns)
                    updateAttributes(forIndexPath: rowsForUpdate, byContentValue: .fixed(value: layoutAttributesForItem.frame.origin.y))
                    invalidationIndexes.append(contentsOf: rowsForUpdate)
                }
            }else{
                updateColumnsAttributes(layoutAttributesForItem, attributes, indexPath, contentHeightAdjustment, &invalidationIndexes)
            }
        }
        print("Invalidation: invalidated indexes: \(invalidationIndexes)")
        context.invalidateItems(at: invalidationIndexes)
    }
    /**
     - Returns Inline columns that are not present in *normal* layout. This it useful for telling that fields to update its value.
     */
    private func getInlineColumsIndexes(forRow row: Int, includeFirst: Bool = false) -> [IndexPath]{
        
        return layoutAttributesForItems.filter { (item) -> Bool in
            let shouldInclude: Bool = includeFirst ? true : item.key.row % inlineColumns != 0
            
            return item.value.collectionPosition == .inline && item.key.row/inlineColumns == row && shouldInclude
        }.map { (keyValue) -> IndexPath in
            return keyValue.key
        }
    }
    
    private func updateAttributes(forIndexPath indexPaths: [IndexPath], byContentValue contentValue: ContentUpdateValue){
        
        for indexPath in indexPaths{
            switch contentValue {
            case .fixed(value: let value):
                layoutAttributesForItems[indexPath]?.frame.origin.y = value
            case .offset(value: let value):
                layoutAttributesForItems[indexPath]?.frame.origin.y += value
            }
        }
    }
    
    fileprivate func updateColumnsAttributes(_ layoutAttributesForItem: CustomCollectionViewLayoutAttribute, _ attributes: CustomCollectionViewLayoutAttribute, _ indexPath: IndexPath, _ contentHeightAdjustment: CGFloat, _ invalidationIndexes: inout [IndexPath]) {
        
        if layoutAttributesForItem.column == attributes.column{
            
            updateAttributes(forIndexPath: [indexPath], byContentValue: .offset(value: contentHeightAdjustment))
            invalidationIndexes.append(indexPath)
        }
    }
}