From 4ec2b2e3117b0b51a1e29090c0c304fbd04e0c3c Mon Sep 17 00:00:00 2001
From: John Estropia <rommel.estropia@gmail.com>
Date: Mon, 5 Oct 2020 23:12:17 +0900
Subject: [PATCH] Optimize ListSnapshot collection implementation

---
 ...Internals.DiffableDataSourceSnapshot.swift | 63 +++++++++++++++-
 Sources/ListSnapshot.swift                    | 74 ++++++++++++++++---
 2 files changed, 125 insertions(+), 12 deletions(-)

diff --git a/Sources/Internals.DiffableDataSourceSnapshot.swift b/Sources/Internals.DiffableDataSourceSnapshot.swift
index e4388c6..c46bc38 100644
--- a/Sources/Internals.DiffableDataSourceSnapshot.swift
+++ b/Sources/Internals.DiffableDataSourceSnapshot.swift
@@ -78,7 +78,7 @@ extension Internals {
 
         var numberOfItems: Int {
 
-            return self.structure.allItemIDs.count
+            return self.structure.allItemsCount
         }
 
         var numberOfSections: Int {
@@ -106,6 +106,59 @@ extension Internals {
             return self.itemIdentifiers(inSection: identifier).count
         }
 
+        func itemIdentifier(atAllItemsIndex index: Int) -> NSManagedObjectID? {
+
+            guard index >= 0 else {
+
+                return nil
+            }
+            var remainingIndex = index
+            for section in self.structure.sections {
+
+                let elements = section.elements
+                let sectionCount = elements.count
+                if remainingIndex < sectionCount {
+
+                    return elements[remainingIndex].differenceIdentifier
+                }
+
+                remainingIndex -= sectionCount
+            }
+            return nil
+        }
+
+        func itemIdentifiers(atAllItemsBounds bounds: Range<Int>) -> [NSManagedObjectID] {
+
+            var remainingIndex = bounds.lowerBound
+            var itemIdentifiers: [NSManagedObjectID] = []
+            for section in self.structure.sections {
+
+                let elements = section.elements
+                let sectionCount = elements.count
+                if remainingIndex < sectionCount {
+
+                    itemIdentifiers.append(
+                        contentsOf: elements[remainingIndex..<min(sectionCount, bounds.count)]
+                            .map({ $0.differenceIdentifier })
+                    )
+                }
+                else if !itemIdentifiers.isEmpty {
+
+                    itemIdentifiers.append(
+                        contentsOf: elements.prefix(bounds.count - itemIdentifiers.count)
+                            .map({ $0.differenceIdentifier })
+                    )
+                }
+                if itemIdentifiers.count >= bounds.count {
+
+                    return itemIdentifiers
+                }
+
+                remainingIndex -= sectionCount
+            }
+            return itemIdentifiers
+        }
+
         func itemIdentifiers(inSection identifier: String) -> [NSManagedObjectID] {
 
             return self.structure.items(in: identifier)
@@ -332,6 +385,14 @@ extension Internals {
                 return self.sections.map({ $0.differenceIdentifier })
             }
 
+            var allItemsCount: Int {
+
+                return self.sections.reduce(into: 0) { (result, section) in
+
+                    result += section.elements.count
+                }
+            }
+
             var allItemIDs: [NSManagedObjectID] {
 
                 return self.sections.lazy.flatMap({ $0.elements }).map({ $0.differenceIdentifier })
diff --git a/Sources/ListSnapshot.swift b/Sources/ListSnapshot.swift
index 9ed1f92..d1bea9f 100644
--- a/Sources/ListSnapshot.swift
+++ b/Sources/ListSnapshot.swift
@@ -71,7 +71,7 @@ public struct ListSnapshot<O: DynamicObject>: RandomAccessCollection, Hashable {
     public subscript(index: Index) -> ObjectPublisher<O> {
 
         let context = self.context!
-        let itemID = self.diffableSnapshot.itemIdentifiers[index]
+        let itemID = self.diffableSnapshot.itemIdentifier(atAllItemsIndex: index)!
         return context.objectPublisher(objectID: itemID)
     }
 
@@ -83,16 +83,13 @@ public struct ListSnapshot<O: DynamicObject>: RandomAccessCollection, Hashable {
      */
     public subscript(safeIndex index: Index) -> ObjectPublisher<O>? {
 
-        guard let context = self.context else {
+        guard
+            let context = self.context,
+            let itemID = self.diffableSnapshot.itemIdentifier(atAllItemsIndex: index)
+        else {
 
             return nil
         }
-        let itemIDs = self.diffableSnapshot.itemIdentifiers
-        guard itemIDs.indices.contains(index) else {
-
-            return nil
-        }
-        let itemID = itemIDs[index]
         return context.objectPublisher(objectID: itemID)
     }
 
@@ -603,14 +600,69 @@ public struct ListSnapshot<O: DynamicObject>: RandomAccessCollection, Hashable {
     
     public var startIndex: Index {
         
-        return self.diffableSnapshot.itemIdentifiers.startIndex
+        return 0
     }
     
     public var endIndex: Index {
         
-        return self.diffableSnapshot.itemIdentifiers.endIndex
+        return self.diffableSnapshot.numberOfItems
     }
-    
+
+    public func index(after i: Index) -> Index {
+
+        return i + 1
+    }
+
+    public func formIndex(after i: inout Index) {
+
+        return i += 1
+    }
+
+    public func index(before i: Index) -> Index {
+
+        return i - 1
+    }
+
+    public func formIndex(before i: inout Index) {
+
+        return i -= 1
+    }
+
+
+    // MARK: BidirectionalCollection
+
+    public func index(_ i: Index, offsetBy distance: Int) -> Index {
+
+        return i + distance
+    }
+
+    public func index(_ i: Index, offsetBy distance: Int, limitedBy limit: Int) -> Index? {
+
+        let length = limit - i
+        if distance > 0
+            ? length >= 0 && length < distance
+            : length <= 0 && length > distance {
+
+            return nil
+        }
+        return i + distance
+    }
+
+    public func distance(from start: Index, to end: Index) -> Int {
+
+        return end - start
+    }
+
+    public subscript(bounds: Range<Index>) -> ArraySlice<Element> {
+
+        guard let context = self.context else {
+
+            return .init()
+        }
+        let itemIDs = self.diffableSnapshot.itemIdentifiers(atAllItemsBounds: bounds)
+        return ArraySlice(itemIDs.map(context.objectPublisher(objectID:)))
+    }
+
     
     // MARK: Sequence