Implemented logic to retain and release cells. Used the same for focu…

…ssed cells

Proton/Sources/Swift/Table/TableCell.swift

@@ -49,7 +49,10 @@ public class TableCell {
 
     public var onEditorInitialized: ((TableCell, EditorView) -> Void)?
 
+    @MainActor
     private var retainCount = 0
+    /// Returns `true` if the cell has one or more retain invocations.
+    @MainActor
     public var isRetained: Bool {
         retainCount > 0
     }
@@ -189,10 +192,18 @@ public class TableCell {
         contentView?.removeFocus()
     }
 
+    /// Retains the cell to prevent it from getting recycled when viewport changes and cell is scrolled off-screen
+    /// Cell with focus is automatically retained and released.
+    /// - Note: A cell may be retained as many time as needed but needs to have a corresponding release for every retain.
+    /// Failing to release would mean that the cell will never participate in virtualization and will always be kept alive even when off-screen.
+    /// - Important: It is responsibility of consumer to ensure that the cells are correctly released after being retained.
+    @MainActor
     public func retain() {
         retainCount += 1
     }
 
+    /// Releases a retained cell. Calling this function on a non-retained cell is a no-op.
+    @MainActor
     public func release() {
         retainCount = max(0, retainCount - 1)
     }

Proton/Sources/Swift/Table/TableView.swift

@@ -416,6 +416,8 @@ public class TableView: UIView {
 
     var cellsInViewport: [TableCell] = [] {
         didSet {
+            reclaimReleasedCells()
+
             guard oldValue != cellsInViewport else { return }
 
             let oldCells = Set(oldValue)
@@ -424,7 +426,8 @@ public class TableView: UIView {
             let toReclaim = oldCells.subtracting(newCells)
 
             toReclaim.forEach { [weak self] cell in
-                // Ignore reclaiming the cell if it has focus
+                // Ignore reclaiming the cell if these are retained
+                // Cell having focus is always retained and released on lost focus
                 if cell.isRetained == false {
                     self?.repository.enqueue(cell: cell)
                 } else {
@@ -433,19 +436,12 @@ public class TableView: UIView {
             }
 
             toGenerate.forEach { [weak self] cell in
-                // Ignore generating the cell if it has focus. The focussed cell is not reclaimed, hence need not be regenerated.
-                // In absence of this check, there may be cases where the focussed cell gets duplicated
+                // Ignore generating the cell if it is already retained. The retained cell is not reclaimed, hence need not be regenerated.
+                // In absence of this check, there may be cases where the retained cell gets duplicated
                 if cell.isRetained == false {
                     self?.repository.dequeue(for: cell)
                 }
             }
-
-            retainedCells.forEach {
-                if $0.isRetained == false {
-                    self.repository.enqueue(cell: $0)
-                    retainedCells.remove($0)
-                }
-            }
         }
     }
 
@@ -500,6 +496,15 @@ public class TableView: UIView {
             .intersects(adjustedViewport) }
     }
 
+    private func reclaimReleasedCells() {
+        retainedCells.forEach {
+            if $0.isRetained == false {
+                self.repository.enqueue(cell: $0)
+                retainedCells.remove($0)
+            }
+        }
+    }
+
     func cellBelow(_ cell: TableCell) -> TableCell? {
         guard let row = cell.rowSpan.max(),
               let column = cell.columnSpan.min() else {

Proton/Tests/Table/TableViewAttachmentSnapshotTests.swift

@@ -1094,6 +1094,46 @@ class TableViewAttachmentSnapshotTests: SnapshotTestCase {
         assertSnapshot(matching: viewController.view, as: .image, record: recordMode)
     }
 
+    func FLAKY_testRecyclesRetainedCell() {
+        var viewport = CGRect(x: 0, y: 100, width: 350, height: 200)
+        delegate.viewport = viewport
+
+        Utility.drawRect(rect: viewport, color: .red, in: editor)
+
+        let attachment = AttachmentGenerator.makeTableViewAttachment(id: 1, numRows: 20, numColumns: 5)
+        attachment.view.delegate = delegate
+        editor.replaceCharacters(in: .zero, with: "Some text in editor")
+        editor.insertAttachment(in: editor.textEndRange, attachment: attachment)
+        editor.replaceCharacters(in: editor.textEndRange, with: "Text after grid")
+
+        XCTAssertEqual(attachment.view.containerAttachment, attachment)
+
+        viewController.render(size: CGSize(width: 400, height: 700))
+        assertSnapshot(matching: viewController.view, as: .image, record: recordMode)
+
+        let cell10 = attachment.view.cellAt(rowIndex: 1, columnIndex: 0)
+        cell10?.editor?.setFocus()
+
+        viewport = CGRect(x: 0, y: 300, width: 350, height: 200)
+        delegate.viewport = viewport
+        attachment.view.scrollViewDidScroll(editor.scrollView)
+
+        Utility.drawRect(rect: viewport, color: .red, in: editor)
+        viewController.render(size: CGSize(width: 400, height: 700))
+        assertSnapshot(matching: viewController.view, as: .image, record: recordMode)
+
+        let cell40 = attachment.view.cellAt(rowIndex: 4, columnIndex: 0)
+        cell40?.editor?.setFocus()
+
+        viewport = CGRect(x: 0, y: 320, width: 350, height: 200)
+        delegate.viewport = viewport
+        attachment.view.scrollViewDidScroll(editor.scrollView)
+
+        Utility.drawRect(rect: viewport, color: .red, in: editor)
+        viewController.render(size: CGSize(width: 400, height: 700))
+        assertSnapshot(matching: viewController.view, as: .image, record: recordMode)
+    }
+
     func FLAKY_testPreventsRecyclingNestedEditorFocussedCell() {
         var viewport = CGRect(x: 0, y: 100, width: 350, height: 200)
         delegate.viewport = viewport