Properly track coroutine-ness of comprehensions

Author: msimacek

graalpython/com.oracle.graal.python.pegparser/src/com/oracle/graal/python/pegparser/scope/ScopeEnvironment.java

@@ -393,6 +393,7 @@ private void handleComprehension(ExprTy e, String scopeName, ComprehensionTy[] g
             outermost.iter.accept(this);
             currentScope.comprehensionIterExpression--;
             enterBlock(scopeName, Scope.ScopeType.Function, e);
+            boolean isAsync;
             try {
                 currentScope.comprehensionType = comprehensionType;
                 if (outermost.isAsync) {
@@ -414,9 +415,13 @@ private void handleComprehension(ExprTy e, String scopeName, ComprehensionTy[] g
                 if (isGenerator) {
                     currentScope.flags.add(ScopeFlags.IsGenerator);
                 }
+                isAsync = currentScope.isCoroutine() && !isGenerator;
             } finally {
                 exitBlock();
             }
+            if (isAsync) {
+                currentScope.flags.add(ScopeFlags.IsCoroutine);
+            }
         }
 
         private void raiseIfComprehensionBlock(ExprTy node) {
@@ -545,6 +550,7 @@ public Void visit(ExprTy.Attribute node) {
         public Void visit(ExprTy.Await node) {
             raiseIfAnnotationBlock("await expression", node);
             node.value.accept(this);
+            currentScope.flags.add(ScopeFlags.IsCoroutine);
             return null;
         }
 

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/compiler/Compiler.java

@@ -526,6 +526,7 @@ private Void addOp(OpCodes code, int arg, byte[] followingArgs) {
     }
 
     private void addOp(OpCodes code, int arg, byte[] followingArgs, SourceRange location) {
+        assert code != YIELD_VALUE || unit.scope.isGenerator() || unit.scope.isCoroutine();
         Block b = unit.currentBlock;
         Instruction insn = new Instruction(code, arg, followingArgs, null, location);
         b.instr.add(insn);
@@ -1570,10 +1571,15 @@ private Void visitComprehension(ExprTy node, String name, ComprehensionTy[] gene
         SourceRange savedLocation = setLocation(node);
         try {
             enterScope(name, CompilationScope.Comprehension, node, 1, 0, 0, false, false);
+            boolean isAsyncGenerator = unit.scope.isCoroutine();
             if (type != ComprehensionType.GENEXPR) {
                 // The result accumulator, empty at the beginning
                 addOp(COLLECTION_FROM_STACK, type.typeBits);
             }
+            // TODO allow top-level await
+            if (isAsyncGenerator && type != ComprehensionType.GENEXPR && unit.scopeType != CompilationScope.AsyncFunction && unit.scopeType != CompilationScope.Comprehension) {
+                errorCallback.onError(ErrorType.Syntax, unit.currentLocation, "asynchronous comprehension outside of an asynchronous function");
+            }
             visitComprehensionGenerator(generators, 0, element, value, type);
             if (type != ComprehensionType.GENEXPR) {
                 addOp(RETURN_VALUE);
@@ -1590,15 +1596,12 @@ private Void visitComprehension(ExprTy node, String name, ComprehensionTy[] gene
             addOp(CALL_COMPREHENSION);
             // a genexpr will create an asyncgen, which we cannot await
             if (type != ComprehensionType.GENEXPR) {
-                for (ComprehensionTy gen : generators) {
-                    // if we have a non-genexpr async comprehension, the call will produce a
-                    // coroutine which we need to await
-                    if (gen.isAsync) {
-                        addOp(GET_AWAITABLE);
-                        addOp(LOAD_NONE);
-                        addYieldFrom();
-                        break;
-                    }
+                // if we have a non-genexpr async comprehension, the call will produce a
+                // coroutine which we need to await
+                if (isAsyncGenerator) {
+                    addOp(GET_AWAITABLE);
+                    addOp(LOAD_NONE);
+                    addYieldFrom();
                 }
             }
             return null;