From 2dac7f5c4751bf772bf573ad7e3665e1482a5bfd Mon Sep 17 00:00:00 2001
From: David Lacroix <lacroix.david1@gmail.com>
Date: Wed, 22 May 2024 11:42:02 -0400
Subject: [PATCH] [lexical-selection] Bug Fix  /  Fixes text formatting with
 segmented and token nodes #6059 (#6062)

---
 .../__tests__/unit/LexicalSelection.test.tsx  | 106 ++++++++++++++++++
 packages/lexical/src/LexicalSelection.ts      |  20 ++--
 2 files changed, 118 insertions(+), 8 deletions(-)

diff --git a/packages/lexical-selection/src/__tests__/unit/LexicalSelection.test.tsx b/packages/lexical-selection/src/__tests__/unit/LexicalSelection.test.tsx
index 2a9930a4276..1d8caa48066 100644
--- a/packages/lexical-selection/src/__tests__/unit/LexicalSelection.test.tsx
+++ b/packages/lexical-selection/src/__tests__/unit/LexicalSelection.test.tsx
@@ -625,6 +625,112 @@ describe('LexicalSelection tests', () => {
       ],
       name: 'Format selection that starts on element and ends on text and retain selection',
     },
+
+    {
+      expectedHTML:
+        '<div contenteditable="true" style="user-select: text; white-space: pre-wrap; word-break: break-word;" data-lexical-editor="true">' +
+        '<p class="editor-paragraph"><br></p>' +
+        '<p class="editor-paragraph" dir="ltr">' +
+        '<strong class="editor-text-bold" data-lexical-text="true">Hello</strong><strong class="editor-text-bold" data-lexical-text="true"> world</strong>' +
+        '</p>' +
+        '<p class="editor-paragraph"><br></p>' +
+        '</div>',
+      expectedSelection: {
+        anchorOffset: 2,
+        anchorPath: [1, 0, 0],
+        focusOffset: 0,
+        focusPath: [2],
+      },
+      inputs: [
+        insertParagraph(),
+        insertTokenNode('Hello'),
+        insertText(' world'),
+        insertParagraph(),
+        moveNativeSelection([1, 0, 0], 2, [2], 0),
+        formatBold(),
+      ],
+      name: 'Format selection that starts on middle of token node should format complete node',
+    },
+
+    {
+      expectedHTML:
+        '<div contenteditable="true" style="user-select: text; white-space: pre-wrap; word-break: break-word;" data-lexical-editor="true">' +
+        '<p class="editor-paragraph"><br></p>' +
+        '<p class="editor-paragraph" dir="ltr">' +
+        '<strong class="editor-text-bold" data-lexical-text="true">Hello </strong><strong class="editor-text-bold" data-lexical-text="true">world</strong>' +
+        '</p>' +
+        '<p class="editor-paragraph"><br></p>' +
+        '</div>',
+      expectedSelection: {
+        anchorOffset: 0,
+        anchorPath: [0],
+        focusOffset: 2,
+        focusPath: [1, 1, 0],
+      },
+      inputs: [
+        insertParagraph(),
+        insertText('Hello '),
+        insertTokenNode('world'),
+        insertParagraph(),
+        moveNativeSelection([0], 0, [1, 1, 0], 2),
+        formatBold(),
+      ],
+      name: 'Format selection that ends on middle of token node should format complete node',
+    },
+
+    {
+      expectedHTML:
+        '<div contenteditable="true" style="user-select: text; white-space: pre-wrap; word-break: break-word;" data-lexical-editor="true">' +
+        '<p class="editor-paragraph"><br></p>' +
+        '<p class="editor-paragraph" dir="ltr">' +
+        '<strong class="editor-text-bold" data-lexical-text="true">Hello</strong><span data-lexical-text="true"> world</span>' +
+        '</p>' +
+        '<p class="editor-paragraph"><br></p>' +
+        '</div>',
+      expectedSelection: {
+        anchorOffset: 2,
+        anchorPath: [1, 0, 0],
+        focusOffset: 3,
+        focusPath: [1, 0, 0],
+      },
+      inputs: [
+        insertParagraph(),
+        insertTokenNode('Hello'),
+        insertText(' world'),
+        insertParagraph(),
+        moveNativeSelection([1, 0, 0], 2, [1, 0, 0], 3),
+        formatBold(),
+      ],
+      name: 'Format token node if it is the single one selected',
+    },
+
+    {
+      expectedHTML:
+        '<div contenteditable="true" style="user-select: text; white-space: pre-wrap; word-break: break-word;" data-lexical-editor="true">' +
+        '<p class="editor-paragraph"><br></p>' +
+        '<p class="editor-paragraph" dir="ltr">' +
+        '<strong class="editor-text-bold" data-lexical-text="true">Hello </strong><strong class="editor-text-bold" data-lexical-text="true">beautiful</strong><strong class="editor-text-bold" data-lexical-text="true"> world</strong>' +
+        '</p>' +
+        '<p class="editor-paragraph"><br></p>' +
+        '</div>',
+      expectedSelection: {
+        anchorOffset: 0,
+        anchorPath: [0],
+        focusOffset: 0,
+        focusPath: [2],
+      },
+      inputs: [
+        insertParagraph(),
+        insertText('Hello '),
+        insertTokenNode('beautiful'),
+        insertText(' world'),
+        insertParagraph(),
+        moveNativeSelection([0], 0, [2], 0),
+        formatBold(),
+      ],
+      name: 'Format selection that contains a token node in the middle should format the token node',
+    },
+
     // Tests need fixing:
     // ...GRAPHEME_SCENARIOS.flatMap(({description, grapheme}) => [
     //   {
diff --git a/packages/lexical/src/LexicalSelection.ts b/packages/lexical/src/LexicalSelection.ts
index a62cb36e00b..44db3020f3a 100644
--- a/packages/lexical/src/LexicalSelection.ts
+++ b/packages/lexical/src/LexicalSelection.ts
@@ -1132,8 +1132,11 @@ export class RangeSelection implements BaseSelection {
       if (startOffset === endOffset) {
         return;
       }
-      // The entire node is selected, so just format it
-      if (startOffset === 0 && endOffset === firstNode.getTextContentSize()) {
+      // The entire node is selected or it is token, so just format it
+      if (
+        $isTokenOrSegmented(firstNode) ||
+        (startOffset === 0 && endOffset === firstNode.getTextContentSize())
+      ) {
         firstNode.setFormat(firstNextFormat);
       } else {
         // Node is partially selected, so split it into two nodes
@@ -1157,7 +1160,7 @@ export class RangeSelection implements BaseSelection {
     }
     // Multiple nodes selected
     // The entire first node isn't selected, so split it
-    if (startOffset !== 0) {
+    if (startOffset !== 0 && !$isTokenOrSegmented(firstNode)) {
       [, firstNode as TextNode] = firstNode.splitText(startOffset);
       startOffset = 0;
     }
@@ -1167,7 +1170,10 @@ export class RangeSelection implements BaseSelection {
     // If the offset is 0, it means no actual characters are selected,
     // so we skip formatting the last node altogether.
     if (endOffset > 0) {
-      if (endOffset !== lastNode.getTextContentSize()) {
+      if (
+        endOffset !== lastNode.getTextContentSize() &&
+        !$isTokenOrSegmented(lastNode)
+      ) {
         [lastNode as TextNode] = lastNode.splitText(endOffset);
       }
       lastNode.setFormat(lastNextFormat);
@@ -1176,10 +1182,8 @@ export class RangeSelection implements BaseSelection {
     // Process all text nodes in between
     for (let i = firstIndex + 1; i < lastIndex; i++) {
       const textNode = selectedTextNodes[i];
-      if (!textNode.isToken()) {
-        const nextFormat = textNode.getFormatFlags(formatType, lastNextFormat);
-        textNode.setFormat(nextFormat);
-      }
+      const nextFormat = textNode.getFormatFlags(formatType, lastNextFormat);
+      textNode.setFormat(nextFormat);
     }
 
     // Update selection only if starts/ends on text node