Skip to content

Commit cf3e0e3

Browse files
auvredandrewbranch
auvred
and
andrewbranch
authored
6.4x speedup of AST materialization in JS API (microsoft/typescript-go#2835)
Co-authored-by: Andrew Branch <andrew@wheream.io>
1 parent d7238a7 commit cf3e0e3

4 files changed

Lines changed: 91 additions & 41 deletions

File tree

‎tsc/_packages/api/src/node.ts‎

Lines changed: 78 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -189,11 +189,9 @@ const KIND_NODE_LIST = 2 ** 32 - 1;
189189

190190
export class RemoteNodeBase {
191191
parent: RemoteNode;
192-
protected view: DataView;
193-
protected decoder: TextDecoder;
192+
view: DataView;
193+
decoder: TextDecoder;
194194
protected index: number;
195-
/** Keys are positions */
196-
protected _children: Map<number, RemoteNode | RemoteNodeList> | undefined;
197195

198196
constructor(view: DataView, decoder: TextDecoder, index: number, parent: RemoteNode) {
199197
this.view = view;
@@ -267,8 +265,6 @@ export class RemoteNodeList extends Array<RemoteNode> implements NodeArray<Remot
267265
protected view: DataView;
268266
protected decoder: TextDecoder;
269267
protected index: number;
270-
/** Keys are positions */
271-
protected _children: Map<number, RemoteNode | RemoteNodeList> | undefined;
272268

273269
get pos(): number {
274270
return this.view.getUint32(this.byteIndex + NODE_OFFSET_POS, true);
@@ -293,24 +289,74 @@ export class RemoteNodeList extends Array<RemoteNode> implements NodeArray<Remot
293289
private get byteIndex(): number {
294290
return this.offsetNodes + this.index * NODE_LEN;
295291
}
292+
private sourceFile: RemoteSourceFile;
296293

297-
constructor(view: DataView, decoder: TextDecoder, index: number, parent: RemoteNode) {
294+
constructor(view: DataView, decoder: TextDecoder, index: number, parent: RemoteNode, sourceFile: RemoteSourceFile) {
298295
super();
299296
this.view = view;
300297
this.decoder = decoder;
301298
this.index = index;
302299
this.parent = parent;
303300
this.length = this.data;
301+
this.sourceFile = sourceFile;
304302

305303
const length = this.length;
306-
for (let i = 0; i < length; i++) {
304+
for (let i = 16; i < length; i++) {
307305
Object.defineProperty(this, i, {
308306
get() {
309307
return this.at(i);
310308
},
311309
});
312310
}
313311
}
312+
get 0(): RemoteNode {
313+
return this.at(0);
314+
}
315+
get 1(): RemoteNode {
316+
return this.at(1);
317+
}
318+
get 2(): RemoteNode {
319+
return this.at(2);
320+
}
321+
get 3(): RemoteNode {
322+
return this.at(3);
323+
}
324+
get 4(): RemoteNode {
325+
return this.at(4);
326+
}
327+
get 5(): RemoteNode {
328+
return this.at(5);
329+
}
330+
get 6(): RemoteNode {
331+
return this.at(6);
332+
}
333+
get 7(): RemoteNode {
334+
return this.at(7);
335+
}
336+
get 8(): RemoteNode {
337+
return this.at(8);
338+
}
339+
get 9(): RemoteNode {
340+
return this.at(9);
341+
}
342+
get 10(): RemoteNode {
343+
return this.at(10);
344+
}
345+
get 11(): RemoteNode {
346+
return this.at(11);
347+
}
348+
get 12(): RemoteNode {
349+
return this.at(12);
350+
}
351+
get 13(): RemoteNode {
352+
return this.at(13);
353+
}
354+
get 14(): RemoteNode {
355+
return this.at(14);
356+
}
357+
get 15(): RemoteNode {
358+
return this.at(15);
359+
}
314360

315361
*[Symbol.iterator](): ArrayIterator<RemoteNode> {
316362
let next = this.index + 1;
@@ -325,8 +371,11 @@ export class RemoteNodeList extends Array<RemoteNode> implements NodeArray<Remot
325371
if (!Number.isInteger(index)) {
326372
return undefined!;
327373
}
374+
if (index >= this.data || (index < 0 && -index > this.data)) {
375+
return undefined!;
376+
}
328377
if (index < 0) {
329-
index = this.length - index;
378+
index = this.length + index;
330379
}
331380
let next = this.index + 1;
332381
for (let i = 0; i < index; i++) {
@@ -337,15 +386,14 @@ export class RemoteNodeList extends Array<RemoteNode> implements NodeArray<Remot
337386
}
338387

339388
private getOrCreateChildAtNodeIndex(index: number): RemoteNode | RemoteNodeList {
340-
const pos = this.view.getUint32(this.offsetNodes + index * NODE_LEN + NODE_OFFSET_POS, true);
341-
let child = (this._children ??= new Map()).get(pos);
389+
let child = this.sourceFile.nodes[index];
342390
if (!child) {
343391
const kind = this.view.getUint32(this.offsetNodes + index * NODE_LEN + NODE_OFFSET_KIND, true);
344392
if (kind === KIND_NODE_LIST) {
345393
throw new Error("NodeList cannot directly contain another NodeList");
346394
}
347-
child = new RemoteNode(this.view, this.decoder, index, this.parent);
348-
this._children.set(pos, child);
395+
child = new RemoteNode(this.view, this.decoder, index, this.parent, this.sourceFile);
396+
this.sourceFile.nodes[index] = child;
349397
}
350398
return child;
351399
}
@@ -362,21 +410,15 @@ export class RemoteNodeList extends Array<RemoteNode> implements NodeArray<Remot
362410

363411
export class RemoteNode extends RemoteNodeBase implements Node {
364412
protected static NODE_LEN: number = NODE_LEN;
365-
private sourceFile: SourceFile;
366-
id: string;
413+
protected sourceFile: RemoteSourceFile;
414+
get id(): string {
415+
return `${this.pos}.${this.end}.${this.kind}.${this.sourceFile.path}`;
416+
}
367417

368-
constructor(view: DataView, decoder: TextDecoder, index: number, parent: RemoteNode) {
418+
constructor(view: DataView, decoder: TextDecoder, index: number, parent: RemoteNode, sourceFile: RemoteSourceFile) {
369419
super(view, decoder, index, parent);
370-
let sourceFile: RemoteNode = this;
371-
while (sourceFile && sourceFile.kind !== SyntaxKind.SourceFile) {
372-
sourceFile = sourceFile.parent;
373-
}
374-
if (!sourceFile) {
375-
throw new Error("SourceFile not found");
376-
}
377-
this.sourceFile = sourceFile as unknown as SourceFile;
420+
this.sourceFile = sourceFile;
378421
// Node handle format: pos.end.kind.path
379-
this.id = `${this.pos}.${this.end}.${this.kind}.${this.sourceFile.path}`;
380422
}
381423

382424
forEachChild<T>(visitNode: (node: Node) => T, visitList?: (list: NodeArray<Node>) => T): T | undefined {
@@ -411,7 +453,7 @@ export class RemoteNode extends RemoteNodeBase implements Node {
411453
}
412454

413455
getSourceFile(): SourceFile {
414-
return this.sourceFile;
456+
return this.sourceFile as unknown as SourceFile;
415457
}
416458

417459
protected getString(index: number): string {
@@ -422,22 +464,18 @@ export class RemoteNode extends RemoteNodeBase implements Node {
422464
}
423465

424466
private getOrCreateChildAtNodeIndex(index: number): RemoteNode | RemoteNodeList {
425-
const pos = this.view.getUint32(this.offsetNodes + index * NODE_LEN + NODE_OFFSET_POS, true);
426-
let child = (this._children ??= new Map()).get(pos);
467+
let child = this.sourceFile.nodes[index];
427468
if (!child) {
428469
const kind = this.view.getUint32(this.offsetNodes + index * NODE_LEN + NODE_OFFSET_KIND, true);
429470
child = kind === KIND_NODE_LIST
430-
? new RemoteNodeList(this.view, this.decoder, index, this)
431-
: new RemoteNode(this.view, this.decoder, index, this);
432-
this._children.set(pos, child);
471+
? new RemoteNodeList(this.view, this.decoder, index, this, this.sourceFile)
472+
: new RemoteNode(this.view, this.decoder, index, this, this.sourceFile);
473+
this.sourceFile.nodes[index] = child;
433474
}
434475
return child;
435476
}
436477

437478
private hasChildren(): boolean {
438-
if (this._children) {
439-
return true;
440-
}
441479
if (this.byteIndex >= this.view.byteLength - NODE_LEN) {
442480
return false;
443481
}
@@ -994,9 +1032,14 @@ export class RemoteNode extends RemoteNodeBase implements Node {
9941032
}
9951033

9961034
export class RemoteSourceFile extends RemoteNode {
1035+
readonly nodes: (RemoteNode | RemoteNodeList)[];
1036+
9971037
constructor(data: Uint8Array, decoder: TextDecoder) {
9981038
const view = new DataView(data.buffer, data.byteOffset, data.byteLength);
999-
super(view, decoder, 1, undefined!);
1039+
super(view, decoder, 1, undefined!, {} as unknown as RemoteSourceFile);
1040+
this.sourceFile = this;
1041+
this.nodes = Array((this.view.byteLength - this.offsetNodes) / NODE_LEN);
1042+
this.nodes[1] = this;
10001043
}
10011044
}
10021045

‎tsc/_packages/api/test/async/api.bench.ts‎

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -8,11 +8,12 @@ import {
88
type SourceFile,
99
SyntaxKind,
1010
} from "@typescript/ast";
11-
import fs, { existsSync } from "node:fs";
11+
import { existsSync } from "node:fs";
1212
import path from "node:path";
1313
import { fileURLToPath } from "node:url";
1414
import { Bench } from "tinybench";
1515
import ts from "typescript";
16+
import { RemoteSourceFile } from "../../src/node.ts";
1617

1718
const isMain = process.argv[1] === fileURLToPath(import.meta.url);
1819
if (isMain) {
@@ -84,12 +85,14 @@ export async function runBenchmarks(singleIteration?: boolean) {
8485
await getCheckerTS();
8586
}, { beforeAll: all(spawnAPI, loadSnapshot) })
8687
.add("materialize program.ts", async () => {
87-
file.forEachChild(function visit(node) {
88+
const { view, decoder } = file as unknown as RemoteSourceFile;
89+
new RemoteSourceFile(new Uint8Array(view.buffer, view.byteOffset, view.byteLength), decoder).forEachChild(function visit(node) {
8890
node.forEachChild(visit);
8991
});
9092
}, { beforeAll: all(spawnAPI, loadSnapshot, getProgramTS) })
9193
.add("materialize checker.ts", async () => {
92-
file.forEachChild(function visit(node) {
94+
const { view, decoder } = file as unknown as RemoteSourceFile;
95+
new RemoteSourceFile(new Uint8Array(view.buffer, view.byteOffset, view.byteLength), decoder).forEachChild(function visit(node) {
9396
node.forEachChild(visit);
9497
});
9598
}, { beforeAll: all(spawnAPI, loadSnapshot, getCheckerTS) })

‎tsc/_packages/api/test/sync/api.bench.ts‎

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -16,11 +16,12 @@ import {
1616
type SourceFile,
1717
SyntaxKind,
1818
} from "@typescript/ast";
19-
import fs, { existsSync } from "node:fs";
19+
import { existsSync } from "node:fs";
2020
import path from "node:path";
2121
import { fileURLToPath } from "node:url";
2222
import { Bench } from "tinybench";
2323
import ts from "typescript";
24+
import { RemoteSourceFile } from "../../src/node.ts";
2425

2526
const isMain = process.argv[1] === fileURLToPath(import.meta.url);
2627
if (isMain) {
@@ -92,12 +93,14 @@ export function runBenchmarks(singleIteration?: boolean) {
9293
getCheckerTS();
9394
}, { beforeAll: all(spawnAPI, loadSnapshot) })
9495
.add("materialize program.ts", () => {
95-
file.forEachChild(function visit(node) {
96+
const { view, decoder } = file as unknown as RemoteSourceFile;
97+
new RemoteSourceFile(new Uint8Array(view.buffer, view.byteOffset, view.byteLength), decoder).forEachChild(function visit(node) {
9698
node.forEachChild(visit);
9799
});
98100
}, { beforeAll: all(spawnAPI, loadSnapshot, getProgramTS) })
99101
.add("materialize checker.ts", () => {
100-
file.forEachChild(function visit(node) {
102+
const { view, decoder } = file as unknown as RemoteSourceFile;
103+
new RemoteSourceFile(new Uint8Array(view.buffer, view.byteOffset, view.byteLength), decoder).forEachChild(function visit(node) {
101104
node.forEachChild(visit);
102105
});
103106
}, { beforeAll: all(spawnAPI, loadSnapshot, getCheckerTS) })

‎tsc/_packages/api/tsconfig.dev.json‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
{
22
"extends": "./tsconfig.base.json",
33
"compilerOptions": {
4+
"rewriteRelativeImportExtensions": false,
45
"module": "nodenext",
56
"strict": true,
67
"noEmit": true,

0 commit comments

Comments
 (0)