From 4f89e6afbda3b352c902e77d6d312f7efb287e59 Mon Sep 17 00:00:00 2001
From: Erik Demaine <edemaine@mit.edu>
Date: Thu, 26 Dec 2024 08:58:53 -0500
Subject: [PATCH] Fix top-level await/yield detection in `iife` mode

---
 source/parser.hera      |  2 --
 source/parser/lib.civet | 29 +++++++++++++++++------------
 test/iife.civet         | 31 +++++++++++++++++++++++++++++++
 3 files changed, 48 insertions(+), 14 deletions(-)

diff --git a/source/parser.hera b/source/parser.hera
index 3155a6c7..a15f9b6d 100644
--- a/source/parser.hera
+++ b/source/parser.hera
@@ -221,8 +221,6 @@ Program
       children: [reset, init, ws1, statements, ws2],
       bare: true,
       root: true,
-      topLevelAwait: hasAwait(statements),
-      topLevelYield: hasYield(statements),
     }
     processProgram(program)
     return program
diff --git a/source/parser/lib.civet b/source/parser/lib.civet
index 2c9639ab..969f513d 100644
--- a/source/parser/lib.civet
+++ b/source/parser/lib.civet
@@ -1546,20 +1546,9 @@ function processProgram(root: BlockStatement): void
   assert.equal state.forbidTrailingMemberProperty#, 1, "forbidTrailingMemberProperty"
   assert.equal state.JSXTagStack#, 1, "JSXTagStack"
 
-  let rootIIFE: ASTNode
-  if config.iife or config.repl
-    // Avoid top-level await if code is async (same as CoffeeScript),
-    // and avoid top-level yield* if code is a generator,
-    // so that code works in CommonJS environment and so REPL can eval it
-    rootIIFE = wrapIIFE root.expressions, root.topLevelAwait,
-      if root.topLevelYield then "*"
-    newExpressions: StatementTuple[] := [['', rootIIFE]]
-    root.children = root.children.map & is root.expressions ? newExpressions : &
-    root.expressions = newExpressions
-
   addParentPointers(root)
 
-  { expressions: statements } := root
+  { expressions: statements } .= root
 
   processPlaceholders(statements)
   processNegativeIndexAccess(statements)
@@ -1574,6 +1563,21 @@ function processProgram(root: BlockStatement): void
   processFinallyClauses(statements)
   processBreaksContinues(statements)
 
+  // Now that e.g. `do*` and `async do` have been processed
+  // (processIterationExpressions), we can detect top-level await/yield
+  root.topLevelAwait = hasAwait statements
+  root.topLevelYield = hasYield statements
+  let rootIIFE: ASTNode
+  if config.iife or config.repl
+    // Avoid top-level await if code is async (same as CoffeeScript),
+    // and avoid top-level yield* if code is a generator,
+    // so that code works in CommonJS environment and so REPL can eval it
+    rootIIFE = wrapIIFE root.expressions, root.topLevelAwait,
+      if root.topLevelYield then "*"
+    statements = [['', rootIIFE]]
+    root.children = root.children.map & is root.expressions ? statements : &
+    root.expressions = statements
+
   // Hoist hoistDec attributes to actual declarations.
   // NOTE: This should come after iteration expressions get processed
   // into IIFEs.
@@ -1581,6 +1585,7 @@ function processProgram(root: BlockStatement): void
 
   // Adding implicit returns should happen after hoisting any ref declarations
   // so their target node can be found in the block without being inside a return
+  // and after IIFE wrapper
   processFunctions(statements, config)
 
   processCoffeeClasses(statements) if config.coffeeClasses
diff --git a/test/iife.civet b/test/iife.civet
index 5f29d3d4..e0d9f173 100644
--- a/test/iife.civet
+++ b/test/iife.civet
@@ -1,4 +1,6 @@
 {testCase} from ./helper.civet
+{compile} from ../source/main.civet
+assert from assert
 
 describe "iife directive", ->
   testCase """
@@ -47,6 +49,35 @@ describe "iife directive", ->
     (function*(){yield 5})()
   """
 
+  describe "topLevel", ->
+    function ast(code: string)
+      compile ```
+        "civet iife"
+        ${code}
+      ```, ast: true
+
+    it "simple no await", =>
+      assert.equal (ast("5") |> await |> .topLevelAwait), false
+    it "simple await", =>
+      assert.equal (ast("await 5") |> await |> .topLevelAwait), true
+    it "await in function", =>
+      assert.equal (ast("=> await 5") |> await |> .topLevelAwait), false
+    it "await in do", =>
+      assert.equal (ast("do await 5") |> await |> .topLevelAwait), true
+    it "await in async do", =>
+      assert.equal (ast("async do await 5") |> await |> .topLevelAwait), false
+
+    it "simple no yield", =>
+      assert.equal (ast("5") |> await |> .topLevelYield), false
+    it "simple yield", =>
+      assert.equal (ast("yield 5") |> await |> .topLevelYield), true
+    it "yield in function", =>
+      assert.equal (ast("-> yield 5") |> await |> .topLevelYield), false
+    it "yield in do", =>
+      assert.equal (ast("do yield 5") |> await |> .topLevelYield), true
+    it "yield in do*", =>
+      assert.equal (ast("do* yield 5") |> await |> .topLevelYield), false
+
 describe "repl directive", ->
   testCase """
     top-level declarations