ForLoopsLowering: Handle Sequence<*>.withIndex().

compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/ir/Ir.kt

@@ -95,8 +95,8 @@ open class BuiltinSymbolsBase(protected val irBuiltIns: IrBuiltIns, protected va
     private fun progression(name: String) = getClass(Name.identifier(name), "kotlin", "ranges")
     private fun progressionOrNull(name: String) = getClassOrNull(Name.identifier(name), "kotlin", "ranges")
 
-    // The "...OrNull" variants are used for unsigned (and progressions) because the minimal stdlib used in tests do not include those
+    // The "...OrNull" variants are used for the classes below because the minimal stdlib used in tests do not include those classes.
-    // those classes. It was not feasible to add them to the JS reduced runtime because all its transitive dependencies also need to be
+    // It was not feasible to add them to the JS reduced runtime because all its transitive dependencies also need to be
     // added, which would include a lot of the full stdlib.
     open val uByte = getClassOrNull(Name.identifier("UByte"), "kotlin")
     open val uShort = getClassOrNull(Name.identifier("UShort"), "kotlin")
@@ -104,6 +104,7 @@ open class BuiltinSymbolsBase(protected val irBuiltIns: IrBuiltIns, protected va
     open val uLong = getClassOrNull(Name.identifier("ULong"), "kotlin")
     val uIntProgression = progressionOrNull("UIntProgression")
     val uLongProgression = progressionOrNull("ULongProgression")
+    val sequence = getClassOrNull(Name.identifier("Sequence"), "kotlin", "sequences")
 
     val charProgression = progression("CharProgression")
     val intProgression = progression("IntProgression")

compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/loops/HeaderInfo.kt

@@ -320,7 +320,7 @@ internal class DefaultHeaderInfoBuilder(context: CommonBackendContext, scopeOwne
 // DefaultHeaderInfoBuilder. The differences between the two are that NestedHeaderInfoBuilderForWithIndex:
 //
 //   - Has NO WithIndexHandler. We do not attempt to optimize `*.withIndex().withIndex()`.
-//   - Has DefaultIterableHandler. This allows us to optimize `Iterable<*>.withIndex()` and `Sequence<*>.withIndex()`.
+//   - Has Default(Iterable|Sequence)Handler. This allows us to optimize `Iterable<*>.withIndex()` and `Sequence<*>.withIndex()`.
 internal class NestedHeaderInfoBuilderForWithIndex(context: CommonBackendContext, scopeOwnerSymbol: () -> IrSymbol) :
     HeaderInfoBuilder(context, scopeOwnerSymbol) {
     // NOTE: No WithIndexHandler; we cannot lower `iterable.withIndex().withIndex()`.
@@ -330,12 +330,13 @@ internal class NestedHeaderInfoBuilderForWithIndex(context: CommonBackendContext
 
     // NOTE: StringIterationHandler MUST come before CharSequenceIterationHandler.
     // String is subtype of CharSequence and therefore its handler is more specialized.
-    // DefaultIterableHandler must come last as it is handles iterables not handled by more specialized handlers.
+    // Default(Iterable|Sequence)Handler must come last as they handle iterables not handled by more specialized handlers.
     override val expressionHandlers = listOf(
         ArrayIterationHandler(context),
         DefaultProgressionHandler(context),
         StringIterationHandler(context),
         CharSequenceIterationHandler(context),
-        DefaultIterableHandler(context)
+        DefaultIterableHandler(context),
+        DefaultSequenceHandler(context),
     )
 }

compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/loops/handlers/DefaultIterableHandler.kt

@@ -19,10 +19,9 @@ import org.jetbrains.kotlin.ir.util.getSimpleFunction
 import org.jetbrains.kotlin.util.OperatorNameConventions
 
 /** Builds a [HeaderInfo] for Iterables not handled by more specialized handlers. */
-internal open class DefaultIterableHandler(private val context: CommonBackendContext) :
+internal class DefaultIterableHandler(private val context: CommonBackendContext) : ExpressionHandler {
-    ExpressionHandler {
 
-    protected open val iterableClassSymbol = context.ir.symbols.iterable
+    private val iterableClassSymbol = context.ir.symbols.iterable
 
     override fun matchIterable(expression: IrExpression) = expression.type.isSubtypeOfClass(iterableClassSymbol)
 
@@ -34,6 +33,4 @@ internal open class DefaultIterableHandler(private val context: CommonBackendCon
                 scope.createTmpVariable(irCall(iteratorFun).apply { dispatchReceiver = expression }, nameHint = "iterator")
             )
         }
-}
+}
-
-// TODO: Handle Sequences by extending DefaultIterableHandler.

compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/loops/handlers/DefaultSequenceHandler.kt

@@ -0,0 +1,37 @@
+/*
+ * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.backend.common.lower.loops.handlers
+
+import org.jetbrains.kotlin.backend.common.CommonBackendContext
+import org.jetbrains.kotlin.backend.common.lower.createIrBuilder
+import org.jetbrains.kotlin.backend.common.lower.loops.ExpressionHandler
+import org.jetbrains.kotlin.backend.common.lower.loops.HeaderInfo
+import org.jetbrains.kotlin.backend.common.lower.loops.IterableHeaderInfo
+import org.jetbrains.kotlin.ir.builders.createTmpVariable
+import org.jetbrains.kotlin.ir.builders.irCall
+import org.jetbrains.kotlin.ir.expressions.IrExpression
+import org.jetbrains.kotlin.ir.symbols.IrSymbol
+import org.jetbrains.kotlin.ir.types.isSubtypeOfClass
+import org.jetbrains.kotlin.ir.util.getSimpleFunction
+import org.jetbrains.kotlin.util.OperatorNameConventions
+
+/** Builds a [HeaderInfo] for Sequences not handled by more specialized handlers. */
+internal class DefaultSequenceHandler(private val context: CommonBackendContext) : ExpressionHandler {
+
+    private val sequenceClassSymbol = context.ir.symbols.sequence
+
+    override fun matchIterable(expression: IrExpression) =
+        sequenceClassSymbol != null && expression.type.isSubtypeOfClass(sequenceClassSymbol)
+
+    override fun build(expression: IrExpression, scopeOwner: IrSymbol): HeaderInfo? =
+        with(context.createIrBuilder(scopeOwner, expression.startOffset, expression.endOffset)) {
+            val iteratorFun =
+                sequenceClassSymbol!!.getSimpleFunction(OperatorNameConventions.ITERATOR.asString())!!.owner
+            IterableHeaderInfo(
+                scope.createTmpVariable(irCall(iteratorFun).apply { dispatchReceiver = expression }, nameHint = "iterator")
+            )
+        }
+}

compiler/testData/codegen/bytecodeText/forLoop/forInSequenceWithIndex/forInEmptySequenceWithIndex.kt

@@ -1,6 +1,3 @@
-// IGNORE_BACKEND: JVM_IR
-// TODO KT-36774 Generate optimized loops over 'xs.withIndex()' in JVM_IR
-// TODO: Handle Sequences by extending DefaultIterableHandler.
 val xs = listOf<Any>().asSequence()
 
 fun box(): String {

compiler/testData/codegen/bytecodeText/forLoop/forInSequenceWithIndex/forInSequenceTypeParameterWithIndex.kt

@@ -1,6 +1,3 @@
-// IGNORE_BACKEND: JVM_IR
-// TODO KT-36774 Generate optimized loops over 'xs.withIndex()' in JVM_IR
-// TODO: Handle Sequences by extending DefaultIterableHandler.
 fun <T : Sequence<*>> test(sequence: T): String {
     val s = StringBuilder()
 

compiler/testData/codegen/bytecodeText/forLoop/forInSequenceWithIndex/forInSequenceWithIndex.kt

@@ -1,6 +1,3 @@
-// IGNORE_BACKEND: JVM_IR
-// TODO KT-36774 Generate optimized loops over 'xs.withIndex()' in JVM_IR
-// TODO: Handle Sequences by extending DefaultIterableHandler.
 val xs = listOf("a", "b", "c", "d").asSequence()
 
 fun box(): String {

compiler/testData/codegen/bytecodeText/forLoop/forInSequenceWithIndex/forInSequenceWithIndexNoElementVar.kt

@@ -1,6 +1,3 @@
-// IGNORE_BACKEND: JVM_IR
-// TODO KT-36774 Generate optimized loops over 'xs.withIndex()' in JVM_IR
-// TODO: Handle Sequences by extending DefaultIterableHandler.
 val xs = listOf("a", "b", "c", "d").asSequence()
 
 fun box(): String {

compiler/testData/codegen/bytecodeText/forLoop/forInSequenceWithIndex/forInSequenceWithIndexNoIndexVar.kt

@@ -1,6 +1,3 @@
-// IGNORE_BACKEND: JVM_IR
-// TODO KT-36774 Generate optimized loops over 'xs.withIndex()' in JVM_IR
-// TODO: Handle Sequences by extending DefaultIterableHandler.
 val xs = listOf("a", "b", "c", "d").asSequence()
 
 fun box(): String {

compiler/testData/codegen/bytecodeText/forLoop/forInSequenceWithIndex/forInSequenceWithIndexThrowsCME.kt

@@ -1,6 +1,3 @@
-// IGNORE_BACKEND: JVM_IR
-// TODO KT-36774 Generate optimized loops over 'xs.withIndex()' in JVM_IR
-// TODO: Handle Sequences by extending DefaultIterableHandler.
 // FULL_JDK
 
 val xsl = arrayListOf("a", "b", "c", "d")

compiler/testData/codegen/bytecodeText/forLoop/forInSequenceWithIndex/forInSequenceWithIndexWithExplicitlyTypedIndexVariable.kt

@@ -1,6 +1,3 @@
-// IGNORE_BACKEND: JVM_IR
-// TODO KT-36774 Generate optimized loops over 'xs.withIndex()' in JVM_IR
-// TODO: Handle Sequences by extending DefaultIterableHandler.
 val xs = listOf("a", "b", "c", "d").asSequence()
 
 fun useAny(x: Any) {}