Add APIs to the JS bindings to generate MLIR

Author: robertoraggi

packages/cxx-frontend/src/Parser.ts

@@ -51,6 +51,7 @@ interface ParserParams {
 export class Parser {
   #unit: Unit | undefined;
   #ast: AST | undefined;
+  #pendingAST: Promise<AST> | undefined;
 
   static async init({
     wasm,
@@ -97,24 +98,54 @@ export class Parser {
     this.#unit = cxx.createUnit(source, path, { resolve, readFile });
   }
 
-  async parse() {
+  async parse(): Promise<AST> {
+    return await this.getASTAsync();
+  }
+
+  async emitIR(): Promise<string> {
+    const _ = await this.getASTAsync();
+    return this.#unit?.emitIR() ?? "";
+  }
+
+  async #parseHelper(): Promise<AST> {
+    if (this.#pendingAST) {
+      return await this.#pendingAST;
+    }
+
     if (!this.#unit) {
-      return;
+      throw new Error("Parser has been disposed");
     }
+
     await this.#unit.parse();
-    this.#ast = AST.from(this.#unit.getHandle(), this);
+
+    const ast = AST.from(this.#unit.getHandle(), this);
+
+    if (!ast) {
+      throw new Error("Failed to create AST");
+    }
+
+    this.#ast = ast;
+
+    return ast;
   }
 
   dispose() {
     this.#unit?.delete();
     this.#unit = undefined;
     this.#ast = undefined;
+    this.#pendingAST = undefined;
   }
 
   getUnitHandle(): number {
     return this.#unit?.getUnitHandle() ?? 0;
   }
 
+  async getASTAsync(): Promise<AST> {
+    this.#pendingAST ??= this.#parseHelper();
+    return await this.#pendingAST;
+  }
+
+  // internal
   getAST(): AST | undefined {
     return this.#ast;
   }

packages/cxx-frontend/src/Unit.ts

@@ -26,4 +26,5 @@ export interface Unit {
   getHandle(): number;
   getUnitHandle(): number;
   getDiagnostics(): Diagnostic[];
+  emitIR(): string;
 }

packages/cxx-frontend/src/cxx-js.d.ts

@@ -101,6 +101,7 @@ export type CXX = {
   TranslationUnit: TranslationUnit;
 
   createUnit(source: string, path: string, api: Api): Unit;
+  emitIR(handle: number): string;
   getASTKind(handle: number): number;
   getASTSlot(handle: number, slot: number): number;
   getASTSlotKind(handle: number, slot: number): ASTSlotKind;

src/frontend/CMakeLists.txt

@@ -26,7 +26,6 @@ add_executable(cxx
 target_link_libraries(cxx PRIVATE cxx-lsp)
 
 if (CXX_ENABLE_MLIR)
-    target_compile_definitions(cxx PRIVATE CXX_WITH_MLIR)
     target_link_libraries(cxx PRIVATE cxx-mlir)
 endif()
 

src/js/CMakeLists.txt

@@ -27,6 +27,10 @@ add_executable(cxx-js ${SOURCES})
 
 target_link_libraries(cxx-js cxx-parser embind)
 
+if (CXX_ENABLE_MLIR)
+    target_link_libraries(cxx-js cxx-mlir)
+endif()
+
 target_link_options(cxx-js PUBLIC
     "SHELL:-s ENVIRONMENT=web"
     "SHELL:-s MINIMAL_RUNTIME"

src/js/cxx/api.cc

@@ -37,6 +37,13 @@
 #include <optional>
 #include <sstream>
 
+#ifdef CXX_WITH_MLIR
+#include <cxx/mlir/codegen.h>
+#include <cxx/mlir/cxx_dialect.h>
+#include <llvm/Support/raw_os_ostream.h>
+#include <mlir/IR/MLIRContext.h>
+#endif
+
 using namespace emscripten;
 
 namespace {
@@ -179,6 +186,31 @@ struct WrappedUnit {
 
     co_return val{true};
   }
+
+  auto emitIR() -> std::string {
+#ifdef CXX_WITH_MLIR
+    mlir::MLIRContext context;
+    context.loadDialect<mlir::cxx::CxxDialect>();
+
+    cxx::Codegen codegen(context, unit.get());
+
+    auto ir = codegen(unit->ast());
+
+    mlir::OpPrintingFlags flags;
+    flags.enableDebugInfo(true, true);
+
+    std::ostringstream out;
+    llvm::raw_os_ostream os(out);
+    ir.module->print(os, flags);
+    os.flush();
+
+    auto code = out.str();
+
+    return code;
+#else
+    return {};
+#endif
+  }
 };
 
 auto getTokenText(std::intptr_t handle, std::intptr_t unitHandle)
@@ -381,7 +413,8 @@ EMSCRIPTEN_BINDINGS(cxx) {
       .function("parse", &WrappedUnit::parse)
       .function("getHandle", &WrappedUnit::getHandle)
       .function("getUnitHandle", &WrappedUnit::getUnitHandle)
-      .function("getDiagnostics", &WrappedUnit::getDiagnostics);
+      .function("getDiagnostics", &WrappedUnit::getDiagnostics)
+      .function("emitIR", &WrappedUnit::emitIR);
 
   function("createUnit", &createUnit, allow_raw_pointers());
   function("getASTKind", &getASTKind);

src/mlir/cxx/mlir/CMakeLists.txt

@@ -20,6 +20,9 @@ target_link_libraries(cxx-mlir PUBLIC
     MLIRControlFlowDialect
     MLIRSCFDialect
 )
+
+target_compile_definitions(cxx-mlir PUBLIC CXX_WITH_MLIR)
+
 add_dependencies(cxx-mlir MLIRCxxOpsIncGen)
 
 if (CMAKE_CXX_COMPILER_ID MATCHES "Clang|GNU")

templates/cxx-browser-esm-vite/src/main.ts

@@ -44,12 +44,10 @@ async function main() {
     source,
   });
 
-  parser.parse();
+  const ast = await parser.parse();
 
   const rows: string[] = [];
 
-  const ast = parser.getAST();
-
   for (const { node, depth } of ast?.walk().preVisit() ?? []) {
     if (!(node instanceof AST)) {
       continue;

templates/cxx-parse-esm/index.js

@@ -51,18 +51,16 @@ async function main() {
 
   const parser = new Parser({ source, path: "source.cc" });
 
-  await parser.parse();
+  const ast = await parser.parse();
 
   const diagnostics = parser.getDiagnostics();
 
   if (diagnostics.length > 0) {
     console.log("diagnostics", diagnostics);
   }
 
-  const ast = parser.getAST();
-
   for (const { node, slot, depth } of ast?.walk().preVisit() ?? []) {
-    if (!node instanceof AST) continue;
+    if (!(node instanceof AST)) continue;
     const ind = " ".repeat(depth * 2);
     const kind = ASTKind[node.getKind()];
     const member = slot !== undefined ? `${ASTSlot[slot]}: ` : "";

templates/cxx-parse/index.js

@@ -59,7 +59,7 @@ async function main() {
   const ast = parser.getAST();
 
   for (const { node, slot, depth } of ast?.walk().preVisit() ?? []) {
-    if (!node instanceof AST) continue;
+    if (!(node instanceof AST)) continue;
     const ind = " ".repeat(depth * 2);
     const kind = ASTKind[node.getKind()];
     const member = slot !== undefined ? `${ASTSlot[slot]}: ` : "";