diff --git a/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/loops/HeaderInfo.kt b/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/loops/HeaderInfo.kt index 198472df9cc..8fc6b4583d8 100644 --- a/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/loops/HeaderInfo.kt +++ b/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 { - /** Returns true if the handler can build a [HeaderInfo] from the expression. */ - fun match(expression: E): Boolean + /** Returns true if the handler can build a [HeaderInfo] from the iterable expression. */ + 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)) { - build(expression, data, scopeOwner) - } else { - null - } + fun handle(expression: E, iteratorCall: IrCall?, data: D, scopeOwner: IrSymbol) = + if ((iteratorCall == null || iteratorCallMatcher == null || iteratorCallMatcher!!(iteratorCall)) && matchIterable(expression)) { + build(expression, data, scopeOwner) + } else { + null + } } internal interface ExpressionHandler : HeaderInfoHandler { @@ -225,13 +233,13 @@ internal interface ExpressionHandler : HeaderInfoHandler internal interface HeaderInfoFromCallHandler : HeaderInfoHandler { val matcher: IrCallMatcher - override fun match(expression: IrCall) = matcher(expression) + override fun matchIterable(expression: IrCall) = matcher(expression) } internal typealias ProgressionHandler = HeaderInfoFromCallHandler internal class HeaderInfoBuilder(context: CommonBackendContext, private val scopeOwnerSymbol: () -> IrSymbol) : - IrElementVisitor { + IrElementVisitor { 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? { - return expressionHandlers.firstNotNullResult { it.handle(expression, null, scopeOwnerSymbol()) } - ?: super.visitExpression(expression, data) + override fun visitExpression(iterable: IrExpression, iteratorCall: IrCall?): HeaderInfo? { + return expressionHandlers.firstNotNullResult { it.handle(iterable, iteratorCall, null, scopeOwnerSymbol()) } + ?: super.visitExpression(iterable, iteratorCall) } } \ No newline at end of file diff --git a/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/loops/HeaderProcessor.kt b/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/loops/HeaderProcessor.kt index 56e10ff58a1..395e5c691f9 100644 --- a/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/loops/HeaderProcessor.kt +++ b/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() - // - // 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 + val iteratorCall = variable.initializer as? IrCall + val iterable = iteratorCall?.run { 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 - } + extensionReceiver } 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) diff --git a/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/loops/ProgressionHandlers.kt b/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/loops/ProgressionHandlers.kt index 769703cc8f8..f411753112b 100644 --- a/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/loops/ProgressionHandlers.kt +++ b/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 + 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.") } 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) : +internal class ReversedHandler(context: CommonBackendContext, private val visitor: HeaderInfoBuilder) : HeaderInfoFromCallHandler { 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 -) : - ExpressionHandler { + private val canCacheLast: Boolean +) : 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) { - override fun match(expression: IrExpression) = expression.type.run { isArray() || isPrimitiveArray() } +internal class ArrayIterationHandler(context: CommonBackendContext) : IndexedGetIterationHandler(context, canCacheLast = true) { + 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) { - override fun match(expression: IrExpression) = expression.type.isSubtypeOfClass(context.ir.symbols.charSequence) +internal open class CharSequenceIterationHandler(context: CommonBackendContext, canCacheLast: Boolean = false) : + 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() +} diff --git a/compiler/testData/codegen/box/ranges/forInCharSequenceWithCustomIterator.kt b/compiler/testData/codegen/box/ranges/forInCharSequenceWithCustomIterator.kt new file mode 100644 index 00000000000..cbe68b5f5ee --- /dev/null +++ b/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 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" +} \ No newline at end of file diff --git a/compiler/tests/org/jetbrains/kotlin/codegen/BlackBoxCodegenTestGenerated.java b/compiler/tests/org/jetbrains/kotlin/codegen/BlackBoxCodegenTestGenerated.java index 55e50afcdc3..a653595a71b 100644 --- a/compiler/tests/org/jetbrains/kotlin/codegen/BlackBoxCodegenTestGenerated.java +++ b/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"); diff --git a/compiler/tests/org/jetbrains/kotlin/codegen/LightAnalysisModeTestGenerated.java b/compiler/tests/org/jetbrains/kotlin/codegen/LightAnalysisModeTestGenerated.java index 54f10c5d202..bbd6560b840 100644 --- a/compiler/tests/org/jetbrains/kotlin/codegen/LightAnalysisModeTestGenerated.java +++ b/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"); diff --git a/compiler/tests/org/jetbrains/kotlin/codegen/ir/IrBlackBoxCodegenTestGenerated.java b/compiler/tests/org/jetbrains/kotlin/codegen/ir/IrBlackBoxCodegenTestGenerated.java index 23d634c8d8b..8c59b6e992b 100644 --- a/compiler/tests/org/jetbrains/kotlin/codegen/ir/IrBlackBoxCodegenTestGenerated.java +++ b/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"); diff --git a/js/js.tests/test/org/jetbrains/kotlin/js/test/ir/semantics/IrJsCodegenBoxTestGenerated.java b/js/js.tests/test/org/jetbrains/kotlin/js/test/ir/semantics/IrJsCodegenBoxTestGenerated.java index 6e395d33c51..c294230c190 100644 --- a/js/js.tests/test/org/jetbrains/kotlin/js/test/ir/semantics/IrJsCodegenBoxTestGenerated.java +++ b/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"); diff --git a/js/js.tests/test/org/jetbrains/kotlin/js/test/semantics/JsCodegenBoxTestGenerated.java b/js/js.tests/test/org/jetbrains/kotlin/js/test/semantics/JsCodegenBoxTestGenerated.java index 9ef1fbb055c..fd01c2d120e 100644 --- a/js/js.tests/test/org/jetbrains/kotlin/js/test/semantics/JsCodegenBoxTestGenerated.java +++ b/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");