Support async iterator case with coroutines

Not all the `hasNext` operators return types is exactly Z,
suspend operators return boxed versions.

So the fix is just coercing result value after invoked function to Z

compiler/backend/src/org/jetbrains/kotlin/codegen/ExpressionCodegen.java

@@ -954,14 +954,8 @@ public class ExpressionCodegen extends KtVisitor<StackValue, StackValue> impleme
                                                                       "No hasNext() function " + DiagnosticUtils.atLocation(loopRange));
             @SuppressWarnings("ConstantConditions") Call fakeCall = makeFakeCall(new TransientReceiver(iteratorCall.getResultingDescriptor().getReturnType()));
             StackValue result = invokeFunction(fakeCall, hasNextCall, StackValue.local(iteratorVarIndex, asmTypeForIterator));
-            result.put(result.type, v);
+            result.put(Type.BOOLEAN_TYPE, v);
 
-            FunctionDescriptor hasNext = hasNextCall.getResultingDescriptor();
-            KotlinType type = hasNext.getReturnType();
-            assert type != null && KotlinTypeChecker.DEFAULT.isSubtypeOf(type, DescriptorUtilsKt.getBuiltIns(hasNext).getBooleanType());
-
-            Type asmType = asmType(type);
-            StackValue.coerce(asmType, Type.BOOLEAN_TYPE, v);
             v.ifeq(loopExit);
         }
 

compiler/testData/codegen/box/coroutines/asyncIterator.kt

@@ -0,0 +1,122 @@
+// WITH_RUNTIME
+// WITH_COROUTINES
+// TARGET_BACKEND: JVM
+import kotlin.coroutines.*
+
+interface AsyncGenerator<in T> {
+    suspend fun yield(value: T)
+}
+
+interface AsyncSequence<out T> {
+    operator fun iterator(): AsyncIterator<T>
+}
+
+interface AsyncIterator<out T> {
+    operator suspend fun hasNext(): Boolean
+    operator suspend fun next(): T
+}
+
+fun <T> asyncGenerate(block: suspend AsyncGenerator<T>.() -> Unit): AsyncSequence<T> = object : AsyncSequence<T> {
+    override fun iterator(): AsyncIterator<T> {
+        val iterator = AsyncGeneratorIterator<T>()
+        iterator.nextStep = block.createCoroutine(receiver = iterator, completion = iterator)
+        return iterator
+    }
+}
+
+class AsyncGeneratorIterator<T>: AsyncIterator<T>, AsyncGenerator<T>, Continuation<Unit> {
+    var computedNext = false
+    var nextValue: T? = null
+    var nextStep: Continuation<Unit>? = null
+
+    // if (computesNext) computeContinuation is Continuation<T>
+    // if (!computesNext) computeContinuation is Continuation<Boolean>
+    var computesNext = false
+    var computeContinuation: Continuation<*>? = null
+
+    suspend fun computeHasNext(): Boolean = CoroutineIntrinsics.suspendCoroutineOrReturn { c ->
+        computesNext = false
+        computeContinuation = c
+        nextStep!!.resume(Unit)
+        CoroutineIntrinsics.SUSPENDED
+    }
+
+    suspend fun computeNext(): T = CoroutineIntrinsics.suspendCoroutineOrReturn { c ->
+        computesNext = true
+        computeContinuation = c
+        nextStep!!.resume(Unit)
+        CoroutineIntrinsics.SUSPENDED
+    }
+
+    @Suppress("UNCHECKED_CAST")
+    fun resumeIterator(exception: Throwable?) {
+        if (exception != null) {
+            done()
+            computeContinuation!!.resumeWithException(exception)
+            return
+        }
+        if (computesNext) {
+            computedNext = false
+            (computeContinuation as Continuation<T>).resume(nextValue as T)
+        } else {
+            (computeContinuation as Continuation<Boolean>).resume(nextStep != null)
+        }
+    }
+
+    override suspend fun hasNext(): Boolean {
+        if (!computedNext) return computeHasNext()
+        return nextStep != null
+    }
+
+    override suspend fun next(): T {
+        if (!computedNext) return computeNext()
+        computedNext = false
+        return nextValue as T
+    }
+
+    private fun done() {
+        computedNext = true
+        nextStep = null
+    }
+
+    // Completion continuation implementation
+    override fun resume(value: Unit) {
+        done()
+        resumeIterator(null)
+    }
+
+    override fun resumeWithException(exception: Throwable) {
+        done()
+        resumeIterator(exception)
+    }
+
+    // Generator implementation
+    override suspend fun yield(value: T): Unit = CoroutineIntrinsics.suspendCoroutineOrReturn { c ->
+        computedNext = true
+        nextValue = value
+        nextStep = c
+        resumeIterator(null)
+        CoroutineIntrinsics.SUSPENDED
+    }
+}
+
+fun builder(c: suspend () -> Unit) {
+    c.startCoroutine(EmptyContinuation)
+}
+
+fun box(): String {
+    val seq = asyncGenerate {
+        yield("O")
+        yield("K")
+    }
+
+    var res = ""
+
+    builder {
+        for (i in seq) {
+            res += i
+        }
+    }
+
+    return res
+}

compiler/testData/codegen/light-analysis/coroutines/asyncIterator.txt

@@ -0,0 +1,65 @@
+public interface AsyncGenerator {
+    public abstract @org.jetbrains.annotations.Nullable method yield(p0: java.lang.Object, @org.jetbrains.annotations.NotNull p1: kotlin.coroutines.Continuation): java.lang.Object
+}
+
+
+public final class AsyncGeneratorIterator {
+    private @org.jetbrains.annotations.Nullable field computeContinuation: kotlin.coroutines.Continuation
+    private field computedNext: boolean
+    private field computesNext: boolean
+    private @org.jetbrains.annotations.Nullable field nextStep: kotlin.coroutines.Continuation
+    private @org.jetbrains.annotations.Nullable field nextValue: java.lang.Object
+    public method <init>(): void
+    public final @org.jetbrains.annotations.Nullable method computeHasNext(@org.jetbrains.annotations.NotNull p0: kotlin.coroutines.Continuation): java.lang.Object
+    public final @org.jetbrains.annotations.Nullable method computeNext(@org.jetbrains.annotations.NotNull p0: kotlin.coroutines.Continuation): java.lang.Object
+    private final method done(): void
+    public final @org.jetbrains.annotations.Nullable method getComputeContinuation(): kotlin.coroutines.Continuation
+    public final method getComputedNext(): boolean
+    public final method getComputesNext(): boolean
+    public final @org.jetbrains.annotations.Nullable method getNextStep(): kotlin.coroutines.Continuation
+    public final @org.jetbrains.annotations.Nullable method getNextValue(): java.lang.Object
+    public @org.jetbrains.annotations.Nullable method hasNext(@org.jetbrains.annotations.NotNull p0: kotlin.coroutines.Continuation): java.lang.Object
+    public @org.jetbrains.annotations.Nullable method next(@org.jetbrains.annotations.NotNull p0: kotlin.coroutines.Continuation): java.lang.Object
+    public method resume(@org.jetbrains.annotations.NotNull p0: kotlin.Unit): void
+    public synthetic method resume(p0: java.lang.Object): void
+    public final @kotlin.Suppress method resumeIterator(@org.jetbrains.annotations.Nullable p0: java.lang.Throwable): void
+    public method resumeWithException(@org.jetbrains.annotations.NotNull p0: java.lang.Throwable): void
+    public final method setComputeContinuation(@org.jetbrains.annotations.Nullable p0: kotlin.coroutines.Continuation): void
+    public final method setComputedNext(p0: boolean): void
+    public final method setComputesNext(p0: boolean): void
+    public final method setNextStep(@org.jetbrains.annotations.Nullable p0: kotlin.coroutines.Continuation): void
+    public final method setNextValue(@org.jetbrains.annotations.Nullable p0: java.lang.Object): void
+    public @org.jetbrains.annotations.Nullable method yield(p0: java.lang.Object, @org.jetbrains.annotations.NotNull p1: kotlin.coroutines.Continuation): java.lang.Object
+}
+
+
+public interface AsyncIterator {
+    public abstract @org.jetbrains.annotations.Nullable method hasNext(@org.jetbrains.annotations.NotNull p0: kotlin.coroutines.Continuation): java.lang.Object
+    public abstract @org.jetbrains.annotations.Nullable method next(@org.jetbrains.annotations.NotNull p0: kotlin.coroutines.Continuation): java.lang.Object
+}
+
+
+public final class AsyncIteratorKt {
+    public final static @org.jetbrains.annotations.NotNull method asyncGenerate(@org.jetbrains.annotations.NotNull p0: kotlin.jvm.functions.Function2): AsyncSequence
+    public final static @org.jetbrains.annotations.NotNull method box(): java.lang.String
+    public final static method builder(@org.jetbrains.annotations.NotNull p0: kotlin.jvm.functions.Function1): void
+}
+
+
+public interface AsyncSequence {
+    public abstract @org.jetbrains.annotations.NotNull method iterator(): AsyncIterator
+}
+
+
+public final class CoroutineUtilKt {
+    public final static @org.jetbrains.annotations.NotNull method handleExceptionContinuation(@org.jetbrains.annotations.NotNull p0: kotlin.jvm.functions.Function1): kotlin.coroutines.Continuation
+    public final static @org.jetbrains.annotations.NotNull method handleResultContinuation(@org.jetbrains.annotations.NotNull p0: kotlin.jvm.functions.Function1): kotlin.coroutines.Continuation
+}
+
+
+public final class EmptyContinuation {
+    public final static field INSTANCE: EmptyContinuation
+    private method <init>(): void
+    public method resume(@org.jetbrains.annotations.Nullable p0: java.lang.Object): void
+    public method resumeWithException(@org.jetbrains.annotations.NotNull p0: java.lang.Throwable): void
+}

compiler/tests-ir-jvm/tests/org/jetbrains/kotlin/codegen/ir/IrBlackBoxCodegenTestGenerated.java

@@ -4493,6 +4493,12 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
             KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("compiler/testData/codegen/box/coroutines"), Pattern.compile("^(.+)\\.kt$"), TargetBackend.JVM, true);
         }
 
+        @TestMetadata("asyncIterator.kt")
+        public void testAsyncIterator() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/asyncIterator.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("await.kt")
         public void testAwait() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/await.kt");

compiler/tests/org/jetbrains/kotlin/codegen/BlackBoxCodegenTestGenerated.java

@@ -4493,6 +4493,12 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
             KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("compiler/testData/codegen/box/coroutines"), Pattern.compile("^(.+)\\.kt$"), TargetBackend.JVM, true);
         }
 
+        @TestMetadata("asyncIterator.kt")
+        public void testAsyncIterator() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/asyncIterator.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("await.kt")
         public void testAwait() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/await.kt");

compiler/tests/org/jetbrains/kotlin/codegen/LightAnalysisModeCodegenTestGenerated.java

@@ -4493,6 +4493,12 @@ public class LightAnalysisModeCodegenTestGenerated extends AbstractLightAnalysis
             KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("compiler/testData/codegen/box/coroutines"), Pattern.compile("^(.+)\\.kt$"), TargetBackend.JVM, true);
         }
 
+        @TestMetadata("asyncIterator.kt")
+        public void testAsyncIterator() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/asyncIterator.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("await.kt")
         public void testAwait() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/await.kt");