WIP - custom semantic tokens

Author: radeksimko

package.json

@@ -85,6 +85,40 @@
         "path": "./syntaxes/terraform.tmGrammar.json"
       }
     ],
+    "semanticTokenTypes": [
+      {"id": "hcl-AttrName", "superType": "property"},
+      {"id": "hcl-BlockType", "superType": "type"},
+      {"id": "hcl-BlockLabel", "superType": "enumMember"},
+
+      {"id": "hcl-Bool", "superType": "keyword"},
+	    {"id": "hcl-String", "superType": "string"},
+	    {"id": "hcl-Number", "superType": "number"},
+	    {"id": "hcl-ObjectKey", "superType": "parameter"},
+	    {"id": "hcl-MapKey", "superType": "parameter"},
+	    {"id": "hcl-Keyword", "superType": "variable"},
+	    {"id": "hcl-TraversalStep", "superType": "variable"},
+	    {"id": "hcl-TypeCapsule"},
+	    {"id": "hcl-TypePrimitive"}
+    ],
+    "semanticTokenModifiers": [
+      {"id": "hcl-dependent"},
+      {"id": "terraform-block-type-data"},
+      {"id": "terraform-block-label-data-type"},
+      {"id": "terraform-block-label-data-name"},
+      {"id": "terraform-block-type-locals"},
+      {"id": "terraform-block-type-module"},
+      {"id": "terraform-block-label-module-name"},
+      {"id": "terraform-block-type-output"},
+      {"id": "terraform-block-label-output-name"},
+      {"id": "terraform-block-type-provider"},
+      {"id": "terraform-block-label-provider-name"},
+      {"id": "terraform-block-type-resource"},
+      {"id": "terraform-block-label-resource-type"},
+      {"id": "terraform-block-label-resource-name"},
+      {"id": "terraform-block-type-variable"},
+      {"id": "terraform-block-label-variable-name"},
+      {"id": "terraform-block-type-terraform"}
+    ],
     "snippets": [
       {
         "language": "terraform",

src/clientHandler.ts

@@ -14,6 +14,7 @@ import { ServerPath } from './serverPath';
 import { ShowReferencesFeature } from './showReferences';
 import { TelemetryFeature } from './telemetry';
 import { config } from './vscodeUtils';
+import { CustomSemanticTokens } from './semanticTokens';
 
 export interface TerraformLanguageClient {
   commandPrefix: string;
@@ -29,6 +30,9 @@ export class ClientHandler {
   private tfClient: TerraformLanguageClient | undefined;
   private commands: string[] = [];
 
+  public extSemanticTokenTypes: string[] = [];
+  public extSemanticTokenModifiers: string[] = [];
+
   constructor(
     private lsPath: ServerPath,
     private outputChannel: vscode.OutputChannel,
@@ -86,6 +90,10 @@ export class ClientHandler {
     const id = `terraform`;
     const client = new LanguageClient(id, serverOptions, clientOptions);
 
+    client.registerFeature(
+      new CustomSemanticTokens(client, this.extSemanticTokenTypes, this.extSemanticTokenModifiers),
+    );
+
     const codeLensReferenceCount = config('terraform').get<boolean>('codelens.referenceCount');
     if (codeLensReferenceCount) {
       client.registerFeature(new ShowReferencesFeature(client));

src/extension.ts

@@ -1,6 +1,6 @@
 import * as vscode from 'vscode';
 import TelemetryReporter from '@vscode/extension-telemetry';
-import { ExecuteCommandParams, ExecuteCommandRequest } from 'vscode-languageclient';
+import { CodeLensResolveRequest, ExecuteCommandParams, ExecuteCommandRequest } from 'vscode-languageclient';
 import { LanguageClient } from 'vscode-languageclient/node';
 import { Utils } from 'vscode-uri';
 import { ClientHandler, TerraformLanguageClient } from './clientHandler';
@@ -25,6 +25,8 @@ export async function activate(context: vscode.ExtensionContext): Promise<void>
 
   const lsPath = new ServerPath(context);
   clientHandler = new ClientHandler(lsPath, outputChannel, reporter);
+  clientHandler.extSemanticTokenTypes = tokenTypesFromExtManifest(context.extension);
+  clientHandler.extSemanticTokenModifiers = tokenModifiersFromExtManifest(context.extension);
 
   // get rid of pre-2.0.0 settings
   if (config('terraform').has('languageServer.enabled')) {
@@ -293,6 +295,25 @@ async function terraformCommand(command: string, languageServerExec = true): Pro
   }
 }
 
+interface ObjectWithId {
+  id: string;
+}
+
+function tokenTypesFromExtManifest(ext: vscode.Extension<unknown>): string[] {
+  if (!ext.packageJSON.contributes.semanticTokenTypes) {
+    return [];
+  }
+  return ext.packageJSON.contributes.semanticTokenTypes.map((token: ObjectWithId) => token.id);
+}
+
+function tokenModifiersFromExtManifest(ext: vscode.Extension<unknown>): string[] {
+  if (!ext.packageJSON.contributes.semanticTokenModifiers) {
+    return [];
+  }
+
+  return ext.packageJSON.contributes.semanticTokenModifiers.map((modifier: ObjectWithId) => modifier.id);
+}
+
 function enabled(): boolean {
   return config('terraform').get('languageServer.external', false);
 }

src/semanticTokens.ts

@@ -0,0 +1,30 @@
+import { BaseLanguageClient, ClientCapabilities, ServerCapabilities, StaticFeature } from 'vscode-languageclient';
+
+export class CustomSemanticTokens implements StaticFeature {
+  constructor(
+    private _client: BaseLanguageClient,
+    private extTokenTypes: string[],
+    private extTokenModifiers: string[],
+  ) {}
+
+  public fillClientCapabilities(capabilities: ClientCapabilities): void {
+    if (!capabilities.textDocument || !capabilities.textDocument.semanticTokens) {
+      return;
+    }
+
+    const tokenTypes = capabilities.textDocument.semanticTokens.tokenTypes;
+    capabilities.textDocument.semanticTokens.tokenTypes = tokenTypes.concat(this.extTokenTypes);
+
+    const tokenModifiers = capabilities.textDocument.semanticTokens.tokenModifiers;
+    capabilities.textDocument.semanticTokens.tokenModifiers = tokenModifiers.concat(this.extTokenModifiers);
+  }
+
+  // eslint-disable-next-line @typescript-eslint/no-unused-vars
+  public initialize(capabilities: ServerCapabilities): void {
+    return;
+  }
+
+  public dispose(): void {
+    return;
+  }
+}

src/vscodeUtils.ts

@@ -42,6 +42,7 @@ export function isTerraformFile(document?: vscode.TextDocument): boolean {
     return false;
   }
 
+  // TODO: check for supported language IDs here instead
   if (document.fileName.endsWith('tf')) {
     // For the purposes of this extension, anything with the tf file
     // extension is a Terraform file