1import { ASSERT, UNWRAP } from "@clo/lib/assert.ts";
2import { type JSX, type Key, type ReactNode } from "react";
3import { type Components } from "rehype-react";
4import remarkRehype from "remark-rehype";
5import { Fragment, jsx } from "react/jsx-runtime";
6import { type Processor, unified } from "unified";
7import type { Root as HastRoot, Root } from "hast";
8import type { Literal, Node, Parent } from "unist";
9import { memoizedHastToReact, nodeDeepEquals, type RenderState } from "./hast.ts";
10import remarkParse from "remark-parse";
11import { Predict } from "./Predict.ts";
12
13export const defaultProcessor = unified().use(remarkParse);
14
15export type BaseProcessor = Processor<
16 Node,
17 undefined,
18 undefined,
19 undefined,
20 undefined | JSX.Element
21>;
22type AstProcessor = Processor<Parent, Parent, Parent, Parent, Parent>;
23
24/**
25 * The Memoizer is an incremental markdown compiler that is bolted on top of a
26 * unified markdown processor. It works by splitting documents into blocks, then
27 * using diffing to tell where in the document new text was added or removed. At
28 * the end, render states are diffed to memoize as many React nodes as possible.
29 */
30export class Memoizer {
31 // options incrementally updated via `reconfigure`
32 #baseProcessor: BaseProcessor = defaultProcessor;
33 #astProcessor: AstProcessor | null = null;
34 #components: Partial<Components> = {};
35
36 // incremental parsing graph, structure of arrays
37 #predict: Predict | null = null;
38 #content: string = "";
39 #positions: number[] = [];
40 #parsed: Node[][] = [];
41 // React block identity must survive index shifts when blocks are inserted above.
42 #keys: string[] = [];
43 #nextKey = 0;
44 #renderStates: (RenderState | null)[] = [];
45 #reactNodes: readonly ReactNode[] = [];
46
47 reconfigure(
48 processor: Processor<Node, undefined, undefined, undefined, undefined | JSX.Element>,
49 predict: boolean,
50 components: Partial<Components>,
51 ) {
52 if (
53 this.#astProcessor !== null &&
54 this.#baseProcessor === processor &&
55 componentsAreEqual(this.#components, components) &&
56 !!this.#predict === predict
57 ) {
58 return;
59 }
60
61 this.#baseProcessor = processor;
62 if (!processor.attachers.some((plugin) => plugin[0] === remarkRehype)) {
63 processor = processor();
64 processor.use(remarkRehype);
65 processor.freeze();
66 }
67
68 this.#astProcessor = processor as unknown as AstProcessor;
69 this.#components = components;
70 this.#content = "";
71 this.#positions = [];
72 this.#parsed = [];
73 this.#keys = [];
74 this.#nextKey = 0;
75 this.#renderStates = [];
76 this.#predict = predict ? new Predict() : null;
77 // NOTE: do not reset reactNodes so that pipeline changes can preserve the react node state
78 }
79
80 update(content: string): readonly ReactNode[] {
81 if (this.#predict) content = this.#predict.update(content);
82 let previous = this.#content;
83 if (content === previous) return this.#reactNodes;
84
85 // Locate the last block that has changed content. This is done quickly by finding
86 // the first changed character, and snapping to the nearest cached block.
87 const positions = this.#positions;
88 const firstCharDifference = indexOfDiff(content, previous);
89 let blockStart = positions.length - 1;
90 for (; blockStart >= 0; blockStart -= 1) {
91 if (firstCharDifference >= positions[blockStart]!) break;
92 }
93 let parseOffset = positions[blockStart] ?? 0;
94 if (blockStart === -1) blockStart = 0;
95
96 // Compute and replace changed block regions. This behavior relys on the
97 // fact that blocks generally do not affect each other. To account for
98 // things like reflinks, we cause re-parses with these nodes to perform a
99 // full transform. This detection is not fine grained (that would be too
100 // much effort), but any extra re-renders get eliminated with AST diffing.
101 const processor = UNWRAP(this.#astProcessor);
102 const parsed = this.#parsed;
103 const replacedTailAffectedDocument = parsed
104 .slice(blockStart)
105 .some((block) => block.some(nodeAffectsDocument));
106 let parsedTree = processor.parse(content.slice(parseOffset));
107 let newPositions = parsedTree.children.map(
108 (child) => UNWRAP(UNWRAP(child.position).start.offset) + parseOffset,
109 );
110
111 // Index incremental state with new results
112 this.#content = content;
113 ASSERT(positions.length === parsed.length);
114 positions.splice(blockStart, positions.length - blockStart, ...newPositions);
115 {
116 const newParsed = extractBlocks(parsedTree, newPositions, parseOffset);
117 ASSERT(newParsed.length === newPositions.length);
118 parsed.splice(blockStart, parsed.length - blockStart, ...newParsed);
119 }
120
121 // If a reflink is involved, fall back to a full reparse. Rehydrating a
122 // synthetic root from cached blocks can lose source offsets on transformed
123 // top-level nodes, which would scramble block bucketing and key identity.
124 if (replacedTailAffectedDocument || nodeAffectsDocument(parsedTree)) {
125 blockStart = 0;
126 parseOffset = 0;
127 parsedTree = processor.parse(content);
128 newPositions = parsedTree.children.map(
129 (child) => UNWRAP(UNWRAP(child.position).start.offset) + parseOffset,
130 );
131 positions.splice(0, positions.length, ...newPositions);
132 const newParsed = extractBlocks(parsedTree, newPositions, parseOffset);
133 ASSERT(newParsed.length === newPositions.length);
134 parsed.splice(0, parsed.length, ...newParsed);
135 }
136
137 // For all affected blocks, transform their ASTs. This is done in
138 // a group to preserve context from clobbers, such as reference links. Since
139 // editing a clobbering block force-transforms the document and changing
140 // text re-transforms all text after it, this perfectly implements reference
141 // links. However, the latter of these assumptions is not perfectly ideal.
142 const newTransformed = extractBlocks(processor.runSync(parsedTree), newPositions, parseOffset);
143
144 // React
145 const keys = this.#keys;
146 const renderStates = this.#renderStates;
147 const previousReactNodes = this.#reactNodes;
148 const previousLength = renderStates.length;
149 let reactNodes: ReactNode[] | null = null;
150
151 // Insertions and removals above an unchanged tail shift indices without
152 // changing block content. Matching the shared suffix is enough to keep the
153 // unaffected tail bound to its previous keys and render states.
154 let suffixLength = 0;
155 for (; suffixLength < newTransformed.length; suffixLength += 1) {
156 const nextIndex = positions.length - 1 - suffixLength;
157 const previousIndex = previousLength - 1 - suffixLength;
158 if (nextIndex < blockStart || previousIndex < blockStart) break;
159
160 const next = UNWRAP(newTransformed[nextIndex - blockStart]);
161 const state = renderStates[previousIndex];
162 const prev = state?.node.type === "root" ? (state.node as Root) : null;
163 if (!prev || !nodeDeepEquals(next, prev.children)) break;
164 }
165
166 // (If realignment is going to occurs, the array must be cloned)
167 const previousMiddleEnd = previousLength - suffixLength;
168 const nextMiddleEnd = positions.length - suffixLength;
169 const realign = previousLength !== positions.length || previousMiddleEnd !== nextMiddleEnd;
170 if (realign) reactNodes = previousReactNodes.slice(0, positions.length);
171
172 if (positions.length > previousLength) {
173 keys.length = renderStates.length = positions.length;
174 }
175 // Realignment reads from the previous layout and writes the next layout
176 // into the same backing arrays. Traversal direction prevents later reads
177 // from observing values that were already overwritten earlier in the pass.
178 for (
179 let direction = positions.length < previousLength ? 1 : -1,
180 i = direction === 1 ? 0 : positions.length - 1,
181 end = direction === 1 ? positions.length : -1;
182 i !== end;
183 i += direction
184 ) {
185 let previousIndex: number = -1;
186 if (i < blockStart) {
187 previousIndex = i;
188 } else if (i >= nextMiddleEnd) {
189 previousIndex = previousLength - (positions.length - i);
190 } else if (i < previousMiddleEnd) {
191 previousIndex = i;
192 }
193 if (previousIndex < 0 || previousIndex >= previousLength) {
194 keys[i] = String(this.#nextKey);
195 this.#nextKey += 1;
196 renderStates[i] = null;
197 if (reactNodes) reactNodes[i] = undefined as ReactNode | undefined;
198 continue;
199 }
200
201 keys[i] = keys[previousIndex] ?? String(this.#nextKey++);
202 renderStates[i] = renderStates[previousIndex] ?? null;
203 if (reactNodes) reactNodes[i] = previousReactNodes[previousIndex];
204 }
205
206 // Create and splice the React elements
207 if (positions.length < previousLength) {
208 keys.length = renderStates.length = positions.length;
209 }
210 for (let i = blockStart, len = nextMiddleEnd; i < len; i += 1) {
211 const ast = UNWRAP(newTransformed[i - blockStart]);
212 const previousState = renderStates[i] ?? null;
213
214 // `memoizedHastToReact` handles diffing and incrementally updating the AST.
215 let { react: result, state } = memoizedHastToReact(
216 { type: "root", children: ast } as HastRoot,
217 previousState,
218 this.#components,
219 );
220 if (previousState === state) continue;
221 renderStates[i] = state;
222
223 // Do a little trolling and unwrap the fragment to clean up the Virtual DOM
224 const key = keys[i]!;
225 const rendered = result as JSX.Element;
226 const children = rendered.props.children;
227 if (rendered.type === Fragment && children?.type) {
228 result = setReactKey(children as JSX.Element, key);
229 } else {
230 result = setReactKey(rendered, key);
231 }
232 reactNodes ??= previousReactNodes.slice(0, positions.length);
233 reactNodes[i] = result;
234 }
235
236 if (reactNodes) return (this.#reactNodes = reactNodes);
237 return this.#reactNodes;
238 }
239}
240
241export function indexOfDiff(a: string, b: string) {
242 var i = 0;
243 while (a[i] === b[i]) i++;
244 return i;
245}
246
247export function componentsAreEqual(left: Partial<Components>, right: Partial<Components>) {
248 const leftRecord = left as Record<string, unknown>;
249 const rightRecord = right as Record<string, unknown>;
250 const leftKeys = Object.keys(leftRecord);
251 if (leftKeys.length !== Object.keys(rightRecord).length) return false;
252 for (const key of leftKeys) {
253 if (leftRecord[key] !== rightRecord[key]) return false;
254 }
255 return true;
256}
257
258function setReactKey(element: JSX.Element, key: Key) {
259 const { ...props } = element.props;
260 return jsx(element.type, props, key);
261}
262
263const affectsDocument = new Set(["definition", "footnoteDefinition"]);
264
265export function nodeAffectsDocument(node: Node) {
266 if (affectsDocument.has(node.type)) return true;
267 if (!("children" in node) || !Array.isArray(node.children)) return false;
268 for (const child of (node as Parent).children) {
269 if (nodeAffectsDocument(child)) return true;
270 }
271 return false;
272}
273
274export function extractBlocks(tree: Parent, positions: number[], offset: number) {
275 const blocks = Array.from(positions, (): Node[] => []);
276
277 let blockIndex = 0;
278 for (const child of tree.children) {
279 const startOffset = child.position?.start.offset;
280 if (startOffset !== undefined) {
281 for (; blockIndex + 1 < positions.length; blockIndex += 1) {
282 if (startOffset + offset < positions[blockIndex + 1]!) break;
283 }
284 blocks[blockIndex]!.push(child);
285 continue;
286 }
287
288 if (child.type === "text" && /^\s*$/.test(String((child as Literal).value))) continue;
289 blocks[blockIndex]!.push(child);
290 }
291
292 return blocks;
293}