Add an easy way to run performance tests for tree-sitter (#233072)

* Add an easy way to run performance tests for tree-sitter
Part os #210475

* Compile build folder

.vscode-test.js

@@ -42,6 +42,11 @@ const extensions = [
 		workspaceFolder: `extensions/vscode-colorize-tests/test`,
 		mocha: { timeout: 60_000 }
 	},
+	{
+		label: 'vscode-colorize-perf-tests',
+		workspaceFolder: `extensions/vscode-colorize-perf-tests/test`,
+		mocha: { timeout: 6000_000 }
+	},
 	{
 		label: 'configuration-editing',
 		workspaceFolder: path.join(os.tmpdir(), `confeditout-${Math.floor(Math.random() * 100000)}`),

.vscode/launch.json

@@ -202,6 +202,24 @@
 				"order": 5
 			}
 		},
+		{
+			"type": "extensionHost",
+			"request": "launch",
+			"name": "VS Code Tokenizer Performance Tests",
+			"runtimeExecutable": "${execPath}",
+			"args": [
+				"${workspaceFolder}/extensions/vscode-colorize-perf-tests/test",
+				"--extensionDevelopmentPath=${workspaceFolder}/extensions/vscode-colorize-perf-tests",
+				"--extensionTestsPath=${workspaceFolder}/extensions/vscode-colorize-perf-tests/out"
+			],
+			"outFiles": [
+				"${workspaceFolder}/out/**/*.js"
+			],
+			"presentation": {
+				"group": "5_tests",
+				"order": 6
+			}
+		},
 		{
 			"type": "chrome",
 			"request": "attach",

build/azure-pipelines/darwin/product-build-darwin-test.yml

@@ -61,6 +61,7 @@ steps:
           compile-extension:typescript-language-features \
           compile-extension:vscode-api-tests \
           compile-extension:vscode-colorize-tests \
+          compile-extension:vscode-colorize-perf-tests \
           compile-extension:vscode-test-resolver
       displayName: Build integration tests
 

build/azure-pipelines/linux/product-build-linux-test.yml

@@ -80,6 +80,7 @@ steps:
           compile-extension:typescript-language-features \
           compile-extension:vscode-api-tests \
           compile-extension:vscode-colorize-tests \
+          compile-extension:vscode-colorize-perf-tests \
           compile-extension:vscode-test-resolver
       displayName: Build integration tests
 

build/azure-pipelines/win32/product-build-win32-test.yml

@@ -63,6 +63,7 @@ steps:
           compile-extension:typescript-language-features `
           compile-extension:vscode-api-tests `
           compile-extension:vscode-colorize-tests `
+          compile-extension:vscode-colorize-perf-tests `
           compile-extension:vscode-test-resolver `
         }
       displayName: Build integration tests

build/gulpfile.extensions.js

@@ -65,6 +65,7 @@ const compilations = [
 	'extensions/typescript-language-features/tsconfig.json',
 	'extensions/vscode-api-tests/tsconfig.json',
 	'extensions/vscode-colorize-tests/tsconfig.json',
+	'extensions/vscode-colorize-perf-tests/tsconfig.json',
 	'extensions/vscode-test-resolver/tsconfig.json',
 
 	'.vscode/extensions/vscode-selfhost-test-provider/tsconfig.json',

build/lib/extensions.js

@@ -246,6 +246,7 @@ function fromGithub({ name, version, repo, sha256, metadata }) {
 const excludedExtensions = [
     'vscode-api-tests',
     'vscode-colorize-tests',
+    'vscode-colorize-perf-tests',
     'vscode-test-resolver',
     'ms-vscode.node-debug',
     'ms-vscode.node-debug2',

build/lib/extensions.ts

@@ -278,6 +278,7 @@ export function fromGithub({ name, version, repo, sha256, metadata }: IExtension
 const excludedExtensions = [
 	'vscode-api-tests',
 	'vscode-colorize-tests',
+	'vscode-colorize-perf-tests',
 	'vscode-test-resolver',
 	'ms-vscode.node-debug',
 	'ms-vscode.node-debug2',

build/npm/dirs.js

@@ -44,6 +44,7 @@ const dirs = [
 	'extensions/typescript-language-features',
 	'extensions/vscode-api-tests',
 	'extensions/vscode-colorize-tests',
+	'extensions/vscode-colorize-perf-tests',
 	'extensions/vscode-test-resolver',
 	'remote',
 	'remote/web',

extensions/vscode-colorize-perf-tests/.gitignore

@@ -0,0 +1,2 @@
+out
+node_modules

extensions/vscode-colorize-perf-tests/.npmrc

@@ -0,0 +1,2 @@
+legacy-peer-deps="true"
+timeout=180000

extensions/vscode-colorize-perf-tests/.vscode/launch.json

@@ -0,0 +1,17 @@
+// A launch configuration that compiles the extension and then opens it inside a new window
+{
+	"version": "0.1.0",
+	"configurations": [
+		{
+			"name": "Launch Tests",
+			"type": "extensionHost",
+			"request": "launch",
+			"runtimeExecutable": "${execPath}",
+			"args": ["${workspaceFolder}/../../", "${workspaceFolder}/test", "--extensionDevelopmentPath=${workspaceFolder}", "--extensionTestsPath=${workspaceFolder}/out" ],
+			"stopOnEntry": false,
+			"sourceMaps": true,
+			"outDir": "${workspaceFolder}/out",
+			"preLaunchTask": "npm"
+		}
+	]
+}

extensions/vscode-colorize-perf-tests/.vscode/tasks.json

@@ -0,0 +1,11 @@
+{
+	"version": "2.0.0",
+	"command": "npm",
+	"type": "shell",
+	"presentation": {
+		"reveal": "silent"
+	},
+	"args": ["run", "compile"],
+	"isBackground": true,
+	"problemMatcher": "$tsc-watch"
+}

extensions/vscode-colorize-perf-tests/package-lock.json

@@ -0,0 +1,42 @@
+{
+  "name": "vscode-colorize-perf-tests",
+  "version": "0.0.1",
+  "lockfileVersion": 3,
+  "requires": true,
+  "packages": {
+    "": {
+      "name": "vscode-colorize-perf-tests",
+      "version": "0.0.1",
+      "license": "MIT",
+      "dependencies": {
+        "jsonc-parser": "^3.2.0"
+      },
+      "devDependencies": {
+        "@types/node": "20.x"
+      },
+      "engines": {
+        "vscode": "*"
+      }
+    },
+    "node_modules/@types/node": {
+      "version": "20.11.24",
+      "resolved": "https://registry.npmjs.org/@types/node/-/node-20.11.24.tgz",
+      "integrity": "sha512-Kza43ewS3xoLgCEpQrsT+xRo/EJej1y0kVYGiLFE1NEODXGzTfwiC6tXTLMQskn1X4/Rjlh0MQUvx9W+L9long==",
+      "dev": true,
+      "dependencies": {
+        "undici-types": "~5.26.4"
+      }
+    },
+    "node_modules/jsonc-parser": {
+      "version": "3.2.0",
+      "resolved": "https://registry.npmjs.org/jsonc-parser/-/jsonc-parser-3.2.0.tgz",
+      "integrity": "sha512-gfFQZrcTc8CnKXp6Y4/CBT3fTc0OVuDofpre4aEeEpSBPV5X5v4+Vmx+8snU7RLPrNHPKSgLxGo9YuQzz20o+w=="
+    },
+    "node_modules/undici-types": {
+      "version": "5.26.5",
+      "resolved": "https://registry.npmjs.org/undici-types/-/undici-types-5.26.5.tgz",
+      "integrity": "sha512-JlCMO+ehdEIKqlFxk6IfVoAUVmgz7cU7zD/h9XZ0qzeosSHmUJVOzSQvvYSYWXkFXC+IfLKSIffhv0sVZup6pA==",
+      "dev": true
+    }
+  }
+}

extensions/vscode-colorize-perf-tests/package.json

@@ -0,0 +1,31 @@
+{
+  "name": "vscode-colorize-perf-tests",
+  "description": "Colorize performance tests for VS Code",
+  "version": "0.0.1",
+  "publisher": "vscode",
+  "license": "MIT",
+  "private": true,
+  "activationEvents": [
+    "onLanguage:json"
+  ],
+  "main": "./out/colorizerTestMain",
+  "engines": {
+    "vscode": "*"
+  },
+  "icon": "media/icon.png",
+  "scripts": {
+    "vscode:prepublish": "node ../../node_modules/gulp/bin/gulp.js --gulpfile ../../build/gulpfile.extensions.js compile-extension:vscode-colorize-perf-tests ./tsconfig.json",
+    "watch": "gulp watch-extension:vscode-colorize-perf-tests",
+    "compile": "gulp compile-extension:vscode-colorize-perf-tests"
+  },
+  "dependencies": {
+    "jsonc-parser": "^3.2.0"
+  },
+  "devDependencies": {
+    "@types/node": "20.x"
+  },
+  "repository": {
+    "type": "git",
+    "url": "https://github.com/microsoft/vscode.git"
+  }
+}

extensions/vscode-colorize-perf-tests/src/colorizer.test.ts

@@ -0,0 +1,146 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import * as fs from 'fs';
+import 'mocha';
+import { basename, join, normalize } from 'path';
+import { commands, ConfigurationTarget, Uri, workspace } from 'vscode';
+
+interface BestsAndWorsts {
+	bestParse?: number;
+	bestCapture?: number;
+	bestMetadata?: number;
+	bestCombined: number;
+	worstParse?: number;
+	worstCapture?: number;
+	worstMetadata?: number;
+	worstCombined: number;
+}
+
+function findBestsAndWorsts(results: { parseTime?: number; captureTime?: number; metadataTime?: number; tokenizeTime?: number }[]): BestsAndWorsts {
+	let bestParse: number | undefined;
+	let bestCapture: number | undefined;
+	let bestMetadata: number | undefined;
+	let bestCombined: number | undefined;
+	let worstParse: number | undefined;
+	let worstCapture: number | undefined;
+	let worstMetadata: number | undefined;
+	let worstCombined: number | undefined;
+
+	for (let i = 0; i < results.length; i++) {
+		const result = results[i];
+		if (result.parseTime && result.captureTime && result.metadataTime) {
+			// Tree Sitter
+			const combined = result.parseTime + result.captureTime + result.metadataTime;
+			if (bestParse === undefined || result.parseTime < bestParse) {
+				bestParse = result.parseTime;
+			}
+			if (bestCapture === undefined || result.captureTime < bestCapture) {
+				bestCapture = result.captureTime;
+			}
+			if (bestMetadata === undefined || result.metadataTime < bestMetadata) {
+				bestMetadata = result.metadataTime;
+			}
+			if (bestCombined === undefined || combined < bestCombined) {
+				bestCombined = combined;
+			}
+			if (i !== 0) {
+				if (worstParse === undefined || result.parseTime > worstParse) {
+					worstParse = result.parseTime;
+				}
+				if (worstCapture === undefined || result.captureTime > worstCapture) {
+					worstCapture = result.captureTime;
+				}
+				if (worstMetadata === undefined || result.metadataTime > worstMetadata) {
+					worstMetadata = result.metadataTime;
+				}
+				if (worstCombined === undefined || combined > worstCombined) {
+					worstCombined = combined;
+				}
+			}
+		} else if (result.tokenizeTime) {
+			// TextMate
+			if (bestCombined === undefined || result.tokenizeTime < bestCombined) {
+				bestCombined = result.tokenizeTime;
+			}
+			if (i !== 0 && (worstCombined === undefined || result.tokenizeTime > worstCombined)) {
+				worstCombined = result.tokenizeTime;
+			}
+		}
+	}
+	return {
+		bestParse,
+		bestCapture,
+		bestMetadata,
+		bestCombined: bestCombined!,
+		worstParse,
+		worstCapture,
+		worstMetadata,
+		worstCombined: worstCombined!,
+	};
+}
+
+interface TreeSitterTimes {
+	parseTime: number;
+	captureTime: number;
+	metadataTime: number;
+}
+
+interface TextMateTimes {
+	tokenizeTime: number;
+}
+
+async function runCommand<TimesType = TreeSitterTimes | TextMateTimes>(command: string, file: Uri, times: number): Promise<TimesType[]> {
+	const results: TimesType[] = [];
+	for (let i = 0; i < times; i++) {
+		results.push(await commands.executeCommand(command, file));
+	}
+	return results;
+}
+
+async function doTest(file: Uri, times: number) {
+	const treeSitterResults = await runCommand<TreeSitterTimes>('_workbench.colorizeTreeSitterTokens', file, times);
+
+	const { bestParse, bestCapture, bestMetadata, bestCombined, worstParse, worstCapture, worstMetadata, worstCombined } = findBestsAndWorsts(treeSitterResults);
+	const textMateResults = await runCommand<TextMateTimes>('_workbench.colorizeTextMateTokens', file, times);
+	const textMateBestWorst = findBestsAndWorsts(textMateResults);
+
+	const toString = (time: number, charLength: number) => {
+		// truncate time to charLength characters
+		return time.toString().slice(0, charLength).padEnd(charLength, ' ');
+	};
+	const numLength = 7;
+	const resultString = `                        | First   | Best    | Worst   |
+| --------------------- | ------- | ------- | ------- |
+| TreeSitter (parse)    | ${toString(treeSitterResults[0].parseTime, numLength)} | ${toString(bestParse!, numLength)} | ${toString(worstParse!, numLength)} |
+| TreeSitter (capture)  | ${toString(treeSitterResults[0].captureTime, numLength)} | ${toString(bestCapture!, numLength)} | ${toString(worstCapture!, numLength)} |
+| TreeSitter (metadata) | ${toString(treeSitterResults[0].metadataTime, numLength)} | ${toString(bestMetadata!, numLength)} | ${toString(worstMetadata!, numLength)} |
+| TreeSitter (total)    | ${toString(treeSitterResults[0].parseTime + treeSitterResults[0].captureTime + treeSitterResults[0].metadataTime, numLength)} | ${toString(bestCombined, numLength)} | ${toString(worstCombined, numLength)} |
+| TextMate              | ${toString(textMateResults[0].tokenizeTime, numLength)} | ${toString(textMateBestWorst.bestCombined, numLength)} | ${toString(textMateBestWorst.worstCombined, numLength)} |
+`;
+	console.log(`File ${basename(file.fsPath)}:`);
+	console.log(resultString);
+}
+
+suite('Tokenization Performance', () => {
+	const testPath = normalize(join(__dirname, '../test'));
+	const fixturesPath = join(testPath, 'colorize-fixtures');
+	let originalSettingValue: any;
+
+	suiteSetup(async function () {
+		originalSettingValue = workspace.getConfiguration('editor').get('experimental.preferTreeSitter');
+		await workspace.getConfiguration('editor').update('experimental.preferTreeSitter', ["typescript"], ConfigurationTarget.Global);
+	});
+	suiteTeardown(async function () {
+		await workspace.getConfiguration('editor').update('experimental.preferTreeSitter', originalSettingValue, ConfigurationTarget.Global);
+	});
+
+	for (const fixture of fs.readdirSync(fixturesPath)) {
+		test(`Full file colorize: ${fixture}`, async function () {
+			await commands.executeCommand('workbench.action.closeAllEditors');
+			await doTest(Uri.file(join(fixturesPath, fixture)), 6);
+		});
+	}
+});

extensions/vscode-colorize-perf-tests/src/colorizerTestMain.ts

@@ -0,0 +1,10 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import * as vscode from 'vscode';
+
+export function activate(_context: vscode.ExtensionContext): any {
+
+}

extensions/vscode-colorize-perf-tests/src/index.ts

@@ -0,0 +1,30 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import * as path from 'path';
+import * as testRunner from '../../../test/integration/electron/testrunner';
+
+const suite = 'Performance Colorize Tests';
+
+const options: import('mocha').MochaOptions = {
+	ui: 'tdd',
+	color: true,
+	timeout: 60000
+};
+
+if (process.env.BUILD_ARTIFACTSTAGINGDIRECTORY) {
+	options.reporter = 'mocha-multi-reporters';
+	options.reporterOptions = {
+		reporterEnabled: 'spec, mocha-junit-reporter',
+		mochaJunitReporterReporterOptions: {
+			testsuitesTitle: `${suite} ${process.platform}`,
+			mochaFile: path.join(process.env.BUILD_ARTIFACTSTAGINGDIRECTORY, `test-results/${process.platform}-${process.arch}-${suite.toLowerCase().replace(/[^\w]/g, '-')}-results.xml`)
+		}
+	};
+}
+
+testRunner.configure(options);
+
+export = testRunner;

extensions/vscode-colorize-perf-tests/test/colorize-fixtures/test.ts

@@ -0,0 +1,111 @@
+/* Game of Life
+ * Implemented in TypeScript
+ * To learn more about TypeScript, please visit http://www.typescriptlang.org/
+ */
+
+module Conway {
+
+	export class Cell {
+		public row: number;
+		public col: number;
+		public live: boolean;
+
+		constructor(row: number, col: number, live: boolean) {
+			this.row = row;
+			this.col = col;
+			this.live = live
+		}
+	}
+
+	export class GameOfLife {
+		private gridSize: number;
+		private canvasSize: number;
+		private lineColor: string;
+		private liveColor: string;
+		private deadColor: string;
+		private initialLifeProbability: number;
+		private animationRate: number;
+		private cellSize: number;
+		private world;
+
+
+		constructor() {
+			this.gridSize = 50;
+			this.canvasSize = 600;
+			this.lineColor = '#cdcdcd';
+			this.liveColor = '#666';
+			this.deadColor = '#eee';
+			this.initialLifeProbability = 0.5;
+			this.animationRate = 60;
+			this.cellSize = 0;
+			this.world = this.createWorld();
+			this.circleOfLife();
+		}
+
+		public createWorld() {
+			return this.travelWorld( (cell : Cell) =>  {
+				cell.live = Math.random() < this.initialLifeProbability;
+				return cell;
+			});
+		}
+
+		public circleOfLife() : void {
+			this.world = this.travelWorld( (cell: Cell) => {
+				cell = this.world[cell.row][cell.col];
+				this.draw(cell);
+				return this.resolveNextGeneration(cell);
+			});
+			setTimeout( () => {this.circleOfLife()}, this.animationRate);
+		}
+
+		public resolveNextGeneration(cell : Cell) {
+			var count = this.countNeighbors(cell);
+			var newCell = new Cell(cell.row, cell.col, cell.live);
+			if(count < 2 || count > 3) newCell.live = false;
+			else if(count == 3) newCell.live = true;
+			return newCell;
+		}
+
+		public countNeighbors(cell : Cell) {
+			var neighbors = 0;
+			for(var row = -1; row <=1; row++) {
+				for(var col = -1; col <= 1; col++) {
+					if(row == 0 && col == 0) continue;
+					if(this.isAlive(cell.row + row, cell.col + col)) {
+						neighbors++;
+					}
+				}
+			}
+			return neighbors;
+		}
+
+		public isAlive(row : number, col : number) {
+			if(row < 0 || col < 0 || row >= this.gridSize || col >= this.gridSize) return false;
+			return this.world[row][col].live;
+		}
+
+		public travelWorld(callback) {
+			var result = [];
+			for(var row = 0; row < this.gridSize; row++) {
+				var rowData = [];
+				for(var col = 0; col < this.gridSize; col++) {
+					rowData.push(callback(new Cell(row, col, false)));
+				}
+				result.push(rowData);
+			}
+			return result;
+		}
+
+		public draw(cell : Cell) {
+			if(this.cellSize == 0) this.cellSize = this.canvasSize/this.gridSize;
+
+			this.context.strokeStyle = this.lineColor;
+			this.context.strokeRect(cell.row * this.cellSize, cell.col*this.cellSize, this.cellSize, this.cellSize);
+			this.context.fillStyle = cell.live ? this.liveColor : this.deadColor;
+			this.context.fillRect(cell.row * this.cellSize, cell.col*this.cellSize, this.cellSize, this.cellSize);
+		}
+
+	}
+}
+
+var game = new Conway.GameOfLife();

extensions/vscode-colorize-perf-tests/tsconfig.json

@@ -0,0 +1,13 @@
+{
+	"extends": "../tsconfig.base.json",
+	"compilerOptions": {
+		"outDir": "./out",
+		"types": [
+			"node"
+		]
+	},
+	"include": [
+		"src/**/*",
+		"../../src/vscode-dts/vscode.d.ts"
+	]
+}

src/vs/editor/browser/services/treeSitter/treeSitterParserService.ts

@@ -5,7 +5,7 @@
 
 import type { Parser } from '@vscode/tree-sitter-wasm';
 import { AppResourcePath, FileAccess, nodeModulesAsarUnpackedPath, nodeModulesPath } from '../../../../base/common/network.js';
-import { EDITOR_EXPERIMENTAL_PREFER_TREESITTER, ITreeSitterParserService, ITreeSitterParseResult } from '../../../common/services/treeSitterParserService.js';
+import { EDITOR_EXPERIMENTAL_PREFER_TREESITTER, ITreeSitterParserService, ITreeSitterParseResult, ITextModelTreeSitter } from '../../../common/services/treeSitterParserService.js';
 import { IModelService } from '../../../common/services/model.js';
 import { Disposable, DisposableMap, DisposableStore, dispose, IDisposable } from '../../../../base/common/lifecycle.js';
 import { ITextModel } from '../../../common/model.js';
@@ -31,7 +31,7 @@ function getModuleLocation(environmentService: IEnvironmentService): AppResource
 	return `${(canASAR && environmentService.isBuilt) ? nodeModulesAsarUnpackedPath : nodeModulesPath}/${MODULE_LOCATION_SUBPATH}`;
 }
 
-export class TextModelTreeSitter extends Disposable {
+export class TextModelTreeSitter extends Disposable implements ITextModelTreeSitter {
 	private _onDidChangeParseResult: Emitter<Range[]> = this._register(new Emitter<Range[]>());
 	public readonly onDidChangeParseResult: Event<Range[]> = this._onDidChangeParseResult.event;
 	private _parseResult: TreeSitterParseResult | undefined;
@@ -42,10 +42,15 @@ export class TextModelTreeSitter extends Disposable {
 		private readonly _treeSitterLanguages: TreeSitterLanguages,
 		private readonly _treeSitterImporter: TreeSitterImporter,
 		private readonly _logService: ILogService,
-		private readonly _telemetryService: ITelemetryService
+		private readonly _telemetryService: ITelemetryService,
+		parseImmediately: boolean = true
 	) {
 		super();
-		this._register(Event.runAndSubscribe(this.model.onDidChangeLanguage, (e => this._onDidChangeLanguage(e ? e.newLanguage : this.model.getLanguageId()))));
+		if (parseImmediately) {
+			this._register(Event.runAndSubscribe(this.model.onDidChangeLanguage, (e => this._onDidChangeLanguage(e ? e.newLanguage : this.model.getLanguageId()))));
+		} else {
+			this._register(this.model.onDidChangeLanguage(e => this._onDidChangeLanguage(e ? e.newLanguage : this.model.getLanguageId())));
+		}
 	}
 
 	private readonly _languageSessionDisposables = this._register(new DisposableStore());
@@ -53,6 +58,10 @@ export class TextModelTreeSitter extends Disposable {
 	 * Be very careful when making changes to this method as it is easy to introduce race conditions.
 	 */
 	private async _onDidChangeLanguage(languageId: string) {
+		this.parse(languageId);
+	}
+
+	public async parse(languageId: string = this.model.getLanguageId()): Promise<ITreeSitterParseResult | undefined> {
 		this._languageSessionDisposables.clear();
 		this._parseResult = undefined;
 
@@ -80,6 +89,7 @@ export class TextModelTreeSitter extends Disposable {
 		}
 
 		this._parseResult = treeSitterTree;
+		return this._parseResult;
 	}
 
 	private _getLanguage(languageId: string, token: CancellationToken): Promise<Parser.Language> {
@@ -459,6 +469,10 @@ export class TreeSitterTextModelService extends Disposable implements ITreeSitte
 		this._modelService.getModels().forEach(model => this._createTextModelTreeSitter(model));
 	}
 
+	public getTextModelTreeSitter(model: ITextModel): ITextModelTreeSitter {
+		return new TextModelTreeSitter(model, this._treeSitterLanguages, this._treeSitterImporter, this._logService, this._telemetryService, false);
+	}
+
 	private _createTextModelTreeSitter(model: ITextModel) {
 		const textModelTreeSitter = new TextModelTreeSitter(model, this._treeSitterLanguages, this._treeSitterImporter, this._logService, this._telemetryService);
 		const disposables = new DisposableStore();

src/vs/editor/common/languages.ts

@@ -93,6 +93,7 @@ export interface ITreeSitterTokenizationSupport {
 	captureAtPosition(lineNumber: number, column: number, textModel: model.ITextModel): Parser.QueryCapture[];
 	captureAtPositionTree(lineNumber: number, column: number, tree: Parser.Tree): Parser.QueryCapture[];
 	onDidChangeTokens: Event<{ textModel: model.ITextModel; changes: IModelTokensChangedEvent }>;
+	tokenizeEncodedInstrumented(lineNumber: number, textModel: model.ITextModel): { result: Uint32Array; captureTime: number; metadataTime: number } | undefined;
 }
 
 /**

src/vs/editor/common/services/treeSitterParserService.ts

@@ -20,9 +20,21 @@ export interface ITreeSitterParserService {
 	getParseResult(textModel: ITextModel): ITreeSitterParseResult | undefined;
 	getTree(content: string, languageId: string): Promise<Parser.Tree | undefined>;
 	onDidUpdateTree: Event<{ textModel: ITextModel; ranges: Range[] }>;
+	/**
+	 * For testing purposes so that the time to parse can be measured.
+	*/
+	getTextModelTreeSitter(textModel: ITextModel): ITextModelTreeSitter | undefined;
 }
 
 export interface ITreeSitterParseResult {
 	readonly tree: Parser.Tree | undefined;
 	readonly language: Parser.Language;
 }
+
+export interface ITextModelTreeSitter {
+	/**
+	 * For testing purposes so that the time to parse can be measured.
+	 */
+	parse(languageId?: string): Promise<ITreeSitterParseResult | undefined>;
+	dispose(): void;
+}

src/vs/editor/standalone/browser/standaloneTreeSitterService.ts

@@ -6,7 +6,7 @@
 import type { Parser } from '@vscode/tree-sitter-wasm';
 import { Event } from '../../../base/common/event.js';
 import { ITextModel } from '../../common/model.js';
-import { ITreeSitterParseResult, ITreeSitterParserService } from '../../common/services/treeSitterParserService.js';
+import { ITextModelTreeSitter, ITreeSitterParseResult, ITreeSitterParserService } from '../../common/services/treeSitterParserService.js';
 import { Range } from '../../common/core/range.js';
 
 /**
@@ -14,6 +14,9 @@ import { Range } from '../../common/core/range.js';
  * We use a dummy sertive here to make the build happy.
  */
 export class StandaloneTreeSitterParserService implements ITreeSitterParserService {
+	getTextModelTreeSitter(textModel: ITextModel): ITextModelTreeSitter | undefined {
+		return undefined;
+	}
 	async getTree(content: string, languageId: string): Promise<Parser.Tree | undefined> {
 		return undefined;
 	}

src/vs/editor/test/common/services/testTreeSitterService.ts

@@ -6,10 +6,13 @@
 import type { Parser } from '@vscode/tree-sitter-wasm';
 import { Event } from '../../../../base/common/event.js';
 import { ITextModel } from '../../../common/model.js';
-import { ITreeSitterParserService, ITreeSitterParseResult } from '../../../common/services/treeSitterParserService.js';
+import { ITreeSitterParserService, ITreeSitterParseResult, ITextModelTreeSitter } from '../../../common/services/treeSitterParserService.js';
 import { Range } from '../../../common/core/range.js';
 
 export class TestTreeSitterParserService implements ITreeSitterParserService {
+	getTextModelTreeSitter(textModel: ITextModel): ITextModelTreeSitter | undefined {
+		throw new Error('Method not implemented.');
+	}
 	getTree(content: string, languageId: string): Promise<Parser.Tree | undefined> {
 		throw new Error('Method not implemented.');
 	}

src/vs/workbench/services/textMate/browser/textMateTokenizationFeature.contribution.ts

@@ -7,6 +7,12 @@ import { registerSingleton, InstantiationType } from '../../../../platform/insta
 import { ITextMateTokenizationService } from './textMateTokenizationFeature.js';
 import { TextMateTokenizationFeature } from './textMateTokenizationFeatureImpl.js';
 import { IWorkbenchContribution, WorkbenchPhase, registerWorkbenchContribution2 } from '../../../common/contributions.js';
+import { CommandsRegistry } from '../../../../platform/commands/common/commands.js';
+import { ServicesAccessor } from '../../../../editor/browser/editorExtensions.js';
+import { URI } from '../../../../base/common/uri.js';
+import { TokenizationRegistry } from '../../../../editor/common/languages.js';
+import { ITextFileService } from '../../textfile/common/textfiles.js';
+import { StopWatch } from '../../../../base/common/stopwatch.js';
 
 /**
  * Makes sure the ITextMateTokenizationService is instantiated
@@ -23,3 +29,24 @@ class TextMateTokenizationInstantiator implements IWorkbenchContribution {
 registerSingleton(ITextMateTokenizationService, TextMateTokenizationFeature, InstantiationType.Eager);
 
 registerWorkbenchContribution2(TextMateTokenizationInstantiator.ID, TextMateTokenizationInstantiator, WorkbenchPhase.BlockRestore);
+
+CommandsRegistry.registerCommand('_workbench.colorizeTextMateTokens', async (accessor: ServicesAccessor, resource?: URI): Promise<{ tokenizeTime: number }> => {
+	const textModelService = accessor.get(ITextFileService);
+	const textModel = resource ? (await textModelService.files.resolve(resource)).textEditorModel : undefined;
+	if (!textModel) {
+		throw new Error(`Cannot resolve text model for resource ${resource}`);
+	}
+
+	const tokenizer = await TokenizationRegistry.getOrCreate(textModel.getLanguageId());
+	if (!tokenizer) {
+		throw new Error(`Cannot resolve tokenizer for language ${textModel.getLanguageId()}`);
+	}
+
+	const stopwatch = new StopWatch();
+	let state = tokenizer.getInitialState();
+	for (let i = 1; i <= textModel.getLineCount(); i++) {
+		state = tokenizer.tokenizeEncoded(textModel.getLineContent(i), true, state).endState;
+	}
+	stopwatch.stop();
+	return { tokenizeTime: stopwatch.elapsed() };
+});

src/vs/workbench/services/treeSitter/browser/treeSitterTokenizationFeature.contribution.ts

@@ -8,6 +8,12 @@ import { IWorkbenchContribution, WorkbenchPhase, registerWorkbenchContribution2
 import { TreeSitterTextModelService } from '../../../../editor/browser/services/treeSitter/treeSitterParserService.js';
 import { ITreeSitterParserService } from '../../../../editor/common/services/treeSitterParserService.js';
 import { ITreeSitterTokenizationFeature } from './treeSitterTokenizationFeature.js';
+import { ServicesAccessor } from '../../../../platform/instantiation/common/instantiation.js';
+import { CommandsRegistry } from '../../../../platform/commands/common/commands.js';
+import { URI } from '../../../../base/common/uri.js';
+import { TreeSitterTokenizationRegistry } from '../../../../editor/common/languages.js';
+import { ITextFileService } from '../../textfile/common/textfiles.js';
+import { StopWatch } from '../../../../base/common/stopwatch.js';
 
 /**
  * Makes sure the ITreeSitterTokenizationService is instantiated
@@ -25,3 +31,38 @@ class TreeSitterTokenizationInstantiator implements IWorkbenchContribution {
 registerSingleton(ITreeSitterParserService, TreeSitterTextModelService, InstantiationType.Eager);
 
 registerWorkbenchContribution2(TreeSitterTokenizationInstantiator.ID, TreeSitterTokenizationInstantiator, WorkbenchPhase.BlockRestore);
+
+CommandsRegistry.registerCommand('_workbench.colorizeTreeSitterTokens', async (accessor: ServicesAccessor, resource?: URI): Promise<{ parseTime: number; captureTime: number; metadataTime: number }> => {
+	const treeSitterParserService = accessor.get(ITreeSitterParserService);
+	const textModelService = accessor.get(ITextFileService);
+	const textModel = resource ? (await textModelService.files.resolve(resource)).textEditorModel : undefined;
+	if (!textModel) {
+		throw new Error(`Cannot resolve text model for resource ${resource}`);
+	}
+
+	const tokenizer = await TreeSitterTokenizationRegistry.getOrCreate(textModel.getLanguageId());
+	if (!tokenizer) {
+		throw new Error(`Cannot resolve tokenizer for language ${textModel.getLanguageId()}`);
+	}
+
+	const textModelTreeSitter = treeSitterParserService.getTextModelTreeSitter(textModel);
+	if (!textModelTreeSitter) {
+		throw new Error(`Cannot resolve tree sitter parser for language ${textModel.getLanguageId()}`);
+	}
+	const stopwatch = new StopWatch();
+	await textModelTreeSitter.parse();
+	stopwatch.stop();
+
+	let captureTime = 0;
+	let metadataTime = 0;
+	for (let i = 1; i <= textModel.getLineCount(); i++) {
+		const result = tokenizer.tokenizeEncodedInstrumented(i, textModel);
+		if (result) {
+			captureTime += result.captureTime;
+			metadataTime += result.metadataTime;
+		}
+	}
+	textModelTreeSitter.dispose();
+	textModel.dispose();
+	return { parseTime: stopwatch.elapsed(), captureTime, metadataTime };
+});

src/vs/workbench/services/treeSitter/browser/treeSitterTokenizationFeature.ts

@@ -19,6 +19,7 @@ import { createDecorator, IInstantiationService } from '../../../../platform/ins
 import { IThemeService } from '../../../../platform/theme/common/themeService.js';
 import { ColorThemeData, findMetadata } from '../../themes/common/colorThemeData.js';
 import { ILanguageService } from '../../../../editor/common/languages/language.js';
+import { StopWatch } from '../../../../base/common/stopwatch.js';
 
 const ALLOWED_SUPPORT = ['typescript'];
 type TreeSitterQueries = string;
@@ -166,6 +167,15 @@ class TreeSitterTokenizationSupport extends Disposable implements ITreeSitterTok
 	 * @returns
 	 */
 	public tokenizeEncoded(lineNumber: number, textModel: ITextModel): Uint32Array | undefined {
+		return this._tokenizeEncoded(lineNumber, textModel)?.result;
+	}
+
+	public tokenizeEncodedInstrumented(lineNumber: number, textModel: ITextModel): { result: Uint32Array; captureTime: number; metadataTime: number } | undefined {
+		return this._tokenizeEncoded(lineNumber, textModel);
+	}
+
+	private _tokenizeEncoded(lineNumber: number, textModel: ITextModel): { result: Uint32Array; captureTime: number; metadataTime: number } | undefined {
+		const stopwatch = StopWatch.create();
 		const lineLength = textModel.getLineMaxColumn(lineNumber);
 		const tree = this._getTree(textModel);
 		const captures = this._captureAtRange(lineNumber, new ColumnRange(1, lineLength), tree?.tree);
@@ -245,6 +255,8 @@ class TreeSitterTokenizationSupport extends Disposable implements ITreeSitterTok
 			endOffsetsAndScopes[tokenIndex].endOffset = lineLength - 1;
 			tokenIndex++;
 		}
+		const captureTime = stopwatch.elapsed();
+		stopwatch.reset();
 
 		const tokens: Uint32Array = new Uint32Array((tokenIndex) * 2);
 		for (let i = 0; i < tokenIndex; i++) {
@@ -255,8 +267,8 @@ class TreeSitterTokenizationSupport extends Disposable implements ITreeSitterTok
 			tokens[i * 2] = token.endOffset;
 			tokens[i * 2 + 1] = findMetadata(this._colorThemeData, token.scopes, encodedLanguageId);
 		}
-
+		const metadataTime = stopwatch.elapsed();
-		return tokens;
+		return { result: tokens, captureTime, metadataTime };
 	}
 
 	override dispose() {