From ff10477c164af61579eb6d5d4f477fa4c524adc1 Mon Sep 17 00:00:00 2001
From: Eugene Kazaev <eugene.kazaev@gmail.com>
Date: Fri, 17 May 2024 12:52:50 +0100
Subject: [PATCH] Improvements in reloading cells when using
 UICollectionViewDiffableDataSource

---
 ChatLayout.podspec                            |  2 +-
 .../Classes/Core/Model/ChangeItem.swift       |  2 +-
 .../Classes/Core/Model/StateController.swift  | 28 +++++++++----------
 3 files changed, 16 insertions(+), 16 deletions(-)

diff --git a/ChatLayout.podspec b/ChatLayout.podspec
index 20da9334..84235aff 100644
--- a/ChatLayout.podspec
+++ b/ChatLayout.podspec
@@ -1,6 +1,6 @@
 Pod::Spec.new do |s|
   s.name             = 'ChatLayout'
-  s.version          = '2.0.8'
+  s.version          = '2.0.9'
   s.summary          = 'Chat UI Library. It uses custom UICollectionViewLayout to provide you full control over the presentation.'
   s.swift_version    = '5.8'
 
diff --git a/ChatLayout/Classes/Core/Model/ChangeItem.swift b/ChatLayout/Classes/Core/Model/ChangeItem.swift
index 94133d98..d772b733 100644
--- a/ChatLayout/Classes/Core/Model/ChangeItem.swift
+++ b/ChatLayout/Classes/Core/Model/ChangeItem.swift
@@ -83,7 +83,7 @@ enum ChangeItem: Equatable {
                 self = .itemDelete(itemIndexPath: indexPath)
             }
         case .reload:
-            guard let indexPath = indexPathAfterUpdate else {
+            guard let indexPath = indexPathBeforeUpdate else {
                 assertionFailure("`indexPathAfterUpdate` cannot be `nil` for a `.reload` update action.")
                 return nil
             }
diff --git a/ChatLayout/Classes/Core/Model/StateController.swift b/ChatLayout/Classes/Core/Model/StateController.swift
index 4b0d7548..e733b167 100644
--- a/ChatLayout/Classes/Core/Model/StateController.swift
+++ b/ChatLayout/Classes/Core/Model/StateController.swift
@@ -559,8 +559,6 @@ final class StateController<Layout: ChatLayoutRepresentation> {
 
                 reloadedSectionsIndexesArray.append(sectionIndex)
             case let .itemReload(itemIndexPath: indexPath):
-                reloadedIndexes.insert(indexPath)
-
                 reloadedItemsIndexesArray.append(indexPath)
             case let .itemReconfigure(itemIndexPath: indexPath):
                 reconfiguredItemsIndexesArray.append(indexPath)
@@ -741,18 +739,6 @@ final class StateController<Layout: ChatLayoutRepresentation> {
             }
         }
 
-        for indexPath in reloadedItemsIndexesArray {
-            guard var item = item(for: indexPath.itemPath, kind: .cell, at: .beforeUpdate) else {
-                assertionFailure("Item at index path (\(indexPath.section) - \(indexPath.item)) does not exist.")
-                continue
-            }
-            let oldHeight = item.frame.height
-            let configuration = layoutRepresentation.configuration(for: .cell, at: indexPath)
-            applyConfiguration(configuration, to: &item)
-            afterUpdateModel.replaceItem(item, at: indexPath)
-            visibleBoundsBeforeUpdate.offsettingBy(dx: 0, dy: item.frame.height - oldHeight)
-        }
-
         for indexPath in deletedItemsIndexesArray {
             guard let itemId = itemIdentifier(for: indexPath.itemPath, kind: .cell, at: .beforeUpdate) else {
                 assertionFailure("Item at index path (\(indexPath.section) - \(indexPath.item)) does not exist.")
@@ -795,6 +781,20 @@ final class StateController<Layout: ChatLayoutRepresentation> {
             }
         }
 
+        for indexPath in reloadedItemsIndexesArray {
+            guard var item = item(for: indexPath.itemPath, kind: .cell, at: .beforeUpdate),
+                  let indexPathAfterUpdate = afterUpdateModel.itemPath(by: item.id, kind: .cell)?.indexPath else {
+                assertionFailure("Item at index path (\(indexPath.section) - \(indexPath.item)) does not exist.")
+                continue
+            }
+            reloadedIndexes.insert(indexPathAfterUpdate)
+            let oldHeight = item.frame.height
+            let configuration = layoutRepresentation.configuration(for: .cell, at: indexPathAfterUpdate)
+            applyConfiguration(configuration, to: &item)
+            afterUpdateModel.replaceItem(item, at: indexPathAfterUpdate)
+            visibleBoundsBeforeUpdate.offsettingBy(dx: 0, dy: item.frame.height - oldHeight)
+        }
+
         for indexPath in reconfiguredItemsIndexesArray {
             guard var item = item(for: indexPath.itemPath, kind: .cell, at: .beforeUpdate),
                   let indexPathAfterUpdate = afterUpdateModel.itemPath(by: item.id, kind: .cell)?.indexPath else {