From e0dafb8df7b95959b38ea69d10dda6cf554a8ee2 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Germ=C3=A1n=20Jablo=C3=B1ski?=
 <43938777+GermanJablo@users.noreply.github.com>
Date: Sat, 21 Dec 2024 13:59:01 -0300
Subject: [PATCH] feature: expose forEachSelectedTextNode (#6981)

Co-authored-by: Bob Ippolito <bob@redivi.com>
---
 packages/lexical-selection/src/index.ts       |  2 +
 .../lexical-selection/src/lexical-node.ts     | 38 +++++++++++--------
 packages/lexical/src/LexicalSelection.ts      |  1 +
 3 files changed, 26 insertions(+), 15 deletions(-)

diff --git a/packages/lexical-selection/src/index.ts b/packages/lexical-selection/src/index.ts
index 8d9d47ce635..2fd6becf09b 100644
--- a/packages/lexical-selection/src/index.ts
+++ b/packages/lexical-selection/src/index.ts
@@ -8,6 +8,7 @@
 
 import {
   $addNodeStyle,
+  $forEachSelectedTextNode,
   $isAtNodeEnd,
   $patchStyleText,
   $sliceSelectedTextNodeContent,
@@ -36,6 +37,7 @@ export {
 
 export {
   $addNodeStyle,
+  $forEachSelectedTextNode,
   $isAtNodeEnd,
   $patchStyleText,
   $sliceSelectedTextNodeContent,
diff --git a/packages/lexical-selection/src/lexical-node.ts b/packages/lexical-selection/src/lexical-node.ts
index 32b7b0b802e..22074d1fd6e 100644
--- a/packages/lexical-selection/src/lexical-node.ts
+++ b/packages/lexical-selection/src/lexical-node.ts
@@ -10,6 +10,7 @@ import {
   $getCharacterOffsets,
   $getNodeByKey,
   $getPreviousSelection,
+  $getSelection,
   $isElementNode,
   $isRangeSelection,
   $isRootNode,
@@ -288,23 +289,30 @@ export function $patchStyleText(
       ) => string)
   >,
 ): void {
-  const selectedNodes = selection.getNodes();
-  const selectedNodesLength = selectedNodes.length;
-  const anchorAndFocus = selection.getStartEndPoints();
-  if (anchorAndFocus === null) {
+  if (selection.isCollapsed() && $isRangeSelection(selection)) {
+    $patchStyle(selection, patch);
+  } else {
+    $forEachSelectedTextNode((textNode) => {
+      $patchStyle(textNode, patch);
+    });
+  }
+}
+
+export function $forEachSelectedTextNode(
+  fn: (textNode: TextNode) => void,
+): void {
+  const selection = $getSelection();
+  if (!$isRangeSelection(selection)) {
     return;
   }
-  const [anchor, focus] = anchorAndFocus;
+  const selectedNodes = selection.getNodes();
+  const selectedNodesLength = selectedNodes.length;
+  const {anchor, focus} = selection;
 
   const lastIndex = selectedNodesLength - 1;
   let firstNode = selectedNodes[0];
   let lastNode = selectedNodes[lastIndex];
 
-  if (selection.isCollapsed() && $isRangeSelection(selection)) {
-    $patchStyle(selection, patch);
-    return;
-  }
-
   const firstNodeText = firstNode.getTextContent();
   const firstNodeTextLength = firstNodeText.length;
   const focusOffset = focus.offset;
@@ -355,14 +363,14 @@ export function $patchStyleText(
         $isTokenOrSegmented(firstNode) ||
         (startOffset === 0 && endOffset === firstNodeTextLength)
       ) {
-        $patchStyle(firstNode, patch);
+        fn(firstNode);
         firstNode.select(startOffset, endOffset);
       } else {
         // The node is partially selected, so split it into two nodes
         // and style the selected one.
         const splitNodes = firstNode.splitText(startOffset, endOffset);
         const replacement = startOffset === 0 ? splitNodes[0] : splitNodes[1];
-        $patchStyle(replacement, patch);
+        fn(replacement);
         replacement.select(0, endOffset - startOffset);
       }
     } // multiple nodes selected.
@@ -383,7 +391,7 @@ export function $patchStyleText(
         }
       }
 
-      $patchStyle(firstNode as TextNode, patch);
+      fn(firstNode as TextNode);
     }
 
     if ($isTextNode(lastNode) && lastNode.canHaveFormat()) {
@@ -404,7 +412,7 @@ export function $patchStyleText(
       }
 
       if (endOffset !== 0 || endType === 'element') {
-        $patchStyle(lastNode as TextNode, patch);
+        fn(lastNode as TextNode);
       }
     }
 
@@ -420,7 +428,7 @@ export function $patchStyleText(
         selectedNodeKey !== lastNode.getKey() &&
         !selectedNode.isToken()
       ) {
-        $patchStyle(selectedNode, patch);
+        fn(selectedNode as TextNode);
       }
     }
   }
diff --git a/packages/lexical/src/LexicalSelection.ts b/packages/lexical/src/LexicalSelection.ts
index fb7a62dbea0..1de9a01ffaa 100644
--- a/packages/lexical/src/LexicalSelection.ts
+++ b/packages/lexical/src/LexicalSelection.ts
@@ -1180,6 +1180,7 @@ export class RangeSelection implements BaseSelection {
     }
   }
 
+  // TO-DO: Migrate this method to the new utility function $forEachSelectedTextNode (share similar logic)
   /**
    * Applies the provided format to the TextNodes in the Selection, splitting or
    * merging nodes as necessary.