A ViewModel driven DiffableDataSource implementation

Introduction

We are going to explore a way you can construct datasource’s for your view in a very simple manor. This leads to a very testable design, that requires minimal code, and to never need to use an IndexPath ever again. Let’s get started.

In this example we are going to use a UICollectionView that uses a NSDiffableDataSourceSnapshot to drive the data source.

Let’s begin

First we create two enums. One to represent the sections, and one to represent an item (or cell) inside the collection view with it’s associated value being that cells model.

Here is an enum to represent the sections of our view. Two have no associated value as there is nothing special about them. The third, has a section header that shows a title and a star if it’s featured. As the enum, and all of it’s associated values conform to Hashable we don’t need to manually implement the hash(into: Hasher) function. Note that if you do implement this function, hash all the properties visible to the user, otherwise if you don’t, updates won’t be applied if the non hashed values change. These hash’s also need to be unique otherwise your app will crash so using id’s in them can be a good idea, I would refrain from using index’s.

enum HomeSection: Hashable {
    case banner
    case forYou
    case suggested(SuggestedSectionModel)
}

struct SuggestedSectionModel {
    let title: String
    let isFeatured: Bool
}

Now we can add an enum to represent all the Items in the Section. As you can see, it’s similar to above.

enum HomeItem: Hashable {
    case banner(String)
    case videoContent(URL)
    case imageContent(UIImage)
}

Constructing our Snapshot

The snapshot is what represents the data shown on screen. Every time the data changes, we create a new snapshot, apply it, and have it automatically differentiate and apply the differences.

In this example, our snapshot is bound to the ViewModel, which drives the the UI, so we want to create a Snapshot to represent the data using the two generic types being the HomeSection and HomeItem enums we defined above. You can do this in your Presenter if you like, it just comes down to the architecture you are using in the app.

ViewModel

Here is our basic ViewModel that has three Published properties on it. As this data is fetched from an API it can be dynamically set and the view can update incrementally without having to wait for all of it to load before showing anything.

class HomeViewModel {

    @Published
    var banner: HomeBanner?

    @Published
    var suggested: [MediaItem] = []

    @Published
    var reccomended: [MediaItem] = []
}

Constructing the snapshot

We can add a function to the ViewModel that combines all three properties and turns them into a new snapshot when the data changes.

typealias HomeSnapshot = NSDiffableDataSourceSnapshot<HomeSection, HomeItem>

var dataSourcePublisher: AnyPublisher<HomeSnapshot, Never> {
    Publishers.CombineLatest3($banner, $suggested, $reccomended)
        .map { self.createDataSource(banner: $0, suggested: $1, reccomended: $2) }
        .eraseToAnyPublisher()
}

Now we can implement createDataSource(..)

private func createDataSource(
    banner: Banner?,
    suggested: [MediaItem],
    reccomended: [MediaItem]
) -> HomeSnapshot {

    let snapshot = HomeSnapshot()

    // 1. Add banner
    if let banner {
        // 1.1 Add banner section
        snapshot.appendSections([.banner])
        // 1.2 Add banner item
        snapshot.appendItems(.banner(banner.title), toSection: .banner)
    }

    // 2. Add suggested items
    if !suggested.isEmpty {
        snapshot.appendSections([.suggested])
        let suggestedModels = suggested
            .map { HomeItem.imageContent($0.imageUrl) } 
        snapshot.appendItems(suggestedModels, toSection: .suggested)
    }
    
    // 3. Add reccomendation items
    if !reccomendations.isEmpty {
        // 3.1 Add section model
        let sectionModel = SuggestedSectionModel(
            title: "Suggestions for you",
            isFeatured: true
        )
        let section = .reccomenendations(sectionsModel)
        snapshot.appendSections([section])

        // 3.2 Add suggested models 
        let suggestedModels = suggested
            .map { HomeItem.imageContent($0.imageUrl) } 
        snapshot.appendItems(suggestedModels, toSection: section)
    }
}

Constructing our DataSource on the ViewController

We will first create a UICollectionView followed by a DataSource that is set on it. Our ViewController also has a reference to the HomeViewModel we created just before. Don’t forget to register your cells on viewDidLoad too.

lazy var collectionView = UICollectionView(frame: .zero, collectionViewLayout: layout())
lazy var dataSource: UICollectionViewDiffableDataSource<HomeSection, HomeItem> = createDataSource()

private var viewModel = HomeViewModel()

override func viewDidLoad() {
     collectionView.dataSource = dataSource
     subscribeToSnapshot()
 }

Now we need to implement the function that creates the DataSource. Because our item identifier is an enum, we can simply switch on the enum cases and create a cell, pass it a model to set the cell up, and then return the configured cell.

private func createDataSource() -> UICollectionViewDiffableDataSource<HomeViewSection, HomeViewItem> {
  .init(collectionView: collectionView) { collectionView, indexPath, itemIdentifier in
      switch itemIdentifier {
      case .banner(let string):
          let cell: BannerCell = collectionView.dequeueCell(for: indexPath)
          return cell.setup(string)
          
      case .imageContent(let image):
          let cell: ImageCell = collectionView.dequeueCell(for: indexPath)
          return cell.setup(image)
          
      case .videoContent(let video):
          let cell: VideoCell = collectionView.dequeueCell(for: indexPath)
          return cell.setup(video)
      }
  }
}

And finally we can subscribe to our DataSource snapshot that changes when the ViewModel’s properties do.

private func subscribeToSnapshot() {
    viewModel.dataSourcePublisher
        .receive(on: DispatchQueue.main)
        .sink(receiveValue: { [weak self] snapshot in
            self?.dataSource?.apply(snapshot, animatingDifferences: true)
        })
        .store(in: &cancellables)
}

Summary

While there are many ways to implement diffable data sources, this is one that has worked well for me. Like everything, it’s not perfect and has it’s limitations.

It’s important to remember that as all data that hashed as part of the view models, as it changes, it will cause that cell to reload.

Limitations

If you are using UITextFields in a cell and the value is hashed entering text and updating the view model immediately will cause the text field to be dismissed as the first responder. In these instances, I would suggest storing the values separately and using the latest values on cell setup to avoid this issue.

Thanks for reading!

Hopefully you learn something from this and can use it to help efficiently generate view models and data sources in your application.