Update plugin system

codon/app/main.cpp

@@ -108,11 +108,6 @@ ProcessResult processSource(const std::vector<const char *> &args) {
     defmap.emplace(name, value);
   }
 
-  auto *module = codon::parse(args[0], input.c_str(), /*code=*/"", /*isCode=*/false,
-                              /*isTest=*/false, /*startLine=*/0, defmap);
-  if (!module)
-    return {{}, {}};
-
   const bool isDebug = (optMode == OptMode::Debug);
   auto t = std::chrono::high_resolution_clock::now();
 
@@ -147,6 +142,12 @@ ProcessResult processSource(const std::vector<const char *> &args) {
     }
   }
   t = std::chrono::high_resolution_clock::now();
+
+  auto *module = codon::parse(args[0], input.c_str(), /*code=*/"", /*isCode=*/false,
+                              /*isTest=*/false, /*startLine=*/0, defmap, plm.get());
+  if (!module)
+    return {{}, {}};
+
   pm->run(module);
   LOG_TIME("[T] ir-opt = {:.1f}", std::chrono::duration_cast<std::chrono::milliseconds>(
                                       std::chrono::high_resolution_clock::now() - t)

codon/dsl/dsl.h

@@ -1,5 +1,6 @@
 #pragma once
 
+#include "codon/parser/cache.h"
 #include "codon/sir/sir.h"
 #include "codon/sir/transform/manager.h"
 #include "codon/sir/transform/pass.h"
@@ -15,26 +16,19 @@ namespace codon {
 /// the DSL, like keywords and IR passes.
 class DSL {
 public:
-  /// Represents a keyword, consisting of the keyword
-  /// name (e.g. "if", "while", etc.) and a callback
-  /// to generate the resulting IR node.
-  template <typename Callback> struct Keyword {
-    /// keyword name
-    std::string name;
-    /// callback to produce IR node
-    Callback callback;
+  using KeywordCallback =
+      std::function<ast::StmtPtr(ast::SimplifyVisitor *, ast::CustomStmt *)>;
+
+  struct ExprKeyword {
+    std::string keyword;
+    KeywordCallback callback;
   };
 
-  using ExprKeywordCallback =
-      std::function<ir::Node *(ir::Module *M, const std::vector<ir::Value *> &values)>;
-  using BlockKeywordCallback = std::function<ir::Node *(
-      ir::Module *M, const std::vector<ir::Value *> &values, ir::SeriesFlow *block)>;
-  using BinaryKeywordCallback =
-      std::function<ir::Node *(ir::Module *M, ir::Value *lhs, ir::Value *rhs)>;
-
-  using ExprKeyword = Keyword<ExprKeywordCallback>;
-  using BlockKeyword = Keyword<BlockKeywordCallback>;
-  using BinaryKeyword = Keyword<BinaryKeywordCallback>;
+  struct BlockKeyword {
+    std::string keyword;
+    KeywordCallback callback;
+    bool hasExpr;
+  };
 
   virtual ~DSL() noexcept = default;
 
@@ -60,19 +54,14 @@ public:
   virtual void addLLVMPasses(llvm::PassManagerBuilder *pmb, bool debug) {}
 
   /// Returns a vector of "expression keywords", defined as keywords of
-  /// the form "keyword <expr1> ... <exprN>".
+  /// the form "keyword <expr>".
   /// @return this DSL's expression keywords
   virtual std::vector<ExprKeyword> getExprKeywords() { return {}; }
 
   /// Returns a vector of "block keywords", defined as keywords of the
-  /// form "keyword <expr1> ... <exprN>: <block of code>".
+  /// form "keyword <expr>: <block of code>".
   /// @return this DSL's block keywords
   virtual std::vector<BlockKeyword> getBlockKeywords() { return {}; }
-
-  /// Returns a vector of "binary keywords", defined as keywords of the
-  /// form "<expr1> keyword <expr2>".
-  /// @return this DSL's binary keywords
-  virtual std::vector<BinaryKeyword> getBinaryKeywords() { return {}; }
 };
 
 } // namespace codon

codon/dsl/plugins.cpp

@@ -1,21 +1,17 @@
-#include "codon/dsl/plugins.h"
+#include "plugins.h"
+
 #include "codon/util/common.h"
-#include <dlfcn.h>
 
 namespace codon {
 
-PluginManager::~PluginManager() {
-  for (auto &plugin : plugins) {
-    dlclose(plugin->handle);
-  }
-}
-
 PluginManager::Error PluginManager::load(const std::string &path) {
-  void *handle = dlopen(path.c_str(), RTLD_LAZY);
-  if (!handle)
+  std::string errMsg;
+  auto handle = llvm::sys::DynamicLibrary::getPermanentLibrary(path.c_str(), &errMsg);
+
+  if (!handle.isValid())
     return Error::NOT_FOUND;
 
-  auto *entry = (LoadFunc *)dlsym(handle, "load");
+  auto *entry = (LoadFunc *)handle.getAddressOfSymbol("load");
   if (!entry)
     return Error::NO_ENTRYPOINT;
 
@@ -28,10 +24,7 @@ PluginManager::Error PluginManager::load(DSL *dsl) {
   if (!dsl || !dsl->isVersionSupported(CODON_VERSION_MAJOR, CODON_VERSION_MINOR,
                                        CODON_VERSION_PATCH))
     return Error::UNSUPPORTED_VERSION;
-
   dsl->addIRPasses(pm, debug);
-  // TODO: register new keywords
-
   return Error::NONE;
 }
 

codon/dsl/plugins.h

@@ -1,12 +1,14 @@
 #pragma once
 
-#include "codon/dsl/dsl.h"
-#include "codon/sir/util/iterators.h"
 #include <functional>
 #include <memory>
 #include <string>
 #include <vector>
 
+#include "codon/dsl/dsl.h"
+#include "codon/sir/util/iterators.h"
+#include "llvm/Support/DynamicLibrary.h"
+
 namespace codon {
 
 /// Plugin metadata
@@ -15,10 +17,11 @@ struct Plugin {
   std::unique_ptr<DSL> dsl;
   /// plugin load path
   std::string path;
-  /// plugin dlopen handle
-  void *handle;
+  /// library handle
+  llvm::sys::DynamicLibrary handle;
 
-  Plugin(std::unique_ptr<DSL> dsl, const std::string &path, void *handle)
+  Plugin(std::unique_ptr<DSL> dsl, const std::string &path,
+         const llvm::sys::DynamicLibrary &handle)
       : dsl(std::move(dsl)), path(path), handle(handle) {}
 };
 
@@ -44,8 +47,6 @@ public:
   explicit PluginManager(ir::transform::PassManager *pm, bool debug = false)
       : pm(pm), plugins(), debug(debug) {}
 
-  ~PluginManager();
-
   /// @return iterator to the first plugin
   auto begin() { return ir::util::raw_ptr_adaptor(plugins.begin()); }
   /// @return iterator beyond the last plugin

codon/parser/parser.cpp

@@ -27,7 +27,8 @@ namespace codon {
 
 ir::Module *parse(const std::string &argv0, const std::string &file,
                   const std::string &code, bool isCode, int isTest, int startLine,
-                  const std::unordered_map<std::string, std::string> &defines) {
+                  const std::unordered_map<std::string, std::string> &defines,
+                  PluginManager *plm) {
   try {
     auto d = getenv("CODON_DEBUG");
     if (d) {
@@ -46,6 +47,17 @@ ir::Module *parse(const std::string &argv0, const std::string &file,
       realpath(file.c_str(), abs);
 
     auto cache = std::make_shared<ast::Cache>(argv0);
+    if (plm) {
+      for (auto *plugin : *plm) {
+        for (auto &kw : plugin->dsl->getExprKeywords()) {
+          cache->customExprStmts[kw.keyword] = kw.callback;
+        }
+        for (auto &kw : plugin->dsl->getBlockKeywords()) {
+          cache->customBlockStmts[kw.keyword] = {kw.hasExpr, kw.callback};
+        }
+      }
+    }
+
     ast::StmtPtr codeStmt = isCode ? ast::parseCode(cache, abs, code, startLine)
                                    : ast::parseFile(cache, abs);
     if (_dbg_level) {

codon/parser/parser.h

@@ -4,6 +4,7 @@
 #include <unordered_map>
 #include <vector>
 
+#include "codon/dsl/plugins.h"
 #include "codon/sir/sir.h"
 #include "codon/util/common.h"
 
@@ -13,7 +14,8 @@ codon::ir::Module *parse(const std::string &argv0, const std::string &file,
                          const std::string &code = "", bool isCode = false,
                          int isTest = 0, int startLine = 0,
                          const std::unordered_map<std::string, std::string> &defines =
-                             std::unordered_map<std::string, std::string>{});
+                             std::unordered_map<std::string, std::string>{},
+                         PluginManager *plm = nullptr);
 
 void generateDocstr(const std::string &argv0);