Change the current namespace lookup behavior

* Change the way namespace creation works. We should create new
  namespaces via the context so we can insert them in the contenxt

* Change the way setting the current namespace works. `withCurrentNS`
  method gets a function as input param and a ns name and runs the
  function with the ns set as the current NS and restore the state
  after that

* Change the way Dylibs register with the context and namespaces
  and lookup the dylibs will use the context to use the most recent
  dylib of a namespace.

include/serene/context.h

@@ -42,8 +42,10 @@
 #include <memory>
 
 #define DEFAULT_NS_NAME "serene.user"
+#define INTERNAL_NS     "serene.internal"
 
 namespace serene {
+class SereneContext;
 
 namespace reader {
 class LocationRange;
@@ -66,6 +68,8 @@ enum class CompilationPhase {
   O2,
   O3,
 };
+/// Terminates the serene compiler process in a thread safe manner
+SERENE_EXPORT void terminate(SereneContext &ctx, int exitCode);
 
 class SERENE_EXPORT SereneContext {
   struct Options {
@@ -79,6 +83,9 @@ class SERENE_EXPORT SereneContext {
   };
 
 public:
+  template <typename T>
+  using CurrentNSFn = std::function<T()>;
+
   // --------------------------------------------------------------------------
   // IMPORTANT:
   // These two contextes have to be the very first members of the class in
@@ -113,22 +120,32 @@ public:
   /// namespace with the same name.
   void insertNS(NSPtr &ns);
 
-  /// Sets the n ame of the current namespace in the context and return
-  /// a boolean indicating the status of this operation. The operation
-  /// will fail if the namespace does not exist in the namespace table.
-  bool setCurrentNS(llvm::StringRef ns_name);
+  /// Execute the given function \p f by setting the `currentNS`
+  /// to the given \p nsName. It will restore the value of `currentNS`
+  /// after \p f returned. It also passes the old value of `currentNS`
+  /// to \p f.
+  template <typename T>
+  T withCurrentNS(llvm::StringRef nsName, CurrentNSFn<T> f) {
+    assert(!currentNS.empty() && "The currentNS is not initialized!");
+    auto tmp        = this->currentNS;
+    this->currentNS = nsName.str();
+
+    T res           = f();
+    this->currentNS = tmp;
+    return res;
+  };
 
   /// Return the current namespace that is being processed at the moment
   Namespace &getCurrentNS();
 
   /// Lookup the namespace with the give name in the current context and
   /// return a pointer to it or a `nullptr` in it doesn't exist.
-  Namespace *getNS(llvm::StringRef ns_name);
+  Namespace *getNS(llvm::StringRef nsName);
 
   /// Lookup and return a shared pointer to the given \p ns_name. This
   /// method should be used only if you need to own the namespace as well
   /// and want to keep it long term (like the JIT).
-  NSPtr getSharedPtrToNS(llvm::StringRef ns_name);
+  NSPtr getSharedPtrToNS(llvm::StringRef nsName);
 
   SereneContext()
       : pm(&mlirContext), diagEngine(makeDiagnosticEngine(*this)),
@@ -138,9 +155,11 @@ public:
 
     // We need to create one empty namespace, so that the JIT can
     // start it's operation.
-    auto ns = makeNamespace(*this, DEFAULT_NS_NAME, llvm::None);
+    auto ns = Namespace::make(*this, DEFAULT_NS_NAME, llvm::None);
 
     insertNS(ns);
+    currentNS = ns->name;
+
     // TODO: Get the crash report path dynamically from the cli
     // pm.enableCrashReproducerGeneration("/home/lxsameer/mlir.mlir");
 
@@ -155,6 +174,13 @@ public:
   CompilationPhase getTargetPhase() { return targetPhase; };
   int getOptimizatioLevel();
 
+  // Namespace stuff ---
+
+  /// Create an empty namespace with the given \p name and optional \p filename
+  /// and then insert it into the context
+  NSPtr makeNamespace(llvm::StringRef name,
+                      llvm::Optional<llvm::StringRef> filename);
+
   /// Read a namespace with the given \p name and returne a share pointer
   /// to the name or an Error tree.
   ///
@@ -169,6 +195,7 @@ public:
   /// It will \r a shared pointer to the namespace or an error tree.
   MaybeNS importNamespace(const std::string &name);
   MaybeNS importNamespace(const std::string &name, reader::LocationRange loc);
+  // ---
 
   static std::unique_ptr<llvm::LLVMContext> genLLVMContext() {
     return std::make_unique<llvm::LLVMContext>();
@@ -189,6 +216,17 @@ public:
 
     ctx->jit.swap(*maybeJIT);
 
+    // Make serene.user which is the defult NS available on the
+    // JIT
+    auto loc = reader::LocationRange::UnknownLocation(INTERNAL_NS);
+    auto err = ctx->jit->addNS(*ns, loc);
+
+    // TODO: Fix this by calling to the diag engine
+    if (err) {
+      llvm::errs() << err.getValue().back()->getMessage() << "\n";
+      serene::terminate(*ctx, 1);
+      return nullptr;
+    }
     return ctx;
   };
 
@@ -213,6 +251,12 @@ private:
   CompilationPhase targetPhase;
 
   // TODO: Change it to a LLVM::StringMap
+  // TODO: We need to keep different instances of the namespace
+  //       because if any one of them gets cleaned up via reference
+  //       count (if we are still using shared ptr for namespaces if not
+  //       remove this todo) then we will end up with dangling references
+  //       it the JIT
+
   // The namespace table. Every namespace that needs to be compiled has
   // to register itself with the context and appear on this table.
   // This table acts as a cache as well.
@@ -220,7 +264,7 @@ private:
 
   // Why string vs pointer? We might rewrite the namespace and
   // holding a pointer means that it might point to the old version
-  std::string current_ns;
+  std::string currentNS;
 
   /// A vector of pointers to all the jitDylibs for namespaces. Usually
   /// There will be only one pre NS but in case of forceful reloads of a
@@ -232,9 +276,6 @@ private:
 /// process to store the state
 SERENE_EXPORT std::unique_ptr<SereneContext> makeSereneContext();
 
-/// Terminates the serene compiler process in a thread safe manner
-SERENE_EXPORT void terminate(SereneContext &ctx, int exitCode);
-
 } // namespace serene
 
 #endif

include/serene/namespace.h

@@ -63,6 +63,7 @@
 
 namespace serene {
 class SereneContext;
+class Namespace;
 
 namespace exprs {
 class Expression;
@@ -70,6 +71,8 @@ using Node = std::shared_ptr<Expression>;
 using Ast  = std::vector<Node>;
 } // namespace exprs
 
+using NSPtr                = std::shared_ptr<Namespace>;
+using MaybeNS              = Result<NSPtr, errors::ErrorTree>;
 using MaybeModule          = llvm::Optional<llvm::orc::ThreadSafeModule>;
 using MaybeModuleOp        = llvm::Optional<mlir::OwningOpRef<mlir::ModuleOp>>;
 using SemanticEnv          = Environment<std::string, exprs::Node>;
@@ -80,8 +83,10 @@ using Forms                = std::vector<Form>;
 
 /// Serene's namespaces are the unit of compilation. Any code that needs to be
 /// compiled has to be in a namespace. The official way to create a new
-/// namespace is to use the `makeNamespace` function.
+/// namespace is to use the `makeNamespace` member function of `SereneContext`.
 class SERENE_EXPORT Namespace {
+  friend SereneContext;
+
 private:
   SereneContext &ctx;
 
@@ -98,6 +103,11 @@ private:
 
   std::vector<llvm::StringRef> symbolList;
 
+  /// Create a naw namespace with the given `name` and optional `filename` and
+  /// return a shared pointer to it in the given Serene context.
+  static NSPtr make(SereneContext &ctx, llvm::StringRef name,
+                    llvm::Optional<llvm::StringRef> filename);
+
 public:
   std::string name;
   llvm::Optional<std::string> filename;
@@ -161,17 +171,6 @@ public:
   ~Namespace();
 };
 
-using NSPtr = std::shared_ptr<Namespace>;
-
-using MaybeNS = Result<NSPtr, errors::ErrorTree>;
-/// Create a naw namespace with the given `name` and optional `filename` and
-/// return a shared pointer to it in the given Serene context. If the
-/// `setCurrent` argument is set to true, the created NS will become the
-/// curret namespace in the context
-SERENE_EXPORT NSPtr makeNamespace(SereneContext &ctx, llvm::StringRef name,
-                                  llvm::Optional<llvm::StringRef> filename,
-                                  bool setCurrent = true);
-
 } // namespace serene
 
 #endif

src/libserene/context.cpp

@@ -34,39 +34,31 @@ void SereneContext::insertNS(NSPtr &ns) {
   namespaces[nsName] = ns;
 };
 
-Namespace *SereneContext::getNS(llvm::StringRef ns_name) {
-  if (namespaces.count(ns_name.str()) != 0) {
-    return namespaces[ns_name.str()].get();
+Namespace *SereneContext::getNS(llvm::StringRef nsName) {
+  if (namespaces.count(nsName.str()) != 0) {
+    return namespaces[nsName.str()].get();
   }
 
   return nullptr;
 };
 
-bool SereneContext::setCurrentNS(llvm::StringRef ns_name) {
-  if (namespaces.count(ns_name.str()) != 0) {
-    this->current_ns = ns_name;
-    return true;
-  }
-
-  return false;
-};
-
 Namespace &SereneContext::getCurrentNS() {
-  if (this->current_ns.empty() || (namespaces.count(this->current_ns) == 0)) {
+  llvm::outs() << this->currentNS << "\n";
+  if (this->currentNS.empty() || (namespaces.count(this->currentNS) == 0)) {
     panic(*this, llvm::formatv("getCurrentNS: Namespace '{0}' does not exist",
-                               this->current_ns)
+                               this->currentNS)
                      .str());
   }
 
-  return *namespaces[this->current_ns];
+  return *namespaces[this->currentNS];
 };
 
-NSPtr SereneContext::getSharedPtrToNS(llvm::StringRef ns_name) {
-  if (namespaces.count(ns_name.str()) == 0) {
+NSPtr SereneContext::getSharedPtrToNS(llvm::StringRef nsName) {
+  if (namespaces.count(nsName.str()) == 0) {
     return nullptr;
   }
 
-  return namespaces[ns_name.str()];
+  return namespaces[nsName.str()];
 };
 
 void SereneContext::setOperationPhase(CompilationPhase phase) {
@@ -99,6 +91,17 @@ int SereneContext::getOptimizatioLevel() {
   return 3;
 }
 
+NSPtr SereneContext::makeNamespace(llvm::StringRef name,
+                                   llvm::Optional<llvm::StringRef> filename) {
+  auto ns = Namespace::make(*this, name, filename);
+
+  if (ns != nullptr) {
+    insertNS(ns);
+  }
+
+  return ns;
+};
+
 MaybeNS SereneContext::readNamespace(const std::string &name) {
   auto loc = reader::LocationRange::UnknownLocation(name);
 
@@ -108,7 +111,9 @@ MaybeNS SereneContext::readNamespace(const std::string &name) {
 MaybeNS SereneContext::readNamespace(const std::string &name,
                                      reader::LocationRange loc) {
 
-  return sourceManager.readNamespace(*this, name, loc);
+  return withCurrentNS<MaybeNS>(name, [&]() -> MaybeNS {
+    return this->sourceManager.readNamespace(*this, name, loc);
+  });
 }
 
 MaybeNS SereneContext::importNamespace(const std::string &name) {
@@ -138,6 +143,8 @@ llvm::orc::JITDylib *SereneContext::getLatestJITDylib(Namespace &ns) {
   }
 
   auto vec = jitDylibs[ns.name];
+  // TODO: Make sure that the returning Dylib still exists in the JIT
+  //       by calling jit->engine->getJITDylibByName(dylib_name);
   return vec.empty() ? nullptr : vec.back();
 };
 

src/libserene/jit/halley.cpp

@@ -198,7 +198,17 @@ Halley::Halley(serene::SereneContext &ctx,
 };
 
 MaybeJITPtr Halley::lookup(exprs::Symbol &sym) const {
-  auto *dylib = engine->getJITDylibByName(sym.nsName);
+  HALLEY_LOG("Looking up: " << sym.toString());
+  auto *ns = ctx.getNS(sym.nsName);
+
+  if (ns == nullptr) {
+    return MaybeJITPtr::error(errors::makeErrorTree(
+        sym.location, errors::CantResolveSymbol,
+        "Can't find the namespace in the context: " + sym.nsName));
+  }
+
+  auto *dylib = ctx.getLatestJITDylib(*ns);
+  //
 
   if (dylib == nullptr) {
     return MaybeJITPtr::error(
@@ -274,9 +284,10 @@ void Halley::registerSymbols(
 llvm::Optional<errors::ErrorTree> Halley::addNS(Namespace &ns,
                                                 reader::LocationRange &loc) {
 
-  llvm::outs() << llvm::formatv("{0}#{1}", ns.name,
-                                ctx.getNumberOfJITDylibs(ns) + 1)
-               << "\n";
+  HALLEY_LOG(llvm::formatv("Creating Dylib {0}#{1}", ns.name,
+                           ctx.getNumberOfJITDylibs(ns) + 1)
+             << "\n");
+
   auto newDylib = engine->createJITDylib(
       llvm::formatv("{0}#{1}", ns.name, ctx.getNumberOfJITDylibs(ns) + 1));
 
@@ -287,10 +298,6 @@ llvm::Optional<errors::ErrorTree> Halley::addNS(Namespace &ns,
 
   ctx.pushJITDylib(ns, &(*newDylib));
 
-  llvm::outs() << llvm::formatv("{0}#{1}", ns.name,
-                                ctx.getNumberOfJITDylibs(ns) + 1)
-               << "\n";
-
   // TODO: Fix compileToLLVM to return proper errors
   auto maybeModule = ns.compileToLLVM();
 

src/libserene/namespace.cpp

@@ -31,6 +31,8 @@
 #include "serene/semantics.h"
 #include "serene/slir/slir.h"
 
+#include <serene/export.h>
+
 #include <llvm/ADT/StringRef.h>
 #include <llvm/Support/FormatVariadic.h>
 #include <llvm/Support/raw_ostream.h>
@@ -201,19 +203,11 @@ MaybeModule Namespace::compileToLLVMFromOffset(unsigned offset) {
   return llvm::None;
 };
 
+NSPtr Namespace::make(SereneContext &ctx, llvm::StringRef name,
+                      llvm::Optional<llvm::StringRef> filename) {
+  return std::make_shared<Namespace>(ctx, name, filename);
+};
+
 Namespace::~Namespace(){};
 
-NSPtr makeNamespace(SereneContext &ctx, llvm::StringRef name,
-                    llvm::Optional<llvm::StringRef> filename, bool setCurrent) {
-  auto nsPtr = std::make_shared<Namespace>(ctx, name, filename);
-
-  ctx.insertNS(nsPtr);
-
-  if (setCurrent) {
-    if (!ctx.setCurrentNS(nsPtr->name)) {
-      throw std::runtime_error("Couldn't set the current NS");
-    }
-  }
-  return nsPtr;
-};
 } // namespace serene

src/libserene/source_mgr.cpp

@@ -122,7 +122,7 @@ MaybeNS SourceMgr::readNamespace(SereneContext &ctx, std::string name,
 
   // Create the NS and set the AST
   auto ns =
-      makeNamespace(ctx, name, llvm::Optional(llvm::StringRef(importedFile)));
+      ctx.makeNamespace(name, llvm::Optional(llvm::StringRef(importedFile)));
 
   auto errs = ns->addTree(maybeAst.getValue());
   if (errs) {

src/serenec/serenec.cpp

@@ -261,8 +261,7 @@ int main(int argc, char *argv[]) {
 
   cl::ParseCommandLineOptions(argc, argv, banner);
 
-  auto ctx    = makeSereneContext();
-  auto userNS = makeNamespace(*ctx, "user", llvm::None);
+  auto ctx = makeSereneContext();
 
   applySereneCLOptions(*ctx);