[lexical-table] Bug Fix: Shift+Arrow no longer throws when the anchor is off the rect corner

packages/lexical-table/src/LexicalTableSelectionHelpers.ts

@@ -1775,17 +1775,29 @@ function getCorner(
   return [colName, rowName];
 }
 
-function getCornerOrThrow(
+/**
+ * Resolve the corner of `rect` that the anchor sits on.
+ *
+ * `$computeTableCellRectBoundary` grows the rect until it contains every
+ * merged cell that straddles an edge, so the anchor is not guaranteed to be at
+ * a corner of the result — a cell merged across the rect's edge pushes that
+ * edge past the anchor. Fall back the same way {@link $extractRectCorners}
+ * does: to the corner opposite the focus, and finally to the top-left.
+ */
+function getAnchorCorner(
   rect: TableCellRectBoundary,
-  cellValue: TableMapValueType,
+  anchorCellValue: TableMapValueType,
+  focusCellValue: TableMapValueType,
 ): Corner {
-  const corner = getCorner(rect, cellValue);
-  invariant(
-    corner !== null,
-    'getCornerOrThrow: cell %s is not at a corner of rect',
-    cellValue.cell.getKey(),
-  );
-  return corner;
+  const anchorCorner = getCorner(rect, anchorCellValue);
+  if (anchorCorner) {
+    return anchorCorner;
+  }
+  const focusCorner = getCorner(rect, focusCellValue);
+  if (focusCorner) {
+    return oppositeCorner(focusCorner);
+  }
+  return ['minColumn', 'minRow'];
 }
 
 function oppositeCorner([colName, rowName]: Corner): Corner {
@@ -1868,7 +1880,7 @@ function $adjustFocusInDirection(
   );
   const spans = $computeTableCellRectSpans(tableMap, rect);
   const {topSpan, leftSpan, bottomSpan, rightSpan} = spans;
-  const anchorCorner = getCornerOrThrow(rect, anchorCellValue);
+  const anchorCorner = getAnchorCorner(rect, anchorCellValue, focusCellValue);
   const [focusColumn, focusRow] = oppositeCorner(anchorCorner);
   let fCol = rect[focusColumn];
   let fRow = rect[focusRow];

packages/lexical-table/src/__tests__/unit/ShiftArrowOffCornerAnchor.test.ts

@@ -0,0 +1,103 @@
+/**
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ *
+ */
+
+import {buildEditorFromExtensions} from '@lexical/extension';
+import {
+  $computeTableMapSkipCellCheck,
+  $createTableNodeWithDimensions,
+  $createTableSelectionFrom,
+  $isTableNode,
+  $isTableSelection,
+  $mergeCells,
+  TableExtension,
+} from '@lexical/table';
+import {
+  $getRoot,
+  $getSelection,
+  $setSelection,
+  defineExtension,
+  KEY_ARROW_DOWN_COMMAND,
+} from 'lexical';
+import {afterEach, assert, beforeEach, describe, expect, test} from 'vitest';
+
+let container: HTMLDivElement;
+let editor: ReturnType<typeof buildEditorFromExtensions>;
+
+beforeEach(() => {
+  container = document.createElement('div');
+  document.body.appendChild(container);
+  editor = buildEditorFromExtensions(
+    defineExtension({
+      dependencies: [TableExtension],
+      name: 'shift-arrow-corner-host',
+    }),
+  );
+  // The arrow key handlers come from the table selection observer, which needs
+  // the editor to have a root element.
+  editor.setRootElement(container);
+});
+
+afterEach(() => {
+  editor.dispose();
+  document.body.removeChild(container);
+});
+
+function $table() {
+  const table = $getRoot().getFirstChild();
+  assert($isTableNode(table), 'expected a TableNode at the root');
+  return table;
+}
+
+describe('Shift+Arrow with an anchor that is not on a rect corner', () => {
+  test('extends the selection instead of throwing', () => {
+    editor.update(
+      () => {
+        const table = $createTableNodeWithDimensions(3, 3, false);
+        $getRoot().clear().append(table);
+        const [map] = $computeTableMapSkipCellCheck(table, null, null);
+        // Merge grid columns 0 and 1 of row 1. This cell straddles the left
+        // edge of the rect selected below, so $computeTableCellRectBoundary
+        // grows the rect out to column 0 — past the anchor, which then sits on
+        // no corner of it.
+        const merged = $mergeCells([map[1][0].cell, map[1][1].cell]);
+        assert(merged !== null, 'expected the cells to merge');
+      },
+      {discrete: true},
+    );
+
+    editor.update(
+      () => {
+        const table = $table();
+        const [map] = $computeTableMapSkipCellCheck(table, null, null);
+        // Anchor at (row 0, column 1), focus at (row 1, column 2).
+        $setSelection(
+          $createTableSelectionFrom(table, map[0][1].cell, map[1][2].cell),
+        );
+      },
+      {discrete: true},
+    );
+
+    expect(() =>
+      editor.update(
+        () => {
+          editor.dispatchCommand(
+            KEY_ARROW_DOWN_COMMAND,
+            new KeyboardEvent('keydown', {key: 'ArrowDown', shiftKey: true}),
+          );
+        },
+        {discrete: true},
+      ),
+    ).not.toThrow();
+
+    editor.read('latest', () => {
+      const selection = $getSelection();
+      assert($isTableSelection(selection), 'expected a TableSelection');
+      expect(selection.isValid()).toBe(true);
+    });
+  });
+});