(feat) Semantic tokens (#752)

#71

This implements the semantic token for TypeScript. Because of TypeScript's implementation, this feature doesn't work with JavaScript.

Author: jasonlyu123

packages/language-server/src/lib/semanticToken/semanticTokenLegend.ts

@@ -0,0 +1,78 @@
+import {
+    SemanticTokensLegend,
+    SemanticTokenModifiers,
+    SemanticTokenTypes
+} from 'vscode-languageserver';
+
+/**
+ * extended from https://github.com/microsoft/TypeScript/blob/35c8df04ad959224fad9037e340c1e50f0540a49/src/services/classifier2020.ts#L9
+ * so that we don't have to map it into our own legend
+ */
+export const enum TokenType {
+    class,
+    enum,
+    interface,
+    namespace,
+    typeParameter,
+    type,
+    parameter,
+    variable,
+    enumMember,
+    property,
+    function,
+    member,
+
+    // svelte
+    event
+}
+
+/**
+ * adopted from https://github.com/microsoft/TypeScript/blob/35c8df04ad959224fad9037e340c1e50f0540a49/src/services/classifier2020.ts#L13
+ * so that we don't have to map it into our own legend
+ */
+export const enum TokenModifier {
+    declaration,
+    static,
+    async,
+    readonly,
+    defaultLibrary,
+    local
+}
+
+export function getSemanticTokenLegends(): SemanticTokensLegend {
+    const tokenModifiers: string[] = [];
+
+    ([
+        [TokenModifier.declaration, SemanticTokenModifiers.declaration],
+        [TokenModifier.static, SemanticTokenModifiers.static],
+        [TokenModifier.async, SemanticTokenModifiers.async],
+        [TokenModifier.readonly, SemanticTokenModifiers.readonly],
+        [TokenModifier.defaultLibrary, SemanticTokenModifiers.defaultLibrary],
+        [TokenModifier.local, 'local']
+    ] as const).forEach(([tsModifier, legend]) => (tokenModifiers[tsModifier] = legend));
+
+    const tokenTypes: string[] = [];
+
+    ([
+        [TokenType.class, SemanticTokenTypes.class],
+        [TokenType.enum, SemanticTokenTypes.enum],
+        [TokenType.interface, SemanticTokenTypes.interface],
+        [TokenType.namespace, SemanticTokenTypes.namespace],
+        [TokenType.typeParameter, SemanticTokenTypes.typeParameter],
+        [TokenType.type, SemanticTokenTypes.type],
+        [TokenType.parameter, SemanticTokenTypes.parameter],
+        [TokenType.variable, SemanticTokenTypes.variable],
+        [TokenType.enumMember, SemanticTokenTypes.enumMember],
+        [TokenType.property, SemanticTokenTypes.property],
+        [TokenType.function, SemanticTokenTypes.function],
+
+        // member is renamed to method in vscode codebase to match LSP default
+        [TokenType.member, SemanticTokenTypes.method],
+        [TokenType.event, SemanticTokenTypes.event]
+    ] as const).forEach(([tokenType, legend]) => (tokenTypes[tokenType] = legend));
+
+    return {
+        tokenModifiers,
+        tokenTypes
+    };
+}

packages/language-server/src/ls-config.ts

@@ -17,7 +17,8 @@ const defaultLSConfig: LSConfig = {
         codeActions: { enable: true },
         rename: { enable: true },
         selectionRange: { enable: true },
-        signatureHelp: { enable: true }
+        signatureHelp: { enable: true },
+        semanticTokens: { enable: true }
     },
     css: {
         enable: true,
@@ -93,6 +94,9 @@ export interface LSTypescriptConfig {
     signatureHelp: {
         enable: boolean;
     };
+    semanticTokens: {
+        enable: boolean;
+    }
 }
 
 export interface LSCSSConfig {

packages/language-server/src/plugins/PluginHost.ts

@@ -17,6 +17,7 @@ import {
     Range,
     ReferenceContext,
     SelectionRange,
+    SemanticTokens,
     SignatureHelp,
     SignatureHelpContext,
     SymbolInformation,
@@ -400,6 +401,23 @@ export class PluginHost implements LSProvider, OnWatchFileChanges {
         }
     }
 
+    async getSemanticTokens(textDocument: TextDocumentIdentifier, range?: Range) {
+        const document = this.getDocument(textDocument.uri);
+        if (!document) {
+            throw new Error('Cannot call methods on an unopened document');
+        }
+
+        return (
+            (await this.execute<SemanticTokens>(
+                'getSemanticTokens',
+                [document, range],
+                ExecuteMode.FirstNonNull
+            )) ?? {
+                data: []
+            }
+        );
+    }
+
     onWatchFileChanges(onWatchFileChangesParas: OnWatchFileChangesPara[]): void {
         for (const support of this.plugins) {
             support.onWatchFileChanges?.(onWatchFileChangesParas);

packages/language-server/src/plugins/interfaces.ts

@@ -1,4 +1,4 @@
-import { CompletionContext, FileChangeType, SignatureHelpContext } from 'vscode-languageserver';
+import { CompletionContext, FileChangeType, SemanticTokens, SignatureHelpContext } from 'vscode-languageserver';
 import {
     CodeAction,
     CodeActionContext,
@@ -141,6 +141,13 @@ export interface SelectionRangeProvider {
     getSelectionRange(document: Document, position: Position): Resolvable<SelectionRange | null>;
 }
 
+export interface SemanticTokensProvider {
+    getSemanticTokens(
+        textDocument: Document,
+        range?: Range
+    ): Resolvable<SemanticTokens>
+}
+
 export interface OnWatchFileChangesPara {
     fileName: string;
     changeType: FileChangeType;
@@ -162,7 +169,8 @@ type ProviderBase = DiagnosticsProvider &
     CodeActionsProvider &
     FindReferencesProvider &
     RenameProvider &
-    SignatureHelpProvider;
+    SignatureHelpProvider &
+    SemanticTokensProvider;
 
 export type LSProvider = ProviderBase & BackwardsCompatibleDefinitionsProvider;
 

packages/language-server/src/plugins/typescript/TypeScriptPlugin.ts

@@ -17,7 +17,8 @@ import {
     CompletionList,
     SelectionRange,
     SignatureHelp,
-    SignatureHelpContext
+    SignatureHelpContext,
+    SemanticTokens
 } from 'vscode-languageserver';
 import {
     Document,
@@ -43,7 +44,8 @@ import {
     SelectionRangeProvider,
     SignatureHelpProvider,
     UpdateImportsProvider,
-    OnWatchFileChangesPara
+    OnWatchFileChangesPara,
+    SemanticTokensProvider
 } from '../interfaces';
 import { SnapshotFragment } from './DocumentSnapshot';
 import { CodeActionsProviderImpl } from './features/CodeActionsProvider';
@@ -62,6 +64,7 @@ import { FindReferencesProviderImpl } from './features/FindReferencesProvider';
 import { SelectionRangeProviderImpl } from './features/SelectionRangeProvider';
 import { SignatureHelpProviderImpl } from './features/SignatureHelpProvider';
 import { SnapshotManager } from './SnapshotManager';
+import { SemanticTokensProviderImpl } from './features/SemanticTokensProvider';
 
 export class TypeScriptPlugin
     implements
@@ -75,6 +78,7 @@ export class TypeScriptPlugin
         FindReferencesProvider,
         SelectionRangeProvider,
         SignatureHelpProvider,
+        SemanticTokensProvider,
         OnWatchFileChanges,
         CompletionsProvider<CompletionEntryWithIdentifer> {
     private readonly configManager: LSConfigManager;
@@ -88,6 +92,7 @@ export class TypeScriptPlugin
     private readonly findReferencesProvider: FindReferencesProviderImpl;
     private readonly selectionRangeProvider: SelectionRangeProviderImpl;
     private readonly signatureHelpProvider: SignatureHelpProviderImpl;
+    private readonly semanticTokensProvider: SemanticTokensProviderImpl;
 
     constructor(
         docManager: DocumentManager,
@@ -112,6 +117,7 @@ export class TypeScriptPlugin
         this.findReferencesProvider = new FindReferencesProviderImpl(this.lsAndTsDocResolver);
         this.selectionRangeProvider = new SelectionRangeProviderImpl(this.lsAndTsDocResolver);
         this.signatureHelpProvider = new SignatureHelpProviderImpl(this.lsAndTsDocResolver);
+        this.semanticTokensProvider = new SemanticTokensProviderImpl(this.lsAndTsDocResolver);
     }
 
     async getDiagnostics(document: Document): Promise<Diagnostic[]> {
@@ -401,6 +407,16 @@ export class TypeScriptPlugin
         return this.signatureHelpProvider.getSignatureHelp(document, position, context);
     }
 
+    async getSemanticTokens(textDocument: Document, range?: Range): Promise<SemanticTokens> {
+        if (!this.featureEnabled('semanticTokens')) {
+            return {
+                data: []
+            };
+        }
+
+        return this.semanticTokensProvider.getSemanticTokens(textDocument, range);
+    }
+
     private getLSAndTSDoc(document: Document) {
         return this.lsAndTsDocResolver.getLSAndTSDoc(document);
     }

packages/language-server/src/plugins/typescript/features/SemanticTokensProvider.ts

@@ -0,0 +1,110 @@
+import ts from 'typescript';
+import {
+    Range,
+    SemanticTokens,
+    SemanticTokensBuilder
+} from 'vscode-languageserver';
+import { Document } from '../../../lib/documents';
+import { SemanticTokensProvider } from '../../interfaces';
+import { SnapshotFragment } from '../DocumentSnapshot';
+import { LSAndTSDocResolver } from '../LSAndTSDocResolver';
+import { convertToTextSpan } from '../utils';
+
+export class SemanticTokensProviderImpl implements SemanticTokensProvider {
+    constructor(private readonly lsAndTsDocResolver: LSAndTSDocResolver) {}
+
+    async getSemanticTokens(textDocument: Document, range?: Range): Promise<SemanticTokens> {
+        const { lang, tsDoc } = this.lsAndTsDocResolver.getLSAndTSDoc(textDocument);
+        const fragment = await tsDoc.getFragment();
+        const textSpan = range
+            ? convertToTextSpan(range, fragment)
+            : {
+                  start: 0,
+                  length: tsDoc.parserError
+                      ? fragment.text.length
+                      : // This is appended by svelte2tsx, there's nothing mappable afterwards
+                        fragment.text.lastIndexOf('return { props:') || fragment.text.length
+              };
+
+        const { spans } = lang.getEncodedSemanticClassifications(
+            tsDoc.filePath,
+            textSpan,
+            ts.SemanticClassificationFormat.TwentyTwenty
+        );
+
+        const builder = new SemanticTokensBuilder();
+        let index = 0;
+
+        while (index < spans.length) {
+            // [start, length, encodedClassification, start2, length2, encodedClassification2]
+            const generatedOffset = spans[index++];
+            const generatedLength = spans[index++];
+            const encodedClassification = spans[index++];
+            const classificationType = this.getTokenTypeFromClassification(encodedClassification);
+            if (classificationType < 0) {
+                continue;
+            }
+
+            const originalPosition = this.mapToOrigin(
+                textDocument,
+                fragment,
+                generatedOffset,
+                generatedLength
+            );
+            if (!originalPosition) {
+                continue;
+            }
+
+            const [line, character, length] = originalPosition;
+
+            // remove identifers whose start and end mapped to the same location
+            // like the svelte2tsx inserted render function
+            if (!length) {
+                continue;
+            }
+
+            const modifier = this.getTokenModifierFromClassification(encodedClassification);
+
+            builder.push(line, character, length, classificationType , modifier);
+        }
+
+        return builder.build();
+    }
+
+    private mapToOrigin(
+        document: Document,
+        fragment: SnapshotFragment,
+        generatedOffset: number,
+        generatedLength: number
+    ): [line: number, character: number, length: number] | undefined {
+        const startPosition = fragment.getOriginalPosition(fragment.positionAt(generatedOffset));
+
+        if (startPosition.line < 0) {
+            return;
+        }
+
+        const endPosition = fragment.getOriginalPosition(
+            fragment.positionAt(generatedOffset + generatedLength)
+        );
+        const startOffset = document.offsetAt(startPosition);
+        const endOffset = document.offsetAt(endPosition);
+
+        return [startPosition.line, startPosition.character, endOffset - startOffset];
+    }
+
+    /**
+     *  TSClassification = (TokenType + 1) << TokenEncodingConsts.typeOffset + TokenModifier
+     */
+    private getTokenTypeFromClassification(tsClassification: number): number {
+        return (tsClassification >> TokenEncodingConsts.typeOffset) - 1;
+    }
+
+    private getTokenModifierFromClassification(tsClassification: number) {
+        return tsClassification & TokenEncodingConsts.modifierMask;
+    }
+}
+
+const enum TokenEncodingConsts {
+    typeOffset = 8,
+    modifierMask = (1 << typeOffset) - 1
+}

packages/language-server/src/plugins/typescript/utils.ts

@@ -286,3 +286,13 @@ export function getTsCheckComment(str = ''): string | undefined {
         }
     }
 }
+
+export function convertToTextSpan(range: Range, fragment: SnapshotFragment): ts.TextSpan {
+    const start = fragment.offsetAt(fragment.getGeneratedPosition(range.start));
+    const end = fragment.offsetAt(fragment.getGeneratedPosition(range.end));
+
+    return {
+        start,
+        length: end - start
+    };
+}