[Heavy] Fix OpGen unwinding from LibrarySpec error

Author: ricejasonf

heavy/include/heavy/OpGen.h

@@ -151,6 +151,8 @@ class OpGen : public ValueVisitor<OpGen, mlir::Value> {
   // Err - The stored error to indicate that the compiler is in an error state.
   //       (Use CheckError())
   Value Err = nullptr;
+  // Track depth of calls to RunSync to allow unwinding errors properly.
+  size_t RunSyncDepth = 0;
 
   bool IsLocalDefineAllowed = false;
   heavy::Symbol* ModulePrefix = nullptr;
@@ -330,26 +332,27 @@ class OpGen : public ValueVisitor<OpGen, mlir::Value> {
 
   void createLoadModule(SourceLocation Loc, Symbol* MangledName);
 
+  mlir::Value SetError(heavy::Error* NewErr) {
+    Err = NewErr;
+    Context.setLoc(Err.getSourceLocation());
+    if (RunSyncDepth == 0)
+      Context.Raise(Err);
+    else
+      Context.Yield(Err);
+    return Error();
+  }
+
   template <typename T>
   mlir::Value SetError(SourceLocation Loc, T Str, Value V = Undefined()) {
-    Err = Context.CreateError(Loc, Str, Context.CreatePair(V));
-    Context.setLoc(Loc);
-    Context.Raise(Err);
-    return Error();
+    heavy::Error* E = Context.CreateError(Loc, Str, Context.CreatePair(V));
+    return SetError(E);
   }
 
   template <typename T>
   mlir::Value SetError(T Str, Value V = Undefined()) {
     return SetError(V.getSourceLocation(), Str, V);
   }
 
-  mlir::Value SetError(heavy::Error* NewErr) {
-    Err = NewErr;
-    Context.Raise(Err);
-    return Error();
-  }
-
-
   mlir::Value Error() {
     return createUndefined();
   }

heavy/lib/OpGen.cpp

@@ -259,8 +259,11 @@ void OpGen::VisitLibrarySpec(Value LibSpec) {
   Context.PushCont([](heavy::Context& C, ValueRefs) {
     Value LibSpec = C.getCapture(0);
     bool DidCallSyntax = false;
-    if (Pair* P = dyn_cast<Pair>(LibSpec))
+    if (Pair* P = dyn_cast<Pair>(LibSpec)) {
       C.OpGen->MaybeCallSyntax(P, DidCallSyntax);
+      if (C.OpGen->CheckError())
+        return;
+    }
 
     if (!DidCallSyntax) {
       C.OpGen->SetError("expecting library spec", LibSpec);
@@ -1098,7 +1101,9 @@ mlir::Value OpGen::MaybeCallSyntax(Value Operator, Pair* P,
       Value Input = P;
       heavy::OpGen* OpGen = Context.OpGen;
       assert(OpGen == this && "sanity check");
+      ++RunSyncDepth;
       Value Result = Context.RunSync(Operator, Input);
+      --RunSyncDepth;
       assert(OpGen == Context.OpGen && "OpGen should not unwind itself");
       if (auto ResultErr = dyn_cast_or_null<heavy::Error>(Result))
         SetError(ResultErr);