195 lines · typescript
1import * as base64 from 'base64-js'2import * as vscode from 'vscode'3 4import {MLIRContext} from '../mlirContext';5 6/**7 * The parameters to the mlir/convert(To|From)Bytecode commands. These8 * parameters are:9 * - `uri`: The URI of the file to convert.10 */11type ConvertBytecodeParams = Partial<{uri : string}>;12 13/**14 * The output of the mlir/convert(To|From)Bytecode commands:15 * - `output`: The output buffer of the command, e.g. a .mlir or bytecode16 * buffer.17 */18type ConvertBytecodeResult = Partial<{output : string}>;19 20/**21 * A custom filesystem that is used to convert MLIR bytecode files to text for22 * use in the editor, but still use bytecode on disk.23 */24class BytecodeFS implements vscode.FileSystemProvider {25 mlirContext: MLIRContext;26 27 constructor(mlirContext: MLIRContext) { this.mlirContext = mlirContext; }28 29 /*30 * Forward to the default filesystem for the various methods that don't need31 * to understand the bytecode <-> text translation.32 */33 readDirectory(uri: vscode.Uri): Thenable<[ string, vscode.FileType ][]> {34 return vscode.workspace.fs.readDirectory(uri);35 }36 delete(uri: vscode.Uri): void {37 vscode.workspace.fs.delete(uri.with({scheme : "file"}));38 }39 stat(uri: vscode.Uri): Thenable<vscode.FileStat> {40 return vscode.workspace.fs.stat(uri.with({scheme : "file"}));41 }42 rename(oldUri: vscode.Uri, newUri: vscode.Uri,43 options: {overwrite: boolean}): void {44 vscode.workspace.fs.rename(oldUri.with({scheme : "file"}),45 newUri.with({scheme : "file"}), options);46 }47 createDirectory(uri: vscode.Uri): void {48 vscode.workspace.fs.createDirectory(uri.with({scheme : "file"}));49 }50 watch(_uri: vscode.Uri, _options: {51 readonly recursive: boolean; readonly excludes : readonly string[]52 }): vscode.Disposable {53 return new vscode.Disposable(() => {});54 }55 56 private _emitter = new vscode.EventEmitter<vscode.FileChangeEvent[]>();57 readonly onDidChangeFile: vscode.Event<vscode.FileChangeEvent[]> =58 this._emitter.event;59 60 /*61 * Read in a bytecode file, converting it to text before returning it to the62 * caller.63 */64 async readFile(uri: vscode.Uri): Promise<Uint8Array> {65 // Try to start a language client for this file so that we can parse66 // it.67 const client =68 await this.mlirContext.getOrActivateLanguageClient(uri, 'mlir');69 if (!client) {70 throw new Error(71 'Failed to activate mlir language server to read bytecode');72 }73 74 // Ask the client to do the conversion.75 let result: ConvertBytecodeResult;76 try {77 let params: ConvertBytecodeParams = {uri : uri.toString()};78 result = await client.sendRequest('mlir/convertFromBytecode', params);79 } catch (e) {80 vscode.window.showErrorMessage(e.message);81 throw new Error(`Failed to read bytecode file: ${e}`);82 }83 let resultBuffer = new TextEncoder().encode(result.output);84 85 // NOTE: VSCode does not allow for extensions to manage files above 50mb.86 // Detect that here and if our result is too large for us to manage, alert87 // the user and open it as a new temporary .mlir file.88 if (resultBuffer.length > (50 * 1024 * 1024)) {89 const openAsTempInstead: vscode.MessageItem = {90 title : 'Open as temporary .mlir instead',91 };92 const message: string = `Failed to open bytecode file "${93 uri.toString()}". Cannot edit converted bytecode files larger than 50MB.`;94 const errorResult: vscode.MessageItem|undefined =95 await vscode.window.showErrorMessage(message, openAsTempInstead);96 if (errorResult === openAsTempInstead) {97 let tempFile = await vscode.workspace.openTextDocument({98 language : 'mlir',99 content : result.output,100 });101 await vscode.window.showTextDocument(tempFile);102 }103 throw new Error(message);104 }105 106 return resultBuffer;107 }108 109 /*110 * Save the provided content, which contains MLIR text, as bytecode.111 */112 async writeFile(uri: vscode.Uri, content: Uint8Array,113 _options: {create: boolean, overwrite: boolean}) {114 // Get the language client managing this file.115 let client = this.mlirContext.getLanguageClient(uri, 'mlir');116 if (!client) {117 throw new Error(118 'Failed to activate mlir language server to write bytecode');119 }120 121 // Ask the client to do the conversion.122 let convertParams: ConvertBytecodeParams = {123 uri : uri.toString(),124 };125 const result: ConvertBytecodeResult =126 await client.sendRequest('mlir/convertToBytecode', convertParams);127 await vscode.workspace.fs.writeFile(uri.with({scheme : "file"}),128 base64.toByteArray(result.output));129 }130}131 132/**133 * A custom bytecode document for use by the custom editor provider below.134 */135class BytecodeDocument implements vscode.CustomDocument {136 readonly uri: vscode.Uri;137 138 constructor(uri: vscode.Uri) { this.uri = uri; }139 dispose(): void {}140}141 142/**143 * A custom editor provider for MLIR bytecode that allows for non-binary144 * interpretation.145 */146class BytecodeEditorProvider implements147 vscode.CustomReadonlyEditorProvider<BytecodeDocument> {148 public async openCustomDocument(uri: vscode.Uri, _openContext: any,149 _token: vscode.CancellationToken):150 Promise<BytecodeDocument> {151 return new BytecodeDocument(uri);152 }153 154 public async resolveCustomEditor(document: BytecodeDocument,155 _webviewPanel: vscode.WebviewPanel,156 _token: vscode.CancellationToken):157 Promise<void> {158 // Ask the user for the desired view type.159 const editType = await vscode.window.showQuickPick(160 [ {label : '.mlir', description : "Edit as a .mlir text file"} ],161 {title : 'Select an editor for the bytecode.'},162 );163 164 // If we don't have a valid view type, just bail.165 if (!editType) {166 await vscode.commands.executeCommand(167 'workbench.action.closeActiveEditor');168 return;169 }170 171 // TODO: We should also provide a non-`.mlir` way of viewing the172 // bytecode, which should also ideally have some support for invalid173 // bytecode files.174 175 // Close the active editor given that we aren't using it.176 await vscode.commands.executeCommand('workbench.action.closeActiveEditor');177 178 // Display the file using a .mlir format.179 await vscode.window.showTextDocument(180 document.uri.with({scheme : "mlir.bytecode-mlir"}),181 {preview : true, preserveFocus : false});182 }183}184 185/**186 * Register the necessary providers for supporting MLIR bytecode.187 */188export function registerMLIRBytecodeExtensions(context: vscode.ExtensionContext,189 mlirContext: MLIRContext) {190 vscode.workspace.registerFileSystemProvider("mlir.bytecode-mlir",191 new BytecodeFS(mlirContext));192 vscode.window.registerCustomEditorProvider('mlir.bytecode',193 new BytecodeEditorProvider());194}195