Efficiently loading images in table views and collection views

When your app shows images from the network in a table view or collection view, you need to load the images asynchronously to make sure your list scrolls smoothly. More importantly, you’ll need to somehow connect the image that you’re loading to the correct cell in your list (instead of table view or collection view, I’m going to say list from now on). And if the cell goes out of view and is reused to display new data with a new image, you’ll want to cancel the in-progress image load to make sure new images are loaded immediately. And, to make sure we don’t go to the network more often than needed, we’ll need some way to cache images in memory, or on disk, so we can use the local version of an image if we’ve already fetched it in the past. Based on this, we can identify three core problems that we need to solve when loading images for our cells asynchronously:

1. Setting the loaded image on the correct cell.
2. Canceling an in-progress load when a cell is reused.
3. Caching loaded images to avoid unneeded network calls.

In this post, I’m going to show you how to build a simple image loader class, and write a table view cell that will help us solve all these problems. I will also show you how you can use the same image loader to enhance UIImage with some fancy helpers.

The loader and techniques in this post focus on UIKit do not use Swift's async/await. If you're interested in building an image loader that leverages async/await make sure you check out this post alongside this one.

Building a simple image loader

When you make a GET request using URLSession, you typically do so through a data task. Normally, you don’t hold on to that data task because you don’t need it. But if you keep a reference to your data task around, you can cancel it at a later time. I’m going to use a dictionary of [UUID: URLSessionDataTask] in the image loader we’re building because that will allow me to keep track of running downloads and cancel them later. I’m also going to use a dictionary of [URL: UIImage] as a simple in-memory cache for loaded images. Based on this, we can begin writing the image loader:

class ImageLoader {
  private var loadedImages = [URL: UIImage]()
  private var runningRequests = [UUID: URLSessionDataTask]() 
}

We can also implement a loadImage(_:completion:) method. This method will accept a URL and a completion handler, and it’s going to return a UUID that’s used to uniquely identify each data task later on. The implementation looks as follows:

func loadImage(_ url: URL, _ completion: @escaping (Result<UIImage, Error>) -> Void) -> UUID? {

  // 1
  if let image = loadedImages[url] {
    completion(.success(image))
    return nil
  }

  // 2
  let uuid = UUID()

  let task = URLSession.shared.dataTask(with: url) { data, response, error in
    // 3
    defer {self.runningRequests.removeValue(forKey: uuid) }

    // 4
    if let data = data, let image = UIImage(data: data) {
      self.loadedImages[url] = image
      completion(.success(image))
      return
    }

    // 5
    guard let error = error else {
      // without an image or an error, we'll just ignore this for now
      // you could add your own special error cases for this scenario
      return
    }

    guard (error as NSError).code == NSURLErrorCancelled else {
      completion(.failure(error))
      return
    }

    // the request was cancelled, no need to call the callback
  }
  task.resume()

  // 6
  runningRequests[uuid] = task
  return uuid
}

Let’s go over the preceding code step by step, following the numbered comments.

1. If the URL already exists as a key in our in-memory cache, we can immediately call the completion handler. Since there is no active task and nothing to cancel later, we can return nil instead of a UUID instance.
2. We create a UUID instance that is used to identify the data task that we’re about to create.
3. When the data task completed, it should be removed from the running requests dictionary. We use a defer statement here to remove the running task before we leave the scope of the data task’s completion handler.
4. When the data task completes and we can extract an image from the result of the data task, it is cached in the in-memory cache and the completion handler is called with the loaded image. After this, we can return from the data task’s completion handler.
5. If we receive an error, we check whether the error is due to the task being canceled. If the error is anything other than canceling the task, we forward that to the caller of loadImage(_:completion:).
6. The data task is stored in the running requests dictionary using the UUID that was created in step 2. This UUID is then returned to the caller.

Note that steps 3 through 5 all take place in the data task’s completion handler. This means that the order in which the listed steps execute isn’t linear. Step 1 and 2 are executed first, then step 6 and then 3 through 5.

Now that we have logic to load our images, let’s at some logic that allows us to cancel in-progress image downloads too:

func cancelLoad(_ uuid: UUID) {
  runningRequests[uuid]?.cancel()
  runningRequests.removeValue(forKey: uuid)
}

This method receives a UUID, uses it to find a running data task and cancels that task. It also removes the task from the running tasks dictionary, if it exists. Fairly straightforward, right?

Let’s see how you would use this loader in a table view’s cellForRowAtIndexPath method:

// 1
let token = loader.loadImage(imageUrl) { result in
  do {
    // 2
    let image = try result.get()
    // 3
    DispatchQueue.main.async {
      cell.cellImageView.image = image
    }
  } catch {
    // 4
    print(error)
  }
}

// 5
cell.onReuse = {
  if let token = token {
    self.loader.cancelLoad(token)
  }
}

Let’s go through the preceding code step by step again:

1. The image loader’s loadImage(_:completion:) method is called, and the UUID returned by the loader is stored in a constant.
2. In the completion handler for loadImage(_:completion:), we extract the result from the completion’s result argument.
3. If we successfully extracted an image, we dispatch to the main queue and set the fetched image on the cell’s imageView property. Not sure what dispatching to the main queue is? Read more in this post
4. If something went wrong, print the error. You might want to do something else here in your app.
5. I’ll show you an example of my cell subclass shortly. The important bit is that we use the UUID that we received from loadImage(_:completion:) to cancel the loader’s load operation for that UUID.

Note that we do this in the cellForRowAt method. This means that every time we’re asked for a cell to show in our list, this method is called for that cell. So the load and cancel are pretty tightly coupled to the cell’s life cycle which is exactly what we want in this case. Let’s see what onReuse is in a sample cell:

class ImageCell: UITableViewCell {
  @IBOutlet var cellImageView: UIImageView!
  var onReuse: () -> Void = {}

  override func prepareForReuse() {
    super.prepareForReuse()
    onReuse()
    cellImageView.image = nil
  }
}

The onReuse property is a closure that we call when the cell’s prepareForReuse method is called. We also remove the current image from the cell in prepareForReuse so it doesn’t show an old image while loading a new one. Cells are reused quite often so doing the appropriate cleanup in prepareForReuse is crucial to prevent artifacts from old data on a cell from showing up when you don’t want to.

If you implement all of this in your app, you’ll have a decent strategy for loading images. You would probably want to add a listener for memory warnings that are emitted by your app’s Notification Center, and maybe you would want to cache images to disk as well as memory too, but I don’t think that fits nicely into the scope of this article for now. Keep these two features in mind though if you want to implement your own image loader. Especially listening for memory warnings is important since your app might be killed by the OS if it consumes too much memory by storing images in the in-memory cache.

Enhancing UIImageView to create a beautiful image loading API

Before we implement the fancy helpers, let’s refactor our cell and cellForRowAt method so they already contain the code we want to write. The prepareForReuse method is going to look as follows:

override func prepareForReuse() {
  cellImageView.image = nil
  cellImageView.cancelImageLoad()
}

This will set the current image to nil and tells the image view to stop loading the image it was loading. All of the image loading code in cellForRowAt should be replaced with the following:

cell.cellImageView.loadImage(at: imageUrl)

Yes, all of that code we had before is now a single line.

To make this new way of loading and canceling work, we’re going to implement a special image loader class called UIImageLoader. It will be a singleton object that manages loading for all UIImageView instances in your app which means that you end up using a single cache for your entire app. Normally you might not want that, but in this case, I think it makes sense. The following code outlines the skeleton of the UIImageLoader:

class UIImageLoader {
  static let loader = UIImageLoader()

  private let imageLoader = ImageLoader()
  private var uuidMap = [UIImageView: UUID]()

  private init() {}

  func load(_ url: URL, for imageView: UIImageView) {

  }

  func cancel(for imageView: UIImageView) {

  }
}

The loader itself is a static instance, and it uses the ImageLoader from the previous section to actually load the images and cache them. We also have a dictionary of [UIImageView: UUID] to keep track of currently active image loading tasks. We map these based on the UIImageView so we can connect individual task identifiers to UIImageView instances.

The implementation for the load(_:for:) method looks as follows:

func load(_ url: URL, for imageView: UIImageView) {
  // 1
  let token = imageLoader.loadImage(url) { result in
    // 2
    defer { self.uuidMap.removeValue(forKey: imageView) }
    do {
      // 3
      let image = try result.get()
      DispatchQueue.main.async {
        imageView.image = image
      }
    } catch {
      // handle the error
    }
  }

  // 4
  if let token = token {
    uuidMap[imageView] = token
  }
}

Step by step, this code does the following:

1. We initiate the image load using the URL that was passed too load(_:for:).
2. When the load is completed, we need to clean up the uuidMap by removing the UIImageView for which we’re loading the image from the dictionary.
3. This is similar to what was done in cellForRowAt before. The image is extracted from the result and set on the image view itself.
4. Lastly, if we received a token from the image loader, we keep it around in the [UIImageView: UUID] dictionary so we can reference it later if the load has to be canceled.

The cancel(for:) method has the following implementation:

func cancel(for imageView: UIImageView) {
  if let uuid = uuidMap[imageView] {
    imageLoader.cancelLoad(uuid)
    uuidMap.removeValue(forKey: imageView)
  }
}

If we have an active download for the passed image view, it’s canceled and removed from the uuidMap. Very similar to what you’ve seen before.

All we need to do now is add an extension to UIImageView to add the loadImage(at:) and cancelImageLoad() method you saw earlier:

extension UIImageView {
  func loadImage(at url: URL) {
    UIImageLoader.loader.load(url, for: self)
  }

  func cancelImageLoad() {
    UIImageLoader.loader.cancel(for: self)
  }
}

Both methods pass self to the image loader. Since the extension methods are added to instances of UIImageView, this helps the image loader to connect the URL that it loads to the UIImageView instance that we want to show the image, leaving us with a very simple and easy to use API! Cool stuff, right?

What’s even better is that this new strategy can also be used for images that are not in a table view cell or collection view cell. It can be used for literally any image view in your app!

In summary

Asynchronously loading data can be a tough problem on its own. When paired with the fleeting nature of table view (and collection view) cells, you run into a whole new range of issues. In this post, you saw how you can use URLSession and a very simple in-memory cache to implement a smart mechanism to start, finish and cancel image downloads.

After creating a simple mechanism, you saw how to create an extra loader object and some extensions for UIImageView to create a very straightforward and easy to use API to load images from URLs directly into your image views.

Keep in mind that the implementations I’ve shown you here aren’t production-ready. You’ll need to do some work in terms of memory management and possibly add a disk cache to make these objects ready for prime time.

If you have any questions about this topic, have feedback or anything else, don’t hesitate to shoot me a message on Twitter.