[mlir][PDLL] Add initial support for a PDLL compilation database

The compilation database acts in a similar way to the compilation database
(compile_commands.json) used by clang-tidy, i.e. it provides additional
information about the compilation of project files to help the language
server. The main piece of information provided by the PDLL compilation
database in this commit is the set of include directories used when processing
the input .pdll file. This allows for the server to properly process .pdll files
that use includes anchored by the include directories set up in the build system.

The structure of the textual form of a compilation database is a yaml file
containing documents of the following form:

```
--- !FileInfo:
  filepath: <string> - Absolute file path of the file.
  includes: <string> - Semi-colon delimited list of include directories.
```

This commit also adds support to cmake for automatically generating
a `pdll_compile_commands.yml` file at the top-level of the build
directory.

Differential Revision: https://reviews.llvm.org/D124076

3 files changed, 121 insertions, 48 deletions

diff --git a/mlir/utils/vscode/package.json b/mlir/utils/vscode/package.json
index 230124a6a629..9b4de106991b 100644
--- a/mlir/utils/vscode/package.json
+++ b/mlir/utils/vscode/package.json
@@ -124,6 +124,11 @@
           "type": "string",
           "description": "The file path of the mlir-pdll-lsp-server executable."
         },
+        "mlir.pdll_compilation_databases": {
+          "scope": "resource",
+          "type": "array",
+          "description": "A list of `pdll_compile_commands.yml` database files containing information about .pdll files processed by the server."
+        },
         "mlir.onSettingsChanged": {
           "type": "string",
           "default": "prompt",
diff --git a/mlir/utils/vscode/src/configWatcher.ts b/mlir/utils/vscode/src/configWatcher.ts
index 82b202d9c3c9..68426873c18e 100644
--- a/mlir/utils/vscode/src/configWatcher.ts
+++ b/mlir/utils/vscode/src/configWatcher.ts
@@ -41,41 +41,45 @@ async function promptRestart(settingName: string, promptMessage: string) {
  *  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,
-                               serverSetting: string, serverPath: string) {
+export async function activate(
+    mlirContext: MLIRContext, workspaceFolder: vscode.WorkspaceFolder,
+    serverSettings: string[], serverPaths: string[]) {
   // When a configuration change happens, check to see if we should restart the
   // server.
   mlirContext.subscriptions.push(vscode.workspace.onDidChangeConfiguration(event => {
-    const expandedSetting = `mlir.${serverSetting}`;
-    if (event.affectsConfiguration(expandedSetting, workspaceFolder)) {
-      promptRestart(
-          'onSettingsChanged',
-          `setting '${
-              expandedSetting}' has changed. Do you want to reload the server?`);
+    for (const serverSetting of serverSettings) {
+      const expandedSetting = `mlir.${serverSetting}`;
+      if (event.affectsConfiguration(expandedSetting, workspaceFolder)) {
+        promptRestart(
+            'onSettingsChanged',
+            `setting '${
+                expandedSetting}' has changed. Do you want to reload the server?`);
+      }
     }
   }));
 
-  // If the server path actually exists, track it in case it changes. Check that
-  // the path actually exists.
-  if (serverPath === '') {
-    return;
-  }
-
+  // Setup watchers for the provided server paths.
   const fileWatcherConfig = {
     disableGlobbing : true,
     followSymlinks : true,
     ignoreInitial : true,
     awaitWriteFinish : true,
   };
-  const fileWatcher = chokidar.watch(serverPath, fileWatcherConfig);
-  fileWatcher.on('all', (event, _filename, _details) => {
-    if (event != 'unlink') {
-      promptRestart(
-          'onSettingsChanged',
-          'MLIR language server binary has changed. Do you want to reload the server?');
+  for (const serverPath of serverPaths) {
+    if (serverPath === '') {
+      return;
     }
-  });
-  mlirContext.subscriptions.push(
-      new vscode.Disposable(() => { fileWatcher.close(); }));
+
+    // If the server path actually exists, track it in case it changes.
+    const fileWatcher = chokidar.watch(serverPath, fileWatcherConfig);
+    fileWatcher.on('all', (event, _filename, _details) => {
+      if (event != 'unlink') {
+        promptRestart(
+            'onSettingsChanged',
+            'MLIR language server file has changed. Do you want to reload the server?');
+      }
+    });
+    mlirContext.subscriptions.push(
+        new vscode.Disposable(() => { fileWatcher.close(); }));
+  }
 }
diff --git a/mlir/utils/vscode/src/mlirContext.ts b/mlir/utils/vscode/src/mlirContext.ts
index 3827fcc86d33..76a990c87b6c 100644
--- a/mlir/utils/vscode/src/mlirContext.ts
+++ b/mlir/utils/vscode/src/mlirContext.ts
@@ -86,6 +86,43 @@ export class MLIRContext implements vscode.Disposable {
   }
 
   /**
+   *  Prepare the server options for a PDLL server, e.g. populating any
+   *  accessible compilation databases.
+   */
+  async preparePDLLServerOptions(workspaceFolder: vscode.WorkspaceFolder,
+                                 configsToWatch: string[],
+                                 pathsToWatch: string[],
+                                 additionalServerArgs: string[]) {
+    // Process the compilation databases attached for the workspace folder.
+    let databases =
+        config.get<string[]>('pdll_compilation_databases', workspaceFolder);
+
+    // If no databases were explicitly specified, default to a database in the
+    // 'build' directory within the current workspace.
+    if (databases.length === 0) {
+      if (workspaceFolder) {
+        databases.push(workspaceFolder.uri.fsPath +
+                       '/build/pdll_compile_commands.yml');
+      }
+
+      // Otherwise, try to resolve each of the paths.
+    } else {
+      for await (let database of databases) {
+        database = await this.resolvePath(database, '', workspaceFolder);
+      }
+    }
+
+    configsToWatch.push('pdll_compilation_databases');
+    pathsToWatch.push(...databases);
+
+    // Setup the compilation databases as additional arguments to pass to the
+    // server.
+    databases.filter(database => database !== '');
+    additionalServerArgs.push(...databases.map(
+        (database) => `--pdll-compilation-database=${database}`));
+  }
+
+  /**
    *  Activate the language client for the given language in the given workspace
    *  folder.
    */
@@ -93,12 +130,27 @@ export class MLIRContext implements vscode.Disposable {
                                 serverSettingName: string, languageName: string,
                                 outputChannel: vscode.OutputChannel):
       Promise<vscodelc.LanguageClient> {
+    let configsToWatch: string[] = [];
+    let filepathsToWatch: string[] = [];
+    let additionalServerArgs: string[] = [];
+
+    // Initialize additional configurations for this server.
+    if (languageName === 'pdll') {
+      await this.preparePDLLServerOptions(workspaceFolder, configsToWatch,
+                                          filepathsToWatch,
+                                          additionalServerArgs);
+    }
+
+    // Try to activate the language client.
     const [server, serverPath] = await this.startLanguageClient(
-        workspaceFolder, outputChannel, serverSettingName, languageName);
+        workspaceFolder, outputChannel, serverSettingName, languageName,
+        additionalServerArgs);
+    configsToWatch.push(serverSettingName);
+    filepathsToWatch.push(serverPath);
 
     // Watch for configuration changes on this folder.
-    await configWatcher.activate(this, workspaceFolder, serverSettingName,
-                                 serverPath);
+    await configWatcher.activate(this, workspaceFolder, configsToWatch,
+                                 filepathsToWatch);
     return server;
   }
 
@@ -109,7 +161,8 @@ export class MLIRContext implements vscode.Disposable {
    */
   async startLanguageClient(workspaceFolder: vscode.WorkspaceFolder,
                             outputChannel: vscode.OutputChannel,
-                            serverSettingName: string, languageName: string):
+                            serverSettingName: string, languageName: string,
+                            additionalServerArgs: string[]):
       Promise<[ vscodelc.LanguageClient, string ]> {
     const clientTitle = languageName.toUpperCase() + ' Language Client';
 
@@ -146,12 +199,12 @@ export class MLIRContext implements vscode.Disposable {
       run : {
         command : serverPath,
         transport : vscodelc.TransportKind.stdio,
-        args : []
+        args : additionalServerArgs
       },
       debug : {
         command : serverPath,
         transport : vscodelc.TransportKind.stdio,
-        args : []
+        args : additionalServerArgs
       }
     };
 
@@ -217,45 +270,56 @@ export class MLIRContext implements vscode.Disposable {
   }
 
   /**
-   * Try to resolve the path for the given server setting, with an optional
-   * workspace folder.
+   * Try to resolve the given path, or the default path, with an optional
+   * workspace folder. If a path could not be resolved, just returns the
+   * input filePath.
    */
-  async resolveServerPath(serverSettingName: string,
-                          workspaceFolder: vscode.WorkspaceFolder):
-      Promise<string> {
-    const configServerPath =
-        config.get<string>(serverSettingName, workspaceFolder);
-    let serverPath = configServerPath;
+  async resolvePath(filePath: string, defaultPath: string,
+                    workspaceFolder: vscode.WorkspaceFolder): Promise<string> {
+    const configPath = filePath;
 
     // If the path is already fully resolved, there is nothing to do.
-    if (path.isAbsolute(serverPath)) {
-      return serverPath;
+    if (path.isAbsolute(filePath)) {
+      return filePath;
     }
 
     // If a path hasn't been set, try to use the default path.
-    if (serverPath === '') {
-      serverPath = MLIRContext.getDefaultServerFilename(serverSettingName);
-      if (serverPath === '') {
-        return serverPath;
+    if (filePath === '') {
+      if (defaultPath === '') {
+        return filePath;
       }
+      filePath = defaultPath;
+
       // Fallthrough to try resolving the default path.
     }
 
     // Try to resolve the path relative to the workspace.
-    let filePattern: vscode.GlobPattern = '**/' + serverPath;
+    let filePattern: vscode.GlobPattern = '**/' + filePath;
     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;
+      // If we couldn't resolve it, just return the original path anyways. The
+      // file might not exist yet.
+      return configPath;
     }
     // Otherwise, return the resolved path.
     return foundUris[0].fsPath;
   }
 
+  /**
+   * Try to resolve the path for the given server setting, with an optional
+   * workspace folder.
+   */
+  async resolveServerPath(serverSettingName: string,
+                          workspaceFolder: vscode.WorkspaceFolder):
+      Promise<string> {
+    const serverPath = config.get<string>(serverSettingName, workspaceFolder);
+    const defaultPath = MLIRContext.getDefaultServerFilename(serverSettingName);
+    return this.resolvePath(serverPath, defaultPath, workspaceFolder);
+  }
+
   dispose() {
     this.subscriptions.forEach((d) => { d.dispose(); });
     this.subscriptions = [];