[IR] For loop lowering for iteration over CharacterSequences.

Change-Id: I7aeadfffc80f791ec19a3607e219c1dc4de028db
2019-08-27 08:11:55 +02:00
parent 26fbe3f07b
commit 8ed1317839
6 changed files with 65 additions and 28 deletions
@@ -89,6 +89,8 @@ abstract class Symbols<out T : CommonBackendContext>(val context: T, private val

    val iterator = getClass(Name.identifier("Iterator"), "kotlin", "collections")

+    val charSequence = getClass(Name.identifier("CharSequence"), "kotlin")
+
    val primitiveIteratorsByType = PrimitiveType.values().associate { type ->
        val iteratorClass = getClass(Name.identifier(type.typeName.asString() + "Iterator"), "kotlin", "collections")
        type to iteratorClass
@@ -21,7 +21,6 @@ import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.utils.addToStdlib.firstNotNullResult

 // TODO: Handle withIndex()
-// TODO: Handle Strings/CharSequences
 // TODO: Handle UIntProgression, ULongProgression

 /** Represents a progression type in the Kotlin stdlib. */
@@ -135,12 +134,15 @@ internal class ProgressionHeaderInfo(
    )
 }

-/** Information about a for-loop over an array. The internal induction variable used is an Int. */
-internal class ArrayHeaderInfo(
+/**
+ * Information about a for-loop over an object with an indexed get method (such as arrays or character sequences).
+ * The internal induction variable used is an Int.
+ */
+internal class IndexedGetHeaderInfo(
    first: IrExpression,
    last: IrExpression,
    step: IrExpression,
-    val arrayVariable: IrVariable
+    val objectVariable: IrVariable
 ) : HeaderInfo(
    ProgressionType.INT_PROGRESSION,
    first,
@@ -220,7 +222,8 @@ internal class HeaderInfoBuilder(context: CommonBackendContext, private val scop

    private val expressionHandlers = listOf(
        ArrayIterationHandler(context),
-        DefaultProgressionHandler(context)
+        DefaultProgressionHandler(context),
+        CharSequenceIterationHandler(context)
    )

    override fun visitElement(element: IrElement, data: Nothing?): HeaderInfo? = null
@@ -7,6 +7,7 @@ package org.jetbrains.kotlin.backend.common.lower.loops

 import org.jetbrains.kotlin.backend.common.CommonBackendContext
 import org.jetbrains.kotlin.backend.common.ir.Symbols
+import org.jetbrains.kotlin.backend.common.ir.isTopLevel
 import org.jetbrains.kotlin.backend.common.lower.DeclarationIrBuilder
 import org.jetbrains.kotlin.backend.common.lower.createIrBuilder
 import org.jetbrains.kotlin.backend.common.lower.irIfThen
@@ -22,6 +23,9 @@ import org.jetbrains.kotlin.ir.expressions.impl.IrWhileLoopImpl
 import org.jetbrains.kotlin.ir.symbols.IrSymbol
 import org.jetbrains.kotlin.ir.types.*
 import org.jetbrains.kotlin.ir.util.functions
+import org.jetbrains.kotlin.ir.util.getPackageFragment
+import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.util.OperatorNameConventions

 /**
 * Contains the loop and expression to replace the old loop.
@@ -196,8 +200,8 @@ internal class ProgressionLoopHeader(
        }
 }

-internal class ArrayLoopHeader(
-    override val headerInfo: ArrayHeaderInfo,
+internal class IndexedGetLoopHeader(
+    override val headerInfo: IndexedGetHeaderInfo,
    inductionVariable: IrVariable,
    last: IrVariable,
    step: IrVariable
@@ -205,15 +209,15 @@ internal class ArrayLoopHeader(

    override fun initializeLoopVariable(symbols: Symbols<CommonBackendContext>, builder: DeclarationIrBuilder) = with(builder) {
        // inductionVar = loopVar[inductionVariable]
-        val arrayGetFun = headerInfo.arrayVariable.type.getClass()!!.functions.first { it.name.asString() == "get" }
-        irCall(arrayGetFun).apply {
-            dispatchReceiver = irGet(headerInfo.arrayVariable)
+        val indexedGetFun = headerInfo.objectVariable.type.getClass()!!.functions.first { it.name.asString() == "get" }
+        irCall(indexedGetFun).apply {
+            dispatchReceiver = irGet(headerInfo.objectVariable)
            putValueArgument(0, irGet(inductionVariable))
        }
    }

    override val declarations: List<IrStatement>
-        get() = listOf(headerInfo.arrayVariable, inductionVariable, last, step)
+        get() = listOf(headerInfo.objectVariable, inductionVariable, last, step)

    override fun buildLoop(builder: DeclarationIrBuilder, oldLoop: IrLoop, newBody: IrExpression?): LoopReplacement = with(builder) {
        // Loop is lowered into something like:
@@ -261,8 +265,29 @@ internal class HeaderProcessor(
        }

        // Get the iterable expression, e.g., `someIterable` in the following loop variable declaration:
+        //
        //   val it = someIterable.iterator()
-        val iterable = (variable.initializer as? IrCall)?.dispatchReceiver
+        //
+        // If the `iterator` method is an extension method, make sure that we are calling a known extension
+        // method in the library such as `kotlin.text.StringsKt.iterator` with no value arguments. Other
+        // extension methods could return user-defined iterators.
+        val iterable = (variable.initializer as? IrCall)?.let {
+            val extensionReceiver = it.extensionReceiver
+            if (extensionReceiver != null) {
+                val function = it.symbol.owner
+                if (it.valueArgumentsCount == 0
+                    && function.isTopLevel
+                    && function.getPackageFragment()?.fqName == FqName("kotlin.text")
+                    && function.name == OperatorNameConventions.ITERATOR) {
+                    extensionReceiver
+                } else {
+                    null
+                }
+            } else {
+                it.dispatchReceiver
+            }
+        }
+
        // Collect loop information from the iterable expression.
        val headerInfo = iterable?.accept(headerInfoBuilder, null)
            ?: return null  // If the iterable is not supported.
@@ -316,7 +341,7 @@ internal class HeaderProcessor(
                )

                return when (headerInfo) {
-                    is ArrayHeaderInfo -> ArrayLoopHeader(
+                    is IndexedGetHeaderInfo -> IndexedGetLoopHeader(
                        headerInfo,
                        inductionVariable,
                        lastValue,
@@ -17,10 +17,7 @@ import org.jetbrains.kotlin.ir.declarations.IrVariable
 import org.jetbrains.kotlin.ir.expressions.IrCall
 import org.jetbrains.kotlin.ir.expressions.IrExpression
 import org.jetbrains.kotlin.ir.symbols.IrSymbol
-import org.jetbrains.kotlin.ir.types.IrType
-import org.jetbrains.kotlin.ir.types.getClass
-import org.jetbrains.kotlin.ir.types.isArray
-import org.jetbrains.kotlin.ir.types.toKotlinType
+import org.jetbrains.kotlin.ir.types.*
 import org.jetbrains.kotlin.ir.util.isPrimitiveArray
 import org.jetbrains.kotlin.ir.util.properties
 import org.jetbrains.kotlin.ir.visitors.IrElementVisitor
@@ -310,11 +307,7 @@ internal class DefaultProgressionHandler(private val context: CommonBackendConte
        }
 }

-/** Builds a [HeaderInfo] for arrays. */
-internal class ArrayIterationHandler(private val context: CommonBackendContext) : ExpressionHandler {
-
-    override fun match(expression: IrExpression) = expression.type.run { isArray() || isPrimitiveArray() }
-
+internal abstract class IndexedGetIterationHandler(protected val context: CommonBackendContext) : ExpressionHandler {
    override fun build(expression: IrExpression, scopeOwner: IrSymbol): HeaderInfo? =
        with(context.createIrBuilder(scopeOwner, expression.startOffset, expression.endOffset)) {
            // Consider the case like:
@@ -333,21 +326,37 @@ internal class ArrayIterationHandler(private val context: CommonBackendContext)
            // This also ensures that the semantics of re-assignment of array variables used in the loop is consistent with the semantics
            // proposed in https://youtrack.jetbrains.com/issue/KT-21354.
            val arrayReference = scope.createTemporaryVariable(
-                expression, nameHint = "array",
+                expression, nameHint = "indexedObject",
                origin = IrDeclarationOrigin.FOR_LOOP_IMPLICIT_VARIABLE
            )

            // `last = array.size` (last is exclusive) for the loop `for (i in array.indices)`.
-            val arraySizeProperty = arrayReference.type.getClass()!!.properties.first { it.name.asString() == "size" }
+            val arraySizeProperty = arrayReference.type.getClass()!!.properties.first { it.name.asString() == sizePropertyName() }
            val last = irCall(arraySizeProperty.getter!!).apply {
                dispatchReceiver = irGet(arrayReference)
            }

-            ArrayHeaderInfo(
+            IndexedGetHeaderInfo(
                first = irInt(0),
                last = last,
                step = irInt(1),
-                arrayVariable = arrayReference
+                objectVariable = arrayReference
            )
        }
+
+    abstract fun sizePropertyName() : String
+}
+
+/** Builds a [HeaderInfo] for arrays. */
+internal class ArrayIterationHandler(context: CommonBackendContext) : IndexedGetIterationHandler(context) {
+    override fun match(expression: IrExpression) = expression.type.run { isArray() || isPrimitiveArray() }
+
+    override fun sizePropertyName(): String = "size"
+}
+
+/** Builds a [HeaderInfo] for iteration over characters in a `CharacterSequence`. */
+internal class CharSequenceIterationHandler(context: CommonBackendContext) : IndexedGetIterationHandler(context) {
+    override fun match(expression: IrExpression) = expression.type.isSubtypeOfClass(context.ir.symbols.charSequence)
+
+    override fun sizePropertyName(): String = "length"
 }
@@ -1,4 +1,3 @@
-// IGNORE_BACKEND: JVM_IR
 fun test() {
    var s = ""
    for (c in "testString") {
@@ -1,4 +1,3 @@
-// IGNORE_BACKEND: JVM_IR
 fun test() {
    var s = ""
    for (c in "testString") {