[Heavy] Fix hygiene for define init syntax

Author: ricejasonf

heavy/include/heavy/OpGen.h

@@ -128,6 +128,8 @@ class OpGen : public ValueVisitor<OpGen, mlir::Value> {
   // in LookupEnv.
   SyntaxClosure* CurSyntaxClosure = nullptr;
 
+  struct SyntaxClosureScope;
+
   // Err - The stored error to indicate that the compiler is in an error state.
   //       (Use CheckError())
   Value Err = nullptr;

heavy/include/heavy/base.sld

@@ -34,6 +34,7 @@
     define-syntax
     if
     lambda
+    case-lambda
     quasiquote
     quote
     set!

heavy/lib/Builtins.cpp

@@ -193,9 +193,14 @@ mlir::Value syntax_fn(OpGen& OG, Pair* P) {
 
 mlir::Value lambda(OpGen& OG, Pair* P) {
   Pair* P2 = dyn_cast<Pair>(P->Cdr);
+  if (!P2)
+    return OG.SetError("invalid lambda syntax", P);
   Value Formals = P2->Car;
   Pair* Body = dyn_cast<Pair>(P2->Cdr);
 
+  if (!Body)
+    return OG.SetError("lambda syntax requires body");
+
   return OG.createLambda(Formals, Body, P->getSourceLocation());
 }
 

heavy/lib/OpEval.cpp

@@ -503,7 +503,13 @@ class OpEvalImpl {
   }
 
   BlockItrTy Visit(SetOp Op) {
-    heavy::Binding* B = cast<heavy::Binding>(getBindingOrValue(Op.getBinding()));
+    Value RawValue = getBindingOrValue(Op.getBinding());
+    heavy::Binding* B = dyn_cast<heavy::Binding>(RawValue);
+    if (!B) {
+      Context.RaiseError("invalid set operand: {}", RawValue);
+      return BlockItrTy();
+    }
+
     heavy::Value RHS = getValue(Op.getInput());
     B->setValue(RHS);
     return next(Op);

heavy/lib/OpGen.cpp

@@ -30,6 +30,23 @@
 #include <memory>
 
 using namespace heavy;
+
+struct OpGen::SyntaxClosureScope {
+  OpGen& O;
+  SyntaxClosure* PrevSC;
+
+  SyntaxClosureScope(OpGen& O, SyntaxClosure* SC)
+    : O(O),
+      PrevSC(O.CurSyntaxClosure)
+  {
+    O.CurSyntaxClosure = SC;
+  }
+
+  ~SyntaxClosureScope() {
+    O.CurSyntaxClosure = PrevSC;
+  }
+};
+
 // LambdaScope - RAII object that pushes an operation that is
 //               FunctionLike to the scope stack along with a
 //               BindingScope where we can insert stack local
@@ -884,6 +901,11 @@ bool OpGen::WalkDefineInits(Value Env, IdSet& LocalIds) {
   // Insert the binding initializer.
   auto Undef = cast<heavy::Undefined>(B->getValue());
   Value DefineArgs = Undef.getTracer();
+  // SC may be nullptr.
+  auto* SC = dyn_cast<SyntaxClosure>(Undef.getTracer());
+  SyntaxClosureScope SCScope(*this, SC);
+  if (SC)
+    DefineArgs = SC->Node;
   mlir::Value Init = VisitDefineArgs(DefineArgs);
   mlir::Value BVal = LocalizeValue(BindingTable.lookup(B), B);
   SourceLocation Loc = DefineArgs.getSourceLocation();
@@ -918,14 +940,17 @@ mlir::Value OpGen::createBinding(Binding *B, mlir::Value Init) {
   return BVal;
 }
 
-mlir::Value OpGen::createDefine(Value Id, Value DefineArgs,
-                                           Value OrigCall) {
+mlir::Value OpGen::createDefine(Value Id, Value DefineArgs, Value 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.)
+  // Create the binding with an undefined object that secretly holds
+  // its syntax in the current environment.
+  Value Env;
+
+  if (CurSyntaxClosure)
+    DefineArgs = Context.CreateSyntaxClosure(DefineArgs.getSourceLocation(),
+                                             DefineArgs,
+                                             CurSyntaxClosure->Env);
   Binding* B = Context.CreateBinding(Id, Undefined(DefineArgs));
   // Push to the local environment.
   Context.PushLocalBinding(B);
@@ -1145,11 +1170,8 @@ mlir::Value OpGen::VisitExternName(ExternName* EN) {
 }
 
 mlir::Value OpGen::VisitSyntaxClosure(SyntaxClosure* SC) {
-  SyntaxClosure* PrevSC = CurSyntaxClosure;
-  CurSyntaxClosure = SC;
+  SyntaxClosureScope SCScope(*this, SC);
   mlir::Value Result = Visit(SC->Node);
-  CurSyntaxClosure = PrevSC;
-
   return Result;
 }
 

heavy/test/Evaluate/Inputs/my/lib.sld

@@ -21,11 +21,16 @@
 
     (define-syntax create-op-literal
       (syntax-rules ()
-        ((create-op-literal (Arg) X)
+        ((create-op-literal (EmptyThunk) X)
           ((lambda (Arg)
+            ; Test hygiene of define initializer.
+            (define (MakeAttr Z)
+              (value-attr Z))
+            ; Trigger compiling define initializers with a syntactic closure.
+            (EmptyThunk)
             (create-op "heavy.literal"
               (attributes
-                `("info", (value-attr Arg))))) X))))
+                `("info", (MakeAttr Arg))))) X))))
 
     (write "end of init")
     (newline)

heavy/test/Evaluate/load-module.scm

@@ -5,7 +5,7 @@
 ; CHECK-NEXT: "end of init"
 (import (my lib))
 (import (only (heavy builtins)
-              write newline))
+              write newline lambda))
 
 ; CHECK: "hello module!"
 ; CHECK-NEXT: 5
@@ -16,5 +16,5 @@
 (hello-module-syntax "woof!")
 
 ; CHECK: #op{"heavy.literal"() {info = #heavy<"42">}
-(write (create-op-literal (SomeArg) 42))
+(write (create-op-literal ((lambda () #f)) 42))
 (newline)