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']]); + }); +});