Fix selection started from inside of the table (#5766)

Author: AlexanderReznik

packages/lexical-table/src/LexicalTableObserver.ts

@@ -24,7 +24,6 @@ import {
   $setSelection,
   SELECTION_CHANGE_COMMAND,
 } from 'lexical';
-import {CAN_USE_DOM} from 'shared/canUseDOM';
 import invariant from 'shared/invariant';
 
 import {$isTableCellNode} from './LexicalTableCellNode';
@@ -34,7 +33,12 @@ import {
   $createTableSelection,
   $isTableSelection,
 } from './LexicalTableSelection';
-import {$updateDOMForSelection, getTable} from './LexicalTableSelectionHelpers';
+import {
+  $findTableNode,
+  $updateDOMForSelection,
+  getDOMSelection,
+  getTable,
+} from './LexicalTableSelectionHelpers';
 
 export type TableDOMCell = {
   elem: HTMLElement;
@@ -52,9 +56,6 @@ export type TableDOMTable = {
   rows: number;
 };
 
-const getDOMSelection = (targetWindow: Window | null): Selection | null =>
-  CAN_USE_DOM ? (targetWindow || window).getSelection() : null;
-
 export class TableObserver {
   focusX: number;
   focusY: number;
@@ -283,7 +284,8 @@ export class TableObserver {
         if (
           this.tableSelection != null &&
           this.anchorCellNodeKey != null &&
-          $isTableCellNode(focusTableCellNode)
+          $isTableCellNode(focusTableCellNode) &&
+          tableNode.is($findTableNode(focusTableCellNode))
         ) {
           const focusNodeKey = focusTableCellNode.getKey();
 

packages/lexical-table/src/LexicalTableSelectionHelpers.ts

@@ -23,6 +23,7 @@ import type {
 import {$findMatchingParent} from '@lexical/utils';
 import {
   $createParagraphNode,
+  $createRangeSelectionFromDom,
   $getNearestNodeFromDOMNode,
   $getPreviousSelection,
   $getSelection,
@@ -50,6 +51,7 @@ import {
   SELECTION_CHANGE_COMMAND,
   SELECTION_INSERT_CLIPBOARD_NODES_COMMAND,
 } from 'lexical';
+import {CAN_USE_DOM} from 'shared/canUseDOM';
 import invariant from 'shared/invariant';
 
 import {$isTableCellNode} from './LexicalTableCellNode';
@@ -64,6 +66,11 @@ import {$computeTableMap} from './LexicalTableUtils';
 
 const LEXICAL_ELEMENT_KEY = '__lexicalTableSelection';
 
+export const getDOMSelection = (
+  targetWindow: Window | null,
+): Selection | null =>
+  CAN_USE_DOM ? (targetWindow || window).getSelection() : null;
+
 export function applyTableHandlers(
   tableNode: TableNode,
   tableElement: HTMLTableElementWithWithTableSelectionState,
@@ -684,11 +691,19 @@ export function applyTableHandlers(
 
           if (isPartialyWithinTable) {
             const newSelection = selection.clone();
-            newSelection.focus.set(
-              tableNode.getKey(),
-              isBackward ? 0 : tableNode.getChildrenSize(),
-              'element',
-            );
+            if (isFocusInside) {
+              newSelection.focus.set(
+                tableNode.getKey(),
+                isBackward ? 0 : tableNode.getChildrenSize(),
+                'element',
+              );
+            } else {
+              newSelection.anchor.set(
+                tableNode.getKey(),
+                isBackward ? tableNode.getChildrenSize() : 0,
+                'element',
+              );
+            }
             $setSelection(newSelection);
             $addHighlightStyleToTable(editor, tableObserver);
           } else if (isWithinTable) {
@@ -704,6 +719,51 @@ export function applyTableHandlers(
               );
             }
           }
+        } else if (
+          selection &&
+          $isTableSelection(selection) &&
+          selection.is(prevSelection) &&
+          selection.tableKey === tableNode.getKey()
+        ) {
+          // if selection goes outside of the table we need to change it to Range selection
+          const domSelection = getDOMSelection(editor._window);
+          if (
+            domSelection &&
+            domSelection.anchorNode &&
+            domSelection.focusNode
+          ) {
+            const focusNode = $getNearestNodeFromDOMNode(
+              domSelection.focusNode,
+            );
+            const isFocusOutside =
+              focusNode && !tableNode.is($findTableNode(focusNode));
+
+            const anchorNode = $getNearestNodeFromDOMNode(
+              domSelection.anchorNode,
+            );
+            const isAnchorInside =
+              anchorNode && tableNode.is($findTableNode(anchorNode));
+
+            if (
+              isFocusOutside &&
+              isAnchorInside &&
+              domSelection.rangeCount > 0
+            ) {
+              const newSelection = $createRangeSelectionFromDom(
+                domSelection,
+                editor,
+              );
+              if (newSelection) {
+                newSelection.anchor.set(
+                  tableNode.getKey(),
+                  selection.isBackward() ? tableNode.getChildrenSize() : 0,
+                  'element',
+                );
+                domSelection.removeAllRanges();
+                $setSelection(newSelection);
+              }
+            }
+          }
         }
 
         if (
@@ -1141,7 +1201,7 @@ function $findCellNode(node: LexicalNode): null | TableCellNode {
   return $isTableCellNode(cellNode) ? cellNode : null;
 }
 
-function $findTableNode(node: LexicalNode): null | TableNode {
+export function $findTableNode(node: LexicalNode): null | TableNode {
   const tableNode = $findMatchingParent(node, $isTableNode);
   return $isTableNode(tableNode) ? tableNode : null;
 }

packages/lexical/src/LexicalSelection.ts

@@ -2201,6 +2201,13 @@ export function internalCreateSelection(
   return lastSelection.clone();
 }
 
+export function $createRangeSelectionFromDom(
+  domSelection: Selection | null,
+  editor: LexicalEditor,
+): null | RangeSelection {
+  return internalCreateRangeSelection(null, domSelection, editor, null);
+}
+
 export function internalCreateRangeSelection(
   lastSelection: null | BaseSelection,
   domSelection: Selection | null,

packages/lexical/src/index.ts

@@ -123,6 +123,7 @@ export {
   $createNodeSelection,
   $createPoint,
   $createRangeSelection,
+  $createRangeSelectionFromDom,
   $getCharacterOffsets,
   $getPreviousSelection,
   $getSelection,