[Heavy] Improve rename envs

Author: ricejasonf

heavy/include/heavy/Context.h

@@ -186,6 +186,7 @@ class Context : public ContinuationStack<Context>,
     EnvStack = E;
   }
 
+  EnvFrame* GetLocalEnvFrame(Value Stack, bool IsLambdaScope = false);
   Value& GetLocalEnvStack(Value Stack = nullptr);
   // Provide a hashable value for checking for duplicate identifiers.
   std::pair<uintptr_t, uintptr_t> GetIdentifierUniqueId(Value V);
@@ -232,7 +233,8 @@ 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 = {},
+                         bool IsLambdaScope = false);
   void PopEnvFrame(EnvFrame*);
   void PushLocalBinding(Binding* B);
   void PushEnv(Value V, Value Env);
@@ -355,8 +357,8 @@ class Context : public ContinuationStack<Context>,
   Vector*     CreateVector(ArrayRef<Value> Xs);
   Vector*     CreateVector(unsigned N, Value Default = Undefined());
   ByteVector* CreateByteVector(llvm::ArrayRef<Value> Xs);
-  EnvFrame*   CreateEnvFrame(llvm::ArrayRef<Value> Names);
-  EnvFrame*   CreateEnvFrame(unsigned N);
+  EnvFrame*   CreateEnvFrame(llvm::ArrayRef<Value> Names, bool IsLambdaScope);
+  EnvFrame*   CreateEnvFrame(unsigned N, bool IsLambdaScope = false);
 
   template <typename T>
   Any* CreateAny(T Obj) {

heavy/include/heavy/OpGen.h

@@ -159,7 +159,7 @@ class OpGen : public ValueVisitor<OpGen, mlir::Value> {
 
   void InsertTopLevelCommandOp(SourceLocation Loc);
   bool WalkDefineInits(Value Env, IdSet& LocalNames);
-  bool FinishLocalDefines();
+  void FinishLocalDefines();
 
 public:
   explicit OpGen(heavy::Context& C, heavy::Symbol* ModulePrefix = nullptr);

heavy/include/heavy/Ops.td

@@ -535,6 +535,25 @@ def heavy_RenameGlobalOp : HeavyOp<"rename_global"> {
   ];
 }
 
+def heavy_RenameEnvOp : HeavyOp<"rename_env"> {
+  let description = [{
+    RenameEnvOp creates a new EnvFrame with $bindings
+    which serves as the tail of a new environment stack
+    with the most local EnvFrame from $env on top.
+
+    The local EnvFrame holds the local stack so any syntax closure
+    created will be topped with any newly added local variables
+    or EnvFrames automatically.
+  }];
+
+  let arguments = (ins HeavyValue:$env, Variadic<HeavyValue>:$bindings);
+  let results   = (outs HeavyValue: $result);
+  let builders = [
+    OpBuilder<(ins "::mlir::Value":$env,
+                   "::llvm::ArrayRef<::mlir::Value>":$args)>
+  ];
+}
+
 def heavy_SourceLocOp : HeavyOp<"source-loc"> {
   let summary = "source-loc";
   let description = [{
@@ -640,22 +659,6 @@ def heavy_VectorOp : HeavyOp<"vector"> {
   ];
 }
 
-def heavy_EnvFrameOp : HeavyOp<"env_frame"> {
-  let description = [{
-    EnvFrameOp creates a vector-like object of bindings.
-
-    Since it serves to define lexical scopes in the
-    environment, it exists apart from VectorOp so that
-    the bindings do not get unwrapped during evaluation.
-  }];
-
-  let arguments = (ins Variadic<HeavyValue>:$args);
-  let results   = (outs HeavyValue: $result);
-  let builders = [
-    OpBuilder<(ins "::llvm::ArrayRef<::mlir::Value>":$args)>
-  ];
-}
-
 def heavy_UndefinedOp : HeavyOp<"undefined", [Pure]> {
   let summary = "undefined";
   let description = [{

heavy/include/heavy/Value.h

@@ -1525,6 +1525,7 @@ class EnvFrame final
   friend class Context;
   friend class CopyCollector;
 
+  bool IsLambdaScope;
   unsigned NumBindings;
   size_t numTrailingObjects(OverloadToken<Binding*> const) const {
     return NumBindings;
@@ -1535,12 +1536,15 @@ class EnvFrame final
 
 public:
 
-  EnvFrame(unsigned NumBindings)
+  EnvFrame(unsigned NumBindings, bool IsLambdaScope)
     : ValueBase(ValueKind::EnvFrame),
+      IsLambdaScope(IsLambdaScope),
       NumBindings(NumBindings),
       LocalStack(Empty())
   { }
 
+  bool isLambdaScope() const { return IsLambdaScope; }
+
   llvm::ArrayRef<Binding*> getBindings() const {
     return llvm::ArrayRef<Binding*>(
         getTrailingObjects<Binding*>(), NumBindings);

heavy/lib/Context.cpp

@@ -161,7 +161,7 @@ EnvFrame* Context::PushLambdaFormals(Value Formals,
     }
   }
 
-  return PushEnvFrame(Ids);
+  return PushEnvFrame(Ids, /*IsLambdaScope=*/true);
 }
 
 std::pair<uintptr_t, uintptr_t> Context::GetIdentifierUniqueId(Value V) {
@@ -175,40 +175,38 @@ std::pair<uintptr_t, uintptr_t> Context::GetIdentifierUniqueId(Value V) {
           Env.getOpaqueValue()};
 }
 
+// Get the localiest EnvFrame.
+EnvFrame* Context::GetLocalEnvFrame(Value Stack, bool IsLambdaScope) {
+  EnvFrame* MostLocalEF = nullptr;
+  for (Value EnvNode : Stack) {
+    if (auto* EF = dyn_cast<EnvFrame>(EnvNode)) {
+      if (auto* NestedEF = GetLocalEnvFrame(EF->getLocalStack(),
+                                            IsLambdaScope))
+        return NestedEF;
+      if (!IsLambdaScope || EF->isLambdaScope())
+        MostLocalEF = EF;
+    }
+  }
+  return MostLocalEF;
+}
+
 // Get the localiest EnvFrame->LocalStack or this->EnvStack.
 Value& Context::GetLocalEnvStack(Value Stack) {
   if (!Stack)
     Stack = this->EnvStack;
-  EnvFrame* CurEF = nullptr;
-  if (auto* EnvPair = dyn_cast<Pair>(Stack)) {
-    if (auto* EF = dyn_cast<EnvFrame>(EnvPair->Car)) {
-      Stack = EF->LocalStack;
-      CurEF = EF;
-    }
-  }
-  return CurEF ? CurEF->LocalStack : this->EnvStack;
+  EnvFrame* MostLocalEF = GetLocalEnvFrame(Stack);
+  return MostLocalEF ? MostLocalEF->LocalStack : this->EnvStack;
 }
 
-EnvFrame* Context::PushEnvFrame(llvm::ArrayRef<Value> Ids) {
-  EnvFrame* E = CreateEnvFrame(Ids);
+EnvFrame* Context::PushEnvFrame(llvm::ArrayRef<Value> Ids,
+                                bool IsLambdaScope) {
+  EnvFrame* E = CreateEnvFrame(Ids, IsLambdaScope);
   PushEnv(E, this->EnvStack);
-
-  // Push to any syntax closures not in the current EnvStack.
-  for (Value Id : Ids) {
-    if (auto* SC = dyn_cast<SyntaxClosure>(Id)) {
-      if (this->EnvStack != SC->Env)
-        PushEnv(E, SC->Env);
-    }
-  }
-
   return E;
 }
 
 void Context::PushLocalBinding(Binding* B) {
   PushEnv(B, this->EnvStack);
-  if (auto* SC = dyn_cast<SyntaxClosure>(B->getIdentifier()))
-    if (this->EnvStack != SC->Env)
-      PushEnv(B, SC->Env);
 }
 
 void Context::PushEnv(Value V, Value Env) {
@@ -480,15 +478,18 @@ class Writer : public ValueVisitor<Writer> {
   }
 
   void VisitEnvFrame(EnvFrame* E) {
-    OS << "#<EnvFrame {Bindings: {";
+    OS << "#<EnvFrame {";
+    if (E->isLambdaScope())
+      OS << "IsLambda, ";
+    OS << "Bindings: {";
     // Print the name of each binding.
     llvm::interleaveComma(E->getBindings(), OS,
       [&](Value B) {
         VisitBinding(cast<Binding>(B));
       });
     OS << "} LocalStack: {";
     Visit(E->getLocalStack());
-    OS << "}>";
+    OS << "}}>";
 
   }
 
@@ -1636,14 +1637,14 @@ ByteVector* Context::CreateByteVector(ArrayRef<Value> Xs) {
   return BV;
 }
 
-EnvFrame* Context::CreateEnvFrame(unsigned N) {
+EnvFrame* Context::CreateEnvFrame(unsigned N, bool IsLambdaScope) {
   unsigned MemSize = EnvFrame::sizeToAlloc(N);
   void* Mem = Allocate(MemSize, alignof(EnvFrame));
-  return new (Mem) EnvFrame(N);
+  return new (Mem) EnvFrame(N, IsLambdaScope);
 }
 
-EnvFrame* Context::CreateEnvFrame(llvm::ArrayRef<Value> Ids) {
-  EnvFrame* E = CreateEnvFrame(Ids.size());
+EnvFrame* Context::CreateEnvFrame(llvm::ArrayRef<Value> Ids, bool IsLambdaScope) {
+  EnvFrame* E = CreateEnvFrame(Ids.size(), IsLambdaScope);
   auto Bindings = E->getBindings();
   for (unsigned i = 0; i < Bindings.size(); i++)
     Bindings[i] = CreateBinding(Ids[i], CreateUndefined());

heavy/lib/Dialect.cpp

@@ -263,10 +263,10 @@ void VectorOp::build(mlir::OpBuilder& B, mlir::OperationState& OpState,
   VectorOp::build(B, OpState, HeavyValueT, args);
 }
 
-void EnvFrameOp::build(mlir::OpBuilder& B, mlir::OperationState& OpState,
-                            llvm::ArrayRef<mlir::Value> args) {
+void RenameEnvOp::build(mlir::OpBuilder& B, mlir::OperationState& OpState,
+                    mlir::Value Env, llvm::ArrayRef<mlir::Value> Bindings) {
   mlir::Type HeavyValueT = B.getType<HeavyValueTy>();
-  EnvFrameOp::build(B, OpState, HeavyValueT, args);
+  RenameEnvOp::build(B, OpState, HeavyValueT, Env, Bindings);
 }
 
 void UndefinedOp::build(mlir::OpBuilder& B, mlir::OperationState& OpState) {

heavy/lib/Heap.cpp

@@ -103,11 +103,13 @@ class CopyCollector : private ValueVisitor<CopyCollector, heavy::Value> {
   // EnvFrame
   heavy::Value VisitEnvFrame(heavy::EnvFrame* EnvFrame) {
     llvm::ArrayRef<heavy::Binding*> Bindings = EnvFrame->getBindings();
+    bool IsLambdaScope = EnvFrame->isLambdaScope();
     unsigned MemSize = EnvFrame::sizeToAlloc(Bindings.size());
 
     void* Mem = NewHeap.Allocate(MemSize, alignof(heavy::EnvFrame));
 
-    heavy::EnvFrame* NewE = new (Mem) heavy::EnvFrame(Bindings.size());
+    heavy::EnvFrame* NewE = new (Mem) heavy::EnvFrame(Bindings.size(),
+                                                      IsLambdaScope);
     auto NewBindings = NewE->getBindings();
     for (unsigned i = 0; i < Bindings.size(); i++) {
       NewBindings[i] = cast<heavy::Binding>(Visit(Bindings[i]));

heavy/lib/OpEval.cpp

@@ -176,7 +176,7 @@ class OpEvalImpl {
     else if (isa<MatchIdOp>(Op))      return Visit(cast<MatchIdOp>(Op));
     else if (isa<RenameOp>(Op))       return Visit(cast<RenameOp>(Op));
     else if (isa<RenameGlobalOp>(Op)) return Visit(cast<RenameGlobalOp>(Op));
-    else if (isa<EnvFrameOp>(Op))     return Visit(cast<EnvFrameOp>(Op));
+    else if (isa<RenameEnvOp>(Op))    return Visit(cast<RenameEnvOp>(Op));
     else if (isa<ToVectorOp>(Op))     return Visit(cast<ToVectorOp>(Op));
     else if (isa<VectorOp>(Op))       return Visit(cast<VectorOp>(Op));
     else if (isa<LoadModuleOp>(Op))   return Visit(cast<LoadModuleOp>(Op));
@@ -552,15 +552,21 @@ class OpEvalImpl {
     return next(Op);
   }
 
-  BlockItrTy Visit(EnvFrameOp Op) {
-    unsigned Size = Op.getArgs().size();
-    heavy::EnvFrame* EnvFrame = Context.CreateEnvFrame(Size);
+  BlockItrTy Visit(RenameEnvOp Op) {
+    unsigned NumBindings = Op.getBindings().size();
+    heavy::EnvFrame* EnvFrame = Context.CreateEnvFrame(NumBindings);
     llvm::MutableArrayRef<Binding*> Bindings = EnvFrame->getBindings();
-    auto Vals = Op.getArgs();
-    for (unsigned I = 0; I < Size; ++I)
+    auto Vals = Op.getBindings();
+    for (unsigned I = 0; I < NumBindings; ++I)
       Bindings[I] = cast<Binding>(getBindingOrValue(Vals[I]));
 
-    setValue(Op, EnvFrame);
+    // Add the local environment on top of EnvFrame.
+    heavy::Value Env = getValue(Op.getEnv());
+    heavy::EnvFrame* LocalEnvFrame = Context.GetLocalEnvFrame(Env);
+    // EnvFrame is the tail.
+    heavy::Value NewEnv = Context.CreatePair(LocalEnvFrame, EnvFrame);
+
+    setValue(Op, NewEnv);
     return next(Op);
   }
 

heavy/lib/OpGen.cpp

@@ -876,21 +876,27 @@ mlir::Value OpGen::createSequence(SourceLocation Loc, Value Body) {
 //    been inserted by the `define` syntax. They are
 //    initialized to "undefined" and their corresponding
 //    heavy::Bindings placed in the EnvStack.
-//    Walk the EnvStack up to the nearest EnvFrame and
-//    collect these and insert the lazy initializers via SetOp
+//    Walk the EnvStack owned by the EnvFrame of the current lambda scope.
+//    Collect these and insert the lazy initializers via SetOp
 //    in the lexical order that they were defined.
 bool OpGen::WalkDefineInits(Value Env, IdSet& LocalIds) {
   Pair* P = dyn_cast<Pair>(Env);
   if (!P) return false;
-  if (isa<EnvFrame>(P->Car)) return false;
-  if (WalkDefineInits(P->Cdr, LocalIds)) return true;
+  // Walk the defines before this.
+  if (WalkDefineInits(P->Cdr, LocalIds))
+    return true;
+  // Walk any defines in nested scopes.
+  if (auto* EF = dyn_cast<EnvFrame>(P->Car)) {
+    assert(!EF->isLambdaScope() &&
+        "should not have any nested lambda scopes");
+    return WalkDefineInits(EF->getLocalStack(), LocalIds);
+  }
 
   // Check for duplicate names.
   Binding* B = cast<Binding>(P->Car);
   Value Id = B->getIdentifier();
-  bool IsNameInserted;
-  std::tie(std::ignore, IsNameInserted) =
-    LocalIds.insert(Context.GetIdentifierUniqueId(Id));
+  std::pair<uintptr_t, uintptr_t> UniqueId = Context.GetIdentifierUniqueId(Id);
+  auto [_, IsNameInserted] = LocalIds.insert(UniqueId);
   if (!IsNameInserted) {
     SetError("local variable has duplicate definitions", Id);
     return true;
@@ -916,14 +922,18 @@ bool OpGen::WalkDefineInits(Value Env, IdSet& LocalIds) {
   return false;
 }
 
-bool OpGen::FinishLocalDefines() {
+void OpGen::FinishLocalDefines() {
   TailPosScope TPS(*this);
   IsTailPos = false;
   IsLocalDefineAllowed = false;
   IdSet LocalIds;
-  // Get the localiest EnvFrame.
-  Value EnvStack = Context.GetLocalEnvStack();
-  return WalkDefineInits(EnvStack, LocalIds);
+  // Get the most local EnvFrame that is a lambda scope.
+  EnvFrame* EF = Context.GetLocalEnvFrame(Context.EnvStack,
+                                          /*IsLambdaScope=*/true);
+  assert(EF && "expecting a lambda scope env frame");
+  if (!EF)
+    return;
+  WalkDefineInits(EF->getLocalStack(), LocalIds);
 }
 
 mlir::Value OpGen::createBody(SourceLocation Loc, Value Body) {
@@ -1304,14 +1314,14 @@ mlir::Value OpGen::createContinuation(mlir::Operation* CallOp) {
 }
 
 mlir::Value OpGen::CallSyntax(Value Operator, Pair* P) {
-  Context.EnsureEnvFrame();
   switch (Operator.getKind()) {
     case ValueKind::BuiltinSyntax: {
       Context.setLoc(P->getSourceLocation());
       BuiltinSyntax* BS = cast<BuiltinSyntax>(Operator);
       return BS->Fn(*this, P);
     }
     case ValueKind::Syntax: {
+      EnvFrame* EF = Context.PushEnvFrame();
       heavy::Context& Context = this->Context;
       Context.setLoc(P->getSourceLocation());
       Value Input = P;
@@ -1325,6 +1335,7 @@ mlir::Value OpGen::CallSyntax(Value Operator, Pair* P) {
       assert(OpGen == Context.OpGen && "OpGen should not unwind itself");
       if (auto ResultErr = dyn_cast_or_null<heavy::Error>(Result))
         SetError(ResultErr);
+      Context.PopEnvFrame(EF);
       return toValue(Result);
     }
     default: {

heavy/lib/PatternTemplate.h

@@ -86,7 +86,7 @@ class PatternTemplate : ValueVisitor<PatternTemplate, mlir::Value> {
     }
 
     if (!OpGen.CheckError()) {
-      TemplateGen TG(OpGen, Keyword, PatternVars, Ellipsis);
+      TemplateGen TG(OpGen, EnvArg, Keyword, PatternVars, Ellipsis);
       TG.BuildTemplate(Template);
     }