Flow layout grid menu with gap

flow-layout-class

class GridMenuLayout: UICollectionViewFlowLayout {
  private let numberOfItemsPerRow: CGFloat = 2.0
  private let spacing: CGFloat = 16.0
  private var bigButtonIndex = 2

  override func prepare() {
    super.prepare()

    sectionInset = UIEdgeInsets(
      top: spacing,
      left: spacing,
      bottom: spacing,
      right: spacing
    )

    minimumLineSpacing = spacing
    minimumInteritemSpacing = spacing

    collectionView?.delegate = self
  }

  override func layoutAttributesForElements(in rect: CGRect) -> [UICollectionViewLayoutAttributes]? {
    let attributes = super.layoutAttributesForElements(in: rect)?
      .map { $0.copy() } as? [UICollectionViewLayoutAttributes]

    attributes?
      .reduce([CGFloat: (CGFloat, [UICollectionViewLayoutAttributes])]()) {
        guard $1.representedElementCategory == .cell else { return $0 }
        return $0.merging([ceil($1.center.y): ($1.frame.origin.y, [$1])]) {
          ($0.0 < $1.0 ? $0.0 : $1.0, $0.1 + $1.1)
        }
      }
      .values.forEach { minY, line in
        line.forEach {
          $0.frame = $0.frame.offsetBy(
            dx: 0,
            dy: minY - $0.frame.origin.y
          )
        }
      }
    return attributes
  }
}

extension GridMenuLayout: UICollectionViewDelegate, UICollectionViewDelegateFlowLayout {
  func collectionView(
    _ collectionView: UICollectionView,
    layout collectionViewLayout: UICollectionViewLayout,
    sizeForItemAt indexPath: IndexPath
  ) -> CGSize {
    let totalSpacing = (2 * spacing) + ((numberOfItemsPerRow - 1) * spacing)

    guard let collection = self.collectionView else {
      return CGSize(width: 0, height: 0)
    }

    let cellWidth = (collection.bounds.width - totalSpacing) / numberOfItemsPerRow

    var height: CGFloat = 130
    if indexPath.item == bigButtonIndex {
      height = 180
      bigButtonIndex += 5
    }

    return CGSize(width: cellWidth, height: height)
  }
}