Fix top-level await/yield detection in iife mode

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

source/parser/lib.civet

@@ -1,4 +1,4 @@
-/**
+**
  * lib.civet holds functions that are used inside parser.hera
  *
  * The rules inside parser.hera should be simple and short.
@@ -87,6 +87,7 @@ import {
   prepend
   replaceNode
   replaceNodes
+  spliceChild
   stripTrailingImplicitComma
   trimFirstSpace
   wrapIIFE
@@ -1546,20 +1547,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 +1564,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.

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