[Heavy] Fix generating syntax closures

Author: ricejasonf

heavy/include/heavy/Context.h

@@ -112,6 +112,11 @@ class Context : public ContinuationStack<Context>,
   //    and swap it back upon completion (via RAII)
   //  - TODO EnvStack et al. could probably be moved to OpGen.
   Value EnvStack;
+  Environment* GetTopLevelEnvironment(Value Env = nullptr);
+  void EnsureEnvFrame() {
+    if (isa<Environment>(EnvStack))
+      PushEnvFrame();
+  }
 
 public: // Provide access in lib/Mlir bindings.
   std::unique_ptr<mlir::DialectRegistry> DialectRegistry;
@@ -126,6 +131,7 @@ class Context : public ContinuationStack<Context>,
   heavy::OpGen* OpGen = nullptr;
   heavy::OpEvalPtr OpEval;
 
+
   // Work around DidCallContinuation being set with compiler errors.
   bool CheckError();
 
@@ -162,7 +168,6 @@ class Context : public ContinuationStack<Context>,
   void dumpModuleOp();
   void verifyModule();
   bool OutputModule(llvm::StringRef MangledName, llvm::StringRef ModulePath);
-  void PushTopLevel(Value);
 
   // Return true on invalid kind
   bool CheckKind(ValueKind VK, Value V);
@@ -181,7 +186,9 @@ class Context : public ContinuationStack<Context>,
     EnvStack = E;
   }
 
-  Environment* getTopLevelEnvironment();
+  Value& GetLocalEnvStack(Value Stack = nullptr);
+  // Provide a hashable value for checking for duplicate identifiers.
+  std::pair<uintptr_t, uintptr_t> GetIdentifierUniqueId(Value V);
 
   Module* RegisterModule(llvm::StringRef MangledName,
                          heavy::ModuleLoadNamesFn* LoadNames = nullptr);
@@ -225,19 +232,20 @@ class Context : public ContinuationStack<Context>,
 
   // PushEnvFrame - Creates and pushes an EnvFrame to the
   //                current environment (EnvStack)
-  EnvFrame* PushEnvFrame(llvm::ArrayRef<Value> Names);
+  EnvFrame* PushEnvFrame(llvm::ArrayRef<Value> Names = {});
   void PopEnvFrame(EnvFrame*);
   void PushLocalBinding(Binding* B);
+  void PushEnv(Value V, Value Env);
 
   // PushLambdaFormals - Check formals, create an EnvFrame,
   //                     and push it onto the EnvStack
   //                     Return the pushed EnvFrame or nullptr.
-  EnvFrame* PushLambdaFormals(Value Formals, bool& HasRestParam,
-                              SyntaxClosure* SC);
+  EnvFrame* PushLambdaFormals(Value Formals, bool& HasRestParam);
+
 private:
   bool CheckLambdaFormals(Value Formals,
                           llvm::SmallVectorImpl<Value>& Names,
-                          bool& HasRestParam, SyntaxClosure* SC);
+                          bool& HasRestParam);
 public:
 
   void EmitStackSpaceError();

heavy/include/heavy/OpGen.h

@@ -280,8 +280,7 @@ class OpGen : public ValueVisitor<OpGen, mlir::Value> {
                          llvm::MutableArrayRef<mlir::Value> Args);
   mlir::Value createError(heavy::SourceLocation Loc,
                           llvm::MutableArrayRef<mlir::Value> Args);
-  mlir::Value createOpGen(SourceLocation Loc, mlir::Value Input,
-                          mlir::Value Env);
+  mlir::Value createOpGen(SourceLocation Loc, mlir::Value Input);
   mlir::Value createBody(SourceLocation Loc, Value Body);
   mlir::Value createSequence(SourceLocation Loc, Value Body);
   mlir::Value createSyntaxSpec(Pair* SyntaxSpec, Value OrigCall);
@@ -398,7 +397,7 @@ class OpGen : public ValueVisitor<OpGen, mlir::Value> {
 
   mlir::Value VisitExternName(ExternName* EN);
   mlir::Value VisitSyntaxClosure(SyntaxClosure* SC);
-  mlir::Value VisitSymbol(Symbol* S);
+  mlir::Value VisitSymbol(Symbol* S, Value ClosedEnv = nullptr);
   mlir::Value VisitBinding(Binding* B);
 
   mlir::Value VisitPair(Pair* P);
@@ -410,7 +409,7 @@ class OpGen : public ValueVisitor<OpGen, mlir::Value> {
 
   mlir::Operation* LookupSymbol(llvm::StringRef MangledName);
 
-  heavy::EnvEntry LookupEnv(heavy::Value Id);
+  heavy::EnvEntry LookupEnv(heavy::Value Id, heavy::Value ClosedEnv = nullptr);
 private:
   mlir::Value CallSyntax(Value Operator, Pair* P);
   mlir::Value HandleCall(Pair* P, heavy::EnvEntry FnEnvEntry);

heavy/include/heavy/Ops.td

@@ -162,7 +162,7 @@ def heavy_OpGenOp : HeavyOp<"opgen", [Terminator]> {
     This operation is used for syntax code generation.
   }];
 
-  let arguments = (ins HeavyValue:$input, HeavyValue:$env);
+  let arguments = (ins HeavyValue:$input);
   let results   = (outs);
 }
 

heavy/include/heavy/Value.h

@@ -1248,6 +1248,7 @@ class Syntax final
 //                 transformation.
 //                 The SourceLocation is the location of the
 //                 PairWithSource from substitution.
+class EnvFrame;
 class SyntaxClosure : public ValueBase,
                       public ValueWithSource {
 public:
@@ -1261,9 +1262,9 @@ class SyntaxClosure : public ValueBase,
       Node(Node)
   { }
 
-  // We store one on the stack for lookup.
+  // For SC on the stack.
   SyntaxClosure()
-    : SyntaxClosure(SourceLocation(), Value(), Value())
+    : SyntaxClosure({}, Value(), Value())
   { }
 
   using ValueWithSource::getSourceLocation;
@@ -1281,9 +1282,6 @@ inline bool isIdentifier(Value V) {
   return isa<Symbol>(V);
 }
 
-// Provide a hashable value for checking for duplicate identifiers.
-std::pair<uintptr_t, uintptr_t> getIdentifierUniqueId(Value V);
-
 class Vector final
   : public ValueBase,
     private llvm::TrailingObjects<Vector, Value> {
@@ -1532,10 +1530,15 @@ class EnvFrame final
     return NumBindings;
   }
 
+  // Store shadow scopes.
+  Value LocalStack;
+
 public:
+
   EnvFrame(unsigned NumBindings)
     : ValueBase(ValueKind::EnvFrame),
-      NumBindings(NumBindings)
+      NumBindings(NumBindings),
+      LocalStack(Empty())
   { }
 
   llvm::ArrayRef<Binding*> getBindings() const {
@@ -1548,16 +1551,15 @@ class EnvFrame final
         getTrailingObjects<Binding*>(), NumBindings);
   }
 
+  Value getLocalStack() const { return LocalStack; }
+
   static size_t sizeToAlloc(unsigned NumBindings) {
     return totalSizeToAlloc<Binding*>(NumBindings);
   }
 
-  // Return nullptr if not found.
-  EnvEntry Lookup(Value Id) const {
-    for (Value V : getBindings()) {
-      Binding* B = cast<Binding>(V);
-      EnvEntry Result = B->Lookup(Id);
-      if (Result)
+  EnvEntry LookupBindings(Value Id) const {
+    for (Binding* B : getBindings()) {
+      if (EnvEntry Result = B->Lookup(Id))
         return Result;
     }
     return {};
@@ -1858,6 +1860,11 @@ inline bool isSymbol(Value V , llvm::StringRef Str) {
   }
   return false;
 }
+inline bool isSymbolAux(Value V , llvm::StringRef Str) {
+  if (auto* SC = dyn_cast<SyntaxClosure>(V))
+    V = SC->Node;
+  return isSymbol(V, Str);
+}
 
 inline SourceLocation ValueBase::getSourceLocation() {
   Value Self(this);