[Heavy] Fix defines with syntax closures

Author: ricejasonf

heavy/include/heavy/Context.h

@@ -391,9 +391,7 @@ class Context : public ContinuationStack<Context>,
   }
 
   SyntaxClosure* CreateSyntaxClosure(SourceLocation Loc, Value Node,
-                                     Value Env) {
-    return new (*this) SyntaxClosure(Loc, Env, Node);
-  }
+                                     Value Env);
 
   SourceValue* CreateSourceValue(SourceLocation Loc) {
     return new (*this) SourceValue(Loc);

heavy/include/heavy/OpGen.h

@@ -300,8 +300,8 @@ class OpGen : public ValueVisitor<OpGen, mlir::Value> {
 
   mlir::Value createGlobal(SourceLocation Loc, llvm::StringRef MangledName);
   mlir::Value createBinding(Binding *B, mlir::Value Init);
-  mlir::Value createDefine(Symbol* S, Value Args, Value OrigCall);
-  mlir::Value createTopLevelDefine(Symbol* S, Value Args, Value OrigCall);
+  mlir::Value createDefine(Value Id, Value Args, Value OrigCall);
+  mlir::Value createTopLevelDefine(Value Id, Value Args, Value OrigCall);
   mlir::Value createUndefined();
   mlir::Value createSet(SourceLocation Loc, Value LHS, Value RHS);
 

heavy/lib/Builtins.cpp

@@ -113,30 +113,33 @@ void op_eval(Context& C, ValueRefs Args);
 
 mlir::Value define(OpGen& OG, Pair* P) {
   Pair*   P2    = dyn_cast<Pair>(P->Cdr);
-  Symbol* S     = nullptr;
+  Value Id     = nullptr;
 
   if (!P2)
     return OG.SetError("invalid syntax for define", P);
   if (Pair* LambdaSpec = dyn_cast<Pair>(P2->Car))
-    S = dyn_cast<Symbol>(LambdaSpec->Car);
+    Id = LambdaSpec->Car;
   else
-    S = dyn_cast<Symbol>(P2->Car);
+    Id = P2->Car;
 
-  if (!S)
+  if (!isIdentifier(Id))
     return OG.SetError("invalid syntax for define", P);
-  return OG.createDefine(S, P2, P);
+  return OG.createDefine(Id, P2, P);
 }
 
 mlir::Value define_syntax(OpGen& OG, Pair* P) {
   Pair* P2 = dyn_cast<Pair>(P->Cdr);
   if (!P2) return OG.SetError("invalid define-syntax syntax", P);
-  Symbol* S = dyn_cast<Symbol>(P2->Car);
-  if (!S) return OG.SetError("expecting name for define-syntax", P);
+  Value Id = P2->Car;
+  if (!isIdentifier(Id))
+    return OG.SetError("expecting name for define-syntax", P);
 
   return OG.createSyntaxSpec(P2, P);
 }
 
 mlir::Value syntax_rules(OpGen& OG, Pair* P) {
+  // TODO Support SyntaxClosures.
+
   // The input is the <Syntax Spec> (Keyword (syntax-rules ...))
   // <Syntax Spec> has its own checks in createSyntaxSpec
   Symbol* Keyword = dyn_cast<Symbol>(P->Car);

heavy/lib/Context.cpp

@@ -1628,6 +1628,14 @@ EnvFrame* Context::CreateEnvFrame(llvm::ArrayRef<Value> Ids) {
   return E;
 }
 
+SyntaxClosure* Context::CreateSyntaxClosure(SourceLocation Loc, Value Node,
+                                            Value Env) {
+  // Prevent direct nesting of SyntaxClosures.
+  if (auto* SC = dyn_cast<SyntaxClosure>(Node))
+    return SC;
+  return new (*this) SyntaxClosure(Loc, Env, Node);
+}
+
 bool Context::OutputModule(llvm::StringRef MangledName,
                            llvm::StringRef ModulePath) {
 #if 0

heavy/lib/OpGen.cpp

@@ -771,32 +771,44 @@ mlir::Value OpGen::createBinding(Binding *B, mlir::Value Init) {
   return BVal;
 }
 
-mlir::Value OpGen::createDefine(Symbol* S, Value DefineArgs,
+mlir::Value OpGen::createDefine(Value Id, Value DefineArgs,
                                            Value OrigCall) {
-  if (isTopLevel()) return createTopLevelDefine(S, DefineArgs, OrigCall);
+  if (isTopLevel()) return createTopLevelDefine(Id, DefineArgs, OrigCall);
   if (!IsLocalDefineAllowed) return SetError("unexpected define", OrigCall);
   // Create the binding with a lazy init.
   // (Include everything after the define
   //  keyword to visit it later because it could
   //  be a terse lambda syntax.)
-  Binding* B = Context.CreateBinding(S, DefineArgs);
+  Binding* B = Context.CreateBinding(Id, DefineArgs);
   // Push to the local environment.
   Context.PushLocalBinding(B);
   mlir::Value BVal = createBinding(B, mlir::Value());
   assert(IsLocalDefineAllowed && "define should still be allowed");
   return BVal;
 }
 
-mlir::Value OpGen::createTopLevelDefine(Symbol* S, Value DefineArgs,
+mlir::Value OpGen::createTopLevelDefine(Value Id, Value DefineArgs,
                                         Value OrigCall) {
   SourceLocation DefineLoc = OrigCall.getSourceLocation();
   if (LibraryEnvProc) {
     return SetError("unexpected define", OrigCall);
   }
 
+  assert(isTopLevel() && "expecting top level");
+  Environment* Env = nullptr;
+  if (auto* SC = dyn_cast<SyntaxClosure>(Id))
+    Env = cast<Environment>(SC->Env);
+  else if (CurSyntaxClosure)
+    Env = cast<Environment>(CurSyntaxClosure->Env);
+  else
+    Env = cast<Environment>(Context.EnvStack);
+
+  // Unwrap SyntaxClosures.
+  assert(isIdentifier(Id) && "expecting identifier");
+  Symbol* S = cast<Symbol>(Context.RebuildLiteral(Id));
+
   // If EnvStack isn't an Environment then there is local
   // scope information on top of it
-  Environment* Env = cast<Environment>(Context.EnvStack);
 
   EnvEntry Entry = Env->Lookup(Context, S);
   if (Entry.Value && Entry.MangledName) {
@@ -1336,32 +1348,30 @@ void OpGen::Export(Value NameList) {
   return Context.Cont();
 }
 
-// Expect a Symbol or SyntaxClosure wrapping a Symbol.
 heavy::EnvEntry OpGen::LookupEnv(heavy::Value Id) {
+  assert(isIdentifier(Id) && "expecting an identifier");
+  // For any given SyntaxClosure, use the object as the lookup
+  // in the current environment, and then use the raw Symbol
+  // in the closed environment.
   heavy::EnvEntry Result;
-  if (CurSyntaxClosure) {
-    SyntaxClosure StackSC;
-    SyntaxClosure* SC = nullptr;
-    // Id could be a Symbol or a closed (wrapped) Symbol.
-    Symbol* S = dyn_cast<Symbol>(Id);
-    if (S) {
-      StackSC.Env = CurSyntaxClosure->Env;
-      StackSC.Node = S;
-      SC = &StackSC;
-    } else {
-      SC = cast<SyntaxClosure>(Id);
-      S = cast<Symbol>(SC->Node);
-    }
+  heavy::Value ClosedEnv;
+  if (auto* SC = dyn_cast<SyntaxClosure>(Id)) {
+    ClosedEnv = SC->Env;
+    Id = SC->Node;
+  } else if (CurSyntaxClosure) {
+    ClosedEnv = CurSyntaxClosure->Env;
+  }
 
-    // Invalid input
-    assert((SC && S) && "expecting identifier");
+  assert(isa<Symbol>(Id) && "syntax closure should be unwrapped");
 
-    // Perform lookup using the SyntaxClosure object in the
-    // primary EnvStack. Then use the raw symbol and look
-    // in the closed environment.
-    Result = Context.Lookup(SC);
+  if (ClosedEnv) {
+    SyntaxClosure StackSC;
+    Symbol* S = cast<Symbol>(Id);
+    StackSC.Env = ClosedEnv;
+    StackSC.Node = S;
+    Result = Context.Lookup(&StackSC);
     if (!Result)
-      Result = Context.Lookup(S, SC->Env);
+      Result = Context.Lookup(S, ClosedEnv);
   } else {
     Symbol* S = cast<Symbol>(Id);
     Result = Context.Lookup(S);

heavy/test/Evaluate/define-syntax.scm

@@ -74,3 +74,19 @@
      '(0 1 i ... 9))))
 (write (ez 0 1 2 3 4 5 6))
 (newline)
+
+(define-syntax my-define
+  (syntax-rules ()
+    ((my-define name x)
+     (define name '(my name x)))))
+
+; CHECK: (my my-tl 42)
+(my-define my-tl 42)
+(write my-tl)
+(newline)
+
+; CHECK: Undefined
+((lambda ()
+  (my-define not-my-local 12)
+  (write not-my-local)
+  (newline)))