Support custom semantic tokens & modifiers (#958)

* Support custom semantic tokens & modifiers

* improve function interfaces

package.json

@@ -92,6 +92,58 @@
         "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", "superType": "keyword"},
+	    {"id": "hcl-typePrimitive", "superType": "keyword"}
+    ],
+    "semanticTokenScopes": [
+      {
+        "scopes": {
+          "hcl-attrName": ["variable.other.property"],
+          "hcl-blockType": ["entity.name.type"],
+          "hcl-blockLabel": ["variable.other.enummember"],
+
+          "hcl-bool": ["keyword.control"],
+          "hcl-string": ["string"],
+          "hcl-number": ["constant.numeric"],
+          "hcl-objectKey": ["variable.parameter"],
+          "hcl-mapKey": ["variable.parameter"],
+          "hcl-keyword": ["keyword.control"],
+          "hcl-traversalStep": ["variable.other.readwrite"],
+          "hcl-typeCapsule": ["keyword.control"],
+          "hcl-typePrimitive": ["keyword.control"]
+        }
+      }
+    ],
+    "semanticTokenModifiers": [
+      {"id": "hcl-dependent"},
+
+      {"id": "terraform-data"},
+      {"id": "terraform-locals"},
+      {"id": "terraform-module"},
+      {"id": "terraform-output"},
+      {"id": "terraform-provider"},
+      {"id": "terraform-resource"},
+      {"id": "terraform-provisioner"},
+      {"id": "terraform-connection"},
+      {"id": "terraform-variable"},
+      {"id": "terraform-terraform"},
+      {"id": "terraform-backend"},
+      {"id": "terraform-name"},
+      {"id": "terraform-type"},
+      {"id": "terraform-requiredProviders"}
+    ],
     "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

@@ -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(manifest);
+  clientHandler.extSemanticTokenModifiers = tokenModifiersFromExtManifest(manifest);
 
   // get rid of pre-2.0.0 settings
   if (config('terraform').has('languageServer.enabled')) {
@@ -293,6 +295,32 @@ async function terraformCommand(command: string, languageServerExec = true): Pro
   }
 }
 
+interface PartialManifest {
+  contributes: {
+    semanticTokenTypes?: ObjectWithId[];
+    semanticTokenModifiers?: ObjectWithId[];
+  };
+}
+
+interface ObjectWithId {
+  id: string;
+}
+
+function tokenTypesFromExtManifest(manifest: PartialManifest): string[] {
+  if (!manifest.contributes.semanticTokenTypes) {
+    return [];
+  }
+  return manifest.contributes.semanticTokenTypes.map((token: ObjectWithId) => token.id);
+}
+
+function tokenModifiersFromExtManifest(manifest: PartialManifest): string[] {
+  if (!manifest.contributes.semanticTokenModifiers) {
+    return [];
+  }
+
+  return manifest.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