Animating UICollectionView updates with self-sizing cells the Rx way

UICollectionView is one of the most flexible and powerful, but also mysterious and tricky components in UIKit toolbelt. I've learned it the hard way when trying to build a simple, Facebook-like feed. Today, I'll share my lessons learned and show how to do it yourself with as little overhead as possible.

The feed has following requirements:

• It should scroll smoothly.
• It has to nicely animate updates.
• It needs to support pagination when fetching data from the API.

The outcome looks as follows:

Collection view is a natural candidate for that kind of requirements. It implements reusing components for maximum performance and it can animate the updates nicely.

Note: although the tutorial relies on RxCocoa and RxSwift, it should still be valuable for all non-reactive folks out there.

Building blocks

The feed does not require building a custom collection view layout.  UICollectionViewFlowLayout does all the heavy work to handle it perfectly.

UICollectionView supports batch updates. It allows to animate the changes to the data source, like adding or removing items, with ease. Let's see how it looks in practice:

Since the feed screen is endlessly paginated, the animations will come in very handy to add new items on the fly. However, unlike in the demo presented above, the feed items are certainly not of the same size:

• the bottom text has a variable length and it is wrapped to fit the cell size,
• the header might be different depending on author's name length.

Fortunately, UICollectionView supports the self-sizing cells. It's as easy as implementing size-constrained UICollectionViewCell and setting the estimatedItemSize property on the UICollectionViewFlowLayout. The outcome might look like this:

The problem

The tricky part comes when combining the animated updates together with self-sizing cells. Let's try it ourselves:

It looks really bad. What's more, when collection view is scrolled down the content offset messes up during the animation, too. It results in badly looking jump when adding new items.

The slow animations mode comes in very handy to figure out what's wrong:

It turns out that during the animation UICollectionViewFlowLayout forgets previously calculated sizes of the cells and reverts to whatever is set as estimatedItemSize. This took me quite a while to figure out, but I've learned that there's a pending radar opened since iOS 8 that's still not fixed.

A quick workaround for the issue is to manually calculate and cache the sizes. It's sounds complicated, but is way easier than one can think. Especially when there's only a single cell type in play.

First, we'll need a cell view. The one presented in animations above is a simple label embedded in a content view:

class FeedCell: UICollectionViewCell {
  let label: UILabel = {
    let label = UILabel(frame: .zero)
    label.numberOfLines = 0
    label.lineBreakMode = .byWordWrapping
    return label
  }()

  var widthConstraint: NSLayoutConstraint?

  override init(frame: CGRect) {
    super.init(frame: frame)

    addSubviews()
    setUpLayout()
  }

  private func addSubviews() {
    contentView.addSubview(label)
  }

  private func setUpLayout() {
    widthConstraint = (contentView.widthAnchor == 0)
    label.edgeAnchors == contentView.edgeAnchors
    label.setContentHuggingPriority(.required, for: .vertical)
    label.setContentCompressionResistancePriority(.required, for: .vertical)
  }

  required init?(coder _: NSCoder) { return nil }
}

The constraints are defined with Anchorage, a very lightweight set of operator extensions.

Then, we'll need a model to display in a collection view.

struct FeedItem {
  let id: Int
  let title: String
}

The model is just for the sake of simplicity. The background colour is disregarded. Also, these are the assumptions made:

• We're dealing with  UICollectionViewController instance.
• We have an array of [FeedItem] stored that drive the data source with a single section and number of rows equal to number of items).
• We've correctly implemented performing batch updates for the animations. (Note: I'll show a very easy way to achieve this later.)

class FeedViewController: UICollectionViewController, UICollectionViewDelegateFlowLayout {
  init() {
    super.init(collectionViewLayout: flowLayout)
  }

  private let flowLayout: UICollectionViewFlowLayout = {
    let layout = UICollectionViewFlowLayout()
    layout.minimumLineSpacing = 20
    layout.sectionInset = UIEdgeInsets(top: 15, left: 0, bottom: 45, right: 0)
    // it's important to NOT set an estimatedItemSize
    //layout.estimatedItemSize = removed
    return layout
  }()

  override func viewDidLoad() {
    super.viewDidLoad()
    collectionView?.delegate = self
    // ... remaining setup code
  }

  private func present(item: FeedItem, in: FeedCell) {
    // ... presentation code
  }

  // MARK: - Private

  private var items: [FeedItem] = []
  private lazy var layoutCell: FeedCell = FeedCell(frame: .zero)
  private var sizes: [Int: CGSize] = [:]

  func collectionView(
    _ collectionView: UICollectionView,
    layout collectionViewLayout: UICollectionViewLayout,
    sizeForItemAt indexPath: IndexPath
  ) -> CGSize {
    let item = items.value[indexPath.row]

    if let size = sizes[item.id] {
      return size
    }

    present(item: item, in: layoutCell)
    let size = layoutCell.systemLayoutSizeFitting(UILayoutFittingCompressedSize)
    sizes[item.id] = size
    return size
  }
}

The snippet above is incomplete, but it presents the idea:

• We need to derive from UICollectionViewDelegateFlowLayout and implement a method to calculate each item size (collectionView(_:layout:sizeForItemAt:)).
• The complete cell presentation code is extracted to a private method and reused between cell creation and size calculation.
• A single cell instance is stored that's not for reusing. It's used to calculate item size and is created lazily (lazy var layoutCell = FeedCell(frame: .zero) ) to make sure we don't have to deal with optionals.
• We're caching cell sizes in a dictionary (var items: [Int: CGSize]) that maps item identifier to a size of a cell. This way, we need to calculate the size only once per every identifier.
• We're using systemLayoutSizeFitting(UILayoutFittingCompressedSize) method to determine the cell size from auto-layout.

The end result is astounding:

Equipped with an idea on how to build animations, we can now proceed to reaching our main goal.

Building the feed

Let's take a quick recap of how the feed should look like:

We'll start off with modelling the item. It should contain:

• poster's details (avatar, name),
• post location and distance to the user,
• creation date,
• post content: image and text,
• whether the item is liked or not.

The FeedItem model containing all of the information is presented below:

struct FeedItem {
  let id: Int
  let title: String
  let image: URL
  let author: String
  let whenCreated: String
  let distance: String
  let address: String
  let avatar: URL
  let liked: Bool
}

Since the items are paginated, we need another model which describes metadata of fetched results page:

struct FeedPage {
  let count: Int
  let results: [FeedItem]
  let nextPageIndex: Int?
  let previousPageIndex: Int?
}

Those pages are fetched from a service that accepts page index as an input and returns a single page of results:

protocol FeedProviding {
  func fetchPage(atIndex index: Int) -> Single<FeedPage>
}

For the tutorial purposes this will generate fake data with randomized delay, but the interface would likely not change when used with production API.

Now, let's figure out how to create a data source for UICollectionView. We want the changes to be animated, including removing, updating and adding the items. If we were to implement it without 3rd party components, every time the update comes we would need to calculate the diff between previous array of items and the updated one. Having the diff calculated, we could invoke performBatchUpdates calling the suitable update methods manually.

Defining the data source

Fortunately, there's a great tool called RxDataSources which greatly simplifies the task. Provided you have a model that is:

• identifiable (has unique, hashable identifier),
• equatable (can determine whether the item has changed)

RxDataSources can animate the changes for you. Let's modify the FeedItem slightly to meet the criteria:

import RxDataSources

struct FeedItem: IdentifiableType, Equatable {
  // ...

  var identity: Int { 
    return id
  }
}

Swift 4.1 can generate Equatable conformance automatically, so there's only a tiny amount of overhead to make the model RxDataSources-compatible.

Next, we'll bind the model to a collection view. In order to do so, we need to create an instance of RxCollectionViewSectionedAnimatedDataSource, which works kinda like UICollectionViewDataSource. It specifies how to configure the cells based on the model.

collectionView?.register(FeedCell.self, forCellWithReuseIdentifier: "cell")

let dataSource = RxCollectionViewSectionedAnimatedDataSource<
  AnimatableSectionModel<String, FeedItem>>(
  configureCell: { [weak self] (_, collectionView, indexPath, item: FeedItem) in
    let cell = collectionView.dequeueReusableCell(withReuseIdentifier: "cell", for: indexPath)
    self?.present(item: item, in: cell as! FeedCell)
    return cell
  },
  configureSupplementaryView: { _, _, _, _ in UICollectionReusableView(frame: .zero) }
)

Notice that RxCollectionViewSectionedAnimatedDataSource is generalised with an odd generic type AnimatableSectionModel<String, FeedItem>. The data source can also animate section updates, that's why it needs to operate on a section model, not an item model. We won't use it in a single section tutorial, but it's mandatory to satisfy the data sources API.

The quickest way to work around that is to use predefined AnimatableSectionModel which requires us to specify:

• The section model that needs to be identifiable. It's straightforward to use it in a conjunction with a simple value type like a string, integer or floating point because those conform to IdentifiableType in RxDataSources by default. I've picked the string type for the sake of the example.
• An array of identifiable and equatable item models. That's our FeedItem.

Following type declaration:

RxDataSources<AnimatableSectionModel<String, FeedItem>>

means that the source of the data must be the observable sequence of type [AnimatableSectionModel<String, FeedItem>].

The presentation is extracted to a private method, as stated in animating self-sizing cells part of the tutorial:

private func present(item: FeedItem, in cell: FeedCell) {
  cell.widthConstraint?.constant = collectionView.frame.width - 40.0
  cell.postLabel.text = item.title
  cell.postImage.af_setImage(withURL: item.image)
  cell.headerView.author.topLabel.text = item.author
  cell.headerView.author.bottomLabel.text = item.whenCreated
  cell.headerView.distance.topLabel.text = item.distance
  cell.headerView.distance.bottomLabel.text = item.address
  cell.headerView.avatar.af_setImage(withURL: item.avatar)
  cell.footerView.likeButton.setTitle(item.liked ? "👎" : "👍", for: .normal)
}

It's worth noting how FeedCell exposes a constraint which describes the width of the content view. It could also be constrained inside a cell using UIScreen.main.bounds.width, but it's less flexible and won't allow us to change interface orientation in future. Also, you can see that images in the example are downloaded using AlamofireImage dependency.

Configuring collection view

With data source instance, it's very easy to bind FeedItem directly into collection view:

let feedProvider = FeedProvider()

feedProvider.fetchPage(atIndex: 1)
  .map { $0.results }
  .asDriver(onErrorJustReturn: [])
  .map { [AnimatableSectionModel(model: "", items: $0]) }
  .drive(collectionView.rx.items(dataSource: dataSource))
  .disposed(by: disposeBag)

Hence we use a constant data source for now, it's not too flashy yet. But a cute fade in animation can already be seen when items first appear in a collection view.

The empty string used as a model for the section doesn't matter at this point. We're using a single section that won't be updated, so it can be an arbitrary string.

Building the pagination

Pagination is a tricky problem to solve in reactive world. It's a very stateful piece of code that's hard to describe as a combination of event streams. The recommended way to implement a pagination by RxSwift author  is to use another dependency called RxFeedback.

RxFeedback is a simplistic architecture based on RxSwift library. It consists of:

• state - a structure that describes current state of the system,
• mutations - formal description of events that can modify the state,
• feedback loops - mapping between state updates and mutations.

We can also try to describe them in type system notation:

• struct State { // ... data } - fields represent the current state of the system.
• enum Mutation { // ... mutations } - cases are mutations that alter the system state.
• (Driver<State>) -> Signal<Mutation> - function(s) describing mapping state updates to mutations.

First, we need to describe a state of our pagination subsystem. We need to store fetched items, an index of the next page and a flag to indicate whether we should already trigger the next page loading:

struct FeedState {
  var items: [FeedItem]
  var nextPage: Int?
  var shouldLoadNext: Bool

  init() {
    items = []
    nextPage = 1
    shouldLoadNext = true
  }
}

The initial state has no items and should fetch the first page. The next step is to define available mutations. We need to support following actions:

• loading next page (both success and failure),
• triggering fetch of a next page,
• liking an item,
• removing an item.

We can describe those actions as FeedMutation enum:

enum FeedMutation {
  case loadNextPage
  case pageLoaded(page: FeedPage)
  case error
  case like(id: Int)
  case drop(id: Int)
}

The next step is to define how the state is updated by the mutation. It's done by the means of reduce function which accepts current state and mutation and returns an updated state:

extension FeedState {
  static func reduce(state: FeedState, mutation: FeedMutation) -> FeedState {
    var copy = state

    switch mutation {
    case .loadNextPage:
      copy.shouldLoadNext = true
    case .error:
      copy.shouldLoadNext = false
    case let .pageLoaded(page):
      copy.shouldLoadNext = false
      copy.nextPage = page.nextPageIndex
      copy.items += page.results
    case let .like(id):
      copy.items = copy.items.map { item in
        item.id == id ? item.copy(liked: !item.liked) : item
      }
    case let .drop(id):
      copy.items = copy.items.filter { $0.id != id }
    }

    return copy
  }
}

Let's take a minute to analyse the reduce method:

• When next page fetch is triggered (.loadNextPage) mutation, we set the shouldLoadNext to true to indicate that we should start the call.
• When page is loaded (.pageLoaded) or the fetch fails (.error) the shouldLoadNext flag is set to false to stop loading the page for now.
• In case of successful page load (.pageLoaded) the fetched items are appended to the state and the nextPage is updated to indicate index we should load next.
• When feed item is liked (.like) the liked flag is toggled on an item that was updated.
• When feed item is removed (.drop), the item with a specified identifier is removed from the list.

Building the feedback loops

Having described the state and possible mutations, we need to implement feedback loops. The first one describes what happens on the UI:

struct FeedState {
  // ...
  typealias Feedback: (Driver<FeedState>) -> Signal<FeedMutation>
}
      
let uiFeedback: FeedState.Feedback = bind(self) { me, state in
  let subscriptions: [Disposable] = [
    me.bindDataSource(state: state, to: me.collectionView)
  ]

  let mutations: [Signal<FeedMutation>] = [
    me.nextPageMutation(state: state, in: me.collectionView),
    me.mutations.asSignal()
  ]

  return Bindings(subscriptions: subscriptions, events: mutations)
}

UI feedback consists of data bindings (used for presentation) and events invoked by user interactions. First, let's cover the presentation. We'll use the RxCollectionViewAnimatedDataSource defined earlier to bind the state items to a collection view:

private func bindDataSource(state: Driver<FeedState>, to collectionView: UICollectionView) -> Disposable {
  let dataSource = // ... defined earlier

  return state
    .map { $0.items }
    .map { [AnimatableSectionModel(model: "section", items: $0)] }
    .drive(collectionView.rx.items(dataSource: dataSource))
}

Now, let's map mutations caused by user interactions. We'll start with triggering .loadNextPage mutation when user scrolls down to the last item of a collection view:

private func nextPageMutation(state: Driver<FeedState>, in collectionView: UICollectionView) -> Signal<FeedMutation> {
  return state.flatMapLatest { state in
    guard !state.shouldLoadNext else {
      return Signal.empty()
    }

    return collectionView.rx.willDisplayCell
      .asSignal()
      .filter { $0.at.item == state.items.count - 1 }
      .map { _ in .loadNextPage }
  }
}
    
private let items: BehaviorRelay<[FeedItem]> = BehaviorRelay(value: [])

If items are not fetched already, we listen for willDisplayCell event and check whether it's the last item in the feed. If so, we emit the .loadNextPage event.

The missing part is attaching the mutations coming from the feed items itself.

private let mutations: PublishRelay<FeedMutation> = PublishRelay()
    
private func present(item: FeedItem, in cell: FeedCell) {
  // ...

  cell.footerView.likeButton.rx.tap
    .map { .like(id: item.id) }
    .bind(to: mutations)
    .disposed(by: cell.disposeBag)

  cell.footerView.dropButton.rx.tap
    .map { .drop(id: item.id) }
    .bind(to: mutations)
    .disposed(by: cell.disposeBag)
}

With that being configured, we can already like and remove items from the feed. The only missing part is to fetch the next page from the service.

Fetching the next page

To fetch the next page, we need other type of a feedback loop. It's called react and it queries the current state. If specific value is present in the state, the side effect is performed.

For the pagination, we need to know an index of the page to fetch. We should only return it if the fetch was requested. A simple extension will come in handy.

extension FeedState {
  var loadPageIndex: Int? {
    return shouldLoadNext ? nextPage : nil
  }
}

Now, let's implement the fetch page feedback:

let feedProvider = FeedProvider()

let loadPageFeedback: FeedState.Feedback = react(
  query: { $0.loadPageIndex },
  effects: { pageIndex -> Signal<FeedMutation> in
    feedProvider.fetchPage(atIndex: pageIndex)
      .map { .pageLoaded(page: $0) }
      .asSignal(onErrorJustReturn: .error)
  }
)

Using the react method, we specify that each occurrence of loadPageIndex value in the state should fetch the page at returned index. If page is successfully loaded, .pageLoaded signal is emitted. If error happens, we emit the .error event.

Having the feedback loops defined, the last missing part is to glue everything together using Driver.system method:

Driver.system(
  initialState: FeedState(),
  reduce: FeedState.reduce,
  feedback: uiFeedback, loadPageFeedback
)
.drive(state)
.disposed(by: disposeBag)

Conclusion

That's it! Using a set of 3rd parties built on top of RxSwift, we've managed to implement a performant feed with animated updates with as little code as possible. The complete source for the tutorial is hosted at turekj/RxFeedExample GitHub.  Happy feeding!