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

Proton/Sources/Swift/Table/TableCell.swift

@@ -49,6 +49,14 @@ 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
+    }
+
     /// Additional attributes that can be stored on Cell to identify various aspects like Header, Numbered etc.
     public var additionalAttributes: [String: Any] = [:]
 
@@ -184,6 +192,22 @@ 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)
+    }
+
     func hideEditor() {
         contentView?.hideEditor()
     }
@@ -192,6 +216,7 @@ public class TableCell {
         contentView?.showEditor()
     }
 
+
     func performWithoutChangingFirstResponder(_ closure: () -> Void) {
         editor?.disableFirstResponder()
         closure()

Proton/Sources/Swift/Table/TableView.swift

@@ -412,8 +412,12 @@ public class TableView: UIView {
         removeScrollObserver()
     }
 
+    private var retainedCells = Set<TableCell>()
+
     var cellsInViewport: [TableCell] = [] {
         didSet {
+            reclaimReleasedCells()
+
             guard oldValue != cellsInViewport else { return }
 
             let oldCells = Set(oldValue)
@@ -422,16 +426,19 @@ public class TableView: UIView {
             let toReclaim = oldCells.subtracting(newCells)
 
             toReclaim.forEach { [weak self] cell in
-                // Ignore reclaiming the cell if it has focus
-                if cell.containsFirstResponder == false {
+                // 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 {
+                    self?.retainedCells.insert(cell)
                 }
             }
 
             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
-                if cell.containsFirstResponder == false {
+                // 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)
                 }
             }
@@ -489,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 {
@@ -854,11 +870,13 @@ extension TableView: TableContentViewDelegate {
     }
 
     func tableContentView(_ tableContentView: TableContentView, didReceiveFocusAt range: NSRange, in cell: TableCell) {
+        cell.retain()
         resetColumnResizingHandles(selectedCell: cell)
         delegate?.tableView(self, didReceiveFocusAt: range, in: cell)
     }
 
     func tableContentView(_ tableContentView: TableContentView, didLoseFocusFrom range: NSRange, in cell: TableCell) {
+        cell.release()
         removeSelectionBorders()
         delegate?.tableView(self, didLoseFocusFrom: range, in: cell)
     }

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