[Exploration] Tree-sitter tokenization exploration (Fixes #161256)

package.json

@@ -57,7 +57,8 @@
     "hygiene": "node --max_old_space_size=4095 ./node_modules/gulp/bin/gulp.js hygiene",
     "core-ci": "node --max_old_space_size=4095 ./node_modules/gulp/bin/gulp.js core-ci",
     "extensions-ci": "node --max_old_space_size=4095 ./node_modules/gulp/bin/gulp.js extensions-ci",
-    "webview-generate-csp-hash": "npx github:apaatsio/csp-hash-from-html csp-hash ./src/vs/workbench/contrib/webview/browser/pre/index.html"
+    "webview-generate-csp-hash": "npx github:apaatsio/csp-hash-from-html csp-hash ./src/vs/workbench/contrib/webview/browser/pre/index.html",
+    "generate-tree-sitter-wasm": "npx tree-sitter build-wasm node_modules/tree-sitter-typescript/typescript"
   },
   "dependencies": {
     "@microsoft/1ds-core-js": "^3.2.2",
@@ -86,6 +87,7 @@
     "vscode-proxy-agent": "^0.12.0",
     "vscode-regexpp": "^3.1.0",
     "vscode-textmate": "7.0.1",
+    "web-tree-sitter": "^0.20.7",
     "xterm": "5.1.0-beta.13",
     "xterm-addon-canvas": "0.3.0-beta.1",
     "xterm-addon-search": "0.11.0-beta.1",
@@ -203,6 +205,8 @@
     "source-map": "0.6.1",
     "source-map-support": "^0.3.2",
     "style-loader": "^1.3.0",
+    "tree-sitter-cli": "^0.20.6",
+    "tree-sitter-typescript": "^0.20.1",
     "ts-loader": "^9.2.7",
     "ts-node": "^10.9.1",
     "tsec": "0.1.4",

src/vs/base/common/stopwatch.ts

@@ -10,15 +10,19 @@ const hasPerformanceNow = (globals.performance && typeof globals.performance.now
 export class StopWatch {
 
 	private _highResolution: boolean;
-	private _startTime: number;
-	private _stopTime: number;
+	private _startTime!: number;
+	private _stopTime!: number;
 
 	public static create(highResolution: boolean = true): StopWatch {
 		return new StopWatch(highResolution);
 	}
 
 	constructor(highResolution: boolean) {
 		this._highResolution = hasPerformanceNow && highResolution;
+		this.reset();
+	}
+
+	public reset(): void {
 		this._startTime = this._now();
 		this._stopTime = -1;
 	}

src/vs/editor/browser/services/treeSitterServices/treeSitterColorizationService.ts

@@ -0,0 +1,123 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+// eslint-disable-next-line local/code-import-patterns
+import { init } from 'web-tree-sitter';
+import { registerSingleton } from 'vs/platform/instantiation/common/extensions';
+import { IModelService } from 'vs/editor/common/services/model';
+import { FileAccess } from 'vs/base/common/network';
+import { createDecorator, ServicesAccessor } from 'vs/platform/instantiation/common/instantiation';
+import { IThemeService } from 'vs/platform/theme/common/themeService';
+import { DisposableStore } from 'vs/base/common/lifecycle';
+import { TreeSitterColorizationTree } from 'vs/editor/browser/services/treeSitterTrees/colorization/treeSitterColorizationTree';
+import { Iterable } from 'vs/base/common/iterator';
+import { Action2, registerAction2 } from 'vs/platform/actions/common/actions';
+import { ITreeSitterService } from 'vs/editor/browser/services/treeSitterServices/treeSitterService';
+import { IFileService } from 'vs/platform/files/common/files';
+
+const ITreeSitterColorizationService = createDecorator<ITreeSitterColorizationService>('ITreeSitterColorizationService');
+
+export interface ITreeSitterColorizationService {
+	readonly _serviceBrand: undefined;
+	registerTreeSittersForColorization(asynchronous: boolean): void;
+	dispose(): void;
+	clearCache(): void;
+}
+
+export class TreeSitterColorizationService implements ITreeSitterColorizationService {
+
+	readonly _serviceBrand: undefined;
+	private readonly _disposableStore: DisposableStore = new DisposableStore();
+	private readonly _treeSittersColorizationTrees: TreeSitterColorizationTree[] = [];
+
+	constructor(
+		@ITreeSitterService private readonly _treeSitterService: ITreeSitterService,
+		@IModelService private readonly _modelService: IModelService,
+		@IThemeService private readonly _themeService: IThemeService,
+		@IFileService private readonly _fileService: IFileService
+	) { }
+
+	registerTreeSittersForColorization(asynchronous: boolean = true) {
+		this._initializeTreeSitterService(asynchronous);
+		const models = this._modelService.getModels();
+		for (const model of models) {
+			if (model.getLanguageId() === 'typescript') {
+				model.tokenization.setTokens([]);
+				this._treeSittersColorizationTrees.push(new TreeSitterColorizationTree(model, this._treeSitterService, this._themeService, this._fileService, asynchronous));
+			}
+		}
+	}
+
+	private _initializeTreeSitterService(asynchronous: boolean = true) {
+		init({
+			locateFile(_file: string, _folder: string) {
+				return FileAccess.asBrowserUri('../../../../../../node_modules/web-tree-sitter/tree-sitter.wasm', require).toString(true);
+			}
+		}).then(async () => {
+			this._disposableStore.add(this._modelService.onModelAdded((model) => {
+				if (model.getLanguageId() === 'typescript') {
+					model.tokenization.setTokens([]);
+					this._treeSittersColorizationTrees.push(new TreeSitterColorizationTree(model, this._treeSitterService, this._themeService, this._fileService, asynchronous));
+				}
+			}));
+			this._disposableStore.add(this._modelService.onModelLanguageChanged((event) => {
+				const model = event.model;
+				if (model.getLanguageId() === 'typescript') {
+					model.tokenization.setTokens([]);
+					this._treeSittersColorizationTrees.push(new TreeSitterColorizationTree(model, this._treeSitterService, this._themeService, this._fileService, asynchronous));
+				}
+			}));
+			this._disposableStore.add(this._modelService.onModelRemoved((model) => {
+				if (model.getLanguageId() === 'typescript') {
+					const treeSitterTreeToDispose = Iterable.find(this._treeSittersColorizationTrees, tree => tree.id === model.id);
+					if (treeSitterTreeToDispose) {
+						treeSitterTreeToDispose.dispose();
+					}
+				}
+			}));
+		});
+	}
+
+	dispose(): void {
+		this._disposableStore.dispose();
+		for (const tree of this._treeSittersColorizationTrees) {
+			tree.dispose();
+		}
+	}
+
+	clearCache(): void {
+		for (const tree of this._treeSittersColorizationTrees) {
+			tree.dispose();
+		}
+		this._treeSittersColorizationTrees.length = 0;
+		this._disposableStore.clear();
+		this._treeSitterService.clearCache();
+	}
+}
+
+registerSingleton(ITreeSitterColorizationService, TreeSitterColorizationService, true);
+
+// Asynchronous colorization that runs when the process is idle
+registerAction2(class extends Action2 {
+	constructor() {
+		super({ id: 'toggleAsynchronousTreeSitterColorization', title: 'Toggle Asynchronous Tree-Sitter Colorization', f1: true });
+	}
+	run(accessor: ServicesAccessor) {
+		const treeSitterTokenizationService = accessor.get(ITreeSitterColorizationService);
+		treeSitterTokenizationService.clearCache();
+		treeSitterTokenizationService.registerTreeSittersForColorization(true);
+	}
+});
+
+// Synchronous colorization for testing the performance
+registerAction2(class extends Action2 {
+	constructor() {
+		super({ id: 'toggleSynchronousTreeSitterColorization', title: 'Toggle Synchronous Tree-Sitter Colorization', f1: true });
+	}
+	run(accessor: ServicesAccessor) {
+		const treeSitterTokenizationService = accessor.get(ITreeSitterColorizationService);
+		treeSitterTokenizationService.clearCache();
+		treeSitterTokenizationService.registerTreeSittersForColorization(false);
+	}
+});

src/vs/editor/browser/services/treeSitterServices/treeSitterService.test.ts

@@ -0,0 +1,107 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+import { ITreeSitterService, TreeSitterService } from 'vs/editor/browser/services/treeSitterServices/treeSitterService';
+import { LanguageFeatureDebounceService } from 'vs/editor/common/services/languageFeatureDebounce';
+import { LanguageFeaturesService } from 'vs/editor/common/services/languageFeaturesService';
+import { LanguageService } from 'vs/editor/common/services/languageService';
+import { ModelService } from 'vs/editor/common/services/modelService';
+import { ServiceCollection } from 'vs/platform/instantiation/common/serviceCollection';
+import { NullLogService } from 'vs/platform/log/common/log';
+import { UndoRedoService } from 'vs/platform/undoRedo/common/undoRedoService';
+import { Range } from 'vs/editor/common/core/range';
+import { Selection } from 'vs/editor/common/core/selection';
+import { FileService } from 'vs/platform/files/common/fileService';
+import { Schemas } from 'vs/base/common/network';
+import { InstantiationService } from 'vs/platform/instantiation/common/instantiationService';
+import { IModelService } from 'vs/editor/common/services/model';
+import { IFileService } from 'vs/platform/files/common/files';
+/* eslint-disable */
+import { TestLanguageConfigurationService } from 'vs/editor/test/common/modes/testLanguageConfigurationService';
+import { createTextModel } from 'vs/editor/test/common/testTextModel';
+import { TestConfigurationService } from 'vs/platform/configuration/test/common/testConfigurationService';
+import { TestDialogService } from 'vs/platform/dialogs/test/common/testDialogService';
+import { TestNotificationService } from 'vs/platform/notification/test/common/testNotificationService';
+import { TestThemeService } from 'vs/platform/theme/test/common/testThemeService';
+import { TestTextResourcePropertiesService } from 'vs/workbench/test/common/workbenchTestServices';
+import { DiskFileSystemProvider } from 'vs/platform/files/node/diskFileSystemProvider';
+/* eslint-enable*/
+
+suite('Testing the Tree-Sitter Service', () => {
+
+	const configService = new TestConfigurationService();
+	configService.setUserConfiguration('editor', { 'detectIndentation': false });
+	const dialogService = new TestDialogService();
+	const notificationService = new TestNotificationService();
+	const logService = new NullLogService();
+	const nullLogService = new NullLogService();
+	const modelService = new ModelService(
+		configService,
+		new TestTextResourcePropertiesService(configService),
+		new TestThemeService(),
+		nullLogService,
+		new UndoRedoService(dialogService, notificationService),
+		new LanguageService(),
+		new TestLanguageConfigurationService(),
+		new LanguageFeatureDebounceService(logService),
+		new LanguageFeaturesService()
+	);
+
+	const diskFileSystemProvider = new DiskFileSystemProvider(logService);
+	const fileService = new FileService(nullLogService);
+	fileService.registerProvider(Schemas.file, diskFileSystemProvider);
+
+	const treeSitterService = new TreeSitterService(modelService, fileService);
+
+	const servicesCollection = new ServiceCollection(
+		[IModelService, modelService],
+		[IFileService, fileService],
+		[ITreeSitterService, treeSitterService]
+	);
+
+	const text = [
+		'function foo() {',
+		'',
+		'}',
+		'/* comment related to TestClass',
+		' end of the comment */',
+		'@classDecorator',
+		'class TestClass {',
+		'// comment related to the function functionOfClass',
+		'functionOfClass(){',
+		'function function1(){',
+		'}',
+		'}}',
+		'function bar() { function insideBar() {}',
+		'}'
+	].join('\n');
+
+	const instantiationService = new InstantiationService(servicesCollection);
+	const treeSitterServiceInstance = instantiationService.createInstance(TreeSitterService);
+
+	// TODO: fix - test retrieves wrong web-tree-sitter bindings
+	test.skip('Checking that parse tree not recomputed if already computed', async () => {
+
+		const model = createTextModel(text, 'typescript');
+		const tree = await treeSitterServiceInstance.getTreeSitterTree(model);
+		await Promise.all([tree.parseTreeAndCountCalls(), tree.parseTreeAndCountCalls()]).then((values) => {
+			console.log('Values returned are : ', values);
+		});
+		model.dispose();
+	});
+
+	// TODO: fix - test retrieves wrong web-tree-sitter bindings
+	test.skip('Checking that parse tree is recomputed when edit is made', async () => {
+		const model = createTextModel(text, 'typescript');
+		const tree = await treeSitterServiceInstance.getTreeSitterTree(model);
+		await tree.parseTreeAndCountCalls().then((value: number) => {
+			console.log('value 1 : ', value);
+		});
+		model.pushEditOperations(null, [{ range: new Range(1, 5, 1, 5), text: '1' }], () => [new Selection(1, 5, 1, 5)]);
+		await tree.parseTreeAndCountCalls().then((value: number) => {
+			console.log('value 2 : ', value);
+		});
+		model.dispose();
+	});
+});