[lexical-code-prism][lexical-code-shiki] Bug Fix: don't restore a negative selection offset when a code line starts with a line break

packages/lexical-code-prism/src/CodeHighlighterPrism.ts

@@ -247,19 +247,30 @@ function $updateAndRetainSelection(
   }
 
   // If it was non-element anchor then we walk through child nodes
-  // and looking for a position of original text offset
-  node.getChildren().some(_node => {
-    const isText = $isTextNode(_node);
-    if (isText || $isLineBreakNode(_node)) {
-      const textContentSize = _node.getTextContentSize();
-      if (isText && textContentSize >= textOffset) {
-        _node.select(textOffset, textOffset);
-        return true;
+  // and looking for a position of original text offset. A LineBreakNode
+  // consumes one unit of the offset but can't host a text point, so when the
+  // offset lands on one we use an element point on the code node instead of
+  // letting the offset go negative and selecting the next text node at an
+  // out-of-range position.
+  const children = node.getChildren();
+  for (let index = 0; index < children.length; index++) {
+    const child = children[index];
+    if ($isTextNode(child)) {
+      const textContentSize = child.getTextContentSize();
+      if (textContentSize >= textOffset) {
+        child.select(textOffset, textOffset);
+        return;
       }
       textOffset -= textContentSize;
+    } else if ($isLineBreakNode(child)) {
+      if (textOffset === 0) {
+        node.select(index, index);
+        return;
+      }
+      textOffset -= 1;
     }
-    return false;
-  });
+  }
+  node.select(children.length, children.length);
 }
 
 // Finds minimal diff range between two nodes lists. It returns from/to range boundaries of prevNodes

packages/lexical-code-prism/src/__tests__/unit/CodeHighlighterPrismRetainSelection.test.ts

@@ -0,0 +1,71 @@
+/**
+ * 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 {$createCodeNode} from '@lexical/code';
+import {registerCodeHighlighting} from '@lexical/code-prism';
+import {
+  $createTextNode,
+  $getRoot,
+  $getSelection,
+  $isRangeSelection,
+} from 'lexical';
+import {initializeUnitTest} from 'lexical/src/__tests__/utils';
+import {describe, expect, test} from 'vitest';
+
+describe('CodeHighlighterPrism $updateAndRetainSelection', () => {
+  initializeUnitTest(testEnv => {
+    test.each([
+      ['\nfoo', 'X\nfoo'],
+      ['\n\nfoo', 'X\n\nfoo'],
+      ['foo\nbar', 'Xfoo\nbar'],
+    ])(
+      'retains a non-negative offset for code content %j (#8943)',
+      async (codeText, expectedText) => {
+        const {editor} = testEnv;
+        registerCodeHighlighting(editor);
+
+        await editor.update(() => {
+          const code = $createCodeNode('javascript');
+          $getRoot().clear().append(code);
+          code.append($createTextNode(codeText));
+          // Caret at the very start of the code block, while its content is
+          // still un-flattened, so the highlighting transform has to restore it.
+          code.select(0, 0);
+        });
+
+        editor.read(() => {
+          const selection = $getSelection();
+          expect($isRangeSelection(selection)).toBe(true);
+          if (!$isRangeSelection(selection)) {
+            return;
+          }
+          for (const point of [selection.anchor, selection.focus]) {
+            expect(point.offset).toBeGreaterThanOrEqual(0);
+            if (point.type === 'text') {
+              expect(point.offset).toBeLessThanOrEqual(
+                point.getNode().getTextContentSize(),
+              );
+            }
+          }
+        });
+
+        await editor.update(() => {
+          const selection = $getSelection();
+          expect($isRangeSelection(selection)).toBe(true);
+          if ($isRangeSelection(selection)) {
+            selection.insertText('X');
+          }
+        });
+
+        expect(
+          editor.read(() => $getRoot().getFirstChild()!.getTextContent()),
+        ).toBe(expectedText);
+      },
+    );
+  });
+});

packages/lexical-code-shiki/src/CodeHighlighterShiki.ts

@@ -272,19 +272,30 @@ function $updateAndRetainSelection(
   }
 
   // If it was non-element anchor then we walk through child nodes
-  // and looking for a position of original text offset
-  node.getChildren().some(_node => {
-    const isText = $isTextNode(_node);
-    if (isText || $isLineBreakNode(_node)) {
-      const textContentSize = _node.getTextContentSize();
-      if (isText && textContentSize >= textOffset) {
-        _node.select(textOffset, textOffset);
-        return true;
+  // and looking for a position of original text offset. A LineBreakNode
+  // consumes one unit of the offset but can't host a text point, so when the
+  // offset lands on one we use an element point on the code node instead of
+  // letting the offset go negative and selecting the next text node at an
+  // out-of-range position.
+  const children = node.getChildren();
+  for (let index = 0; index < children.length; index++) {
+    const child = children[index];
+    if ($isTextNode(child)) {
+      const textContentSize = child.getTextContentSize();
+      if (textContentSize >= textOffset) {
+        child.select(textOffset, textOffset);
+        return;
       }
       textOffset -= textContentSize;
+    } else if ($isLineBreakNode(child)) {
+      if (textOffset === 0) {
+        node.select(index, index);
+        return;
+      }
+      textOffset -= 1;
     }
-    return false;
-  });
+  }
+  node.select(children.length, children.length);
 }
 
 // Finds minimal diff range between two nodes lists. It returns from/to range boundaries of prevNodes

packages/lexical-code-shiki/src/__tests__/unit/CodeShikiRetainSelection.test.ts

@@ -0,0 +1,96 @@
+/**
+ * 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 {$createCodeNode} from '@lexical/code';
+import {
+  CodeShikiExtension,
+  loadCodeLanguage,
+  loadCodeTheme,
+} from '@lexical/code-shiki';
+import {buildEditorFromExtensions} from '@lexical/extension';
+import {RichTextExtension} from '@lexical/rich-text';
+import {
+  $createTextNode,
+  $getRoot,
+  $getSelection,
+  $isRangeSelection,
+  defineExtension,
+} from 'lexical';
+import {beforeAll, describe, expect, test} from 'vitest';
+
+function createEditor() {
+  return buildEditorFromExtensions(
+    defineExtension({
+      dependencies: [RichTextExtension, CodeShikiExtension],
+      name: 'code-shiki-retain-selection-test',
+    }),
+  );
+}
+
+describe('CodeHighlighterShiki $updateAndRetainSelection', () => {
+  beforeAll(async () => {
+    // Shiki defers highlighting until the grammar and theme have loaded, so
+    // load them up front to make the transform run on the first update.
+    await loadCodeLanguage('javascript');
+    await loadCodeTheme('one-light');
+  });
+
+  test.each([
+    ['\nfoo', 'X\nfoo'],
+    ['\n\nfoo', 'X\n\nfoo'],
+    ['foo\nbar', 'Xfoo\nbar'],
+  ])(
+    'retains a non-negative offset for code content %j (#8943)',
+    (codeText, expectedText) => {
+      using editor = createEditor();
+
+      editor.update(
+        () => {
+          const code = $createCodeNode('javascript');
+          $getRoot().clear().append(code);
+          code.append($createTextNode(codeText));
+          // Caret at the very start of the code block, while its content is
+          // still un-flattened, so the highlighting transform has to restore it.
+          code.select(0, 0);
+        },
+        {discrete: true},
+      );
+
+      editor.read(() => {
+        const selection = $getSelection();
+        expect($isRangeSelection(selection)).toBe(true);
+        if (!$isRangeSelection(selection)) {
+          return;
+        }
+        for (const point of [selection.anchor, selection.focus]) {
+          expect(point.offset).toBeGreaterThanOrEqual(0);
+          if (point.type === 'text') {
+            expect(point.offset).toBeLessThanOrEqual(
+              point.getNode().getTextContentSize(),
+            );
+          }
+        }
+      });
+
+      editor.update(
+        () => {
+          const selection = $getSelection();
+          expect($isRangeSelection(selection)).toBe(true);
+          if ($isRangeSelection(selection)) {
+            selection.insertText('X');
+          }
+        },
+        {discrete: true},
+      );
+
+      expect(
+        editor.read(() => $getRoot().getFirstChild()!.getTextContent()),
+      ).toBe(expectedText);
+    },
+  );
+});