[mlir-vscode] Add better support for multiple workspace folders

[rriddle]

Summary

We currently only launch one set of language clients when starting the extension,
but this has the unfortunate effect of applying the same settings to all workspace
folders. This commit adds support for multiple workspace folders by launching
a server for each folder in the workspace. This allows for having different servers
for different workspace folders, e.g. when there are multiple MLIR projects in
the same workspace.

[jpienaar]

LG, I'm just not sure how much the optionality/null is getting you here vs extra paths created.

mlir/utils/vscode/src/mlirContext.ts

44

Is it inefficient to specify workspace folder rather than null in the singleton case?

[rriddle]

In D122793#3419370, @jpienaar wrote:

LG, I'm just not sure how much the optionality/null is getting you here vs extra paths created.

The problem is that a workspace folder isn't actually available in all cases, so to avoid the null checks we'd have to make the extension fail in certain cases (though to be fair, I'm not sure how important those specific cases are).

mlir/utils/vscode/src/mlirContext.ts

44

In this case (the one where the comment is pointing) there is no workspace folder, such as when you just open a .mlir file. In the case where we only have one folder in the workspace, passing in the folder shouldn't be inefficient given that it helps gate various things (such as finding the server). Providing the workspace folder when possible also properly handles when the folder itself overrides the extension configuration (e.g. to specify a specific server).

mlir/utils/vscode/package.json

 {
   "name": "vscode-mlir",
   "displayName": "MLIR",
   "description": "MLIR Language Extension",
-  "version": "0.0.4",
+  "version": "0.0.5",
   "publisher": "llvm-vs-code-extensions",
   "homepage": "https://mlir.llvm.org/",
   "icon": "icon.png",
   "engines": {
     "vscode": "^1.52.0"
   },
   "categories": [
     "Programming Languages"

mlir/utils/vscode/src/config.ts

 import * as vscode from 'vscode';
 
 /**
- *  Gets the config value `mlir.<key>`.
+ *  Gets the config value `mlir.<key>`, with an optional workspace folder.
  */
-export function get<T>(key: string): T {
-  return vscode.workspace.getConfiguration('mlir').get<T>(key);
+export function get<T>(key: string,
+                       workspaceFolder: vscode.WorkspaceFolder = null): T {
+  return vscode.workspace.getConfiguration('mlir', workspaceFolder).get<T>(key);
 }
 
 /**
  *  Sets the config value `mlir.<key>`.
  */
 export function update<T>(key: string, value: T,
                           target?: vscode.ConfigurationTarget) {
   return vscode.workspace.getConfiguration('mlir').update(key, value, target);
 }

mlir/utils/vscode/src/configWatcher.ts

default:
     default:
       break;
     }
     break;
   }
 }
 
 /**
- *  Activate the watchers that track configuration changes which decide when to
- *  restart the server.
+ *  Activate watchers that track configuration changes for the given workspace
+ *  folder, or null if the workspace is top-level.
  */
 export async function activate(mlirContext: MLIRContext,
+                               workspaceFolder: vscode.WorkspaceFolder,
                                serverPathsToWatch: string[]) {
   // When a configuration change happens, check to see if we should restart the
   // server.
   mlirContext.subscriptions.push(vscode.workspace.onDidChangeConfiguration(event => {
     const settings: string[] = [ 'server_path', 'pdll_server_path' ];
     for (const setting of settings) {
       const expandedSetting = `mlir.${setting}`;
-      if (event.affectsConfiguration(expandedSetting)) {
+      if (event.affectsConfiguration(expandedSetting, workspaceFolder)) {
         promptRestart(
             'onSettingsChanged',
             `setting '${
                 expandedSetting}' has changed. Do you want to reload the server?`);
         break;
       }
     }
   }));

mlir/utils/vscode/src/mlirContext.ts

 import * as fs from 'fs';
 import * as path from 'path';
 import * as vscode from 'vscode';
 import * as vscodelc from 'vscode-languageclient';
 
 import * as config from './config';
 import * as configWatcher from './configWatcher';
 
 /**
+ *  This class represents the context of a specific workspace folder.
+ */
+class WorkspaceFolderContext {
+  constructor(mlirServer: vscodelc.LanguageClient,
+              pdllServer: vscodelc.LanguageClient) {
+    this.mlirServer = mlirServer;
+    this.pdllServer = pdllServer;
+  }
+  mlirServer!: vscodelc.LanguageClient;
+  pdllServer!: vscodelc.LanguageClient;
+}
+
+/**
  *  This class manages all of the MLIR extension state,
  *  including the language client.
  */
 export class MLIRContext implements vscode.Disposable {
   subscriptions: vscode.Disposable[] = [];
-  client!: vscodelc.LanguageClient;
-  pdllClient!: vscodelc.LanguageClient;
+  workspaceFolders: WorkspaceFolderContext[] = [];
 
   /**
    *  Activate the MLIR context, and start the language clients.
    */
   async activate(outputChannel: vscode.OutputChannel,
                  warnOnEmptyServerPath: boolean) {
+    // Start clients for each workspace folder.
+    if (vscode.workspace.workspaceFolders &&
+        vscode.workspace.workspaceFolders.length > 0) {
+      for (const workspaceFolder of vscode.workspace.workspaceFolders) {
+        this.workspaceFolders.push(await this.activateWorkspaceFolder(
+            workspaceFolder, outputChannel, warnOnEmptyServerPath));
+      }
+    } else {
+      this.workspaceFolders.push(await this.activateWorkspaceFolder(
+          null, outputChannel, warnOnEmptyServerPath));
+    }
+  }
+
+  /**
+   *  Activate the context for the given workspace folder, and start the
+   *  language clients.
+   */
+  async activateWorkspaceFolder(workspaceFolder: vscode.WorkspaceFolder,
+                                outputChannel: vscode.OutputChannel,
+                                warnOnEmptyServerPath: boolean):
+      Promise<WorkspaceFolderContext> {
     // Create the language clients for mlir and pdll.
-    let mlirServerPath: string, pdllServerPath: string;
-    [this.client, mlirServerPath] = await this.startLanguageClient(
-        outputChannel, warnOnEmptyServerPath, 'server_path', 'mlir');
-    [this.pdllClient, pdllServerPath] = await this.startLanguageClient(
-        outputChannel, warnOnEmptyServerPath, 'pdll_server_path', 'pdll');
+    const [mlirServer, mlirServerPath] = await this.startLanguageClient(
+        workspaceFolder, outputChannel, warnOnEmptyServerPath, 'server_path',
+        'mlir');
+    const [pdllServer, pdllServerPath] = await this.startLanguageClient(
+        workspaceFolder, outputChannel, warnOnEmptyServerPath,
+        'pdll_server_path', 'pdll');
 
-    // Watch for configuration changes.
+    // Watch for configuration changes on this folder.
     const serverPathsToWatch = [ mlirServerPath, pdllServerPath ];
-    await configWatcher.activate(this, serverPathsToWatch);
+    await configWatcher.activate(this, workspaceFolder, serverPathsToWatch);
+
+    return new WorkspaceFolderContext(mlirServer, pdllServer);
   }
 
   /**
    *  Start a new language client for the given language. Returns an array
    *  containing the opened server, or null if the server could not be started,
    *  and the resolved server path.
    */
-  async startLanguageClient(outputChannel: vscode.OutputChannel,
+  async startLanguageClient(workspaceFolder: vscode.WorkspaceFolder,
+                            outputChannel: vscode.OutputChannel,
                             warnOnEmptyServerPath: boolean,
                             serverSettingName: string, languageName: string):
       Promise<[ vscodelc.LanguageClient, string ]> {
     const clientTitle = languageName.toUpperCase() + ' Language Client';
 
     // Get the path of the lsp-server that is used to provide language
     // functionality.
-    var serverPath = await this.resolveServerPath(serverSettingName);
+    var serverPath =
+        await this.resolveServerPath(serverSettingName, workspaceFolder);
 
     // If we aren't emitting warnings on an empty server path, and the server
     // path is empty, bail.
     if (!warnOnEmptyServerPath && serverPath === '') {
       return [ null, serverPath ];
     }
 
     // Check that the file actually exists.
const serverOptions: vscodelc.ServerOptions = {
       },
       debug : {
         command : serverPath,
         transport : vscodelc.TransportKind.stdio,
         args : []
       }
     };
 
+    // Configure file patterns relative to the workspace folder.
+    let filePattern: vscode.GlobPattern = '**/*.' + languageName;
+    let selectorPattern: string = null;
+    if (workspaceFolder) {
+      filePattern = new vscode.RelativePattern(workspaceFolder, filePattern);
+      selectorPattern = `${workspaceFolder.uri.fsPath}/**/*`;
+    }
+
     // Configure the client options.
     const clientOptions: vscodelc.LanguageClientOptions = {
-      documentSelector : [ {scheme : 'file', language : languageName} ],
+      documentSelector : [
+        {scheme : 'file', language : languageName, pattern : selectorPattern}
+      ],
       synchronize : {
         // Notify the server about file changes to language files contained in
         // the workspace.
-        fileEvents :
-            vscode.workspace.createFileSystemWatcher('**/*.' + languageName)
+        fileEvents : vscode.workspace.createFileSystemWatcher(filePattern)
       },
       outputChannel : outputChannel,
+      workspaceFolder : workspaceFolder
     };
 
     // Create the language client and start the client.
     let languageClient = new vscodelc.LanguageClient(
         languageName + '-lsp', clientTitle, serverOptions, clientOptions);
     this.subscriptions.push(languageClient.start());
     return [ languageClient, serverPath ];
   }
 
   /**
    * Given a server setting, return the default server path.
    */
   static getDefaultServerFilename(serverSettingName: string): string {
     if (serverSettingName === 'pdll_server_path') {
       return 'mlir-pdll-lsp-server';
     }
     if (serverSettingName === 'server_path') {
       return 'mlir-lsp-server';
     }
     return '';
   }
 
   /**
-   * Try to resolve the path for the given server setting.
+   * Try to resolve the path for the given server setting, with an optional
+   * workspace folder.
    */
-  async resolveServerPath(serverSettingName: string): Promise<string> {
-    let configServerPath = config.get<string>(serverSettingName);
+  async resolveServerPath(serverSettingName: string,
+                          workspaceFolder: vscode.WorkspaceFolder):
+      Promise<string> {
+    const configServerPath =
+        config.get<string>(serverSettingName, workspaceFolder);
     let serverPath = configServerPath;
 
     // If the path is already fully resolved, there is nothing to do.
     if (path.isAbsolute(serverPath)) {
       return serverPath;
     }
 
     // If a path hasn't been set, try to use the default path.
     if (serverPath === '') {
       serverPath = MLIRContext.getDefaultServerFilename(serverSettingName);
       if (serverPath === '') {
         return serverPath;
       }
       // Fallthrough to try resolving the default path.
     }
 
     // Try to resolve the path relative to the workspace.
-    const foundUris: vscode.Uri[] =
-        await vscode.workspace.findFiles('**/' + serverPath, null, 1);
+    let filePattern: vscode.GlobPattern = '**/' + serverPath;
+    if (workspaceFolder) {
+      filePattern = new vscode.RelativePattern(workspaceFolder, filePattern);
+    }
+    let foundUris = await vscode.workspace.findFiles(filePattern, null, 1);
     if (foundUris.length === 0) {
       // If we couldn't resolve it, just return the current configuration path
       // anyways. The file might not exist yet.
       return configServerPath;
     }
     // Otherwise, return the resolved path.
     return foundUris[0].fsPath;
   }
 
   dispose() {
     this.subscriptions.forEach((d) => { d.dispose(); });
     this.subscriptions = [];
+    this.workspaceFolders = [];
   }
 }