feat(grid): preserve selection highlight after data refresh

apps/desktop/src/components/grid/DataGrid.vue

@@ -172,6 +172,7 @@ import { buildDataGridColumnLookupItems, dataGridColumnCommentFor, filterDataGri
 import { uniqueDataGridColumnOrderKeys } from "@/lib/dataGrid/dataGridColumnOrder";
 import { dataGridColumnLayoutScopeKey, TABLE_DATA_GRID_COLUMN_ORDER_CHANGED_EVENT, tableDataGridColumnOrderScopeKey } from "@/lib/dataGrid/dataGridColumnLayoutStorage";
 import { summarizeSelection } from "@/lib/dataGrid/gridSelection";
+import { captureDataGridSelection, reactivatesCellSelectionAfterRestore, restoreDataGridSelection, type PersistedDataGridSelection } from "@/lib/dataGrid/dataGridSelectionPersistence";
 import { dataGridFrameCoversRow, dataGridSelectionEdgeMask, dataGridSelectionFrameKindAtCell, dataGridSelectionUsesOuterFrame, resolveDataGridSelectionFrames } from "@/lib/dataGrid/dataGridSelectionFrames";
 import {
   createDataGridCellContextMenuItems,
@@ -414,6 +415,7 @@ if (isDebugLoggingEnabled()) {
 const transposeRowIndex = ref<number | null>(null);
 const showTranspose = ref(false);
 const preserveTransposeOnNextResult = ref(false);
+let preservedSelectionOnNextResult: { selection: PersistedDataGridSelection; sourceResult: QueryResult } | null = null;
 
 watch(
   () => props.result,
@@ -2646,6 +2648,7 @@ watch(
     // and the completion was triggered by a refresh/rollback.
     if (!loading && prevLoading && isRefreshingData.value) {
       isRefreshingData.value = false;
+      if (preservedSelectionOnNextResult?.sourceResult === props.result) preservedSelectionOnNextResult = null;
       const total = paginationTotalRowCount.value;
       if (!total || total <= 0) return;
       const lastPageNum = Math.max(1, Math.ceil(total / pageSize.value));
@@ -3386,6 +3389,8 @@ async function onToolbarRefresh() {
   if (infiniteScrollEnabled.value) {
     resetInfiniteScrollState();
   }
+  const selection = captureCurrentSelectionForRefresh();
+  preservedSelectionOnNextResult = selection ? { selection, sourceResult: props.result } : null;
   preserveTransposeOnNextResult.value = showTranspose.value;
   isRefreshingData.value = true;
   beginDataGridNativeSelectionBlock(dataGridNativeSelectionBlockOwner);
@@ -3848,6 +3853,7 @@ const {
   isSelectingAll,
   isSelectingCells,
   selectedRange,
+  selectionAnchor,
   selectedCells,
   selectedCellKeys,
   selectedCellMatrix,
@@ -3877,6 +3883,64 @@ const {
   isRowSelected,
 } = selection;
 
+function captureCurrentSelectionForRefresh(): PersistedDataGridSelection | null {
+  return captureDataGridSelection({
+    columns: props.result.columns,
+    sourceColumns: props.sourceColumns,
+    rows: props.result.rows,
+    primaryKeys: props.tableMeta?.primaryKeys ?? [],
+    visibleColumnIndexes: visibleColumnIndexes.value,
+    displayItems: displayItems.value,
+    selectedRowIds: selectedRowIds.value,
+    selectedColumnIndexes: selectedColumnIndexes.value,
+    selectedCellKeys: selectedCellKeys.value,
+    selectionAnchor: selectionAnchor.value,
+    selectionFocus: selectionFocus.value,
+    selectingAll: isSelectingAll.value,
+    lastClickedRowIndex: selection.lastClickedRowIndex.value,
+  });
+}
+
+function restoreSelectionAfterRefresh(snapshot: PersistedDataGridSelection) {
+  const restored = restoreDataGridSelection({
+    snapshot,
+    columns: props.result.columns,
+    sourceColumns: props.sourceColumns,
+    rows: props.result.rows,
+    visibleColumnIndexes: visibleColumnIndexes.value,
+    displayItems: displayItems.value,
+  });
+  if (!restored) return;
+
+  if (restored.kind === "rows") {
+    selectedRowIds.value = new Set(restored.rowIds);
+    selection.lastClickedRowIndex.value = restored.anchorRowIndex;
+  } else if (restored.kind === "columns") {
+    selectedColumnIndexes.value = new Set(restored.columnIndexes);
+  } else if (restored.kind === "range") {
+    selectionAnchor.value = restored.anchor;
+    selectionFocus.value = restored.focus;
+    isSelectingAll.value = restored.selectingAll;
+  } else {
+    selectedCellKeys.value = restored.cellKeys;
+  }
+
+  // 恢复 range/离散单元格选区后重新激活 cell-selection 拖拽状态机：refresh watcher
+  // 在 restore 前调的 clearCellSelection 已把 isSelectingCells 置 false，而它仅在
+  // beginCellSelection（鼠标按下）时才重新置 true。不在此补偿，则恢复的高亮虽能显示，
+  // 但 extendCellSelection/handleSelectionPointerMove 的 isSelectingCells 守卫会挡住
+  // 拖拽扩展，用户必须重新点一次单元格。行选/列选是独立状态机，不应激活。
+  if (reactivatesCellSelectionAfterRestore(restored)) {
+    isSelectingCells.value = true;
+  }
+
+  if (restored.kind === "columns") return;
+  nextTick(() => {
+    if (restored.kind === "range") scrollCellIntoView(restored.scrollRowIndex, restored.focus.colIndex);
+    else scrollGridRowIntoView(restored.scrollRowIndex);
+  });
+}
+
 const multiRowCount = computed(() => {
   if (hasRowSelection.value) return selectedRowCount.value;
   const range = selectedRange.value;
@@ -7494,6 +7558,8 @@ watch(isTransposeMode, (active) => {
 watch(
   () => props.result,
   (result, previousResult) => {
+    const selectionSnapshot = preservedSelectionOnNextResult?.selection;
+    preservedSelectionOnNextResult = null;
     const shouldPreserveTranspose = preserveTransposeOnNextResult.value;
     preserveTransposeOnNextResult.value = false;
     if (isDataGridPrefixAppend(previousResult, result)) return;
@@ -7512,6 +7578,7 @@ watch(
       closeTranspose(false);
     }
     exitTransaction();
+    if (selectionSnapshot) restoreSelectionAfterRefresh(selectionSnapshot);
   },
 );
 

apps/desktop/src/lib/__tests__/dataGrid/dataGridSelectionPersistence.spec.ts

@@ -0,0 +1,281 @@
+import { describe, expect, it } from "vitest";
+import { captureDataGridSelection, reactivatesCellSelectionAfterRestore, restoreDataGridSelection, type CaptureDataGridSelectionOptions } from "@/lib/dataGrid/dataGridSelectionPersistence";
+
+function baseOptions(overrides: Partial<CaptureDataGridSelectionOptions> = {}): CaptureDataGridSelectionOptions {
+  return {
+    columns: ["id", "tenant_id", "name"],
+    rows: [
+      [1, 10, "Ada"],
+      [2, 10, "Grace"],
+      [3, 20, "Linus"],
+    ],
+    primaryKeys: ["id"],
+    visibleColumnIndexes: [0, 1, 2],
+    displayItems: [
+      { id: 0, sourceIndex: 0 },
+      { id: 1, sourceIndex: 1 },
+      { id: 2, sourceIndex: 2 },
+    ],
+    selectedRowIds: new Set(),
+    selectedColumnIndexes: new Set(),
+    selectedCellKeys: new Set(),
+    selectionAnchor: null,
+    selectionFocus: null,
+    selectingAll: false,
+    ...overrides,
+  };
+}
+
+function restore(snapshot: NonNullable<ReturnType<typeof captureDataGridSelection>>, overrides: Partial<Parameters<typeof restoreDataGridSelection>[0]> = {}) {
+  return restoreDataGridSelection({
+    snapshot,
+    columns: ["id", "tenant_id", "name"],
+    rows: [
+      [1, 10, "Ada"],
+      [2, 10, "Grace"],
+      [3, 20, "Linus"],
+    ],
+    visibleColumnIndexes: [0, 1, 2],
+    displayItems: [
+      { id: 0, sourceIndex: 0 },
+      { id: 1, sourceIndex: 1 },
+      { id: 2, sourceIndex: 2 },
+    ],
+    ...overrides,
+  });
+}
+
+describe("data grid selection persistence", () => {
+  it("restores a selected cell by primary key after rows move", () => {
+    const snapshot = captureDataGridSelection(baseOptions({ selectionAnchor: { rowIndex: 1, colIndex: 2 }, selectionFocus: { rowIndex: 1, colIndex: 2 } }))!;
+
+    expect(
+      restore(snapshot, {
+        rows: [
+          [2, 10, "Grace updated"],
+          [1, 10, "Ada"],
+          [3, 20, "Linus"],
+        ],
+      }),
+    ).toEqual({ kind: "range", anchor: { rowIndex: 0, colIndex: 2 }, focus: { rowIndex: 0, colIndex: 2 }, selectingAll: false, scrollRowIndex: 0 });
+  });
+
+  it("supports composite keys and source-column aliases", () => {
+    const snapshot = captureDataGridSelection(
+      baseOptions({
+        columns: ["tenant", "record", "display_name"],
+        sourceColumns: ["tenant_id", "id", "name"],
+        rows: [
+          [10, 1, "Ada"],
+          [10, 2, "Grace"],
+          [20, 3, "Linus"],
+        ],
+        primaryKeys: ["tenant_id", "id"],
+        selectionAnchor: { rowIndex: 2, colIndex: 2 },
+        selectionFocus: { rowIndex: 2, colIndex: 2 },
+      }),
+    )!;
+
+    expect(
+      restoreDataGridSelection({
+        snapshot,
+        columns: ["tenant", "record", "display_name"],
+        sourceColumns: ["tenant_id", "id", "name"],
+        rows: [
+          [20, 3, "Linus updated"],
+          [10, 1, "Ada"],
+          [10, 2, "Grace"],
+        ],
+        visibleColumnIndexes: [0, 1, 2],
+        displayItems: [
+          { id: 0, sourceIndex: 0 },
+          { id: 1, sourceIndex: 1 },
+          { id: 2, sourceIndex: 2 },
+        ],
+      }),
+    ).toMatchObject({ kind: "range", anchor: { rowIndex: 0, colIndex: 2 } });
+  });
+
+  it("does not mistake an unrelated source column alias for a primary key", () => {
+    const snapshot = captureDataGridSelection(
+      baseOptions({
+        columns: ["id", "name"],
+        sourceColumns: ["external_id", "name"],
+        rows: [
+          [101, "Ada"],
+          [102, "Grace"],
+        ],
+        primaryKeys: ["id"],
+        visibleColumnIndexes: [0, 1],
+        displayItems: [
+          { id: 0, sourceIndex: 0 },
+          { id: 1, sourceIndex: 1 },
+        ],
+        selectionAnchor: { rowIndex: 1, colIndex: 1 },
+        selectionFocus: { rowIndex: 1, colIndex: 1 },
+      }),
+    )!;
+
+    expect(snapshot.identity.mode).toBe("row");
+  });
+
+  it("restores row selection and preserves its Shift-selection anchor", () => {
+    const snapshot = captureDataGridSelection(baseOptions({ selectedRowIds: new Set([0, 2]), lastClickedRowIndex: 2 }))!;
+
+    expect(
+      restore(snapshot, {
+        rows: [
+          [3, 20, "Linus"],
+          [2, 10, "Grace"],
+          [1, 10, "Ada"],
+        ],
+      }),
+    ).toEqual({ kind: "rows", rowIds: [2, 0], anchorRowIndex: 0, scrollRowIndex: 2 });
+  });
+
+  it("restores sparse cells and visible column selections", () => {
+    const cellSnapshot = captureDataGridSelection(baseOptions({ selectedCellKeys: new Set(["0:0", "2:2"]) }))!;
+    const columnSnapshot = captureDataGridSelection(baseOptions({ selectedColumnIndexes: new Set([0, 2]) }))!;
+
+    expect(restore(cellSnapshot)).toEqual({ kind: "cells", cellKeys: new Set(["0:0", "2:2"]), scrollRowIndex: 2 });
+    expect(restore(columnSnapshot, { visibleColumnIndexes: [2, 0] })).toEqual({ kind: "columns", columnIndexes: [1, 0] });
+  });
+
+  it("keeps select-all semantics when the refreshed result gains rows", () => {
+    const snapshot = captureDataGridSelection(baseOptions({ selectionAnchor: { rowIndex: 0, colIndex: 0 }, selectionFocus: { rowIndex: 2, colIndex: 2 }, selectingAll: true }))!;
+
+    expect(
+      restore(snapshot, {
+        rows: [
+          [1, 10, "Ada"],
+          [2, 10, "Grace"],
+          [3, 20, "Linus"],
+          [4, 30, "Margaret"],
+        ],
+        displayItems: [
+          { id: 0, sourceIndex: 0 },
+          { id: 1, sourceIndex: 1 },
+          { id: 2, sourceIndex: 2 },
+          { id: 3, sourceIndex: 3 },
+        ],
+      }),
+    ).toEqual({ kind: "range", anchor: { rowIndex: 0, colIndex: 0 }, focus: { rowIndex: 3, colIndex: 2 }, selectingAll: true, scrollRowIndex: 0 });
+  });
+
+  it("uses the hidden row identifier when it is the declared key", () => {
+    const snapshot = captureDataGridSelection(
+      baseOptions({
+        columns: ["name", "__DBX_ROWID"],
+        rows: [
+          ["Ada", "AAA1"],
+          ["Grace", "AAA2"],
+        ],
+        primaryKeys: ["__DBX_ROWID"],
+        visibleColumnIndexes: [0],
+        displayItems: [
+          { id: 0, sourceIndex: 0 },
+          { id: 1, sourceIndex: 1 },
+        ],
+        selectionAnchor: { rowIndex: 1, colIndex: 0 },
+        selectionFocus: { rowIndex: 1, colIndex: 0 },
+      }),
+    )!;
+
+    expect(
+      restoreDataGridSelection({
+        snapshot,
+        columns: ["name", "__DBX_ROWID"],
+        rows: [
+          ["Grace changed", "AAA2"],
+          ["Ada", "AAA1"],
+        ],
+        visibleColumnIndexes: [0],
+        displayItems: [
+          { id: 0, sourceIndex: 0 },
+          { id: 1, sourceIndex: 1 },
+        ],
+      }),
+    ).toMatchObject({ kind: "range", anchor: { rowIndex: 0, colIndex: 0 } });
+  });
+
+  it("falls back to type-safe full-row matching and distinguishes duplicate occurrences", () => {
+    const snapshot = captureDataGridSelection(
+      baseOptions({
+        columns: ["value"],
+        rows: [[1], ["1"], ["same"], ["same"]],
+        primaryKeys: [],
+        visibleColumnIndexes: [0],
+        displayItems: [
+          { id: 0, sourceIndex: 0 },
+          { id: 1, sourceIndex: 1 },
+          { id: 2, sourceIndex: 2 },
+          { id: 3, sourceIndex: 3 },
+        ],
+        selectionAnchor: { rowIndex: 3, colIndex: 0 },
+        selectionFocus: { rowIndex: 3, colIndex: 0 },
+      }),
+    )!;
+
+    expect(
+      restoreDataGridSelection({
+        snapshot,
+        columns: ["value"],
+        rows: [["same"], [1], ["same"], ["1"]],
+        visibleColumnIndexes: [0],
+        displayItems: [
+          { id: 0, sourceIndex: 0 },
+          { id: 1, sourceIndex: 1 },
+          { id: 2, sourceIndex: 2 },
+          { id: 3, sourceIndex: 3 },
+        ],
+      }),
+    ).toMatchObject({ kind: "range", anchor: { rowIndex: 2, colIndex: 0 } });
+  });
+
+  it("clears safely when the row disappears, is filtered out, or the fallback row shape changes", () => {
+    const keyedSnapshot = captureDataGridSelection(baseOptions({ selectionAnchor: { rowIndex: 1, colIndex: 1 }, selectionFocus: { rowIndex: 1, colIndex: 1 } }))!;
+    const fallbackSnapshot = captureDataGridSelection(baseOptions({ primaryKeys: [], selectionAnchor: { rowIndex: 1, colIndex: 1 }, selectionFocus: { rowIndex: 1, colIndex: 1 } }))!;
+
+    expect(restore(keyedSnapshot, { rows: [[1, 10, "Ada"]], displayItems: [{ id: 0, sourceIndex: 0 }] })).toBeNull();
+    expect(
+      restore(keyedSnapshot, {
+        displayItems: [
+          { id: 0, sourceIndex: 0 },
+          { id: 2, sourceIndex: 2 },
+        ],
+      }),
+    ).toBeNull();
+    expect(
+      restore(fallbackSnapshot, {
+        rows: [
+          [1, 10, "Ada"],
+          [2, 10, "Grace changed"],
+          [3, 20, "Linus"],
+        ],
+      }),
+    ).toBeNull();
+    expect(restore(fallbackSnapshot, { columns: ["id", "name"], rows: [[2, "Grace"]], visibleColumnIndexes: [0, 1], displayItems: [{ id: 0, sourceIndex: 0 }] })).toBeNull();
+  });
+
+  it("reactivates the cell-selection drag state machine only for range and sparse-cell restores", () => {
+    // 刷新后恢复 range/离散单元格选区时，组件层需重新置 isSelectingCells=true，
+    // 否则 clearCellSelection 已将其清零、拖拽扩展会被守卫挡住（见 extendCellSelection/
+    // handleSelectionPointerMove）。行选/列选是独立状态机，不应激活 cell selection。
+    const rangeSnapshot = captureDataGridSelection(baseOptions({ selectionAnchor: { rowIndex: 0, colIndex: 0 }, selectionFocus: { rowIndex: 2, colIndex: 1 } }))!;
+    const cellsSnapshot = captureDataGridSelection(baseOptions({ selectedCellKeys: new Set(["0:0", "2:2"]) }))!;
+    const rowsSnapshot = captureDataGridSelection(baseOptions({ selectedRowIds: new Set([0, 2]), lastClickedRowIndex: 2 }))!;
+    const columnsSnapshot = captureDataGridSelection(baseOptions({ selectedColumnIndexes: new Set([0, 2]) }))!;
+
+    expect(reactivatesCellSelectionAfterRestore(restore(rangeSnapshot)!)).toBe(true);
+    expect(reactivatesCellSelectionAfterRestore(restore(cellsSnapshot)!)).toBe(true);
+    expect(reactivatesCellSelectionAfterRestore(restore(rowsSnapshot)!)).toBe(false);
+    expect(reactivatesCellSelectionAfterRestore(restore(columnsSnapshot)!)).toBe(false);
+  });
+
+  it("does not signal reactivation when the restored range falls back to null", () => {
+    // 行消失/被过滤导致 restore 返回 null 时，组件层不会走到激活分支；
+    // 这里锁定 restore 仍返回 null，确保不会误把失败恢复当作可扩展选区。
+    const snapshot = captureDataGridSelection(baseOptions({ selectionAnchor: { rowIndex: 1, colIndex: 1 }, selectionFocus: { rowIndex: 1, colIndex: 1 } }))!;
+    expect(restore(snapshot, { rows: [[1, 10, "Ada"]], displayItems: [{ id: 0, sourceIndex: 0 }] })).toBeNull();
+  });
+});

apps/desktop/src/lib/dataGrid/dataGridSelectionPersistence.ts

@@ -0,0 +1,291 @@
+import type { CellPosition } from "@/lib/dataGrid/gridSelection";
+
+type CellValue = string | number | boolean | null;
+
+export interface DataGridSelectionRow {
+  id: number;
+  sourceIndex?: number;
+  isNew?: boolean;
+  isDraft?: boolean;
+}
+
+interface RowIdentityToken {
+  key: string;
+  occurrence: number;
+}
+
+interface ColumnIdentityToken {
+  resultName: string;
+  sourceName?: string;
+  occurrence: number;
+}
+
+type PersistedDataGridSelectionState =
+  | { kind: "rows"; rows: RowIdentityToken[]; anchorRow?: RowIdentityToken }
+  | { kind: "columns"; columns: ColumnIdentityToken[] }
+  | { kind: "range"; anchor: { row: RowIdentityToken; column: ColumnIdentityToken }; focus: { row: RowIdentityToken; column: ColumnIdentityToken }; selectingAll: boolean }
+  | { kind: "cells"; cells: Array<{ row: RowIdentityToken; column: ColumnIdentityToken }> };
+
+export interface PersistedDataGridSelection {
+  identity: { mode: "keys"; primaryKeys: readonly string[] } | { mode: "row"; columnSignature: string };
+  state: PersistedDataGridSelectionState;
+}
+
+export interface CaptureDataGridSelectionOptions {
+  columns: readonly string[];
+  sourceColumns?: readonly (string | undefined)[];
+  rows: readonly (readonly CellValue[])[];
+  primaryKeys: readonly string[];
+  visibleColumnIndexes: readonly number[];
+  displayItems: readonly DataGridSelectionRow[];
+  selectedRowIds: ReadonlySet<number>;
+  selectedColumnIndexes: ReadonlySet<number>;
+  selectedCellKeys: ReadonlySet<string>;
+  selectionAnchor: CellPosition | null;
+  selectionFocus: CellPosition | null;
+  selectingAll: boolean;
+  lastClickedRowIndex?: number | null;
+}
+
+export type RestoredDataGridSelection =
+  | { kind: "rows"; rowIds: number[]; anchorRowIndex: number | null; scrollRowIndex: number }
+  | { kind: "columns"; columnIndexes: number[] }
+  | { kind: "range"; anchor: CellPosition; focus: CellPosition; selectingAll: boolean; scrollRowIndex: number }
+  | { kind: "cells"; cellKeys: Set<string>; scrollRowIndex: number };
+
+// 刷新后恢复 range/离散单元格选区时，需要重新激活 cell-selection 拖拽状态机
+// （isSelectingCells）。否则 restore 前调用的 clearCellSelection 已将其置 false，
+// 而 beginCellSelection 才会重新置 true，导致用户必须重新点一次单元格才能拖拽扩展
+// 刚恢复的选区。行选/列选是独立状态机，不应激活 cell selection。
+export function reactivatesCellSelectionAfterRestore(restored: RestoredDataGridSelection): boolean {
+  return restored.kind === "range" || restored.kind === "cells";
+}
+
+export interface RestoreDataGridSelectionOptions {
+  snapshot: PersistedDataGridSelection;
+  columns: readonly string[];
+  sourceColumns?: readonly (string | undefined)[];
+  rows: readonly (readonly CellValue[])[];
+  visibleColumnIndexes: readonly number[];
+  displayItems: readonly DataGridSelectionRow[];
+}
+
+function normalizeColumnName(value: string | undefined): string {
+  return value?.trim().toLocaleLowerCase() ?? "";
+}
+
+function columnSignature(columns: readonly string[], sourceColumns: readonly (string | undefined)[] | undefined): string {
+  return columns.map((column, index) => `${normalizeColumnName(column)}\u0001${normalizeColumnName(sourceColumns?.[index])}`).join("\u0002");
+}
+
+function identityColumnIndexes(columns: readonly string[], sourceColumns: readonly (string | undefined)[] | undefined, primaryKeys: readonly string[]): number[] | null {
+  if (primaryKeys.length === 0) return null;
+  const indexes = primaryKeys.map((primaryKey) => {
+    const target = normalizeColumnName(primaryKey);
+    const sourceIndex = sourceColumns?.findIndex((column) => normalizeColumnName(column) === target) ?? -1;
+    if (sourceIndex >= 0) return sourceIndex;
+    return columns.findIndex((column, index) => !normalizeColumnName(sourceColumns?.[index]) && normalizeColumnName(column) === target);
+  });
+  return indexes.every((index) => index >= 0) ? indexes : null;
+}
+
+function serializeCellValue(value: CellValue): string {
+  if (value === null) return "null";
+  if (typeof value === "string") return `string:${value.length}:${value}`;
+  if (typeof value === "boolean") return value ? "boolean:1" : "boolean:0";
+  if (Number.isNaN(value)) return "number:NaN";
+  if (Object.is(value, -0)) return "number:-0";
+  return `number:${String(value)}`;
+}
+
+function rowIdentityKey(row: readonly CellValue[], indexes: readonly number[]): string {
+  return indexes.map((index) => serializeCellValue(row[index] ?? null)).join("\u0000");
+}
+
+function buildRowTokens(rows: readonly (readonly CellValue[])[], indexes: readonly number[]): RowIdentityToken[] {
+  const occurrences = new Map<string, number>();
+  return rows.map((row) => {
+    const key = rowIdentityKey(row, indexes);
+    const occurrence = occurrences.get(key) ?? 0;
+    occurrences.set(key, occurrence + 1);
+    return { key, occurrence };
+  });
+}
+
+function rowTokenKey(token: RowIdentityToken): string {
+  return `${token.key}\u0003${token.occurrence}`;
+}
+
+function buildColumnTokens(columns: readonly string[], sourceColumns: readonly (string | undefined)[] | undefined): ColumnIdentityToken[] {
+  const occurrences = new Map<string, number>();
+  return columns.map((resultName, index) => {
+    const sourceName = sourceColumns?.[index];
+    const key = `${normalizeColumnName(resultName)}\u0001${normalizeColumnName(sourceName)}`;
+    const occurrence = occurrences.get(key) ?? 0;
+    occurrences.set(key, occurrence + 1);
+    return { resultName, sourceName, occurrence };
+  });
+}
+
+function columnTokenKey(token: ColumnIdentityToken): string {
+  return `${normalizeColumnName(token.resultName)}\u0001${normalizeColumnName(token.sourceName)}\u0003${token.occurrence}`;
+}
+
+function parseCellKey(key: string): CellPosition | null {
+  const [row, column] = key.split(":");
+  const rowIndex = Number(row);
+  const colIndex = Number(column);
+  if (!Number.isInteger(rowIndex) || !Number.isInteger(colIndex) || rowIndex < 0 || colIndex < 0) return null;
+  return { rowIndex, colIndex };
+}
+
+export function captureDataGridSelection(options: CaptureDataGridSelectionOptions): PersistedDataGridSelection | null {
+  const hasSelection = options.selectedRowIds.size > 0 || options.selectedColumnIndexes.size > 0 || options.selectedCellKeys.size > 0 || (!!options.selectionAnchor && !!options.selectionFocus);
+  if (!hasSelection) return null;
+
+  const keyIndexes = identityColumnIndexes(options.columns, options.sourceColumns, options.primaryKeys);
+  const identity = keyIndexes ? ({ mode: "keys", primaryKeys: [...options.primaryKeys] } as const) : ({ mode: "row", columnSignature: columnSignature(options.columns, options.sourceColumns) } as const);
+  const columnTokens = buildColumnTokens(options.columns, options.sourceColumns);
+  if (options.selectedColumnIndexes.size > 0) {
+    const columns = [...options.selectedColumnIndexes]
+      .map((visibleColumnIndex) => {
+        const actualColumnIndex = options.visibleColumnIndexes[visibleColumnIndex];
+        return actualColumnIndex === undefined ? undefined : columnTokens[actualColumnIndex];
+      })
+      .filter((token): token is ColumnIdentityToken => !!token);
+    return columns.length > 0 ? { identity, state: { kind: "columns", columns } } : null;
+  }
+
+  const rowTokens = buildRowTokens(options.rows, keyIndexes ?? options.columns.map((_, index) => index));
+  const displayRowToken = (displayRowIndex: number): RowIdentityToken | undefined => {
+    const sourceIndex = options.displayItems[displayRowIndex]?.sourceIndex;
+    return sourceIndex === undefined ? undefined : rowTokens[sourceIndex];
+  };
+  const visibleColumnToken = (visibleColumnIndex: number): ColumnIdentityToken | undefined => {
+    const actualColumnIndex = options.visibleColumnIndexes[visibleColumnIndex];
+    return actualColumnIndex === undefined ? undefined : columnTokens[actualColumnIndex];
+  };
+
+  if (options.selectedRowIds.size > 0) {
+    const rows = options.displayItems
+      .filter((item) => options.selectedRowIds.has(item.id) && item.sourceIndex !== undefined)
+      .map((item) => rowTokens[item.sourceIndex!])
+      .filter((token): token is RowIdentityToken => !!token);
+    if (rows.length === 0) return null;
+    const anchorSourceIndex = options.lastClickedRowIndex === null || options.lastClickedRowIndex === undefined ? undefined : options.displayItems[options.lastClickedRowIndex]?.sourceIndex;
+    return { identity, state: { kind: "rows", rows, anchorRow: anchorSourceIndex === undefined ? undefined : rowTokens[anchorSourceIndex] } };
+  }
+
+  if (options.selectedCellKeys.size > 0) {
+    const cells = [...options.selectedCellKeys]
+      .map(parseCellKey)
+      .filter((cell): cell is CellPosition => !!cell)
+      .map((cell) => ({ row: displayRowToken(cell.rowIndex), column: visibleColumnToken(cell.colIndex) }))
+      .filter((cell): cell is { row: RowIdentityToken; column: ColumnIdentityToken } => !!cell.row && !!cell.column);
+    return cells.length > 0 ? { identity, state: { kind: "cells", cells } } : null;
+  }
+
+  if (options.selectionAnchor && options.selectionFocus) {
+    const anchorRow = displayRowToken(options.selectionAnchor.rowIndex);
+    const anchorColumn = visibleColumnToken(options.selectionAnchor.colIndex);
+    const focusRow = displayRowToken(options.selectionFocus.rowIndex);
+    const focusColumn = visibleColumnToken(options.selectionFocus.colIndex);
+    if (!anchorRow || !anchorColumn || !focusRow || !focusColumn) return null;
+    return {
+      identity,
+      state: {
+        kind: "range",
+        anchor: { row: anchorRow, column: anchorColumn },
+        focus: { row: focusRow, column: focusColumn },
+        selectingAll: options.selectingAll,
+      },
+    };
+  }
+
+  return null;
+}
+
+export function restoreDataGridSelection(options: RestoreDataGridSelectionOptions): RestoredDataGridSelection | null {
+  const state = options.snapshot.state;
+  const columnTokens = buildColumnTokens(options.columns, options.sourceColumns);
+  const actualColumnIndexByToken = new Map(columnTokens.map((token, index) => [columnTokenKey(token), index]));
+  const visibleColumnIndex = (token: ColumnIdentityToken): number | undefined => {
+    const actualColumnIndex = actualColumnIndexByToken.get(columnTokenKey(token));
+    if (actualColumnIndex === undefined) return undefined;
+    const index = options.visibleColumnIndexes.indexOf(actualColumnIndex);
+    return index >= 0 ? index : undefined;
+  };
+
+  if (state.kind === "columns") {
+    const columnIndexes = state.columns.map(visibleColumnIndex).filter((index): index is number => index !== undefined);
+    return columnIndexes.length > 0 ? { kind: "columns", columnIndexes } : null;
+  }
+
+  const identityIndexes =
+    options.snapshot.identity.mode === "keys"
+      ? identityColumnIndexes(options.columns, options.sourceColumns, options.snapshot.identity.primaryKeys)
+      : options.snapshot.identity.columnSignature === columnSignature(options.columns, options.sourceColumns)
+        ? options.columns.map((_, index) => index)
+        : null;
+  if (!identityIndexes) return null;
+
+  const rowTokens = buildRowTokens(options.rows, identityIndexes);
+  const sourceIndexByToken = new Map(rowTokens.map((token, sourceIndex) => [rowTokenKey(token), sourceIndex]));
+  const displayRowIndexBySourceIndex = new Map<number, number>();
+  options.displayItems.forEach((item, displayRowIndex) => {
+    if (item.sourceIndex !== undefined) displayRowIndexBySourceIndex.set(item.sourceIndex, displayRowIndex);
+  });
+  const displayRowIndex = (token: RowIdentityToken): number | undefined => {
+    const sourceIndex = sourceIndexByToken.get(rowTokenKey(token));
+    return sourceIndex === undefined ? undefined : displayRowIndexBySourceIndex.get(sourceIndex);
+  };
+
+  if (state.kind === "rows") {
+    const restoredRows = state.rows
+      .map((token) => displayRowIndex(token))
+      .filter((index): index is number => index !== undefined)
+      .map((index) => ({ index, id: options.displayItems[index]!.id }));
+    if (restoredRows.length === 0) return null;
+    return {
+      kind: "rows",
+      rowIds: restoredRows.map((row) => row.id),
+      anchorRowIndex: state.anchorRow ? (displayRowIndex(state.anchorRow) ?? restoredRows[0]!.index) : restoredRows[0]!.index,
+      scrollRowIndex: restoredRows[0]!.index,
+    };
+  }
+
+  if (state.kind === "range") {
+    if (state.selectingAll && options.displayItems.length > 0 && options.visibleColumnIndexes.length > 0) {
+      return {
+        kind: "range",
+        anchor: { rowIndex: 0, colIndex: 0 },
+        focus: { rowIndex: options.displayItems.length - 1, colIndex: options.visibleColumnIndexes.length - 1 },
+        selectingAll: true,
+        scrollRowIndex: 0,
+      };
+    }
+    const anchorRow = displayRowIndex(state.anchor.row);
+    const anchorColumn = visibleColumnIndex(state.anchor.column);
+    const focusRow = displayRowIndex(state.focus.row);
+    const focusColumn = visibleColumnIndex(state.focus.column);
+    if (anchorRow === undefined || anchorColumn === undefined || focusRow === undefined || focusColumn === undefined) return null;
+    return {
+      kind: "range",
+      anchor: { rowIndex: anchorRow, colIndex: anchorColumn },
+      focus: { rowIndex: focusRow, colIndex: focusColumn },
+      selectingAll: state.selectingAll,
+      scrollRowIndex: focusRow,
+    };
+  }
+
+  const cellKeys = new Set<string>();
+  let scrollRowIndex = -1;
+  for (const cell of state.cells) {
+    const rowIndex = displayRowIndex(cell.row);
+    const colIndex = visibleColumnIndex(cell.column);
+    if (rowIndex === undefined || colIndex === undefined) continue;
+    cellKeys.add(`${rowIndex}:${colIndex}`);
+    scrollRowIndex = rowIndex;
+  }
+  return cellKeys.size > 0 ? { kind: "cells", cellKeys, scrollRowIndex } : null;
+}