[Heavy] Fix internal define initialization

Author: ricejasonf

heavy/lib/OpGen.cpp

@@ -420,6 +420,10 @@ heavy::FuncOp OpGen::createFunction(SourceLocation Loc,
 mlir::Value OpGen::createLambda(Value Formals, Value Body,
                                 SourceLocation Loc,
                                 llvm::StringRef Name) {
+  // Flush any internal definitions from containing body if any.
+  if (IsLocalDefineAllowed)
+    FinishLocalDefines();
+
   // Ensure we are no longer top level.
   if (!TopLevelOp)
     InsertTopLevelCommandOp(Loc);
@@ -457,9 +461,9 @@ mlir::Value OpGen::createLambda(Value Formals, Value Body,
     // If Result is null then it already
     // has a terminator.
     if (mlir::Value Result = createBody(Loc, Body)) {
-      if (IsLocalDefineAllowed) {
+      // Flush the internal definitions if we have not done so.
+      if (IsLocalDefineAllowed)
         FinishLocalDefines();
-      }
       Result = LocalizeValue(Result);
       Builder.create<ContOp>(Result.getLoc(), Result);
     }
@@ -706,6 +710,9 @@ bool OpGen::FinishLocalDefines() {
 }
 
 mlir::Value OpGen::createBody(SourceLocation Loc, Value Body) {
+  // Flush any internal definitions before starting a new body.
+  if (IsLocalDefineAllowed)
+    FinishLocalDefines();
   IsLocalDefineAllowed = true;
   return createSequence(Loc, Body);
 }

heavy/test/Evaluate/if.scm

@@ -0,0 +1,51 @@
+; RUN: heavy-scheme %s | FileCheck %s
+(import (heavy base))
+
+; CHECK: foo
+(write (if #t 'foo 'bar))
+(newline)
+
+; CHECK: bar
+(write (if #f 'foo 'bar))
+(newline)
+
+; CHECK-NEXT: foo
+(if #t
+  (write 'foo)
+  (write 'bar))
+(newline)
+
+; CHECK-NEXT: bar
+(if #f
+  (write 'foo)
+  (write 'bar))
+(newline)
+
+; CHECK-NEXT: foo
+(if (null? '())
+  (write 'foo)
+  (write 'bar))
+(newline)
+
+; CHECK-NEXT: bar
+(if (null? 42)
+  (write 'foo)
+  (write 'bar))
+(newline)
+
+(define fn
+  ((lambda()
+    (define captured-foo 'foo)
+    (define captured-bar 'bar)
+    (lambda (x)
+      (if (null? x)
+        (write captured-foo)
+        (write captured-bar))))))
+
+; CHECK-NEXT: foo
+(fn '())
+(newline)
+
+; CHECK-NEXT: bar
+(fn 42)
+(newline)

heavy/test/Evaluate/lambda.scm

@@ -1,6 +1,16 @@
 ; RUN: heavy-scheme %s | FileCheck %s
 (import (heavy base))
 
+; CHECK: foobar
+(define fn
+  ((lambda()
+    (define captured-foo 'foo)
+    (define captured-bar 'bar)
+    (lambda ()
+      (write captured-foo)
+      (write captured-bar)))))
+(fn)(newline)
+
 (define global -1)
 
 (set! global 1)