[Heavy] Improve OpEval scope mapping; Fix OpEval ownership

Author: ricejasonf

heavy/include/heavy/Context.h

@@ -54,9 +54,13 @@ class Context;
 void compile(Context&, Value V, Value Env, Value Handler);
 void eval(Context&, Value V, Value Env);
 void write(llvm::raw_ostream&, Value);
+void opEval(mlir::Operation*);
 
 class OpEvalImpl;
-void opEval(mlir::Operation*);
+struct OpEvalDeleter {
+  void operator()(OpEvalImpl*) const;
+};
+using OpEvalPtr = std::unique_ptr<OpEvalImpl, OpEvalDeleter>;
 
 class ContextLocalLookup {
   friend struct ContextLocal;
@@ -120,8 +124,7 @@ class Context : public ContinuationStack<Context>,
 
 public:
   heavy::OpGen* OpGen = nullptr;
-  // FIXME OpEval is not cleaned up or owned by anything.
-  heavy::OpEvalImpl* OpEval = nullptr;
+  heavy::OpEvalPtr OpEval;
 
   // Work around DidCallContinuation being set with compiler errors.
   bool CheckError();

heavy/lib/Context.cpp

@@ -47,6 +47,7 @@
 #include "llvm/Support/raw_ostream.h"
 #include <algorithm>
 #include <cstring>
+#include <memory>
 #include <string>
 
 using namespace heavy;
@@ -901,7 +902,7 @@ namespace {
 
   // Append to Output the path to a module sans any file extension.
   void getModulePath(heavy::Context& C, llvm::StringRef MangledName,
-                     llvm::SmallVectorImpl<char>& Output) { 
+                     llvm::SmallVectorImpl<char>& Output) {
     llvm::StringRef ModulePath =
         HEAVY_BASE_VAR(module_path).get(C).getStringRef();
     if (!ModulePath.empty())
@@ -1280,7 +1281,7 @@ void Context::RaiseError(String* Msg, llvm::ArrayRef<Value> IrrArgs) {
 }
 
 // ManagedObjectWind - Manage the lifetime of a C++ object within a dynamic
-// extent via a provided type-erased desctructor. 
+// extent via a provided type-erased desctructor.
 void Context::ManagedObjectWind(void* Ptr, DestructorTy Destructor,
                        Value Before, Value Thunk, Value After) {
   // Sentinel is referenced by each lambda

heavy/lib/OpEval.cpp

@@ -17,6 +17,7 @@
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/ScopedHashTable.h"
 #include "llvm/Support/Casting.h"
+#include <memory>
 #include <stack>
 
 namespace heavy {
@@ -28,14 +29,23 @@ class OpEvalImpl {
 
   heavy::Context& Context;
   ValueMapTy ValueMap;
-  std::stack<ValueMapScope> ValueMapScopes;
 
   void setValue(mlir::Value M, heavy::Value H) {
     assert(M && "must be set to a valid value");
     assert(H && "must be set to a valid value");
     ValueMap.insert(M, H);
   }
 
+  void setValueParentScope(mlir::Value M, heavy::Value H) {
+    assert(M && "must be set to a valid value");
+    assert(H && "must be set to a valid value");
+    ValueMapScope* CurScope = ValueMap.getCurScope();
+    assert(CurScope && "expecting a scope");
+    ValueMapScope* ParentScope = CurScope->getParentScope();
+    assert(ParentScope && "expecting parent scope");
+    ValueMap.insertIntoScope(ParentScope, M, H);
+  }
+
   heavy::Value getBindingOrValue(mlir::Value M) {
     assert(M && "lookup requires a valid value");
     heavy::Value V = ValueMap.lookup(M);
@@ -76,22 +86,12 @@ class OpEvalImpl {
 public:
   OpEvalImpl(heavy::Context& C)
     : Context(C),
-      ValueMap(),
-      ValueMapScopes()
-  {
-    // there has to be at least one scope on the stack
-    ValueMapScopes.emplace(ValueMap);
-  }
+      ValueMap()
+  { }
 
   // Prevent copy/move since we capture this in a few lambdas.
   OpEvalImpl(OpEvalImpl const&) = delete;
 
-  ~OpEvalImpl() {
-    // Pop the scopes in order.
-    while (!ValueMapScopes.empty())
-      ValueMapScopes.pop();
-  }
-
   void Eval(mlir::Operation* Op) {
     if (isa<GlobalOp, CommandOp, LoadModuleOp>(Op)) {
       Visit(Op);
@@ -119,16 +119,6 @@ class OpEvalImpl {
   }
 
 private:
-  void push_scope() {
-    ValueMapScopes.emplace(ValueMap);
-  }
-
-  void pop_scope() {
-    ValueMapScopes.pop();
-    assert(ValueMapScopes.size() > 0 &&
-        "scope stack must be balanced");
-  }
-
   BlockItrTy next(mlir::Operation* Op) {
     return ++BlockItrTy(Op);
   }
@@ -243,7 +233,7 @@ class OpEvalImpl {
       }
     }
 
-    push_scope();
+    auto Scope = ValueMapScope(ValueMap);
     mlir::Block& Body = F.getBody().front();
     if (NumParams != NumParamsMax)
       LoadArgs(Body, Args);
@@ -253,7 +243,6 @@ class OpEvalImpl {
     BlockItrTy Itr = Body.begin();
     while (Itr != BlockItrTy())
       Itr = Visit(&*Itr);
-    pop_scope();
 
     // The terminator operation should have
     // called Context.Apply() or one of those
@@ -483,20 +472,20 @@ class OpEvalImpl {
       C.Cont(Undefined());
     }, ValueRefs());
 
-    push_scope();
+    auto Scope = ValueMapScope(ValueMap);
     BlockItrTy Itr = Op.getInitializer().front().begin();
     while (Itr != BlockItrTy())
       Itr = Visit(&*Itr);
-    pop_scope();
+
     return BlockItrTy();
   }
 
   BlockItrTy Visit(CommandOp Op) {
-    push_scope();
+    auto Scope = ValueMapScope(ValueMap);
     BlockItrTy Itr = Op.getBody().front().begin();
     while (Itr != BlockItrTy())
       Itr = Visit(&*Itr);
-    pop_scope();
+
     return BlockItrTy();
   }
 
@@ -672,7 +661,7 @@ class OpEvalImpl {
     // Recall that each pack is a list in reverse order.
     while (isa<heavy::Pair>(Packs.front())) {
       // Set each block argument to the cdr of each pack.
-      push_scope();
+      auto Scope = ValueMapScope(ValueMap);
       for (unsigned i = 0; i < Packs.size(); i++) {
         auto* Pair = dyn_cast<heavy::Pair>(Packs[i]);
         if (!Pair)
@@ -683,7 +672,6 @@ class OpEvalImpl {
       BlockItrTy Itr = Op.getBody().front().begin();
       while (Itr != BlockItrTy())
         Itr = Visit(&*Itr);
-      // ResolveOp calls pop_scope.
     }
 
     for (unsigned i = 0; i < Packs.size(); i++) {
@@ -718,13 +706,12 @@ class OpEvalImpl {
     while (auto* Pair = dyn_cast<heavy::Pair>(E)) {
       if (Pair == Tail)
         break;
-      push_scope();
+      auto Scope = ValueMapScope(ValueMap);
       BlockItrTy Itr = Op.getBody().front().begin();
       setValue(BodyArg, Pair->Car);
       while (Itr != BlockItrTy())
         Itr = Visit(&*Itr);
       E = Pair->Cdr;
-      // ResolveOp calls pop_scope.
     }
 
     return next(Op);
@@ -736,21 +723,14 @@ class OpEvalImpl {
       assert(ParentOp.getPacks().size() ==
              Op.getArgs().size() &&
              "expecting value ranges of equal size");
-      // Store packs on the stack until we can pop the scope.
-      llvm::SmallVector<heavy::Value, 4> Packs(
-          ParentOp.getPacks().size(), nullptr);
       mlir::ValueRange ResolveArgs = Op.getArgs();
       mlir::ResultRange MResults = ParentOp.getPacks();
       // Construct each pack as a reversed ordered list.
-      for (unsigned i = 0; i < Packs.size(); i++)
-        Packs[i] = Context.CreatePair(getValue(ResolveArgs[i]),
-                                      getValue(MResults[i]));
-
-      // Set the results in the parent scope.
-      pop_scope();
-      for (unsigned i = 0; i < Packs.size(); i++)
-        setValue(MResults[i], Packs[i]);
-
+      for (unsigned i = 0; i < ResolveArgs.size(); i++) {
+        Value Pack = Context.CreatePair(getValue(ResolveArgs[i]),
+                                        getValue(MResults[i]));
+        setValueParentScope(MResults[i], Pack);
+      }
       return BlockItrTy();
     }
 
@@ -759,11 +739,7 @@ class OpEvalImpl {
     heavy::Value CurrentResult = getValue(Op.getArgs().front());
     heavy::Value Result = Context.CreatePair(CurrentResult,
                                       getValue(ExpandPacksOp.getResult()));
-
-    // Set the result in the parent scope.
-    pop_scope();
-    setValue(ExpandPacksOp.getResult(), Result);
-
+    setValueParentScope(ExpandPacksOp.getResult(), Result);
     return BlockItrTy();
   }
 
@@ -829,6 +805,10 @@ class OpEvalImpl {
   }
 };
 
+void OpEvalDeleter::operator()(OpEvalImpl* Ptr) const {
+  delete Ptr;
+}
+
 namespace base {
 void op_eval(Context& C, ValueRefs Args) {
   if (Args.size() != 1) {
@@ -839,16 +819,9 @@ void op_eval(Context& C, ValueRefs Args) {
     return C.RaiseError("expecting operation");
   }
 
-  // TODO Make the module own the OpEval instance. (maybe?)
   if (!C.OpEval) {
-    C.OpEval = new OpEvalImpl(C);
-    C.PushModuleCleanup(HEAVY_BASE_LIB_STR,
-      C.CreateLambda([](Context& C, ValueRefs) {
-        llvm_unreachable("FIXME find out why we do not get here");
-        assert(C.OpEval && "OpEval should be set.");
-        // Cleanup the OpEval object.
-        delete C.OpEval;
-      }, CaptureList{}));
+    // Still easier than extern template specialization of std::default_delete
+    C.OpEval = OpEvalPtr(new OpEvalImpl(C), OpEvalDeleter());
   }
 
   C.OpEval->Eval(Op);