Ensure ForLoopsLowering only handles the CharSequence.iterator()

extension function from the standard library (kotlin.text.iterator()).

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

@@ -203,17 +203,25 @@ internal class IndexedGetHeaderInfo(
 
 /** Matches an iterable expression and builds a [HeaderInfo] from the expression. */
 internal interface HeaderInfoHandler<E : IrExpression, D> {
-    /** Returns true if the handler can build a [HeaderInfo] from the expression. */
+    /** Returns true if the handler can build a [HeaderInfo] from the iterable expression. */
-    fun match(expression: E): Boolean
+    fun matchIterable(expression: E): Boolean
+
+    /**
+     * Matches the `iterator()` call that produced the iterable; if the call matches (or the matcher is null),
+     * the handler can build a [HeaderInfo] from the iterable.
+     */
+    val iteratorCallMatcher: IrCallMatcher?
+        get() = null
 
     /** Builds a [HeaderInfo] from the expression. */
     fun build(expression: E, data: D, scopeOwner: IrSymbol): HeaderInfo?
 
-    fun handle(expression: E, data: D, scopeOwner: IrSymbol) = if (match(expression)) {
+    fun handle(expression: E, iteratorCall: IrCall?, data: D, scopeOwner: IrSymbol) =
-        build(expression, data, scopeOwner)
+        if ((iteratorCall == null || iteratorCallMatcher == null || iteratorCallMatcher!!(iteratorCall)) && matchIterable(expression)) {
-    } else {
+            build(expression, data, scopeOwner)
-        null
+        } else {
-    }
+            null
+        }
 }
 
 internal interface ExpressionHandler : HeaderInfoHandler<IrExpression, Nothing?> {
@@ -225,13 +233,13 @@ internal interface ExpressionHandler : HeaderInfoHandler<IrExpression, Nothing?>
 internal interface HeaderInfoFromCallHandler<D> : HeaderInfoHandler<IrCall, D> {
     val matcher: IrCallMatcher
 
-    override fun match(expression: IrCall) = matcher(expression)
+    override fun matchIterable(expression: IrCall) = matcher(expression)
 }
 
 internal typealias ProgressionHandler = HeaderInfoFromCallHandler<ProgressionType>
 
 internal class HeaderInfoBuilder(context: CommonBackendContext, private val scopeOwnerSymbol: () -> IrSymbol) :
-    IrElementVisitor<HeaderInfo?, Nothing?> {
+    IrElementVisitor<HeaderInfo?, IrCall?> {
 
     private val symbols = context.ir.symbols
 
@@ -264,26 +272,26 @@ internal class HeaderInfoBuilder(context: CommonBackendContext, private val scop
         CharSequenceIterationHandler(context)
     )
 
-    override fun visitElement(element: IrElement, data: Nothing?): HeaderInfo? = null
+    override fun visitElement(element: IrElement, data: IrCall?): HeaderInfo? = null
 
     /** Builds a [HeaderInfo] for iterable expressions that are calls (e.g., `.reversed()`, `.indices`. */
-    override fun visitCall(expression: IrCall, data: Nothing?): HeaderInfo? {
+    override fun visitCall(iterable: IrCall, iteratorCall: IrCall?): HeaderInfo? {
         // Return the HeaderInfo from the first successful match. First, try to match a `reversed()` call.
-        val reversedHeaderInfo = reversedHandler.handle(expression, null, scopeOwnerSymbol())
+        val reversedHeaderInfo = reversedHandler.handle(iterable, iteratorCall, null, scopeOwnerSymbol())
         if (reversedHeaderInfo != null)
             return reversedHeaderInfo
 
         // Try to match a call to build a progression (e.g., `.indices`, `downTo`).
-        val progressionType = ProgressionType.fromIrType(expression.type, symbols)
+        val progressionType = ProgressionType.fromIrType(iterable.type, symbols)
         val progressionHeaderInfo =
-            progressionType?.run { progressionHandlers.firstNotNullResult { it.handle(expression, this, scopeOwnerSymbol()) } }
+            progressionType?.run { progressionHandlers.firstNotNullResult { it.handle(iterable, iteratorCall, this, scopeOwnerSymbol()) } }
 
-        return progressionHeaderInfo ?: super.visitCall(expression, data)
+        return progressionHeaderInfo ?: super.visitCall(iterable, iteratorCall)
     }
 
     /** Builds a [HeaderInfo] for iterable expressions not handled in [visitCall]. */
-    override fun visitExpression(expression: IrExpression, data: Nothing?): HeaderInfo? {
+    override fun visitExpression(iterable: IrExpression, iteratorCall: IrCall?): HeaderInfo? {
-        return expressionHandlers.firstNotNullResult { it.handle(expression, null, scopeOwnerSymbol()) }
+        return expressionHandlers.firstNotNullResult { it.handle(iterable, iteratorCall, null, scopeOwnerSymbol()) }
-            ?: super.visitExpression(expression, data)
+            ?: super.visitExpression(iterable, iteratorCall)
     }
 }

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

@@ -7,7 +7,6 @@ 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
@@ -24,8 +23,6 @@ import org.jetbrains.kotlin.ir.symbols.IrSymbol
 import org.jetbrains.kotlin.ir.types.*
 import org.jetbrains.kotlin.ir.util.deepCopyWithSymbols
 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
 
 /**
@@ -265,30 +262,17 @@ internal class HeaderProcessor(
         // Get the iterable expression, e.g., `someIterable` in the following loop variable declaration:
         //
         //   val it = someIterable.iterator()
-        //
+        val iteratorCall = variable.initializer as? IrCall
-        // If the `iterator` method is an extension method, make sure that we are calling a known extension
+        val iterable = iteratorCall?.run {
-        // 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
+                extensionReceiver
-                if (it.valueArgumentsCount == 0
-                    && function.isTopLevel
-                    && function.getPackageFragment()?.fqName == FqName("kotlin.text")
-                    && function.name == OperatorNameConventions.ITERATOR
-                ) {
-                    extensionReceiver
-                } else {
-                    null
-                }
             } else {
-                it.dispatchReceiver
+                dispatchReceiver
             }
         }
 
         // Collect loop information from the iterable expression.
-        val headerInfo = iterable?.accept(headerInfoBuilder, null)
+        val headerInfo = iterable?.accept(headerInfoBuilder, iteratorCall)
             ?: return null  // If the iterable is not supported.
 
         val builder = context.createIrBuilder(scopeOwnerSymbol(), variable.startOffset, variable.endOffset)

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

@@ -20,9 +20,7 @@ import org.jetbrains.kotlin.ir.expressions.IrExpression
 import org.jetbrains.kotlin.ir.symbols.IrSymbol
 import org.jetbrains.kotlin.ir.types.*
 import org.jetbrains.kotlin.ir.util.*
-import org.jetbrains.kotlin.ir.visitors.IrElementVisitor
 import org.jetbrains.kotlin.name.FqName
-import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.util.OperatorNameConventions
 import kotlin.math.absoluteValue
 
@@ -224,7 +222,7 @@ internal class UntilHandler(private val context: CommonBackendContext, private v
 /** Builds a [HeaderInfo] for progressions built using the `step` extension function. */
 internal class StepHandler(
     private val context: CommonBackendContext,
-    private val visitor: IrElementVisitor<HeaderInfo?, Nothing?>
+    private val visitor: HeaderInfoBuilder
 ) : ProgressionHandler {
 
     private val symbols = context.ir.symbols
@@ -529,7 +527,7 @@ internal class CollectionIndicesHandler(context: CommonBackendContext) : Indices
 internal class CharSequenceIndicesHandler(context: CommonBackendContext) : IndicesHandler(context) {
 
     override val matcher = SimpleCalleeMatcher {
-        extensionReceiver { it != null && it.type.run { isSubtypeOfClass(context.ir.symbols.charSequence) } }
+        extensionReceiver { it != null && it.type.run { isCharSequence() } }
         fqName { it == FqName("kotlin.text.<get-indices>") }
         parameterCount { it == 0 }
     }
@@ -542,7 +540,7 @@ internal class CharSequenceIndicesHandler(context: CommonBackendContext) : Indic
 }
 
 /** Builds a [HeaderInfo] for calls to reverse an iterable. */
-internal class ReversedHandler(context: CommonBackendContext, private val visitor: IrElementVisitor<HeaderInfo?, Nothing?>) :
+internal class ReversedHandler(context: CommonBackendContext, private val visitor: HeaderInfoBuilder) :
     HeaderInfoFromCallHandler<Nothing?> {
 
     private val symbols = context.ir.symbols
@@ -569,7 +567,7 @@ internal class DefaultProgressionHandler(private val context: CommonBackendConte
 
     private val symbols = context.ir.symbols
 
-    override fun match(expression: IrExpression) = ProgressionType.fromIrType(expression.type, symbols) != null
+    override fun matchIterable(expression: IrExpression) = ProgressionType.fromIrType(expression.type, symbols) != null
 
     override fun build(expression: IrExpression, scopeOwner: IrSymbol): HeaderInfo? =
         with(context.createIrBuilder(scopeOwner, expression.startOffset, expression.endOffset)) {
@@ -599,9 +597,8 @@ internal class DefaultProgressionHandler(private val context: CommonBackendConte
 
 internal abstract class IndexedGetIterationHandler(
     protected val context: CommonBackendContext,
-    val canCacheLast: Boolean = true
+    private val canCacheLast: Boolean
-) :
+) : ExpressionHandler {
-    ExpressionHandler {
     override fun build(expression: IrExpression, scopeOwner: IrSymbol): HeaderInfo? =
         with(context.createIrBuilder(scopeOwner, expression.startOffset, expression.endOffset)) {
             // Consider the case like:
@@ -643,8 +640,8 @@ internal abstract class IndexedGetIterationHandler(
 }
 
 /** Builds a [HeaderInfo] for arrays. */
-internal class ArrayIterationHandler(context: CommonBackendContext) : IndexedGetIterationHandler(context) {
+internal class ArrayIterationHandler(context: CommonBackendContext) : IndexedGetIterationHandler(context, canCacheLast = true) {
-    override fun match(expression: IrExpression) = expression.type.run { isArray() || isPrimitiveArray() }
+    override fun matchIterable(expression: IrExpression) = expression.type.run { isArray() || isPrimitiveArray() }
 
     override val IrType.sizePropertyGetter
         get() = getClass()!!.getPropertyGetter("size")!!.owner
@@ -657,29 +654,24 @@ internal class ArrayIterationHandler(context: CommonBackendContext) : IndexedGet
         }
 }
 
-/** Builds a [HeaderInfo] for iteration over characters in a [String]. */
-internal class StringIterationHandler(context: CommonBackendContext) : IndexedGetIterationHandler(context) {
-    override fun match(expression: IrExpression) = expression.type.isString()
-
-    override val IrType.sizePropertyGetter
-        get() = getClass()!!.getPropertyGetter("length")!!.owner
-
-    override val IrType.getFunction
-        get() = getClass()!!.functions.single {
-            it.name == OperatorNameConventions.GET &&
-                    it.valueParameters.size == 1 &&
-                    it.valueParameters[0].type.isInt()
-        }
-}
-
 /**
  * Builds a [HeaderInfo] for iteration over characters in a [CharSequence].
  *
  * Note: The value for "last" can NOT be cached (i.e., stored in a variable) because the size/length can change within the loop. This means
  * that "last" is re-evaluated with each iteration of the loop.
  */
-internal class CharSequenceIterationHandler(context: CommonBackendContext) : IndexedGetIterationHandler(context, canCacheLast = false) {
+internal open class CharSequenceIterationHandler(context: CommonBackendContext, canCacheLast: Boolean = false) :
-    override fun match(expression: IrExpression) = expression.type.isSubtypeOfClass(context.ir.symbols.charSequence)
+    IndexedGetIterationHandler(context, canCacheLast) {
+    override fun matchIterable(expression: IrExpression) = expression.type.isSubtypeOfClass(context.ir.symbols.charSequence)
+
+    // We only want to handle the known extension function for CharSequence in the standard library (top level `kotlin.text.iterator`).
+    // The behavior of this iterator is well-defined and can be lowered. CharSequences can have their own iterators, either as a member or
+    // extension function, and the behavior of those custom iterators is unknown.
+    override val iteratorCallMatcher = SimpleCalleeMatcher {
+        extensionReceiver { it != null && it.type.run { isCharSequence() } }
+        fqName { it == FqName("kotlin.text.${OperatorNameConventions.ITERATOR}") }
+        parameterCount { it == 0 }
+    }
 
     // The lowering operates on subtypes of CharSequence. Therefore, the IrType could be
     // a type parameter bounded by CharSequence. When that is the case, we cannot get
@@ -694,3 +686,12 @@ internal class CharSequenceIterationHandler(context: CommonBackendContext) : Ind
                     it.valueParameters[0].type.isInt()
         } ?: context.ir.symbols.charSequence.getSimpleFunction(OperatorNameConventions.GET.asString())!!.owner
 }
+
+/**
+ * Builds a [HeaderInfo] for iteration over characters in a [String].
+ *
+ * Note: The value for "last" CAN be cached for Strings as they are immutable and the size/length cannot change.
+ */
+internal class StringIterationHandler(context: CommonBackendContext) : CharSequenceIterationHandler(context, canCacheLast = true) {
+    override fun matchIterable(expression: IrExpression) = expression.type.isString()
+}

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

@@ -0,0 +1,65 @@
+// KJS_WITH_FULL_RUNTIME
+// WITH_RUNTIME
+import kotlin.test.*
+
+open class CharSequenceWithExtensionIterator(val s: String) : CharSequence {
+    fun get(foo: String): Char = TODO("shouldn't be called!")
+    override val length = s.length
+    override fun subSequence(startIndex: Int, endIndex: Int) = s.subSequence(startIndex, endIndex)
+    override fun get(index: Int) = s.get(index)
+}
+
+// Returns no characters
+operator fun CharSequenceWithExtensionIterator.iterator() = object : CharIterator() {
+    public override fun nextChar() = TODO()
+    public override fun hasNext() = false
+}
+
+class CharSequenceWithMemberIterator(s: String) : CharSequenceWithExtensionIterator(s) {
+    // Returns characters in reverse
+    operator fun iterator() = object : CharIterator() {
+        private var index = 0
+        public override fun nextChar() = get(length - ++index)
+        public override fun hasNext() = index < length
+    }
+}
+
+fun collectChars(cs: CharSequence): String {
+    val result = StringBuilder()
+    for (c in cs) {
+        result.append(c)
+    }
+    return result.toString()
+}
+
+fun <T : CharSequence> collectCharsTypeParam(cs: T): String {
+    val result = StringBuilder()
+    for (c in cs) {
+        result.append(c)
+    }
+    return result.toString()
+}
+
+fun box(): String {
+    val csWithExtIt = CharSequenceWithExtensionIterator("1234")
+    val csWithExtItResult = StringBuilder()
+    for (c in csWithExtIt) {
+        csWithExtItResult.append(c)
+    }
+    assertEquals("", csWithExtItResult.toString())
+
+    val csWithMemIt = CharSequenceWithMemberIterator("1234")
+    val csWithMemItResult = StringBuilder()
+    for (c in csWithMemIt) {
+        csWithMemItResult.append(c)
+    }
+    assertEquals("4321", csWithMemItResult.toString())
+
+    // The CharSequence.iterator() extension method should be invoked in all the following calls
+    assertEquals("1234", collectChars(csWithExtIt))
+    assertEquals("1234", collectCharsTypeParam(csWithExtIt))
+    assertEquals("1234", collectChars(csWithMemIt))
+    assertEquals("1234", collectCharsTypeParam(csWithMemIt))
+
+    return "OK"
+}

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

@@ -19305,6 +19305,11 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
             runTest("compiler/testData/codegen/box/ranges/forInCharSequenceLengthIncreasedInLoopBody.kt");
         }
 
+        @TestMetadata("forInCharSequenceWithCustomIterator.kt")
+        public void testForInCharSequenceWithCustomIterator() throws Exception {
+            runTest("compiler/testData/codegen/box/ranges/forInCharSequenceWithCustomIterator.kt");
+        }
+
         @TestMetadata("forInCharSequenceWithMultipleGetFunctions.kt")
         public void testForInCharSequenceWithMultipleGetFunctions() throws Exception {
             runTest("compiler/testData/codegen/box/ranges/forInCharSequenceWithMultipleGetFunctions.kt");

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

@@ -19305,6 +19305,11 @@ public class LightAnalysisModeTestGenerated extends AbstractLightAnalysisModeTes
             runTest("compiler/testData/codegen/box/ranges/forInCharSequenceLengthIncreasedInLoopBody.kt");
         }
 
+        @TestMetadata("forInCharSequenceWithCustomIterator.kt")
+        public void testForInCharSequenceWithCustomIterator() throws Exception {
+            runTest("compiler/testData/codegen/box/ranges/forInCharSequenceWithCustomIterator.kt");
+        }
+
         @TestMetadata("forInCharSequenceWithMultipleGetFunctions.kt")
         public void testForInCharSequenceWithMultipleGetFunctions() throws Exception {
             runTest("compiler/testData/codegen/box/ranges/forInCharSequenceWithMultipleGetFunctions.kt");

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

@@ -17870,6 +17870,11 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
             runTest("compiler/testData/codegen/box/ranges/forInCharSequenceLengthIncreasedInLoopBody.kt");
         }
 
+        @TestMetadata("forInCharSequenceWithCustomIterator.kt")
+        public void testForInCharSequenceWithCustomIterator() throws Exception {
+            runTest("compiler/testData/codegen/box/ranges/forInCharSequenceWithCustomIterator.kt");
+        }
+
         @TestMetadata("forInCharSequenceWithMultipleGetFunctions.kt")
         public void testForInCharSequenceWithMultipleGetFunctions() throws Exception {
             runTest("compiler/testData/codegen/box/ranges/forInCharSequenceWithMultipleGetFunctions.kt");

js/js.tests/test/org/jetbrains/kotlin/js/test/ir/semantics/IrJsCodegenBoxTestGenerated.java

@@ -15070,6 +15070,11 @@ public class IrJsCodegenBoxTestGenerated extends AbstractIrJsCodegenBoxTest {
             runTest("compiler/testData/codegen/box/ranges/forInCharSequenceLengthIncreasedInLoopBody.kt");
         }
 
+        @TestMetadata("forInCharSequenceWithCustomIterator.kt")
+        public void testForInCharSequenceWithCustomIterator() throws Exception {
+            runTest("compiler/testData/codegen/box/ranges/forInCharSequenceWithCustomIterator.kt");
+        }
+
         @TestMetadata("forInCharSequenceWithMultipleGetFunctions.kt")
         public void testForInCharSequenceWithMultipleGetFunctions() throws Exception {
             runTest("compiler/testData/codegen/box/ranges/forInCharSequenceWithMultipleGetFunctions.kt");

js/js.tests/test/org/jetbrains/kotlin/js/test/semantics/JsCodegenBoxTestGenerated.java

@@ -16225,6 +16225,11 @@ public class JsCodegenBoxTestGenerated extends AbstractJsCodegenBoxTest {
             runTest("compiler/testData/codegen/box/ranges/forInCharSequenceLengthIncreasedInLoopBody.kt");
         }
 
+        @TestMetadata("forInCharSequenceWithCustomIterator.kt")
+        public void testForInCharSequenceWithCustomIterator() throws Exception {
+            runTest("compiler/testData/codegen/box/ranges/forInCharSequenceWithCustomIterator.kt");
+        }
+
         @TestMetadata("forInCharSequenceWithMultipleGetFunctions.kt")
         public void testForInCharSequenceWithMultipleGetFunctions() throws Exception {
             runTest("compiler/testData/codegen/box/ranges/forInCharSequenceWithMultipleGetFunctions.kt");