perf(datagrid): cache column index lookup in selection drawing

CHANGELOG.md

@@ -9,6 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Changed
 
+- Drag-selecting many columns in a wide result set scrolls smoothly instead of lagging; the selection overlay and row highlight now look up column positions from a cache that refreshes when columns are added, removed, or reordered.
 - Data grid now serves the row count from its existing cache instead of recomputing it on every layout pass, reducing CPU churn while scrolling large result sets.
 - Typing in the sidebar table search stays responsive on databases with thousands of tables; filtering runs after a short pause instead of on every keystroke.
 

TablePro/Views/Results/DataGridCoordinator.swift

@@ -38,6 +38,27 @@ final class TableViewCoordinator: NSObject, NSTableViewDelegate, NSTableViewData
     private(set) var identitySchema: ColumnIdentitySchema = .empty
     var currentSortState = SortState()
 
+    private var columnIndexByDataIndex: [Int: Int] = [:]
+    private static let selectionCacheLogger = Logger(subsystem: "com.TablePro", category: "DataGrid.ColumnIndexCache")
+
+    func tableColumnIndex(for dataIndex: Int) -> Int? {
+        if let cached = columnIndexByDataIndex[dataIndex] {
+            return cached
+        }
+        guard let tableView,
+              let identifier = identitySchema.identifier(for: dataIndex) else { return nil }
+        let resolved = tableView.column(withIdentifier: identifier)
+        guard resolved >= 0 else { return nil }
+        columnIndexByDataIndex[dataIndex] = resolved
+        return resolved
+    }
+
+    func invalidateColumnIndexCache() {
+        guard !columnIndexByDataIndex.isEmpty else { return }
+        Self.selectionCacheLogger.debug("invalidate column index cache (had \(self.columnIndexByDataIndex.count))")
+        columnIndexByDataIndex.removeAll()
+    }
+
     func columnIdentifier(for dataIndex: Int) -> NSUserInterfaceItemIdentifier? {
         identitySchema.identifier(for: dataIndex)
     }
@@ -224,6 +245,7 @@ final class TableViewCoordinator: NSObject, NSTableViewDelegate, NSTableViewData
         columnDisplayFormats = []
         cachedRowCount = 0
         cachedColumnCount = 0
+        invalidateColumnIndexCache()
         sortedIDs = nil
         lastUpdateSnapshot = nil
         columnPool.detachFromTableView()
@@ -666,6 +688,7 @@ final class TableViewCoordinator: NSObject, NSTableViewDelegate, NSTableViewData
         guard schemaChanged else { return false }
         identitySchema = nextSchema
         displayCache.removeAll()
+        invalidateColumnIndexCache()
         return true
     }
 

TablePro/Views/Results/DataGridRowView.swift

@@ -99,9 +99,10 @@ class DataGridRowView: NSTableRowView {
     }
 
     private func drawCellSelectionFill(in dirtyRect: NSRect) {
-        guard let selection = coordinator?.selectionController.selection,
-              !selection.isEmpty,
-              let tableView = coordinator?.tableView else { return }
+        guard let coordinator,
+              let tableView = coordinator.tableView else { return }
+        let selection = coordinator.selectionController.selection
+        guard !selection.isEmpty else { return }
         let columns = selection.columns(in: rowIndex)
         guard !columns.isEmpty else { return }
 
@@ -110,10 +111,8 @@ class DataGridRowView: NSTableRowView {
             : NSColor.selectedContentBackgroundColor.withAlphaComponent(0.28)
         fillColor.setFill()
 
-        let schema = coordinator?.identitySchema
         for dataColumn in columns {
-            guard let schema,
-                  let tableColumnIndex = DataGridView.tableColumnIndex(for: dataColumn, in: tableView, schema: schema) else { continue }
+            guard let tableColumnIndex = coordinator.tableColumnIndex(for: dataColumn) else { continue }
             let columnRect = tableView.rect(ofColumn: tableColumnIndex)
             let localRect = NSRect(x: columnRect.minX, y: 0, width: columnRect.width, height: bounds.height)
             guard localRect.intersects(dirtyRect) else { continue }

TablePro/Views/Results/DataGridView.swift

@@ -271,6 +271,7 @@ struct DataGridView: NSViewRepresentable {
                 savedLayout: savedLayout
             )
             coordinator.isRebuildingColumns = false
+            coordinator.invalidateColumnIndexCache()
 
             if savedLayout == nil {
                 coordinator.scheduleLayoutPersist()

TablePro/Views/Results/Extensions/DataGridView+Selection.swift

@@ -14,6 +14,7 @@ extension TableViewCoordinator {
 
     func tableViewColumnDidMove(_ notification: Notification) {
         guard !isRebuildingColumns else { return }
+        invalidateColumnIndexCache()
         layoutPersistTask?.cancel()
         persistColumnLayoutToStorage()
     }

TablePro/Views/Results/Selection/GridSelectionOverlay.swift

@@ -34,13 +34,12 @@ final class GridSelectionOverlay: NSView {
 
     override func draw(_ dirtyRect: NSRect) {
         guard let tableView, let coordinator else { return }
-        let schema = coordinator.identitySchema
         let totalRows = tableView.numberOfRows
         let editingCell = activeOverlayCell(in: coordinator)
 
         NSColor.selectedContentBackgroundColor.withAlphaComponent(Self.borderAlpha).setStroke()
         for rect in selection.rectangles {
-            guard let frame = frame(for: rect, in: tableView, schema: schema) else { continue }
+            guard let frame = frame(for: rect, in: tableView, coordinator: coordinator) else { continue }
             guard frame.intersects(dirtyRect) else { continue }
             if isFullHeight(rect, totalRows: totalRows) { continue }
             if let editingCell, rect.contains(editingCell) { continue }
@@ -53,7 +52,7 @@ final class GridSelectionOverlay: NSView {
         if let active = selection.activeCell,
            editingCell != active,
            selection.rectangles.count > 1 || (selection.rectangles.first?.rows.count ?? 0) > 1 || (selection.rectangles.first?.columns.count ?? 0) > 1,
-           let frame = frame(for: GridRect(cell: active), in: tableView, schema: schema),
+           let frame = frame(for: GridRect(cell: active), in: tableView, coordinator: coordinator),
            frame.intersects(dirtyRect) {
             NSColor.controlAccentColor.setStroke()
             let inset = frame.insetBy(dx: Self.activeCellBorderWidth / 2, dy: Self.activeCellBorderWidth / 2)
@@ -78,7 +77,7 @@ final class GridSelectionOverlay: NSView {
         return rect.rows.lowerBound <= 0 && rect.rows.upperBound >= totalRows - 1
     }
 
-    private func frame(for rect: GridRect, in tableView: NSTableView, schema: ColumnIdentitySchema) -> NSRect? {
+    private func frame(for rect: GridRect, in tableView: NSTableView, coordinator: TableViewCoordinator) -> NSRect? {
         guard tableView.numberOfRows > 0, tableView.numberOfColumns > 0 else { return nil }
         let firstRow = max(0, rect.rows.lowerBound)
         let lastRow = min(tableView.numberOfRows - 1, rect.rows.upperBound)
@@ -92,7 +91,7 @@ final class GridSelectionOverlay: NSView {
         var leadingX = CGFloat.infinity
         var trailingX = -CGFloat.infinity
         for dataColumn in rect.columns.lowerBound...rect.columns.upperBound {
-            guard let tableColumnIndex = DataGridView.tableColumnIndex(for: dataColumn, in: tableView, schema: schema) else { continue }
+            guard let tableColumnIndex = coordinator.tableColumnIndex(for: dataColumn) else { continue }
             let columnRect = tableView.rect(ofColumn: tableColumnIndex)
             leadingX = min(leadingX, columnRect.minX)
             trailingX = max(trailingX, columnRect.maxX)

TableProTests/Views/Results/Extensions/ColumnIndexCacheTests.swift

@@ -0,0 +1,149 @@
+import AppKit
+import Foundation
+import SwiftUI
+@testable import TablePro
+import Testing
+
+@MainActor
+private final class StubColumnLayoutPersister: ColumnLayoutPersisting {
+    func load(for tableName: String, connectionId: UUID) -> ColumnLayoutState? { nil }
+    func save(_ layout: ColumnLayoutState, for tableName: String, connectionId: UUID) {}
+    func clear(for tableName: String, connectionId: UUID) {}
+}
+
+@Suite("TableViewCoordinator column index cache")
+@MainActor
+struct ColumnIndexCacheTests {
+    private func makeCoordinator() -> TableViewCoordinator {
+        TableViewCoordinator(
+            changeManager: AnyChangeManager(DataChangeManager()),
+            isEditable: false,
+            selectedRowIndices: .constant([]),
+            delegate: nil,
+            layoutPersister: StubColumnLayoutPersister()
+        )
+    }
+
+    private func attachColumns(_ tableView: NSTableView, count: Int) {
+        tableView.addTableColumn(
+            NSTableColumn(identifier: ColumnIdentitySchema.rowNumberIdentifier)
+        )
+        for slot in 0..<count {
+            tableView.addTableColumn(
+                NSTableColumn(identifier: ColumnIdentitySchema.slotIdentifier(slot))
+            )
+        }
+    }
+
+    @Test("resolved table column index mirrors the schema-backed column order")
+    func resolvesDataColumnToTableColumnIndex() {
+        let coordinator = makeCoordinator()
+        let tableView = NSTableView()
+        attachColumns(tableView, count: 3)
+        coordinator.tableView = tableView
+        _ = coordinator.rebuildColumnMetadataCache(
+            from: TableRows.from(
+                queryRows: [],
+                columns: ["id", "name", "email"],
+                columnTypes: Array(repeating: ColumnType.text(rawType: nil), count: 3)
+            )
+        )
+
+        #expect(coordinator.tableColumnIndex(for: 0) == 1)
+        #expect(coordinator.tableColumnIndex(for: 1) == 2)
+        #expect(coordinator.tableColumnIndex(for: 2) == 3)
+    }
+
+    @Test("repeated lookups keep returning the same value")
+    func lookupsAreStableAcrossCalls() {
+        let coordinator = makeCoordinator()
+        let tableView = NSTableView()
+        attachColumns(tableView, count: 2)
+        coordinator.tableView = tableView
+        _ = coordinator.rebuildColumnMetadataCache(
+            from: TableRows.from(
+                queryRows: [],
+                columns: ["a", "b"],
+                columnTypes: Array(repeating: ColumnType.text(rawType: nil), count: 2)
+            )
+        )
+
+        let first = coordinator.tableColumnIndex(for: 1)
+        for _ in 0..<5 {
+            #expect(coordinator.tableColumnIndex(for: 1) == first)
+        }
+    }
+
+    @Test("invalidate after a column reorder reflects the new layout")
+    func invalidateReflectsReorderedColumns() {
+        let coordinator = makeCoordinator()
+        let tableView = NSTableView()
+        attachColumns(tableView, count: 3)
+        coordinator.tableView = tableView
+        _ = coordinator.rebuildColumnMetadataCache(
+            from: TableRows.from(
+                queryRows: [],
+                columns: ["id", "name", "email"],
+                columnTypes: Array(repeating: ColumnType.text(rawType: nil), count: 3)
+            )
+        )
+        #expect(coordinator.tableColumnIndex(for: 0) == 1)
+
+        tableView.moveColumn(1, toColumn: 3)
+        coordinator.tableViewColumnDidMove(
+            Notification(name: NSTableView.columnDidMoveNotification, object: tableView)
+        )
+
+        #expect(coordinator.tableColumnIndex(for: 0) == 3)
+        #expect(coordinator.tableColumnIndex(for: 1) == 1)
+        #expect(coordinator.tableColumnIndex(for: 2) == 2)
+    }
+
+    @Test("schema rebuild drops stale cached indices when columns shrink")
+    func schemaRebuildInvalidatesCache() {
+        let coordinator = makeCoordinator()
+        let tableView = NSTableView()
+        attachColumns(tableView, count: 2)
+        coordinator.tableView = tableView
+        _ = coordinator.rebuildColumnMetadataCache(
+            from: TableRows.from(
+                queryRows: [],
+                columns: ["a", "b"],
+                columnTypes: Array(repeating: ColumnType.text(rawType: nil), count: 2)
+            )
+        )
+        #expect(coordinator.tableColumnIndex(for: 1) == 2)
+
+        if let second = tableView.tableColumns.last(where: {
+            $0.identifier == ColumnIdentitySchema.slotIdentifier(1)
+        }) {
+            tableView.removeTableColumn(second)
+        }
+        _ = coordinator.rebuildColumnMetadataCache(
+            from: TableRows.from(
+                queryRows: [],
+                columns: ["a"],
+                columnTypes: [ColumnType.text(rawType: nil)]
+            )
+        )
+
+        #expect(coordinator.tableColumnIndex(for: 1) == nil)
+    }
+
+    @Test("an out-of-range data column resolves to nil")
+    func outOfRangeReturnsNil() {
+        let coordinator = makeCoordinator()
+        let tableView = NSTableView()
+        attachColumns(tableView, count: 1)
+        coordinator.tableView = tableView
+        _ = coordinator.rebuildColumnMetadataCache(
+            from: TableRows.from(
+                queryRows: [],
+                columns: ["only"],
+                columnTypes: [ColumnType.text(rawType: nil)]
+            )
+        )
+
+        #expect(coordinator.tableColumnIndex(for: 5) == nil)
+    }
+}