[Heavy] Fix access to global bindings; Expose current-builder

Author: ricejasonf

heavy/include/heavy/Value.h

@@ -2043,6 +2043,7 @@ struct ContextLocal {
   heavy::Value init(heavy::Context& C, heavy::Value Value = nullptr);
   uintptr_t key() const { return reinterpret_cast<uintptr_t>(this); }
   heavy::Value get(heavy::ContextLocalLookup const& C) const;
+  heavy::Value get_binding(heavy::ContextLocalLookup const& C) const;
   void set(heavy::ContextLocalLookup& C, heavy::Value Value);
 };
 

heavy/lib/Context.cpp

@@ -1447,6 +1447,12 @@ heavy::Value ContextLocal::get(heavy::ContextLocalLookup const& C) const {
   return Value;
 }
 
+heavy::Value ContextLocal::get_binding(
+      heavy::ContextLocalLookup const& C) const {
+  heavy::Value Value = C.LookupTable.lookup(key());
+  return dyn_cast<heavy::Binding>(Value);
+}
+
 
 // Module
 

heavy/lib/Mlir.cpp

@@ -424,10 +424,9 @@ void block_arg(Context& C, ValueRefs Args) {
 
 static void with_builder_impl(Context& C, mlir::OpBuilder const& Builder,
                               heavy::Value Thunk) {
-  mlir::MLIRContext* MLIRContext = getCurrentContext(C);
   heavy::Value PrevBuilder = C.CreateBinding(heavy::Empty());
   heavy::Value NewBuilder = CreateTagged(C, kind::mlir_builder,
-                              mlir::OpBuilder(MLIRContext));
+                              Builder);
 
   heavy::Value Before = C.CreateLambda(
     [](heavy::Context& C, heavy::ValueRefs Args) {
@@ -712,6 +711,8 @@ void load_dialect(Context& C, heavy::ValueRefs Args) {
     return C.RaiseError("expecting dialect name");
 
   mlir::MLIRContext* MLIRContext = getCurrentContext(C);
+  // Ensure the registry is up to date.
+  MLIRContext->appendDialectRegistry(*C.DialectRegistry);
   mlir::Dialect* Dialect = MLIRContext->getOrLoadDialect(Name);
   if (Dialect == nullptr)
     return C.RaiseError(C.CreateString("failed to load dialect: ", Name), {});
@@ -817,6 +818,7 @@ void HEAVY_MLIR_LOAD_MODULE(heavy::Context& C) {
   heavy::initModuleNames(C, HEAVY_MLIR_LIB_STR, {
     {"create-op", HEAVY_MLIR_VAR(create_op)},
     {"%create-op", HEAVY_MLIR_VAR(create_op_impl)},
+    {"current-builder", HEAVY_MLIR_VAR(current_builder).get_binding(C)},
     {"region", HEAVY_MLIR_VAR(region)},
     {"entry-block", HEAVY_MLIR_VAR(entry_block)},
     {"results", HEAVY_MLIR_VAR(results)},