[IR] For loop lowering for iteration over CharacterSequences.

Change-Id: I7aeadfffc80f791ec19a3607e219c1dc4de028db

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

@@ -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

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

@@ -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

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

@@ -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(
+internal class IndexedGetLoopHeader(
-    override val headerInfo: ArrayHeaderInfo,
+    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" }
+        val indexedGetFun = headerInfo.objectVariable.type.getClass()!!.functions.first { it.name.asString() == "get" }
-        irCall(arrayGetFun).apply {
+        irCall(indexedGetFun).apply {
-            dispatchReceiver = irGet(headerInfo.arrayVariable)
+            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,

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

@@ -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.*
-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.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 abstract class IndexedGetIterationHandler(protected val context: CommonBackendContext) : ExpressionHandler {
-internal class ArrayIterationHandler(private val context: CommonBackendContext) : ExpressionHandler {
-
-    override fun match(expression: IrExpression) = expression.type.run { isArray() || isPrimitiveArray() }
-
     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"
 }

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

@@ -1,4 +1,3 @@
-// IGNORE_BACKEND: JVM_IR
 fun test() {
     var s = ""
     for (c in "testString") {

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

@@ -1,4 +1,3 @@
-// IGNORE_BACKEND: JVM_IR
 fun test() {
     var s = ""
     for (c in "testString") {