[mlir:PDLL-LSP] Add a custom LSP command for viewing the output of PDLL

[rriddle]

Summary

This commit adds a new PDLL specific LSP command, pdll.viewOutput, that
allows for viewing the intermediate outputs of a given PDLL file. The available
intermediate forms currently mirror those in mlir-pdll, namely: AST, MLIR, CPP.
This is extremely useful for a developer of PDLL, as it simplifies various testing,
and is also quite useful for users as they can easily view what is actually being
generated for their PDLL files.

This new command is added to the vscode client, and is available in the right
client context menu of PDLL files, or via the vscode command palette.

[jpienaar]

A little bit too dev focussed on first impression, but esp as new language this is useful for transparency and could reduce debug cycles quite a lot. (We should have a visualizer for matching at some point ;-))

[rriddle]

In D124783#3544074, @jpienaar wrote:

A little bit too dev focussed on first impression, but esp as new language this is useful for transparency and could reduce debug cycles quite a lot. (We should have a visualizer for matching at some point ;-))

Yeah agreed it's mostly dev focused, but it shouldn't negatively affect normal users at all (I liken it to the "View AST" support in clangd). It's been somewhat useful for me quite a few times, and it should also help users that want to see what's actually being generated (we've seen people get frustrated with this on the tablegen side before).

mlir/lib/Tools/mlir-pdll-lsp-server/CMakeLists.txt

 llvm_add_library(MLIRPdllLspServerLib
   CompilationDatabase.cpp
   LSPServer.cpp
   PDLLServer.cpp
+  Protocol.cpp
   MlirPdllLspServerMain.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Tools/mlir-pdll-lsp-server
 
   LINK_LIBS PUBLIC
+  MLIRPDLLCodeGen
   MLIRPDLLParser
   MLIRLspServerSupportLib
   )

mlir/lib/Tools/mlir-pdll-lsp-server/CompilationDatabase.cpp

 //===- CompilationDatabase.cpp - PDLL Compilation Database ----------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "CompilationDatabase.h"
 #include "../lsp-server-support/Logging.h"
-#include "../lsp-server-support/Protocol.h"
+#include "Protocol.h"
 #include "mlir/Support/FileUtilities.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/YAMLTraits.h"
 
 using namespace mlir;
 using namespace mlir::lsp;
 
 //===----------------------------------------------------------------------===//

mlir/lib/Tools/mlir-pdll-lsp-server/LSPServer.cpp

 //===- LSPServer.cpp - PDLL Language Server -------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "LSPServer.h"
 
 #include "../lsp-server-support/Logging.h"
-#include "../lsp-server-support/Protocol.h"
 #include "../lsp-server-support/Transport.h"
 #include "PDLLServer.h"
+#include "Protocol.h"
 #include "llvm/ADT/FunctionExtras.h"
 #include "llvm/ADT/StringMap.h"
 
 #define DEBUG_TYPE "pdll-lsp-server"
 
 using namespace mlir;
 using namespace mlir::lsp;
 
struct LSPServer {
 
   //===--------------------------------------------------------------------===//
   // Signature Help
 
   void onSignatureHelp(const TextDocumentPositionParams &params,
                        Callback<SignatureHelp> reply);
 
   //===--------------------------------------------------------------------===//
+  // PDLL View Output
+
+  void onPDLLViewOutput(const PDLLViewOutputParams &params,
+                        Callback<Optional<PDLLViewOutputResult>> reply);
+
+  //===--------------------------------------------------------------------===//
   // Fields
   //===--------------------------------------------------------------------===//
 
   PDLLServer &server;
   JSONTransport &transport;
 
   /// An outgoing notification used to send diagnostics to the client when they
   /// are ready to be processed.

 // Signature Help
 
 void LSPServer::onSignatureHelp(const TextDocumentPositionParams &params,
                                 Callback<SignatureHelp> reply) {
   reply(server.getSignatureHelp(params.textDocument.uri, params.position));
 }
 
 //===----------------------------------------------------------------------===//
+// PDLL ViewOutput
+
+void LSPServer::onPDLLViewOutput(
+    const PDLLViewOutputParams &params,
+    Callback<Optional<PDLLViewOutputResult>> reply) {
+  reply(server.getPDLLViewOutput(params.uri, params.kind));
+}
+
+//===----------------------------------------------------------------------===//
 // Entry Point
 //===----------------------------------------------------------------------===//
 
 LogicalResult mlir::lsp::runPdllLSPServer(PDLLServer &server,
                                           JSONTransport &transport) {
   LSPServer lspServer(server, transport);
   MessageHandler messageHandler(transport);
 
LogicalResult mlir::lsp::runPdllLSPServer(PDLLServer &server,
   // Code Completion
   messageHandler.method("textDocument/completion", &lspServer,
                         &LSPServer::onCompletion);
 
   // Signature Help
   messageHandler.method("textDocument/signatureHelp", &lspServer,
                         &LSPServer::onSignatureHelp);
 
+  // PDLL ViewOutput
+  messageHandler.method("pdll/viewOutput", &lspServer,
+                        &LSPServer::onPDLLViewOutput);
+
   // Diagnostics
   lspServer.publishDiagnostics =
       messageHandler.outgoingNotification<PublishDiagnosticsParams>(
           "textDocument/publishDiagnostics");
 
   // Run the main loop of the transport.
   if (llvm::Error error = transport.run(messageHandler)) {
     Logger::error("Transport error: {0}", error);
     llvm::consumeError(std::move(error));
     return failure();
   }
   return success(lspServer.shutdownRequestReceived);
 }

mlir/lib/Tools/mlir-pdll-lsp-server/PDLLServer.h

 #include "llvm/ADT/StringRef.h"
 #include <memory>
 #include <string>
 
 namespace mlir {
 namespace lsp {
 struct Diagnostic;
 class CompilationDatabase;
+struct PDLLViewOutputResult;
+enum class PDLLViewOutputKind;
 struct CompletionList;
 struct DocumentLink;
 struct DocumentSymbol;
 struct Hover;
 struct Location;
 struct Position;
 struct SignatureHelp;
 class URIForFile;
public:
   /// Get the code completion list for the position within the given file.
   CompletionList getCodeCompletion(const URIForFile &uri,
                                    const Position &completePos);
 
   /// Get the signature help for the position within the given file.
   SignatureHelp getSignatureHelp(const URIForFile &uri,
                                  const Position &helpPos);
 
+  /// Get the output of the given PDLL file, or None if there is no valid
+  /// output.
+  Optional<PDLLViewOutputResult> getPDLLViewOutput(const URIForFile &uri,
+                                                   PDLLViewOutputKind kind);
+
 private:
   struct Impl;
   std::unique_ptr<Impl> impl;
 };
 
 } // namespace lsp
 } // namespace mlir
 
 #endif // LIB_MLIR_TOOLS_MLIRPDLLSPSERVER_SERVER_H_

mlir/lib/Tools/mlir-pdll-lsp-server/PDLLServer.cpp

 //===- PDLLServer.cpp - PDLL Language Server ------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "PDLLServer.h"
 
 #include "../lsp-server-support/Logging.h"
-#include "../lsp-server-support/Protocol.h"
 #include "CompilationDatabase.h"
+#include "Protocol.h"
+#include "mlir/IR/BuiltinOps.h"
 #include "mlir/Tools/PDLL/AST/Context.h"
 #include "mlir/Tools/PDLL/AST/Nodes.h"
 #include "mlir/Tools/PDLL/AST/Types.h"
+#include "mlir/Tools/PDLL/CodeGen/CPPGen.h"
+#include "mlir/Tools/PDLL/CodeGen/MLIRGen.h"
 #include "mlir/Tools/PDLL/ODS/Constraint.h"
 #include "mlir/Tools/PDLL/ODS/Context.h"
 #include "mlir/Tools/PDLL/ODS/Dialect.h"
 #include "mlir/Tools/PDLL/ODS/Operation.h"
 #include "mlir/Tools/PDLL/Parser/CodeComplete.h"
 #include "mlir/Tools/PDLL/Parser/Parser.h"
 #include "llvm/ADT/IntervalMap.h"
 #include "llvm/ADT/StringMap.h"
struct PDLDocument {
   //===--------------------------------------------------------------------===//
   // Signature Help
   //===--------------------------------------------------------------------===//
 
   lsp::SignatureHelp getSignatureHelp(const lsp::URIForFile &uri,
                                       const lsp::Position &helpPos);
 
   //===--------------------------------------------------------------------===//
+  // PDLL ViewOutput
+  //===--------------------------------------------------------------------===//
+
+  void getPDLLViewOutput(raw_ostream &os, lsp::PDLLViewOutputKind kind);
+
+  //===--------------------------------------------------------------------===//
   // Fields
   //===--------------------------------------------------------------------===//
 
   /// The include directories for this file.
   std::vector<std::string> includeDirs;
 
   /// The source manager containing the contents of the input file.
   llvm::SourceMgr sourceMgr;
lsp::SignatureHelp PDLDocument::getSignatureHelp(const lsp::URIForFile &uri,
 
   ast::Context tmpContext(tmpODSContext);
   (void)parsePDLAST(tmpContext, sourceMgr, &completeContext);
 
   return signatureHelp;
 }
 
 //===----------------------------------------------------------------------===//
+// PDLL ViewOutput
+//===----------------------------------------------------------------------===//
+
+void PDLDocument::getPDLLViewOutput(raw_ostream &os,
+                                    lsp::PDLLViewOutputKind kind) {
+  if (failed(astModule))
+    return;
+  if (kind == lsp::PDLLViewOutputKind::AST) {
+    (*astModule)->print(os);
+    return;
+  }
+
+  // Generate the MLIR for the ast module. We also capture diagnostics here to
+  // show to the user, which may be useful if PDLL isn't capturing constraints
+  // expected by PDL.
+  MLIRContext mlirContext;
+  SourceMgrDiagnosticHandler diagHandler(sourceMgr, &mlirContext, os);
+  OwningOpRef<ModuleOp> pdlModule =
+      codegenPDLLToMLIR(&mlirContext, astContext, sourceMgr, **astModule);
+  if (!pdlModule)
+    return;
+  if (kind == lsp::PDLLViewOutputKind::MLIR) {
+    pdlModule->print(os, OpPrintingFlags().enableDebugInfo());
+    return;
+  }
+
+  // Otherwise, generate the output for C++.
+  assert(kind == lsp::PDLLViewOutputKind::CPP &&
+         "unexpected PDLLViewOutputKind");
+  codegenPDLLToCPP(**astModule, *pdlModule, os);
+}
+
+//===----------------------------------------------------------------------===//
 // PDLTextFileChunk
 //===----------------------------------------------------------------------===//
 
 namespace {
 /// This class represents a single chunk of an PDL text file.
 struct PDLTextFileChunk {
   PDLTextFileChunk(uint64_t lineOffset, const lsp::URIForFile &uri,
                    StringRef contents,
void getDocumentLinks(const lsp::URIForFile &uri,
                         std::vector<lsp::DocumentLink> &links);
   Optional<lsp::Hover> findHover(const lsp::URIForFile &uri,
                                  lsp::Position hoverPos);
   void findDocumentSymbols(std::vector<lsp::DocumentSymbol> &symbols);
   lsp::CompletionList getCodeCompletion(const lsp::URIForFile &uri,
                                         lsp::Position completePos);
   lsp::SignatureHelp getSignatureHelp(const lsp::URIForFile &uri,
                                       lsp::Position helpPos);
+  lsp::PDLLViewOutputResult getPDLLViewOutput(lsp::PDLLViewOutputKind kind);
 
 private:
   /// Find the PDL document that contains the given position, and update the
   /// position to be anchored at the start of the found chunk instead of the
   /// beginning of the file.
   PDLTextFileChunk &getChunkFor(lsp::Position &pos);
 
   /// The full string contents of the file.
lsp::CompletionList PDLTextFile::getCodeCompletion(const lsp::URIForFile &uri,
   return completionList;
 }
 
 lsp::SignatureHelp PDLTextFile::getSignatureHelp(const lsp::URIForFile &uri,
                                                  lsp::Position helpPos) {
   return getChunkFor(helpPos).document.getSignatureHelp(uri, helpPos);
 }
 
+lsp::PDLLViewOutputResult
+PDLTextFile::getPDLLViewOutput(lsp::PDLLViewOutputKind kind) {
+  lsp::PDLLViewOutputResult result;
+  {
+    llvm::raw_string_ostream outputOS(result.output);
+    llvm::interleave(
+        llvm::make_pointee_range(chunks),
+        [&](PDLTextFileChunk &chunk) {
+          chunk.document.getPDLLViewOutput(outputOS, kind);
+        },
+        [&] { outputOS << "\n// -----\n\n"; });
+  }
+  return result;
+}
+
 PDLTextFileChunk &PDLTextFile::getChunkFor(lsp::Position &pos) {
   if (chunks.size() == 1)
     return *chunks.front();
 
   // Search for the first chunk with a greater line offset, the previous chunk
   // is the one that contains `pos`.
   auto it = llvm::upper_bound(
       chunks, pos, [](const lsp::Position &pos, const auto &chunk) {

 
 lsp::SignatureHelp lsp::PDLLServer::getSignatureHelp(const URIForFile &uri,
                                                      const Position &helpPos) {
   auto fileIt = impl->files.find(uri.file());
   if (fileIt != impl->files.end())
     return fileIt->second->getSignatureHelp(uri, helpPos);
   return SignatureHelp();
 }
+
+Optional<lsp::PDLLViewOutputResult>
+lsp::PDLLServer::getPDLLViewOutput(const URIForFile &uri,
+                                   PDLLViewOutputKind kind) {
+  auto fileIt = impl->files.find(uri.file());
+  if (fileIt != impl->files.end())
+    return fileIt->second->getPDLLViewOutput(kind);
+  return llvm::None;
+}

mlir/lib/Tools/mlir-pdll-lsp-server/Protocol.h

+//===--- Protocol.h - Language Server Protocol Implementation ---*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file contains structs for LSP commands that are specific to the PDLL
+// server.
+//
+// Each struct has a toJSON and fromJSON function, that converts between
+// the struct and a JSON representation. (See JSON.h)
+//
+// Some structs also have operator<< serialization. This is for debugging and
+// tests, and is not generally machine-readable.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LIB_MLIR_TOOLS_MLIRPDLLLSPSERVER_PROTOCOL_H_
+#define LIB_MLIR_TOOLS_MLIRPDLLLSPSERVER_PROTOCOL_H_
+
+#include "../lsp-server-support/Protocol.h"
+
+namespace mlir {
+namespace lsp {
+//===----------------------------------------------------------------------===//
+// PDLLViewOutputParams
+//===----------------------------------------------------------------------===//
+
+/// The type of output to view from PDLL.
+enum class PDLLViewOutputKind {
+  AST,
+  MLIR,
+  CPP,
+};
+
+/// Represents the parameters used when viewing the output of a PDLL file.
+struct PDLLViewOutputParams {
+  /// The URI of the document to view the output of.
+  URIForFile uri;
+
+  /// The kind of output to generate.
+  PDLLViewOutputKind kind;
+};
+
+/// Add support for JSON serialization.
+bool fromJSON(const llvm::json::Value &value, PDLLViewOutputKind &result,
+              llvm::json::Path path);
+bool fromJSON(const llvm::json::Value &value, PDLLViewOutputParams &result,
+              llvm::json::Path path);
+
+//===----------------------------------------------------------------------===//
+// PDLLViewOutputResult
+//===----------------------------------------------------------------------===//
+
+/// Represents the result of viewing the output of a PDLL file.
+struct PDLLViewOutputResult {
+  /// The string representation of the output.
+  std::string output;
+};
+
+/// Add support for JSON serialization.
+llvm::json::Value toJSON(const PDLLViewOutputResult &value);
+
+} // namespace lsp
+} // namespace mlir
+
+#endif

mlir/lib/Tools/mlir-pdll-lsp-server/Protocol.cpp

+//===--- Protocol.cpp - Language Server Protocol Implementation -----------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file contains the serialization code for the PDLL specific LSP structs.
+//
+//===----------------------------------------------------------------------===//
+
+#include "Protocol.h"
+#include "llvm/ADT/Hashing.h"
+#include "llvm/ADT/SmallString.h"
+#include "llvm/ADT/StringSwitch.h"
+#include "llvm/Support/ErrorHandling.h"
+#include "llvm/Support/Format.h"
+#include "llvm/Support/FormatVariadic.h"
+#include "llvm/Support/JSON.h"
+#include "llvm/Support/Path.h"
+#include "llvm/Support/raw_ostream.h"
+
+using namespace mlir;
+using namespace mlir::lsp;
+
+// Helper that doesn't treat `null` and absent fields as failures.
+template <typename T>
+static bool mapOptOrNull(const llvm::json::Value &params,
+                         llvm::StringLiteral prop, T &out,
+                         llvm::json::Path path) {
+  const llvm::json::Object *o = params.getAsObject();
+  assert(o);
+
+  // Field is missing or null.
+  auto *v = o->get(prop);
+  if (!v || v->getAsNull().hasValue())
+    return true;
+  return fromJSON(*v, out, path.field(prop));
+}
+
+//===----------------------------------------------------------------------===//
+// PDLLViewOutputParams
+//===----------------------------------------------------------------------===//
+
+bool mlir::lsp::fromJSON(const llvm::json::Value &value,
+                         PDLLViewOutputKind &result, llvm::json::Path path) {
+  if (Optional<StringRef> str = value.getAsString()) {
+    if (*str == "ast") {
+      result = PDLLViewOutputKind::AST;
+      return true;
+    }
+    if (*str == "mlir") {
+      result = PDLLViewOutputKind::MLIR;
+      return true;
+    }
+    if (*str == "cpp") {
+      result = PDLLViewOutputKind::CPP;
+      return true;
+    }
+  }
+  return false;
+}
+
+bool mlir::lsp::fromJSON(const llvm::json::Value &value,
+                         PDLLViewOutputParams &result, llvm::json::Path path) {
+  llvm::json::ObjectMapper o(value, path);
+  return o && o.map("uri", result.uri) && o.map("kind", result.kind);
+}
+
+//===----------------------------------------------------------------------===//
+// PDLLViewOutputResult
+//===----------------------------------------------------------------------===//
+
+llvm::json::Value mlir::lsp::toJSON(const PDLLViewOutputResult &value) {
+  return llvm::json::Object{{"output", value.output}};
+}

mlir/test/mlir-pdll-lsp-server/view-output.test

+// RUN: mlir-pdll-lsp-server -lit-test < %s | FileCheck -strict-whitespace %s
+{"jsonrpc":"2.0","id":0,"method":"initialize","params":{"processId":123,"rootPath":"pdll","capabilities":{},"trace":"off"}}
+// -----
+{"jsonrpc":"2.0","method":"textDocument/didOpen","params":{"textDocument":{
+  "uri":"test:///foo.pdll",
+  "languageId":"pdll",
+  "version":1,
+  "text":"Pattern TestPat => erase op<test.op>;"
+}}}
+// -----
+{"jsonrpc":"2.0","id":1,"method":"pdll/viewOutput","params":{
+  "uri":"test:///foo.pdll",
+  "kind":"ast"
+}}
+//      CHECK:  "id": 1
+// CHECK-NEXT:  "jsonrpc": "2.0",
+// CHECK-NEXT:  "result": {
+// CHECK-NEXT:    "output": "-Module{{.*}}PatternDecl{{.*}}Name<TestPat>{{.*}}\n"
+// CHECK-NEXT:  }
+// -----
+{"jsonrpc":"2.0","id":2,"method":"pdll/viewOutput","params":{
+  "uri":"test:///foo.pdll",
+  "kind":"mlir"
+}}
+//      CHECK:  "id": 2
+// CHECK-NEXT:  "jsonrpc": "2.0",
+// CHECK-NEXT:  "result": {
+// CHECK-NEXT:    "output": "module {\n  pdl.pattern @TestPat {{.*}}\n"
+// CHECK-NEXT:  }
+// -----
+{"jsonrpc":"2.0","id":3,"method":"pdll/viewOutput","params":{
+  "uri":"test:///foo.pdll",
+  "kind":"cpp"
+}}
+//      CHECK:  "id": 3
+// CHECK-NEXT:  "jsonrpc": "2.0",
+// CHECK-NEXT:  "result": {
+// CHECK-NEXT:    "output": "{{.*}}struct TestPat : ::mlir::PDLPatternModule{{.*}}\n"
+// CHECK-NEXT:  }
+// -----
+{"jsonrpc":"2.0","id":3,"method":"shutdown"}
+// -----
+{"jsonrpc":"2.0","method":"exit"}

mlir/utils/vscode/package.json

"configuration": {
           ]
         }
       }
     },
     "commands": [
       {
         "command": "mlir.restart",
         "title": "mlir: Restart language server"
+      },
+      {
+        "command": "mlir.viewPDLLOutput",
+        "title": "mlir-pdll: View PDLL output"
+      }
+    ],
+    "menus": {
+      "editor/context": [
+        {
+          "command": "mlir.viewPDLLOutput",
+          "group": "z_commands",
+          "when": "editorLangId == pdll"
         }
       ]
     }
   }
+}

mlir/utils/vscode/src/PDLL/commands/viewOutput.ts

+import * as vscode from 'vscode'
+
+import {Command} from '../../command';
+import {MLIRContext} from '../../mlirContext';
+
+/**
+ * The parameters to the pdll/viewOutput command. These parameters are:
+ * - `uri`: The URI of the file to view.
+ * - `kind`: The kind of the output to generate.
+ */
+type ViewOutputParams = Partial<{uri : string, kind : string;}>;
+
+/**
+ * The output of the commands:
+ * - `output`: The output string of the command, e.g. a .mlir PDL string.
+ */
+type ViewOutputResult = Partial<{output : string}>;
+
+/**
+ * A command that displays the output of the current PDLL document.
+ */
+export class ViewPDLLCommand extends Command {
+  constructor(context: MLIRContext) { super('mlir.viewPDLLOutput', context); }
+
+  async execute() {
+    const editor = vscode.window.activeTextEditor;
+    if (editor.document.languageId != 'pdll')
+      return;
+
+    // Check to see if a language client is active for this document.
+    const workspaceFolder =
+        vscode.workspace.getWorkspaceFolder(editor.document.uri);
+    const pdllClient = this.context.getLanguageClient(workspaceFolder, "pdll");
+    if (!pdllClient) {
+      return;
+    }
+
+    // Ask the user for the desired output type.
+    const outputType =
+        await vscode.window.showQuickPick([ 'ast', 'mlir', 'cpp' ]);
+    if (!outputType) {
+      return;
+    }
+
+    // If we have the language client, ask it to try compiling the document.
+    let outputParams: ViewOutputParams = {
+      uri : editor.document.uri.toString(),
+      kind : outputType,
+    };
+    const result: ViewOutputResult|undefined =
+        await pdllClient.sendRequest('pdll/viewOutput', outputParams);
+    if (!result || result.output.length === 0) {
+      return;
+    }
+
+    // Display the output in a new editor.
+    let outputFileType = 'plaintext';
+    if (outputType == 'mlir') {
+      outputFileType = 'mlir';
+    } else if (outputType == 'cpp') {
+      outputFileType = 'cpp';
+    }
+    await vscode.workspace.openTextDocument(
+        {language : outputFileType, content : result.output});
+  }
+}

mlir/utils/vscode/src/PDLL/pdll.ts

+import * as vscode from 'vscode';
+
+import {MLIRContext} from '../mlirContext';
+import {ViewPDLLCommand} from './commands/viewOutput';
+
+/**
+ *  Register the necessary context and commands for PDLL.
+ */
+export function registerPDLLCommands(context: vscode.ExtensionContext,
+                                     mlirContext: MLIRContext) {
+  context.subscriptions.push(new ViewPDLLCommand(mlirContext));
+}

mlir/utils/vscode/src/command.ts

+import * as vscode from 'vscode';
+import {MLIRContext} from './mlirContext';
+
+/**
+ * This class represents a base vscode command. It handles all of the necessary
+ * command registration and disposal boilerplate.
+ */
+export abstract class Command extends vscode.Disposable {
+  private disposable: vscode.Disposable;
+  protected context: MLIRContext;
+
+  constructor(command: string, context: MLIRContext) {
+    super(() => this.dispose());
+    this.disposable =
+        vscode.commands.registerCommand(command, this.execute, this);
+    this.context = context;
+  }
+
+  dispose() { this.disposable && this.disposable.dispose(); }
+
+  /**
+   * The function executed when this command is invoked.
+   */
+  abstract execute(...args: any[]): any;
+}

mlir/utils/vscode/src/extension.ts

 import * as vscode from 'vscode';
 
 import {MLIRContext} from './mlirContext';
+import {registerPDLLCommands} from './PDLL/pdll';
 
 /**
  *  This method is called when the extension is activated. The extension is
  *  activated the very first time a command is executed.
  */
 export function activate(context: vscode.ExtensionContext) {
   const outputChannel = vscode.window.createOutputChannel('MLIR');
   context.subscriptions.push(outputChannel);
 
   const mlirContext = new MLIRContext();
   context.subscriptions.push(mlirContext);
 
   // Initialize the commands of the extension.
   context.subscriptions.push(
       vscode.commands.registerCommand('mlir.restart', async () => {
         // Dispose and reactivate the context.
         mlirContext.dispose();
         await mlirContext.activate(outputChannel);
       }));
+  registerPDLLCommands(context, mlirContext);
 
   mlirContext.activate(outputChannel);
 }

mlir/utils/vscode/src/mlirContext.ts

export class MLIRContext implements vscode.Disposable {
   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);
   }
 
+  /**
+   * Return the language client for the given language and workspace folder, or
+   * null if no client is active.
+   */
+  getLanguageClient(workspaceFolder: vscode.WorkspaceFolder,
+                    languageName: string): vscodelc.LanguageClient {
+    let workspaceFolderStr =
+        workspaceFolder ? workspaceFolder.uri.toString() : "";
+    let folderContext = this.workspaceFolders.get(workspaceFolderStr);
+    if (!folderContext) {
+      return null;
+    }
+    return folderContext.clients.get(languageName);
+  }
+
   dispose() {
     this.subscriptions.forEach((d) => { d.dispose(); });
     this.subscriptions = [];
     this.workspaceFolders.forEach((d) => { d.dispose(); });
     this.workspaceFolders.clear();
   }
 }