Skip to content

Commit d4eaca5

Browse files
[api] Optimize RemoteNodeList child access (microsoft/typescript-go#4549)
Co-authored-by: Copilot Autofix powered by AI <[email protected]>
1 parent 2ef980a commit d4eaca5

4 files changed

Lines changed: 263 additions & 9 deletions

File tree

‎tsc/_packages/native-preview/src/api/node/node.generated.ts‎

Lines changed: 26 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,12 @@ export class RemoteNodeList extends Array<RemoteNode> implements NodeArray<Remot
4848
protected view: DataView;
4949
protected index: number;
5050
private _byteIndex: number;
51+
// Cursor memoizing the last resolved (logical index -> node index) so that
52+
// sequential forward access (index loops and list[i], plus forEach/map/
53+
// reduce/filter) resumes instead of re-walking from the head, turning an
54+
// O(n) pass over the whole list from O(n^2) into O(n).
55+
private _cursorIndex: number = 0;
56+
private _cursorNodeIndex: number = 0;
5157

5258
get pos(): number {
5359
return this.view.getUint32(this._byteIndex + NODE_OFFSET_POS, true);
@@ -75,6 +81,7 @@ export class RemoteNodeList extends Array<RemoteNode> implements NodeArray<Remot
7581
this.sourceFile = sourceFile;
7682
this._byteIndex = offsetNodes + index * NODE_LEN;
7783
this.length = this.data;
84+
this._cursorNodeIndex = index + 1;
7885

7986
const length = this.length;
8087
for (let i = 16; i < length; i++) {
@@ -165,11 +172,26 @@ export class RemoteNodeList extends Array<RemoteNode> implements NodeArray<Remot
165172
if (index < 0) {
166173
index = this.length + index;
167174
}
168-
let next = this.index + 1;
169-
for (let i = 0; i < index; i++) {
170-
const child = this.getOrCreateChildAtNodeIndex(next);
171-
next = child.next;
175+
// Walk the raw buffer following each node's `next` pointer instead of
176+
// materializing every intermediate RemoteNode just to read it. Resume from
177+
// the memoized cursor when possible so sequential forward access is O(1)
178+
// amortized (a full in-order pass is O(n) rather than O(n^2)).
179+
const offsetNodes = this.sourceFile._offsetNodes;
180+
let i: number;
181+
let next: number;
182+
if (index >= this._cursorIndex) {
183+
i = this._cursorIndex;
184+
next = this._cursorNodeIndex;
185+
}
186+
else {
187+
i = 0;
188+
next = this.index + 1;
189+
}
190+
for (; i < index; i++) {
191+
next = this.view.getUint32(offsetNodes + next * NODE_LEN + NODE_OFFSET_NEXT, true);
172192
}
193+
this._cursorIndex = index;
194+
this._cursorNodeIndex = next;
173195
return this.getOrCreateChildAtNodeIndex(next) as RemoteNode;
174196
}
175197

‎tsc/_packages/native-preview/test/encoder.test.ts‎

Lines changed: 54 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,8 +3,12 @@ import type {
33
SourceFile,
44
Statement,
55
} from "@typescript/native-preview/unstable/ast";
6-
import { SyntaxKind } from "@typescript/native-preview/unstable/ast";
76
import {
7+
SyntaxKind,
8+
TokenFlags,
9+
} from "@typescript/native-preview/unstable/ast";
10+
import {
11+
createArrayLiteralExpression,
812
createBlock,
913
createExpressionStatement,
1014
createFunctionDeclaration,
@@ -14,6 +18,7 @@ import {
1418
createJsxClosingElement,
1519
createJsxElement,
1620
createJsxOpeningElement,
21+
createNumericLiteral,
1722
createPostfixUnaryExpression,
1823
createPrefixUnaryExpression,
1924
createSourceFile,
@@ -268,6 +273,54 @@ describe("Encoder", () => {
268273
// Empty properties should return undefined, not throw
269274
assert.strictEqual(jsxElem.typeArguments, undefined);
270275
});
276+
277+
test("NodeList element access is consistent across forward, backward, and random order", () => {
278+
// Build `[0, 1, 2, ..., 9]` and access the NodeList in several orders to
279+
// ensure the memoized cursor in `at()` never returns a stale element.
280+
const count = 10;
281+
const elements = [];
282+
let text = "[";
283+
for (let i = 0; i < count; i++) {
284+
elements.push(createNumericLiteral(String(i), TokenFlags.None));
285+
text += (i === 0 ? "" : ", ") + i;
286+
}
287+
text += "]";
288+
const stmt = createExpressionStatement(createArrayLiteralExpression(elements));
289+
const sf = makeSF(text, "/test.ts", [stmt]);
290+
291+
const decoded = decode(encodeSourceFile(sf));
292+
const arrayExpr = (decoded.statements!.at(0)! as RemoteNode).expression! as RemoteNode;
293+
const list = arrayExpr.elements! as unknown as RemoteNodeList;
294+
assert.strictEqual(list.length, count);
295+
296+
// Independent decode used as an oracle: fresh cursor per lookup.
297+
const posOf = (i: number) => {
298+
const freshExpr = (decode(encodeSourceFile(sf)).statements!.at(0)! as RemoteNode).expression! as RemoteNode;
299+
return (freshExpr.elements! as unknown as RemoteNodeList).at(i)!.pos;
300+
};
301+
const expected = Array.from({ length: count }, (_, i) => posOf(i));
302+
303+
// Forward sequential (exercises cursor resume).
304+
for (let i = 0; i < count; i++) {
305+
assert.strictEqual(list.at(i)!.pos, expected[i], `forward at(${i})`);
306+
}
307+
// Backward sequential (cursor must reset when index < cursor).
308+
for (let i = count - 1; i >= 0; i--) {
309+
assert.strictEqual(list.at(i)!.pos, expected[i], `backward at(${i})`);
310+
}
311+
// Random order, with repeats.
312+
for (const i of [3, 7, 1, 7, 0, 9, 4, 4, 2, 8, 5, 6, 0, 9]) {
313+
assert.strictEqual(list.at(i)!.pos, expected[i], `random at(${i})`);
314+
}
315+
// Negative indices.
316+
assert.strictEqual(list.at(-1)!.pos, expected[count - 1], "at(-1)");
317+
assert.strictEqual(list.at(-count)!.pos, expected[0], `at(-${count})`);
318+
319+
// Indexed getters (0..15) and spread must agree with at().
320+
assert.strictEqual((list[0] as RemoteNode).pos, expected[0]);
321+
assert.strictEqual((list[9] as RemoteNode).pos, expected[9]);
322+
assert.deepStrictEqual([...list].map(n => n.pos), expected);
323+
});
271324
});
272325

273326
describe("UTF-8 vs UTF-16 position encoding", () => {
Lines changed: 157 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,157 @@
1+
// Micro-benchmark for RemoteNodeList access patterns.
2+
//
3+
// Focuses on the cost of the element accessors (`list[i]`), the `.at()`
4+
// method, and the inherited Array iteration methods (map/forEach/reduce),
5+
// all of which walk the encoded linked list of children.
6+
//
7+
// Run with:
8+
// npm run node -- test/nodelist.bench.ts
9+
// npm run node -- test/nodelist.bench.ts --filter map
10+
11+
import {
12+
type Expression,
13+
type Path,
14+
type SourceFile,
15+
type Statement,
16+
SyntaxKind,
17+
TokenFlags,
18+
} from "@typescript/native-preview/unstable/ast";
19+
import {
20+
createArrayLiteralExpression,
21+
createExpressionStatement,
22+
createNumericLiteral,
23+
createSourceFile,
24+
createToken,
25+
} from "@typescript/native-preview/unstable/ast/factory";
26+
import { fileURLToPath } from "node:url";
27+
import { parseArgs } from "node:util";
28+
import { Bench } from "tinybench";
29+
import { encodeSourceFile } from "../src/api/node/encoder.ts";
30+
import type { RemoteNodeList } from "../src/api/node/node.generated.ts";
31+
import { RemoteSourceFile } from "../src/api/node/node.ts";
32+
33+
const isMain = process.argv[1] === fileURLToPath(import.meta.url);
34+
if (isMain) {
35+
const { values } = parseArgs({
36+
options: {
37+
filter: { type: "string" },
38+
size: { type: "string", default: "500" },
39+
},
40+
});
41+
runBenchmarks(values);
42+
}
43+
44+
function makeSourceFileWithLargeList(elementCount: number): Uint8Array {
45+
const elements: Expression[] = [];
46+
let text = "[";
47+
for (let i = 0; i < elementCount; i++) {
48+
elements.push(createNumericLiteral(String(i), TokenFlags.None));
49+
text += (i === 0 ? "" : ",") + i;
50+
}
51+
text += "]";
52+
const array = createArrayLiteralExpression(elements);
53+
const stmt: Statement = createExpressionStatement(array);
54+
const eof = createToken(SyntaxKind.EndOfFile);
55+
const sf: SourceFile = createSourceFile([stmt], eof, text, "/bench.ts", "/bench.ts" as Path);
56+
return encodeSourceFile(sf);
57+
}
58+
59+
function getElementsList(data: Uint8Array): RemoteNodeList {
60+
const sf = new RemoteSourceFile(data, new TextDecoder());
61+
const stmt = (sf.statements as any).at(0);
62+
return stmt.expression.elements as RemoteNodeList;
63+
}
64+
65+
export function runBenchmarks(options?: { filter?: string; size?: string; }): void {
66+
const { filter, size } = options ?? {};
67+
const elementCount = size ? parseInt(size, 10) : 500;
68+
69+
const encoded = makeSourceFileWithLargeList(elementCount);
70+
71+
// Sanity check that access works and produces the expected length.
72+
{
73+
const list = getElementsList(encoded);
74+
if (list.length !== elementCount) {
75+
throw new Error(`Expected ${elementCount} elements, got ${list.length}`);
76+
}
77+
}
78+
79+
const bench = new Bench({
80+
name: `RemoteNodeList (${elementCount} elements)`,
81+
iterations: 30,
82+
warmupIterations: 10,
83+
});
84+
85+
// Each task re-decodes the source file so node materialization caching
86+
// inside `sourceFile.nodes` does not carry over between iterations, which
87+
// would otherwise mask the cost of walking the list.
88+
bench
89+
.add("index access (list[i])", () => {
90+
const list = getElementsList(encoded);
91+
let sum = 0;
92+
for (let i = 0; i < list.length; i++) {
93+
sum += list[i].pos;
94+
}
95+
if (sum < 0) throw new Error("unreachable");
96+
})
97+
.add("at(i)", () => {
98+
const list = getElementsList(encoded);
99+
let sum = 0;
100+
for (let i = 0; i < list.length; i++) {
101+
sum += list.at(i)!.pos;
102+
}
103+
if (sum < 0) throw new Error("unreachable");
104+
})
105+
.add("for..of (iterator)", () => {
106+
const list = getElementsList(encoded);
107+
let sum = 0;
108+
for (const node of list) {
109+
sum += node.pos;
110+
}
111+
if (sum < 0) throw new Error("unreachable");
112+
})
113+
.add("map", () => {
114+
const list = getElementsList(encoded);
115+
const positions = list.map(node => node.pos);
116+
if (positions.length < 0) throw new Error("unreachable");
117+
})
118+
.add("filter", () => {
119+
const list = getElementsList(encoded);
120+
const filtered = list.filter(node => node.pos >= 0);
121+
if (filtered.length < 0) throw new Error("unreachable");
122+
})
123+
.add("forEach", () => {
124+
const list = getElementsList(encoded);
125+
let sum = 0;
126+
list.forEach(node => {
127+
sum += node.pos;
128+
});
129+
if (sum < 0) throw new Error("unreachable");
130+
})
131+
.add("reduce", () => {
132+
const list = getElementsList(encoded);
133+
const sum = list.reduce((acc, node) => acc + node.pos, 0);
134+
if (sum < 0) throw new Error("unreachable");
135+
})
136+
.add("forEachNode", () => {
137+
const list = getElementsList(encoded);
138+
let sum = 0;
139+
list.forEachNode(node => {
140+
sum += node.pos;
141+
return undefined;
142+
});
143+
if (sum < 0) throw new Error("unreachable");
144+
});
145+
146+
if (filter) {
147+
const pattern = filter.toLowerCase();
148+
for (const task of [...bench.tasks]) {
149+
if (!task.name.toLowerCase().includes(pattern)) {
150+
bench.remove(task.name);
151+
}
152+
}
153+
}
154+
155+
bench.runSync();
156+
console.table(bench.table());
157+
}

‎tsc/_scripts/generate-encoder.ts‎

Lines changed: 26 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1523,6 +1523,12 @@ function emitRemoteNodeList(w: CodeWriter) {
15231523
w.write(` protected view: DataView;`);
15241524
w.write(` protected index: number;`);
15251525
w.write(` private _byteIndex: number;`);
1526+
w.write(` // Cursor memoizing the last resolved (logical index -> node index) so that`);
1527+
w.write(` // sequential forward access (index loops and list[i], plus forEach/map/`);
1528+
w.write(` // reduce/filter) resumes instead of re-walking from the head, turning an`);
1529+
w.write(` // O(n) pass over the whole list from O(n^2) into O(n).`);
1530+
w.write(` private _cursorIndex: number = 0;`);
1531+
w.write(` private _cursorNodeIndex: number = 0;`);
15261532
w.write(``);
15271533
w.write(` get pos(): number {`);
15281534
w.write(` return this.view.getUint32(this._byteIndex + NODE_OFFSET_POS, true);`);
@@ -1550,6 +1556,7 @@ function emitRemoteNodeList(w: CodeWriter) {
15501556
w.write(` this.sourceFile = sourceFile;`);
15511557
w.write(` this._byteIndex = offsetNodes + index * NODE_LEN;`);
15521558
w.write(` this.length = this.data;`);
1559+
w.write(` this._cursorNodeIndex = index + 1;`);
15531560
w.write(``);
15541561
w.write(` const length = this.length;`);
15551562
w.write(` for (let i = 16; i < length; i++) {`);
@@ -1600,11 +1607,26 @@ function emitRemoteNodeList(w: CodeWriter) {
16001607
w.write(` if (index < 0) {`);
16011608
w.write(` index = this.length + index;`);
16021609
w.write(` }`);
1603-
w.write(` let next = this.index + 1;`);
1604-
w.write(` for (let i = 0; i < index; i++) {`);
1605-
w.write(` const child = this.getOrCreateChildAtNodeIndex(next);`);
1606-
w.write(` next = child.next;`);
1610+
w.write(` // Walk the raw buffer following each node's \`next\` pointer instead of`);
1611+
w.write(` // materializing every intermediate RemoteNode just to read it. Resume from`);
1612+
w.write(` // the memoized cursor when possible so sequential forward access is O(1)`);
1613+
w.write(` // amortized (a full in-order pass is O(n) rather than O(n^2)).`);
1614+
w.write(` const offsetNodes = this.sourceFile._offsetNodes;`);
1615+
w.write(` let i: number;`);
1616+
w.write(` let next: number;`);
1617+
w.write(` if (index >= this._cursorIndex) {`);
1618+
w.write(` i = this._cursorIndex;`);
1619+
w.write(` next = this._cursorNodeIndex;`);
1620+
w.write(` }`);
1621+
w.write(` else {`);
1622+
w.write(` i = 0;`);
1623+
w.write(` next = this.index + 1;`);
1624+
w.write(` }`);
1625+
w.write(` for (; i < index; i++) {`);
1626+
w.write(` next = this.view.getUint32(offsetNodes + next * NODE_LEN + NODE_OFFSET_NEXT, true);`);
16071627
w.write(` }`);
1628+
w.write(` this._cursorIndex = index;`);
1629+
w.write(` this._cursorNodeIndex = next;`);
16081630
w.write(` return this.getOrCreateChildAtNodeIndex(next) as RemoteNode;`);
16091631
w.write(` }`);
16101632
w.write(``);

0 commit comments

Comments
 (0)