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

2020-07-21 22:33:53 -07:00
parent 132960a695
commit 174576af61
11 changed files with 47 additions and 32 deletions
@@ -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
-    // those classes. It was not feasible to add them to the JS reduced runtime because all its transitive dependencies also need to be
+    // The "...OrNull" variants are used for the classes below because the minimal stdlib used in tests do not include those classes.
+    // 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")
@@ -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),
    )
 }
@@ -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) :
-    ExpressionHandler {
+internal class DefaultIterableHandler(private val context: CommonBackendContext) : 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.
+}
@@ -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")
+            )
+        }
+}
@@ -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 {
@@ -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()

@@ -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 {
@@ -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 {
@@ -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 {
@@ -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")
@@ -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) {}