1import * as fs from 'fs';
2import * as path from 'path';
3import * as vscode from 'vscode';
4import * as vscodelc from 'vscode-languageclient/node';
5
6import * as config from './config';
7import * as configWatcher from './configWatcher';
8
9/**
10 *  This class represents the context of a specific workspace folder.
11 */
12class WorkspaceFolderContext implements vscode.Disposable {
13  dispose() {
14    this.clients.forEach(async client => await client.stop());
15    this.clients.clear();
16  }
17
18  clients: Map<string, vscodelc.LanguageClient> = new Map();
19}
20
21/**
22 *  This class manages all of the MLIR extension state,
23 *  including the language client.
24 */
25export class MLIRContext implements vscode.Disposable {
26  subscriptions: vscode.Disposable[] = [];
27  workspaceFolders: Map<string, WorkspaceFolderContext> = new Map();
28
29  /**
30   *  Activate the MLIR context, and start the language clients.
31   */
32  async activate(outputChannel: vscode.OutputChannel) {
33    // This lambda is used to lazily start language clients for the given
34    // document. It removes the need to pro-actively start language clients for
35    // every folder within the workspace and every language type we provide.
36    const startClientOnOpenDocument = async (document: vscode.TextDocument) => {
37      if (document.uri.scheme !== 'file') {
38        return;
39      }
40      let serverSettingName: string;
41      if (document.languageId === 'mlir') {
42        serverSettingName = 'server_path';
43      } else if (document.languageId === 'pdll') {
44        serverSettingName = 'pdll_server_path';
45      } else if (document.languageId === 'tablegen') {
46        serverSettingName = 'tablegen_server_path';
47      } else {
48        return;
49      }
50
51      // Resolve the workspace folder if this document is in one. We use the
52      // workspace folder when determining if a server needs to be started.
53      const uri = document.uri;
54      let workspaceFolder = vscode.workspace.getWorkspaceFolder(uri);
55      let workspaceFolderStr =
56          workspaceFolder ? workspaceFolder.uri.toString() : "";
57
58      // Get or create a client context for this folder.
59      let folderContext = this.workspaceFolders.get(workspaceFolderStr);
60      if (!folderContext) {
61        folderContext = new WorkspaceFolderContext();
62        this.workspaceFolders.set(workspaceFolderStr, folderContext);
63      }
64      // Start the client for this language if necessary.
65      if (!folderContext.clients.has(document.languageId)) {
66        let client = await this.activateWorkspaceFolder(
67            workspaceFolder, serverSettingName, document.languageId,
68            outputChannel);
69        folderContext.clients.set(document.languageId, client);
70      }
71    };
72    // Process any existing documents.
73    for (const textDoc of vscode.workspace.textDocuments) {
74      await startClientOnOpenDocument(textDoc);
75    }
76
77    // Watch any new documents to spawn servers when necessary.
78    this.subscriptions.push(
79        vscode.workspace.onDidOpenTextDocument(startClientOnOpenDocument));
80    this.subscriptions.push(
81        vscode.workspace.onDidChangeWorkspaceFolders((event) => {
82          for (const folder of event.removed) {
83            const client = this.workspaceFolders.get(folder.uri.toString());
84            if (client) {
85              client.dispose();
86              this.workspaceFolders.delete(folder.uri.toString());
87            }
88          }
89        }));
90  }
91
92  /**
93   *  Prepare a compilation database option for a server.
94   */
95  async prepareCompilationDatabaseServerOptions(
96      languageName: string, workspaceFolder: vscode.WorkspaceFolder,
97      configsToWatch: string[], pathsToWatch: string[],
98      additionalServerArgs: string[]) {
99    // Process the compilation databases attached for the workspace folder.
100    let databases = config.get<string[]>(
101        `${languageName}_compilation_databases`, workspaceFolder, []);
102
103    // If no databases were explicitly specified, default to a database in the
104    // 'build' directory within the current workspace.
105    if (databases.length === 0) {
106      if (workspaceFolder) {
107        databases.push(workspaceFolder.uri.fsPath +
108                       `/build/${languageName}_compile_commands.yml`);
109      }
110
111      // Otherwise, try to resolve each of the paths.
112    } else {
113      for await (let database of databases) {
114        database = await this.resolvePath(database, '', workspaceFolder);
115      }
116    }
117
118    configsToWatch.push(`${languageName}_compilation_databases`);
119    pathsToWatch.push(...databases);
120
121    // Setup the compilation databases as additional arguments to pass to the
122    // server.
123    databases.filter(database => database !== '');
124    additionalServerArgs.push(...databases.map(
125        (database) => `--${languageName}-compilation-database=${database}`));
126  }
127
128  /**
129   *  Prepare the server options for a PDLL server, e.g. populating any
130   *  accessible compilation databases.
131   */
132  async preparePDLLServerOptions(workspaceFolder: vscode.WorkspaceFolder,
133                                 configsToWatch: string[],
134                                 pathsToWatch: string[],
135                                 additionalServerArgs: string[]) {
136    await this.prepareCompilationDatabaseServerOptions(
137        'pdll', workspaceFolder, configsToWatch, pathsToWatch,
138        additionalServerArgs);
139  }
140
141  /**
142   *  Prepare the server options for a TableGen server, e.g. populating any
143   *  accessible compilation databases.
144   */
145  async prepareTableGenServerOptions(workspaceFolder: vscode.WorkspaceFolder,
146                                     configsToWatch: string[],
147                                     pathsToWatch: string[],
148                                     additionalServerArgs: string[]) {
149    await this.prepareCompilationDatabaseServerOptions(
150        'tablegen', workspaceFolder, configsToWatch, pathsToWatch,
151        additionalServerArgs);
152  }
153
154  /**
155   *  Activate the language client for the given language in the given workspace
156   *  folder.
157   */
158  async activateWorkspaceFolder(workspaceFolder: vscode.WorkspaceFolder,
159                                serverSettingName: string, languageName: string,
160                                outputChannel: vscode.OutputChannel):
161      Promise<vscodelc.LanguageClient> {
162    let configsToWatch: string[] = [];
163    let filepathsToWatch: string[] = [];
164    let additionalServerArgs: string[] = [];
165
166    // Initialize additional configurations for this server.
167    if (languageName === 'pdll') {
168      await this.preparePDLLServerOptions(workspaceFolder, configsToWatch,
169                                          filepathsToWatch,
170                                          additionalServerArgs);
171    } else if (languageName == 'tablegen') {
172      await this.prepareTableGenServerOptions(workspaceFolder, configsToWatch,
173                                              filepathsToWatch,
174                                              additionalServerArgs);
175    }
176
177    // Try to activate the language client.
178    const [server, serverPath] = await this.startLanguageClient(
179        workspaceFolder, outputChannel, serverSettingName, languageName,
180        additionalServerArgs);
181    configsToWatch.push(serverSettingName);
182    filepathsToWatch.push(serverPath);
183
184    // Watch for configuration changes on this folder.
185    await configWatcher.activate(this, workspaceFolder, configsToWatch,
186                                 filepathsToWatch);
187    return server;
188  }
189
190  /**
191   *  Start a new language client for the given language. Returns an array
192   *  containing the opened server, or null if the server could not be started,
193   *  and the resolved server path.
194   */
195  async startLanguageClient(workspaceFolder: vscode.WorkspaceFolder,
196                            outputChannel: vscode.OutputChannel,
197                            serverSettingName: string, languageName: string,
198                            additionalServerArgs: string[]):
199      Promise<[ vscodelc.LanguageClient, string ]> {
200    const clientTitle = languageName.toUpperCase() + ' Language Client';
201
202    // Get the path of the lsp-server that is used to provide language
203    // functionality.
204    var serverPath =
205        await this.resolveServerPath(serverSettingName, workspaceFolder);
206
207    // If the server path is empty, bail. We don't emit errors if the user
208    // hasn't explicitly configured the server.
209    if (serverPath === '') {
210      return [ null, serverPath ];
211    }
212
213    // Check that the file actually exists.
214    if (!fs.existsSync(serverPath)) {
215      vscode.window
216          .showErrorMessage(
217              `${clientTitle}: Unable to resolve path for '${
218                  serverSettingName}', please ensure the path is correct`,
219              "Open Setting")
220          .then((value) => {
221            if (value === "Open Setting") {
222              vscode.commands.executeCommand(
223                  'workbench.action.openWorkspaceSettings',
224                  {openToSide : false, query : `mlir.${serverSettingName}`});
225            }
226          });
227      return [ null, serverPath ];
228    }
229
230    // Configure the server options.
231    const serverOptions: vscodelc.ServerOptions = {
232      command : serverPath,
233      args : additionalServerArgs
234    };
235
236    // Configure file patterns relative to the workspace folder.
237    let filePattern: vscode.GlobPattern = '**/*.' + languageName;
238    let selectorPattern: string = null;
239    if (workspaceFolder) {
240      filePattern = new vscode.RelativePattern(workspaceFolder, filePattern);
241      selectorPattern = `${workspaceFolder.uri.fsPath}/**/*`;
242    }
243
244    // Configure the middleware of the client. This is sort of abused to allow
245    // for defining a "fallback" language server that operates on non-workspace
246    // folders. Workspace folder language servers can properly filter out
247    // documents not within the folder, but we can't effectively filter for
248    // documents outside of the workspace. To support this, and avoid having two
249    // servers targeting the same set of files, we use middleware to inject the
250    // dynamic logic for checking if a document is in the workspace.
251    let middleware = {};
252    if (!workspaceFolder) {
253      middleware = {
254        didOpen : (document, next) : Promise<void> => {
255          if (!vscode.workspace.getWorkspaceFolder(document.uri)) {
256            return next(document);
257          }
258          return Promise.resolve();
259        }
260      };
261    }
262
263    // Configure the client options.
264    const clientOptions: vscodelc.LanguageClientOptions = {
265      documentSelector : [
266        {scheme : 'file', language : languageName, pattern : selectorPattern}
267      ],
268      synchronize : {
269        // Notify the server about file changes to language files contained in
270        // the workspace.
271        fileEvents : vscode.workspace.createFileSystemWatcher(filePattern)
272      },
273      outputChannel : outputChannel,
274      workspaceFolder : workspaceFolder,
275      middleware : middleware,
276
277      // Don't switch to output window when the server returns output.
278      revealOutputChannelOn : vscodelc.RevealOutputChannelOn.Never,
279    };
280
281    // Create the language client and start the client.
282    let languageClient = new vscodelc.LanguageClient(
283        languageName + '-lsp', clientTitle, serverOptions, clientOptions);
284    languageClient.start();
285    return [ languageClient, serverPath ];
286  }
287
288  /**
289   * Given a server setting, return the default server path.
290   */
291  static getDefaultServerFilename(serverSettingName: string): string {
292    if (serverSettingName === 'pdll_server_path') {
293      return 'mlir-pdll-lsp-server';
294    }
295    if (serverSettingName === 'server_path') {
296      return 'mlir-lsp-server';
297    }
298    if (serverSettingName === 'tablegen_server_path') {
299      return 'tblgen-lsp-server';
300    }
301    return '';
302  }
303
304  /**
305   * Try to resolve the given path, or the default path, with an optional
306   * workspace folder. If a path could not be resolved, just returns the
307   * input filePath.
308   */
309  async resolvePath(filePath: string, defaultPath: string,
310                    workspaceFolder: vscode.WorkspaceFolder): Promise<string> {
311    const configPath = filePath;
312
313    // If the path is already fully resolved, there is nothing to do.
314    if (path.isAbsolute(filePath)) {
315      return filePath;
316    }
317
318    // If a path hasn't been set, try to use the default path.
319    if (filePath === '') {
320      if (defaultPath === '') {
321        return filePath;
322      }
323      filePath = defaultPath;
324
325      // Fallthrough to try resolving the default path.
326    }
327
328    // Try to resolve the path relative to the workspace.
329    let filePattern: vscode.GlobPattern = '**/' + filePath;
330    if (workspaceFolder) {
331      filePattern = new vscode.RelativePattern(workspaceFolder, filePattern);
332    }
333    let foundUris = await vscode.workspace.findFiles(filePattern, null, 1);
334    if (foundUris.length === 0) {
335      // If we couldn't resolve it, just return the original path anyways. The
336      // file might not exist yet.
337      return configPath;
338    }
339    // Otherwise, return the resolved path.
340    return foundUris[0].fsPath;
341  }
342
343  /**
344   * Try to resolve the path for the given server setting, with an optional
345   * workspace folder.
346   */
347  async resolveServerPath(serverSettingName: string,
348                          workspaceFolder: vscode.WorkspaceFolder):
349      Promise<string> {
350    const serverPath = config.get<string>(serverSettingName, workspaceFolder);
351    const defaultPath = MLIRContext.getDefaultServerFilename(serverSettingName);
352    return this.resolvePath(serverPath, defaultPath, workspaceFolder);
353  }
354
355  /**
356   * Return the language client for the given language and workspace folder, or
357   * null if no client is active.
358   */
359  getLanguageClient(workspaceFolder: vscode.WorkspaceFolder,
360                    languageName: string): vscodelc.LanguageClient {
361    let workspaceFolderStr =
362        workspaceFolder ? workspaceFolder.uri.toString() : "";
363    let folderContext = this.workspaceFolders.get(workspaceFolderStr);
364    if (!folderContext) {
365      return null;
366    }
367    return folderContext.clients.get(languageName);
368  }
369
370  dispose() {
371    this.subscriptions.forEach((d) => { d.dispose(); });
372    this.subscriptions = [];
373    this.workspaceFolders.forEach((d) => { d.dispose(); });
374    this.workspaceFolders.clear();
375  }
376}
377