KT-24156 Do not optimize for-loops over Strings with custom iterator

'fun CharSequence.iterator()' is an extension function, so one can
overload it with custom implementation.
Other "predefined" containers such as arrays and ranges have member
'fun iterator()', so these containers are not affected.

Check that 'iterator' call corresponds to an extension function
'iterator' defined in package 'kotlin.text' with a receiver of type
'kotlin.CharSequence'.

 #KT-24156 Fixed

compiler/backend/src/org/jetbrains/kotlin/codegen/RangeCodegenUtil.kt

@@ -254,3 +254,9 @@ fun getAsmRangeElementTypeForPrimitiveRangeOrProgression(rangeCallee: CallableDe
 
     throw AssertionError("Unexpected range type: $rangeType")
 }
+
+fun isCharSequenceIterator(descriptor: CallableDescriptor) =
+    descriptor.isTopLevelExtensionOnType("iterator", "kotlin.text") {
+        it.constructor.declarationDescriptor?.isTopLevelInPackage("CharSequence", "kotlin")
+                ?: false
+    }

compiler/backend/src/org/jetbrains/kotlin/codegen/range/RangeValues.kt

@@ -47,6 +47,8 @@ fun ExpressionCodegen.createRangeValueForExpression(rangeExpression: KtExpressio
     val rangeType = bindingContext.getType(rangeExpression)!!
     val asmRangeType = asmType(rangeType)
 
+    val loopRangeIteratorResolvedCall = bindingContext[BindingContext.LOOP_RANGE_ITERATOR_RESOLVED_CALL, rangeExpression]
+
     val builtIns = state.module.builtIns
     return when {
         asmRangeType.sort == Type.ARRAY -> {
@@ -62,9 +64,9 @@ fun ExpressionCodegen.createRangeValueForExpression(rangeExpression: KtExpressio
             PrimitiveRangeRangeValue(rangeExpression)
         isPrimitiveProgression(rangeType) ->
             PrimitiveProgressionRangeValue(rangeExpression)
-        isSubtypeOfString(rangeType, builtIns) ->
+        isSubtypeOfString(rangeType, builtIns) && isCharSequenceIteratorCall(loopRangeIteratorResolvedCall) ->
             CharSequenceRangeValue(true, AsmTypes.JAVA_STRING_TYPE)
-        isSubtypeOfCharSequence(rangeType, builtIns) ->
+        isSubtypeOfCharSequence(rangeType, builtIns) && isCharSequenceIteratorCall(loopRangeIteratorResolvedCall) ->
             CharSequenceRangeValue(false, null)
         else ->
             IterableRangeValue()
@@ -147,4 +149,7 @@ private fun ExpressionCodegen.createReversedRangeValueOrNull(rangeCall: Resolved
     val receiver = rangeCall.extensionReceiver as? ExpressionReceiver ?: return null
     val receiverRangeValue = createRangeValueForExpression(receiver.expression) as? ReversableRangeValue ?: return null
     return ReversedRangeValue(receiverRangeValue)
-}
+}
+
+private fun isCharSequenceIteratorCall(iteratorCall: ResolvedCall<*>?) =
+    iteratorCall?.resultingDescriptor?.let { isCharSequenceIterator(it) } ?: false

compiler/testData/codegen/box/ranges/forInStringWithCustomIterator.kt

@@ -0,0 +1,25 @@
+// TARGET_BACKEND: JVM
+// WITH_RUNTIME
+
+operator fun String.iterator(): IntIterator = object : IntIterator() {
+    private var index = 0
+
+    override fun nextInt() = codePointAt(index).apply {
+        index += Character.charCount(this)
+    }
+
+    override fun hasNext(): Boolean = index < length
+}
+
+fun String.collectInts(): List<Int> {
+    val result = ArrayList<Int>()
+    for (c in this) {
+        result.add(c)
+    }
+    return result
+}
+
+fun box(): String {
+    val ints = String(Character.toChars(127849)).collectInts()
+    return if (ints == listOf(127849)) "OK" else "Fail: $ints"
+}

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

@@ -15480,6 +15480,11 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
             runTest("compiler/testData/codegen/box/ranges/forInStringVarUpdatedInLoopBody.kt");
         }
 
+        @TestMetadata("forInStringWithCustomIterator.kt")
+        public void testForInStringWithCustomIterator() throws Exception {
+            runTest("compiler/testData/codegen/box/ranges/forInStringWithCustomIterator.kt");
+        }
+
         @TestMetadata("forIntRange.kt")
         public void testForIntRange() throws Exception {
             runTest("compiler/testData/codegen/box/ranges/forIntRange.kt");

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

@@ -15480,6 +15480,11 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
             runTest("compiler/testData/codegen/box/ranges/forInStringVarUpdatedInLoopBody.kt");
         }
 
+        @TestMetadata("forInStringWithCustomIterator.kt")
+        public void testForInStringWithCustomIterator() throws Exception {
+            runTest("compiler/testData/codegen/box/ranges/forInStringWithCustomIterator.kt");
+        }
+
         @TestMetadata("forIntRange.kt")
         public void testForIntRange() throws Exception {
             runTest("compiler/testData/codegen/box/ranges/forIntRange.kt");

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

@@ -15480,6 +15480,11 @@ public class LightAnalysisModeTestGenerated extends AbstractLightAnalysisModeTes
             runTest("compiler/testData/codegen/box/ranges/forInStringVarUpdatedInLoopBody.kt");
         }
 
+        @TestMetadata("forInStringWithCustomIterator.kt")
+        public void testForInStringWithCustomIterator() throws Exception {
+            runTest("compiler/testData/codegen/box/ranges/forInStringWithCustomIterator.kt");
+        }
+
         @TestMetadata("forIntRange.kt")
         public void testForIntRange() throws Exception {
             runTest("compiler/testData/codegen/box/ranges/forIntRange.kt");