From eaade844d46c1b3b7c77550e7924d0c31369426d Mon Sep 17 00:00:00 2001 From: Gerard Rovira Date: Thu, 5 Nov 2026 22:23:07 -0800 Subject: [PATCH] [lexical-code-core][lexical-code-prism][lexical-code-shiki] Bug Fix: The code highlighter keeps the NodeState of text inside a code block when it tokenizes the block again MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Description Text inside a code block can carry NodeState of its own: comment ids, a suggestion's members, which collaboration stores on the text. When text with NodeState enters a code block (a quote or paragraph merged into it, or a collaborator's change read with `skipTransforms`), the highlighter keeps it as it is until its next pass over the block, which any change that makes the CodeNode dirty runs (an edit inside the block, or a node inserted next to it). That pass tokenizes the block's text again and splices fresh CodeHighlightNodes in, so the NodeState of the text was dropped: a comment or suggestion was lost on text nobody edited. The highlighters already keep NodeState the other way, when code highlight nodes leave a code block. The new tokens now keep it (`$keepCodeNodeState` in @lexical/code-core, used by the Prism and Shiki highlighters): a token is split where the NodeState of the old children changes, and each part gets the NodeState of the node its text came from. A state the old node knows by its StateConfig stays known on the copy, here and where highlight nodes leave a code block: copied as JSON only, a StateConfig not registered for the new node's class is read back as an unknown value, which collaboration stores whole (a set-valued state as one value instead of one entry per member). What the copy has of its own stays as the copy of the JSON leaves it: a state the old node has at its default value (which its JSON leaves out) is not copied over the value the copy's class sets when it is made (a node replacement's own state), and a key the copy's class has a StateConfig of its own for is read with that one, not given the old node's StateConfig (which would put that StateConfig in the place of the class's own for every node of the class, so that a later read with the class's own throws a State key collision). The highlighters' comparison of old and new children also compares NodeState (`$isSameCodeNodeState`), so a token that gained or lost state is replaced. A block with no NodeState (the common case, on every keystroke in code) is tokenized as before: the helper returns the new tokens untouched after one pass over the children. ## Changeset `.changeset/code-retokenize-keeps-state.md`: a patch for `@lexical/code-core`, `@lexical/code-prism` and `@lexical/code-shiki`. ## Test plan `CodePrismRetokenizeKeepsState.test.ts` and `CodeShikiRetokenizeKeepsState.test.ts` (one per highlighter): - a text node with NodeState (a string and a set-valued state) inside a code block, as a collaborator's change is read (no transforms), keeps both when a paragraph inserted right after the block makes the highlighter tokenize it again, in a highlight token of its own, known by its StateConfig; - a quote merged into a code block keeps the NodeState of its text; - a block with no NodeState is tokenized exactly as before; - with a TextNode replacement, and with a code highlight node replacement, that sets a state when it is made and has a StateConfig of its own for a key: the text that takes a token's place when it leaves a code block, and a new token, keep the state their class sets where the old node has that state at its default value, and read a value the old node has under another StateConfig for that key with their class's own (no State key collision). On main, the two cells for text leaving a code block pass (the copy there is JSON only) and the two for new tokens fail (the state is dropped). ### Before ``` $ vitest --run --project unit CodePrismRetokenizeKeepsState CodeShikiRetokenizeKeepsState × a text node with NodeState inside a code block, tokenized again by an insert next to the block (Prism, Shiki) × a quote merged into a code block keeps the NodeState of its text (Prism, Shiki) ✓ without NodeState, the block is tokenized as before (Prism, Shiki) ✓ text leaving a code block keeps the state its class sets, where the token has that state at its default (Prism, Shiki) ✓ text leaving a code block is read with the StateConfig its class has for a key (Prism, Shiki) × a new token keeps the state its class sets, where the text it comes from has that state at its default (Prism, Shiki) × a new token is read with the StateConfig its class has for a key (Prism, Shiki) AssertionError: expected [] to deeply equal [ true ] AssertionError: expected 'const a = 1;t50z' to be 'const a = 1;t50z[members=s1]' // Object.is equality AssertionError: expected [] to deeply equal [ [ true, 'kept', 'initialized' ] ] AssertionError: expected [] to deeply equal [ [ true, 'kept' ] ] Tests 8 failed | 6 passed (14) ``` ### After ``` $ vitest --run --project unit CodePrismRetokenizeKeepsState CodeShikiRetokenizeKeepsState Tests 14 passed (14) ``` The whole suite, the types and Flow: ``` $ LEXICAL_Y_GC=0 vitest run --project unit --project scripts-unit Test Files 920 passed | 6 skipped (926) Tests 37176 passed | 371 expected fail | 160 skipped (37707) $ pnpm run tsc $ pnpm run flow Found 0 errors ``` Files that went over a time or memory bound while the machine was loaded were run again with four workers, and are counted from that run. Highlighting cost on every keystroke in a large block (Prism, a JavaScript block of `const valueN = fetch(N);` lines, an edit of the number in the middle line, median of five runs of 150 updates each, before -> after): ``` 1,000 lines, no NodeState: 49.7 ms -> 50.1 ms 2,000 lines, no NodeState: 111.9 ms -> 110.7 ms 1,000 lines, 20 tokens with state: 49.3 ms -> 52.5 ms 2,000 lines, 40 tokens with state: 110.4 ms -> 118.2 ms ``` The tokens with state are the numbers of every 50th line. Before, they keep their state too (main's comparison leaves a token whose text and type did not change in place); after, each pass over a block with NodeState also splits the new tokens where the state changes and compares the state of old and new children, which is what the last two pay for. --- .changeset/code-retokenize-keeps-state.md | 9 + .../flow/LexicalCodeCore.js.flow | 10 + .../src/FlatStructureUtils.ts | 160 +++++++ packages/lexical-code-core/src/index.ts | 2 + .../src/CodeHighlighterPrism.ts | 31 +- .../CodePrismRetokenizeKeepsState.test.ts | 425 +++++++++++++++++ .../src/CodeHighlighterShiki.ts | 31 +- .../CodeShikiRetokenizeKeepsState.test.ts | 441 ++++++++++++++++++ 8 files changed, 1105 insertions(+), 4 deletions(-) create mode 100644 .changeset/code-retokenize-keeps-state.md create mode 100644 packages/lexical-code-prism/src/__tests__/unit/CodePrismRetokenizeKeepsState.test.ts create mode 100644 packages/lexical-code-shiki/src/__tests__/unit/CodeShikiRetokenizeKeepsState.test.ts diff --git a/.changeset/code-retokenize-keeps-state.md b/.changeset/code-retokenize-keeps-state.md new file mode 100644 index 00000000000..222ea7e1ed5 --- /dev/null +++ b/.changeset/code-retokenize-keeps-state.md @@ -0,0 +1,9 @@ +--- +'@lexical/code-core': patch +'@lexical/code-prism': patch +'@lexical/code-shiki': patch +--- + +[fix] The code highlighter keeps the NodeState of text inside a code block (comment ids, a suggestion's members) when it tokenizes the block again: a token is split where that state changes, and each part keeps the state of the text it came from (text with NodeState that entered the block, by a merge or a collaborator's change, lost it on the next highlight pass) + +Verify: put a text node with a NodeState value inside a code block (or merge a quote whose text has one into it), then make any edit next to the block; the text keeps its value, in its own highlight token. diff --git a/packages/lexical-code-core/flow/LexicalCodeCore.js.flow b/packages/lexical-code-core/flow/LexicalCodeCore.js.flow index c59d07c7cb2..8d43f172dcc 100644 --- a/packages/lexical-code-core/flow/LexicalCodeCore.js.flow +++ b/packages/lexical-code-core/flow/LexicalCodeCore.js.flow @@ -79,6 +79,16 @@ declare export function $outdentLeadingSpaces( declare export function $plainifyCodeContent(text: string): LexicalNode[]; +declare export function $keepCodeNodeState( + prevNodes: Array, + nextNodes: Array, +): Array; + +declare export function $isSameCodeNodeState( + a: LexicalNode, + b: LexicalNode, +): boolean; + /** * CodeNode */ diff --git a/packages/lexical-code-core/src/FlatStructureUtils.ts b/packages/lexical-code-core/src/FlatStructureUtils.ts index c41ea9aa51a..066218b2836 100644 --- a/packages/lexical-code-core/src/FlatStructureUtils.ts +++ b/packages/lexical-code-core/src/FlatStructureUtils.ts @@ -11,13 +11,16 @@ import { $createLineBreakNode, $createTabNode, $getSiblingCaret, + $getWritableNodeState, $isElementNode, $isLineBreakNode, $isTabNode, + $isTextNode, type CaretDirection, getTextDirection, type LexicalNode, type LineBreakNode, + NODE_STATE_KEY, type RangeSelection, type SiblingCaret, type TabNode, @@ -250,6 +253,163 @@ export function $plainifyCodeContent(text: string): LexicalNode[] { return out; } +/** + * Give `to` the NodeState of `from`, flat states included, keeping the state + * that `to` has for the keys `from` does not set. + */ +function $copyNodeState(from: LexicalNode, to: LexicalNode): void { + // @internal access, like `__text`: the NodeState of the node + const json: Record = from.__state + ? from.__state.toJSON() + : {}; + const state = $getWritableNodeState(to); + for (const [key, value] of Object.entries(json)) { + if (key === NODE_STATE_KEY) { + for (const [nestedKey, nestedValue] of Object.entries( + value as Record, + )) { + state.updateFromUnknown(nestedKey, nestedValue); + } + } else { + state.updateFromUnknown(key, value); + } + } + // States the node knows by their StateConfig stay known on the copy (a + // StateConfig not registered for the copy's class would otherwise be + // read back as an unknown JSON value: a set-valued state, for one). As in + // the JSON: not a default value (the copy keeps what it has, such as state + // its replacement class sets), and the value of a key the copy's class + // has another StateConfig for stays parsed with that one (above). + if (from.__state) { + const {sharedConfigMap} = state.sharedNodeState; + for (const [config, value] of from.__state.getInternalState()[1]) { + if (config.isEqual(value, config.defaultValue)) { + continue; + } + const own = sharedConfigMap.get(config.key); + if (own === undefined || own === config) { + state.updateFromKnown(config, value); + } + } + } +} + +/** Whether `node` has NodeState (such as comment ids or a suggestion). */ +function hasNodeState(node: LexicalNode): boolean { + // @internal access, like `__text` + return ( + node.__state !== undefined && Object.keys(node.__state.toJSON()).length > 0 + ); +} + +/** + * @internal + * Whether `a` and `b` have the same NodeState, for a highlighter's + * comparison of a code block's children with their new tokens. + */ +export function $isSameCodeNodeState(a: LexicalNode, b: LexicalNode): boolean { + // @internal access, like `__text` + if (a.__state === b.__state) { + return true; + } + const aHas = hasNodeState(a); + const bHas = hasNodeState(b); + if (!aHas || !bHas) { + return aHas === bHas; + } + return ( + JSON.stringify(a.__state!.toJSON()) === JSON.stringify(b.__state!.toJSON()) + ); +} + +/** + * @internal + * For a highlighter that tokenizes a code block again: `nextNodes`, the new + * tokens of the text of `prevNodes` (the block's children), with the + * NodeState of `prevNodes` kept (such as comment ids or a suggestion, which + * collaboration stores on the text). A token is split where that state + * changes, and each part gets the NodeState of the node its text came from. + * Returns `nextNodes` itself when no node of `prevNodes` has NodeState. + */ +export function $keepCodeNodeState( + prevNodes: LexicalNode[], + nextNodes: LexicalNode[], +): LexicalNode[] { + // The common case, on every keystroke in code: no NodeState at all. + if (!prevNodes.some(hasNodeState)) { + return nextNodes; + } + const spans: {end: number; node: LexicalNode; start: number}[] = []; + let at = 0; + for (const node of prevNodes) { + const size = node.getTextContentSize(); + if ( + size > 0 && + $isTextNode(node) && + !$isTabNode(node) && + hasNodeState(node) + ) { + spans.push({end: at + size, node, start: at}); + } + at += size; + } + if (spans.length === 0) { + return nextNodes; + } + const out: LexicalNode[] = []; + let next = 0; + at = 0; + for (const node of nextNodes) { + const size = node.getTextContentSize(); + const start = at; + const end = at + size; + at = end; + while (next < spans.length && spans[next].end <= start) { + next++; + } + if ( + !$isCodeHighlightNode(node) || + next === spans.length || + spans[next].start >= end + ) { + out.push(node); + continue; + } + const text = node.getTextContent(); + let pos = start; + let k = next; + while (pos < end) { + const span = k < spans.length && spans[k].start < end ? spans[k] : null; + let partEnd = end; + let from: LexicalNode | null = null; + if (span !== null && pos < span.start) { + partEnd = span.start; + } else if (span !== null) { + partEnd = Math.min(end, span.end); + from = span.node; + } + let part: CodeHighlightNode = node; + if (partEnd - pos !== size) { + part = $createCodeHighlightNode( + text.slice(pos - start, partEnd - start), + node.getHighlightType(), + ); + part.setFormat(node.getFormat()); + part.setStyle(node.getStyle()); + } + if (from !== null) { + $copyNodeState(from, part); + if (partEnd === span!.end) { + k++; + } + } + out.push(part); + pos = partEnd; + } + } + return out; +} + /** * Strip up to `tabSize` leading spaces from a {@link CodeHighlightNode} that * starts a code line, to support outdenting space-indented code lines (e.g. diff --git a/packages/lexical-code-core/src/index.ts b/packages/lexical-code-core/src/index.ts index 519820b8eb1..31f221e7739 100644 --- a/packages/lexical-code-core/src/index.ts +++ b/packages/lexical-code-core/src/index.ts @@ -33,6 +33,8 @@ export { $getFirstCodeNodeOfLine, $getLastCodeNodeOfLine, $getStartOfCodeInLine, + $isSameCodeNodeState, + $keepCodeNodeState, $outdentLeadingSpaces, $plainifyCodeContent, } from './FlatStructureUtils'; diff --git a/packages/lexical-code-prism/src/CodeHighlighterPrism.ts b/packages/lexical-code-prism/src/CodeHighlighterPrism.ts index 245a881d6f8..b43a70dadbe 100644 --- a/packages/lexical-code-prism/src/CodeHighlighterPrism.ts +++ b/packages/lexical-code-prism/src/CodeHighlighterPrism.ts @@ -9,6 +9,8 @@ import { $isCodeHighlightNode, $isCodeNode, + $isSameCodeNodeState, + $keepCodeNodeState, $plainifyCodeContent, CodeExtension, CodeHighlightNode, @@ -111,6 +113,24 @@ function $copyNodeState(from: LexicalNode, to: LexicalNode): void { state.updateFromUnknown(key, value); } } + // States the node knows by their StateConfig stay known on the copy (a + // StateConfig not registered for the copy's class would otherwise be + // read back as an unknown JSON value: a set-valued state, for one). As in + // the JSON: not a default value (the copy keeps what it has, such as state + // its replacement class sets), and the value of a key the copy's class + // has another StateConfig for stays parsed with that one (above). + if (from.__state) { + const {sharedConfigMap} = state.sharedNodeState; + for (const [config, value] of from.__state.getInternalState()[1]) { + if (config.isEqual(value, config.defaultValue)) { + continue; + } + const own = sharedConfigMap.get(config.key); + if (own === undefined || own === config) { + state.updateFromKnown(config, value); + } + } + } } function $textNodeTransform( @@ -200,7 +220,13 @@ function $codeNodeTransform( ? $plainifyCodeContent(currentNode.getTextContent()) : tokenizer.$tokenize(currentNode, currentLanguage ?? undefined); - const diffRange = getDiffRange(currentNode.getChildren(), highlightNodes); + // The NodeState of the block's children (comment ids, a suggestion) + // stays on their text. + const children = currentNode.getChildren(); + const diffRange = getDiffRange( + children, + $keepCodeNodeState(children, highlightNodes), + ); const {from, to, nodesForReplacement} = diffRange; if (from !== to || nodesForReplacement.length) { @@ -353,7 +379,8 @@ function isEqual(nodeA: LexicalNode, nodeB: LexicalNode): boolean { ($isCodeHighlightNode(nodeA) && $isCodeHighlightNode(nodeB) && nodeA.__text === nodeB.__text && - nodeA.__highlightType === nodeB.__highlightType) || + nodeA.__highlightType === nodeB.__highlightType && + $isSameCodeNodeState(nodeA, nodeB)) || ($isTabNode(nodeA) && $isTabNode(nodeB)) || ($isLineBreakNode(nodeA) && $isLineBreakNode(nodeB)) ); diff --git a/packages/lexical-code-prism/src/__tests__/unit/CodePrismRetokenizeKeepsState.test.ts b/packages/lexical-code-prism/src/__tests__/unit/CodePrismRetokenizeKeepsState.test.ts new file mode 100644 index 00000000000..8d1e78e38c9 --- /dev/null +++ b/packages/lexical-code-prism/src/__tests__/unit/CodePrismRetokenizeKeepsState.test.ts @@ -0,0 +1,425 @@ +/** + * 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. + * + */ + +// @vitest-environment node + +import { + $createCodeHighlightNode, + $createCodeNode, + $isCodeHighlightNode, + CodeHighlightNode, +} from '@lexical/code-core'; +import {CodePrismExtension} from '@lexical/code-prism'; +import {buildEditorFromExtensions} from '@lexical/extension'; +import {$createQuoteNode, RichTextExtension} from '@lexical/rich-text'; +import { + $create, + $createParagraphNode, + $createTextNode, + $getRoot, + $getState, + $isElementNode, + $setState, + createState, + defineExtension, + type LexicalEditor, + type StateConfig, + TextNode, +} from 'lexical'; +import {describe, expect, test} from 'vitest'; + +// Text inside a code block can carry NodeState of its own (comment ids, a +// suggestion: collaboration stores them on the text). When the highlighter +// tokenizes the block again, the new tokens keep it: a token is split where +// the state changes, and each part has the state of the text it came from. + +const noteState = createState('note', { + parse: (value: unknown) => (typeof value === 'string' ? value : ''), +}); +/** A set-valued state, as a suggestion's or a comment's members are. */ +const membersState = createState('members', { + parse: (value: unknown) => + Array.isArray(value) ? (value as string[]) : ([] as string[]), +}); + +function editorWithCode(): LexicalEditor & Disposable { + return buildEditorFromExtensions( + defineExtension({ + $initialEditorState: null, + dependencies: [RichTextExtension, CodePrismExtension], + name: 'retokenize-keeps-state', + onError: (error: Error) => { + throw error; + }, + }), + ); +} + +/** Whether every child of the first block is a code highlight node. */ +const allTokens = (editor: LexicalEditor) => + editor.read(() => { + const block = $getRoot().getFirstChildOrThrow(); + return ( + $isElementNode(block) && block.getChildren().every($isCodeHighlightNode) + ); + }); + +/** Each text node of the first block: text, then its states if any. */ +const show = (editor: LexicalEditor) => + editor.read(() => { + const block = $getRoot().getFirstChildOrThrow(); + return ($isElementNode(block) ? block.getChildren() : []).map(node => { + const note = $getState(node, noteState); + const members = $getState(node, membersState); + return ( + node.getTextContent() + + (note === '' ? '' : `[note=${note}]`) + + (members.length === 0 ? '' : `[members=${members.join(',')}]`) + ); + }); + }); + +describe('a code block tokenized again keeps the NodeState of its text', () => { + test('a text node with NodeState inside a code block, tokenized again by an insert next to the block', () => { + using editor = editorWithCode(); + // As a collaborator's change is read: with no transforms, so the text + // node stays one node of its own inside the block. + editor.update( + () => { + const word = $createTextNode('t78z'); + $setState(word, noteState, 'kept'); + $setState(word, membersState, ['s1']); + $getRoot().append( + $createCodeNode('javascript').append( + $createCodeHighlightNode('fetch( '), + word, + $createCodeHighlightNode(' t51z'), + ), + $createParagraphNode().append($createTextNode('after')), + ); + }, + {discrete: true, skipTransforms: true}, + ); + // A paragraph inserted right after the block makes the highlighter + // tokenize it again. + editor.update( + () => { + $getRoot() + .getFirstChildOrThrow() + .insertAfter($createParagraphNode().append($createTextNode('x'))); + }, + {discrete: true}, + ); + // Known by its StateConfig on the new token, as on the text it came from + // (checked before anything reads it there): collaboration stores a known + // set-valued state one entry per member, an unknown one whole. + expect( + editor.read(() => + $getRoot() + .getAllTextNodes() + .filter(node => node.getTextContent() === 't78z') + .map(node => node.__state!.getInternalState()[1].has(membersState)), + ), + ).toEqual([true]); + const nodes = show(editor); + expect(nodes.join('')).toBe('fetch( t78z[note=kept][members=s1] t51z'); + expect(nodes).toContain('t78z[note=kept][members=s1]'); + expect(allTokens(editor)).toBe(true); + }); + + test('a quote merged into a code block keeps the NodeState of its text', () => { + using editor = editorWithCode(); + editor.update( + () => { + const word = $createTextNode('t50z'); + $setState(word, membersState, ['s1']); + $getRoot().append( + $createCodeNode('javascript').append( + $createCodeHighlightNode('const a = 1;'), + ), + $createQuoteNode().append(word), + ); + }, + {discrete: true}, + ); + editor.update( + () => { + const [code, quote] = $getRoot().getChildren(); + if ($isElementNode(code) && $isElementNode(quote)) { + code.append(...quote.getChildren()); + quote.remove(); + } + }, + {discrete: true}, + ); + expect(show(editor).join('')).toBe('const a = 1;t50z[members=s1]'); + expect(allTokens(editor)).toBe(true); + }); + + test('without NodeState, the block is tokenized as before', () => { + using editor = editorWithCode(); + editor.update( + () => { + $getRoot().append( + $createCodeNode('javascript').append( + $createCodeHighlightNode('fetch( '), + $createTextNode('t78z'), + $createCodeHighlightNode(' t51z'), + ), + ); + }, + {discrete: true}, + ); + const withoutState = show(editor); + using plain = editorWithCode(); + plain.update( + () => { + $getRoot().append( + $createCodeNode('javascript').append( + $createCodeHighlightNode('fetch( t78z t51z'), + ), + ); + }, + {discrete: true}, + ); + expect(withoutState).toEqual(show(plain)); + }); +}); + +// The copies of NodeState (onto the new tokens, and onto text that leaves a +// code block) keep what the copy has of its own, as the copy of the JSON +// does: where the old node has a state at its default value (which its JSON +// leaves out), the value the copy's class sets when it is made; and where the +// copy's class has a StateConfig of its own for a key, the value read with +// that one. + +const parseString = (value: unknown) => + typeof value === 'string' ? value : ''; +/** A state the replacement classes below set when they are made. */ +const createdState = createState('created', {parse: parseString}); +/** A key the replacement classes below have a StateConfig of their own for. */ +const ownedState = createState('owned', {parse: parseString}); +/** Another StateConfig for the same key, as another module could make. */ +const otherOwnedState = createState('owned', {parse: parseString}); + +/** A TextNode replacement that sets a state of its own when it is made. */ +class CreatedTextNode extends TextNode { + $config() { + return this.config('test-retokenize-created-text', { + extends: TextNode, + stateConfigs: [createdState, ownedState], + }); + } +} + +/** A code highlight node replacement that sets a state of its own too. */ +class CreatedHighlightNode extends CodeHighlightNode { + $config() { + return this.config('test-retokenize-created-highlight', { + extends: CodeHighlightNode, + stateConfigs: [createdState, ownedState], + }); + } +} + +function editorWithCreatedText(): LexicalEditor & Disposable { + return buildEditorFromExtensions( + defineExtension({ + $initialEditorState: null, + dependencies: [RichTextExtension, CodePrismExtension], + name: 'retokenize-keeps-state-created-text', + nodes: () => [ + CreatedTextNode, + { + replace: TextNode, + with: (node: TextNode) => + $setState( + $create(CreatedTextNode).setTextContent(node.getTextContent()), + createdState, + 'initialized', + ), + withKlass: CreatedTextNode, + }, + ], + onError: (error: Error) => { + throw error; + }, + }), + ); +} + +function editorWithCreatedHighlight(): LexicalEditor & Disposable { + return buildEditorFromExtensions( + defineExtension({ + $initialEditorState: null, + dependencies: [RichTextExtension, CodePrismExtension], + name: 'retokenize-keeps-state-created-highlight', + nodes: () => [ + CreatedHighlightNode, + { + replace: CodeHighlightNode, + with: (node: CodeHighlightNode) => + $setState( + $create(CreatedHighlightNode).setTextContent( + node.getTextContent(), + ), + createdState, + 'initialized', + ), + withKlass: CreatedHighlightNode, + }, + ], + onError: (error: Error) => { + throw error; + }, + }), + ); +} + +/** A code block whose text node `t78z` has the given states. */ +function $appendBlockWithState( + ...states: [StateConfig, string][] +): void { + const word = $createTextNode('t78z'); + for (const [stateConfig, value] of states) { + $setState(word, stateConfig, value); + } + $getRoot().append( + $createCodeNode('javascript').append( + $createCodeHighlightNode('fetch( '), + word, + $createCodeHighlightNode(' t51z'), + ), + ); +} + +/** A code block of one token with the given states. */ +function $appendTokenWithState( + ...states: [StateConfig, string][] +): void { + const token = $createCodeHighlightNode('plain'); + for (const [stateConfig, value] of states) { + $setState(token, stateConfig, value); + } + $getRoot().append($createCodeNode().append(token)); +} + +describe('the NodeState copies keep what the copy has of its own', () => { + test('text leaving a code block keeps the state its class sets, where the token has that state at its default', () => { + using editor = editorWithCreatedText(); + // The note makes the token's state JSON not empty. + editor.update( + () => + $appendTokenWithState( + [noteState, 'kept'], + [createdState, createdState.defaultValue], + ), + {discrete: true, skipTransforms: true}, + ); + editor.update( + () => { + $getRoot().getFirstChildOrThrow().replace($createParagraphNode(), true); + }, + {discrete: true}, + ); + expect( + editor.read(() => + $getRoot() + .getAllTextNodes() + .map(node => [ + $isCodeHighlightNode(node), + $getState(node, noteState), + $getState(node, createdState), + ]), + ), + ).toEqual([[false, 'kept', 'initialized']]); + }); + + test('text leaving a code block is read with the StateConfig its class has for a key', () => { + using editor = editorWithCreatedText(); + editor.update(() => $appendTokenWithState([otherOwnedState, 'kept']), { + discrete: true, + skipTransforms: true, + }); + editor.update( + () => { + $getRoot().getFirstChildOrThrow().replace($createParagraphNode(), true); + }, + {discrete: true}, + ); + expect( + editor.read(() => + $getRoot() + .getAllTextNodes() + .map(node => [ + $isCodeHighlightNode(node), + $getState(node, ownedState), + ]), + ), + ).toEqual([[false, 'kept']]); + }); + + test('a new token keeps the state its class sets, where the text it comes from has that state at its default', () => { + using editor = editorWithCreatedHighlight(); + editor.update( + () => + $appendBlockWithState( + [noteState, 'kept'], + [createdState, createdState.defaultValue], + ), + {discrete: true, skipTransforms: true}, + ); + editor.update( + () => { + $getRoot() + .getFirstChildOrThrow() + .insertAfter($createParagraphNode().append($createTextNode('x'))); + }, + {discrete: true}, + ); + expect( + editor.read(() => + $getRoot() + .getAllTextNodes() + .filter(node => node.getTextContent() === 't78z') + .map(node => [ + $isCodeHighlightNode(node), + $getState(node, noteState), + $getState(node, createdState), + ]), + ), + ).toEqual([[true, 'kept', 'initialized']]); + }); + + test('a new token is read with the StateConfig its class has for a key', () => { + using editor = editorWithCreatedHighlight(); + editor.update(() => $appendBlockWithState([otherOwnedState, 'kept']), { + discrete: true, + skipTransforms: true, + }); + editor.update( + () => { + $getRoot() + .getFirstChildOrThrow() + .insertAfter($createParagraphNode().append($createTextNode('x'))); + }, + {discrete: true}, + ); + expect( + editor.read(() => + $getRoot() + .getAllTextNodes() + .filter(node => node.getTextContent() === 't78z') + .map(node => [ + $isCodeHighlightNode(node), + $getState(node, ownedState), + ]), + ), + ).toEqual([[true, 'kept']]); + }); +}); diff --git a/packages/lexical-code-shiki/src/CodeHighlighterShiki.ts b/packages/lexical-code-shiki/src/CodeHighlighterShiki.ts index 6fef71c293e..cf339071e46 100644 --- a/packages/lexical-code-shiki/src/CodeHighlighterShiki.ts +++ b/packages/lexical-code-shiki/src/CodeHighlighterShiki.ts @@ -9,6 +9,8 @@ import { $isCodeHighlightNode, $isCodeNode, + $isSameCodeNodeState, + $keepCodeNodeState, $plainifyCodeContent, CodeExtension, CodeHighlightNode, @@ -107,6 +109,24 @@ function $copyNodeState(from: LexicalNode, to: LexicalNode): void { state.updateFromUnknown(key, value); } } + // States the node knows by their StateConfig stay known on the copy (a + // StateConfig not registered for the copy's class would otherwise be + // read back as an unknown JSON value: a set-valued state, for one). As in + // the JSON: not a default value (the copy keeps what it has, such as state + // its replacement class sets), and the value of a key the copy's class + // has another StateConfig for stays parsed with that one (above). + if (from.__state) { + const {sharedConfigMap} = state.sharedNodeState; + for (const [config, value] of from.__state.getInternalState()[1]) { + if (config.isEqual(value, config.defaultValue)) { + continue; + } + const own = sharedConfigMap.get(config.key); + if (own === undefined || own === config) { + state.updateFromKnown(config, value); + } + } + } } function $textNodeTransform( @@ -230,7 +250,13 @@ function $codeNodeTransform( lang && !isCodeLanguageLoaded(lang) ? $plainifyCodeContent(currentNode.getTextContent()) : tokenizer.$tokenize(currentNode, lang ?? undefined); - const diffRange = getDiffRange(currentNode.getChildren(), highlightNodes); + // The NodeState of the block's children (comment ids, a suggestion) + // stays on their text. + const children = currentNode.getChildren(); + const diffRange = getDiffRange( + children, + $keepCodeNodeState(children, highlightNodes), + ); const {from, to, nodesForReplacement} = diffRange; if (from !== to || nodesForReplacement.length) { @@ -384,7 +410,8 @@ function isEqual(nodeA: LexicalNode, nodeB: LexicalNode): boolean { $isCodeHighlightNode(nodeB) && nodeA.__text === nodeB.__text && nodeA.__highlightType === nodeB.__highlightType && - nodeA.__style === nodeB.__style) || + nodeA.__style === nodeB.__style && + $isSameCodeNodeState(nodeA, nodeB)) || ($isTabNode(nodeA) && $isTabNode(nodeB)) || ($isLineBreakNode(nodeA) && $isLineBreakNode(nodeB)) ); diff --git a/packages/lexical-code-shiki/src/__tests__/unit/CodeShikiRetokenizeKeepsState.test.ts b/packages/lexical-code-shiki/src/__tests__/unit/CodeShikiRetokenizeKeepsState.test.ts new file mode 100644 index 00000000000..4df5dae2122 --- /dev/null +++ b/packages/lexical-code-shiki/src/__tests__/unit/CodeShikiRetokenizeKeepsState.test.ts @@ -0,0 +1,441 @@ +/** + * 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. + * + */ + +// @vitest-environment node + +import { + $createCodeHighlightNode, + $createCodeNode, + $isCodeHighlightNode, + CodeHighlightNode, +} from '@lexical/code-core'; +import { + CodeShikiExtension, + loadCodeLanguage, + loadCodeTheme, +} from '@lexical/code-shiki'; +import {buildEditorFromExtensions} from '@lexical/extension'; +import {$createQuoteNode, RichTextExtension} from '@lexical/rich-text'; +import { + $create, + $createParagraphNode, + $createTextNode, + $getRoot, + $getState, + $isElementNode, + $setState, + createState, + defineExtension, + type LexicalEditor, + type StateConfig, + TextNode, +} from 'lexical'; +import {beforeAll, describe, expect, test} from 'vitest'; + +// Text inside a code block can carry NodeState of its own (comment ids, a +// suggestion: collaboration stores them on the text). When the highlighter +// tokenizes the block again, the new tokens keep it: a token is split where +// the state changes, and each part has the state of the text it came from. + +const noteState = createState('note', { + parse: (value: unknown) => (typeof value === 'string' ? value : ''), +}); +/** A set-valued state, as a suggestion's or a comment's members are. */ +const membersState = createState('members', { + parse: (value: unknown) => + Array.isArray(value) ? (value as string[]) : ([] as string[]), +}); + +function editorWithCode(): LexicalEditor & Disposable { + return buildEditorFromExtensions( + defineExtension({ + $initialEditorState: null, + dependencies: [RichTextExtension, CodeShikiExtension], + name: 'retokenize-keeps-state', + onError: (error: Error) => { + throw error; + }, + }), + ); +} + +/** Whether every child of the first block is a code highlight node. */ +const allTokens = (editor: LexicalEditor) => + editor.read(() => { + const block = $getRoot().getFirstChildOrThrow(); + return ( + $isElementNode(block) && block.getChildren().every($isCodeHighlightNode) + ); + }); + +/** Each text node of the first block: text, then its states if any. */ +const show = (editor: LexicalEditor) => + editor.read(() => { + const block = $getRoot().getFirstChildOrThrow(); + return ($isElementNode(block) ? block.getChildren() : []).map(node => { + const note = $getState(node, noteState); + const members = $getState(node, membersState); + return ( + node.getTextContent() + + (note === '' ? '' : `[note=${note}]`) + + (members.length === 0 ? '' : `[members=${members.join(',')}]`) + ); + }); + }); + +describe('a code block tokenized again keeps the NodeState of its text', () => { + // Shiki loads grammars and themes asynchronously; loaded, it highlights in + // the update that changes the block. + beforeAll(async () => { + await loadCodeLanguage('javascript'); + await loadCodeTheme('one-light'); + }); + + test('a text node with NodeState inside a code block, tokenized again by an insert next to the block', () => { + using editor = editorWithCode(); + // As a collaborator's change is read: with no transforms, so the text + // node stays one node of its own inside the block. + editor.update( + () => { + const word = $createTextNode('t78z'); + $setState(word, noteState, 'kept'); + $setState(word, membersState, ['s1']); + $getRoot().append( + $createCodeNode('javascript').append( + $createCodeHighlightNode('fetch( '), + word, + $createCodeHighlightNode(' t51z'), + ), + $createParagraphNode().append($createTextNode('after')), + ); + }, + {discrete: true, skipTransforms: true}, + ); + // A paragraph inserted right after the block makes the highlighter + // tokenize it again. + editor.update( + () => { + $getRoot() + .getFirstChildOrThrow() + .insertAfter($createParagraphNode().append($createTextNode('x'))); + }, + {discrete: true}, + ); + // Known by its StateConfig on the new token, as on the text it came from + // (checked before anything reads it there): collaboration stores a known + // set-valued state one entry per member, an unknown one whole. + expect( + editor.read(() => + $getRoot() + .getAllTextNodes() + .filter(node => node.getTextContent() === 't78z') + .map(node => node.__state!.getInternalState()[1].has(membersState)), + ), + ).toEqual([true]); + const nodes = show(editor); + expect(nodes.join('')).toBe('fetch( t78z[note=kept][members=s1] t51z'); + expect(nodes).toContain('t78z[note=kept][members=s1]'); + expect(allTokens(editor)).toBe(true); + }); + + test('a quote merged into a code block keeps the NodeState of its text', () => { + using editor = editorWithCode(); + editor.update( + () => { + const word = $createTextNode('t50z'); + $setState(word, membersState, ['s1']); + $getRoot().append( + $createCodeNode('javascript').append( + $createCodeHighlightNode('const a = 1;'), + ), + $createQuoteNode().append(word), + ); + }, + {discrete: true}, + ); + editor.update( + () => { + const [code, quote] = $getRoot().getChildren(); + if ($isElementNode(code) && $isElementNode(quote)) { + code.append(...quote.getChildren()); + quote.remove(); + } + }, + {discrete: true}, + ); + expect(show(editor).join('')).toBe('const a = 1;t50z[members=s1]'); + expect(allTokens(editor)).toBe(true); + }); + + test('without NodeState, the block is tokenized as before', () => { + using editor = editorWithCode(); + editor.update( + () => { + $getRoot().append( + $createCodeNode('javascript').append( + $createCodeHighlightNode('fetch( '), + $createTextNode('t78z'), + $createCodeHighlightNode(' t51z'), + ), + ); + }, + {discrete: true}, + ); + const withoutState = show(editor); + using plain = editorWithCode(); + plain.update( + () => { + $getRoot().append( + $createCodeNode('javascript').append( + $createCodeHighlightNode('fetch( t78z t51z'), + ), + ); + }, + {discrete: true}, + ); + expect(withoutState).toEqual(show(plain)); + }); +}); + +// The copies of NodeState (onto the new tokens, and onto text that leaves a +// code block) keep what the copy has of its own, as the copy of the JSON +// does: where the old node has a state at its default value (which its JSON +// leaves out), the value the copy's class sets when it is made; and where the +// copy's class has a StateConfig of its own for a key, the value read with +// that one. + +const parseString = (value: unknown) => + typeof value === 'string' ? value : ''; +/** A state the replacement classes below set when they are made. */ +const createdState = createState('created', {parse: parseString}); +/** A key the replacement classes below have a StateConfig of their own for. */ +const ownedState = createState('owned', {parse: parseString}); +/** Another StateConfig for the same key, as another module could make. */ +const otherOwnedState = createState('owned', {parse: parseString}); + +/** A TextNode replacement that sets a state of its own when it is made. */ +class CreatedTextNode extends TextNode { + $config() { + return this.config('test-retokenize-created-text', { + extends: TextNode, + stateConfigs: [createdState, ownedState], + }); + } +} + +/** A code highlight node replacement that sets a state of its own too. */ +class CreatedHighlightNode extends CodeHighlightNode { + $config() { + return this.config('test-retokenize-created-highlight', { + extends: CodeHighlightNode, + stateConfigs: [createdState, ownedState], + }); + } +} + +function editorWithCreatedText(): LexicalEditor & Disposable { + return buildEditorFromExtensions( + defineExtension({ + $initialEditorState: null, + dependencies: [RichTextExtension, CodeShikiExtension], + name: 'retokenize-keeps-state-created-text', + nodes: () => [ + CreatedTextNode, + { + replace: TextNode, + with: (node: TextNode) => + $setState( + $create(CreatedTextNode).setTextContent(node.getTextContent()), + createdState, + 'initialized', + ), + withKlass: CreatedTextNode, + }, + ], + onError: (error: Error) => { + throw error; + }, + }), + ); +} + +function editorWithCreatedHighlight(): LexicalEditor & Disposable { + return buildEditorFromExtensions( + defineExtension({ + $initialEditorState: null, + dependencies: [RichTextExtension, CodeShikiExtension], + name: 'retokenize-keeps-state-created-highlight', + nodes: () => [ + CreatedHighlightNode, + { + replace: CodeHighlightNode, + with: (node: CodeHighlightNode) => + $setState( + $create(CreatedHighlightNode).setTextContent( + node.getTextContent(), + ), + createdState, + 'initialized', + ), + withKlass: CreatedHighlightNode, + }, + ], + onError: (error: Error) => { + throw error; + }, + }), + ); +} + +/** A code block whose text node `t78z` has the given states. */ +function $appendBlockWithState( + ...states: [StateConfig, string][] +): void { + const word = $createTextNode('t78z'); + for (const [stateConfig, value] of states) { + $setState(word, stateConfig, value); + } + $getRoot().append( + $createCodeNode('javascript').append( + $createCodeHighlightNode('fetch( '), + word, + $createCodeHighlightNode(' t51z'), + ), + ); +} + +/** A code block of one token with the given states. */ +function $appendTokenWithState( + ...states: [StateConfig, string][] +): void { + const token = $createCodeHighlightNode('plain'); + for (const [stateConfig, value] of states) { + $setState(token, stateConfig, value); + } + $getRoot().append($createCodeNode().append(token)); +} + +describe('the NodeState copies keep what the copy has of its own', () => { + beforeAll(async () => { + await loadCodeLanguage('javascript'); + await loadCodeTheme('one-light'); + }); + + test('text leaving a code block keeps the state its class sets, where the token has that state at its default', () => { + using editor = editorWithCreatedText(); + // The note makes the token's state JSON not empty. + editor.update( + () => + $appendTokenWithState( + [noteState, 'kept'], + [createdState, createdState.defaultValue], + ), + {discrete: true, skipTransforms: true}, + ); + editor.update( + () => { + $getRoot().getFirstChildOrThrow().replace($createParagraphNode(), true); + }, + {discrete: true}, + ); + expect( + editor.read(() => + $getRoot() + .getAllTextNodes() + .map(node => [ + $isCodeHighlightNode(node), + $getState(node, noteState), + $getState(node, createdState), + ]), + ), + ).toEqual([[false, 'kept', 'initialized']]); + }); + + test('text leaving a code block is read with the StateConfig its class has for a key', () => { + using editor = editorWithCreatedText(); + editor.update(() => $appendTokenWithState([otherOwnedState, 'kept']), { + discrete: true, + skipTransforms: true, + }); + editor.update( + () => { + $getRoot().getFirstChildOrThrow().replace($createParagraphNode(), true); + }, + {discrete: true}, + ); + expect( + editor.read(() => + $getRoot() + .getAllTextNodes() + .map(node => [ + $isCodeHighlightNode(node), + $getState(node, ownedState), + ]), + ), + ).toEqual([[false, 'kept']]); + }); + + test('a new token keeps the state its class sets, where the text it comes from has that state at its default', () => { + using editor = editorWithCreatedHighlight(); + editor.update( + () => + $appendBlockWithState( + [noteState, 'kept'], + [createdState, createdState.defaultValue], + ), + {discrete: true, skipTransforms: true}, + ); + editor.update( + () => { + $getRoot() + .getFirstChildOrThrow() + .insertAfter($createParagraphNode().append($createTextNode('x'))); + }, + {discrete: true}, + ); + expect( + editor.read(() => + $getRoot() + .getAllTextNodes() + .filter(node => node.getTextContent() === 't78z') + .map(node => [ + $isCodeHighlightNode(node), + $getState(node, noteState), + $getState(node, createdState), + ]), + ), + ).toEqual([[true, 'kept', 'initialized']]); + }); + + test('a new token is read with the StateConfig its class has for a key', () => { + using editor = editorWithCreatedHighlight(); + editor.update(() => $appendBlockWithState([otherOwnedState, 'kept']), { + discrete: true, + skipTransforms: true, + }); + editor.update( + () => { + $getRoot() + .getFirstChildOrThrow() + .insertAfter($createParagraphNode().append($createTextNode('x'))); + }, + {discrete: true}, + ); + expect( + editor.read(() => + $getRoot() + .getAllTextNodes() + .filter(node => node.getTextContent() === 't78z') + .map(node => [ + $isCodeHighlightNode(node), + $getState(node, ownedState), + ]), + ), + ).toEqual([[true, 'kept']]); + }); +});