¿Cómo descargar y mostrar una lista de imágenes en una vista de colección para evitar el parpadeo y las imágenes en lugares equivocados?

Estoy luchando por resolver el problema muy común de cargar imágenes dentro de una vista de colección dada una lista de URL. He implementado una UIImageViewextensión. Allí definí una variable de caché:

static var imageCache = NSCache<NSString, UIImage>()

Además, he creado un loadImagemétodo que toma como entrada la clave de caché de la imagen, la URL en sí, una imagen de marcador de posición y una imagen de error y funciona de la siguiente manera:

public func loadImage(cacheKey: String?, urlString: String?, placeholderImage: UIImage?, errorImage: UIImage?) {
    
    image = nil
    
    if let cacheKey = cacheKey, let cachedImage = UIImageView.imageCache.object(forKey: cacheKey as NSString) {
        image = cachedImage
        return
    }
    
    if let placeholder = placeholderImage {
        image = placeholder
    }
    
    if let urlString = urlString, let url = URL(string: urlString) {
        self.downloadImage(url: url, errorImage: errorImage)
    } else {
        image = errorImage
    }
}

La función de descargar imagen procede a crear una tarea de datos, descargar la imagen y asignarla a la vista de imagen:

private func downloadImage(url: URL, errorImage: UIImage?) {
    
    let dataTask = URLSession.shared.dataTask(with: url) { [weak self] (data, response, error ) in
        
        if error != nil {
            DispatchQueue.main.async {
                self?.image = errorImage
            }
            return
        }
        
        if let statusCode = (response as? HTTPURLResponse)?.statusCode {
            if !(200...299).contains(statusCode) {
                DispatchQueue.main.async {
                    self?.image = errorImage
                }
                return
            }
        }
        
        if let data = data, let image = UIImage(data: data) {
            UIImageView.imageCache.setObject(image, forKey: url.absoluteString as NSString)
            DispatchQueue.main.async {
                self?.image = image
            }
        }
        
    }
    
    dataTask.resume()
    
}

Llamo al loadImagemétodo inside cellForItemAt, de modo que no tengo que descargar todos los datos a la vez, sino solo las imágenes que se muestran efectivamente en la pantalla. La forma en que llamo a la función es la siguiente:

cell.albumImage.loadImage(
    cacheKey: bestQualityImage,
    urlString: bestQualityImage,
    placeholderImage: UIImage(named: "Undefined"),
    errorImage: UIImage(named: "Undefined")
)

El principal problema al que me enfrento con la implementación actual es que, a veces, las imágenes no se muestran en el lugar correcto. En otras palabras, si tengo tres elementos en pantalla, a veces vería los tres elementos con la misma imagen en lugar de su imagen respectiva.

Creo que el problema es que cuando se completa la descarga para una celda específica, el cellForItemAt para esa celda ya ha finalizado y la celda obtiene la imagen incorrecta en su lugar.

¿Hay alguna forma en que pueda modificar mi función para corregir este error o debo cambiar mi enfoque por completo?

EDITAR. Al principio pensé que el problema era que estaba usando una extensión e intenté usar una función con un cierre que devolvía la imagen descargada, pero eso no resolvió mi problema.

Answer