From 1b776bd5b68e475c4666d9b49eece63019681ef3 Mon Sep 17 00:00:00 2001 From: Xin Wang Date: Tue, 22 Feb 2022 17:51:58 +0800 Subject: [PATCH] JVM_IR: Specialize iterator if possible fixes KT-47171 --- .../FirBlackBoxCodegenTestGenerated.java | 16 ++++ .../common/lower/loops/ForLoopsLowering.kt | 81 +++++++++++++++++-- .../codegen/box/forLoop/forInIterator.kt | 28 +++++++ .../codegen/box/forLoop/forInIterator.txt | 17 ++++ .../codegen/BlackBoxCodegenTestGenerated.java | 16 ++++ .../IrBlackBoxCodegenTestGenerated.java | 16 ++++ .../LightAnalysisModeTestGenerated.java | 18 +++++ .../js/test/JsCodegenBoxTestGenerated.java | 10 +++ .../test/ir/IrJsCodegenBoxTestGenerated.java | 10 +++ .../IrCodegenBoxWasmTestGenerated.java | 13 +++ .../NativeCodegenBoxTestGenerated.java | 12 +++ 11 files changed, 230 insertions(+), 7 deletions(-) create mode 100644 compiler/testData/codegen/box/forLoop/forInIterator.kt create mode 100644 compiler/testData/codegen/box/forLoop/forInIterator.txt diff --git a/compiler/fir/fir2ir/tests-gen/org/jetbrains/kotlin/test/runners/codegen/FirBlackBoxCodegenTestGenerated.java b/compiler/fir/fir2ir/tests-gen/org/jetbrains/kotlin/test/runners/codegen/FirBlackBoxCodegenTestGenerated.java index 97e133eb238..80d07799bab 100644 --- a/compiler/fir/fir2ir/tests-gen/org/jetbrains/kotlin/test/runners/codegen/FirBlackBoxCodegenTestGenerated.java +++ b/compiler/fir/fir2ir/tests-gen/org/jetbrains/kotlin/test/runners/codegen/FirBlackBoxCodegenTestGenerated.java @@ -17266,6 +17266,22 @@ public class FirBlackBoxCodegenTestGenerated extends AbstractFirBlackBoxCodegenT } } + @Nested + @TestMetadata("compiler/testData/codegen/box/forLoop") + @TestDataPath("$PROJECT_ROOT") + public class ForLoop { + @Test + public void testAllFilesPresentInForLoop() throws Exception { + KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/testData/codegen/box/forLoop"), Pattern.compile("^(.+)\\.kt$"), null, TargetBackend.JVM_IR, true); + } + + @Test + @TestMetadata("forInIterator.kt") + public void testForInIterator() throws Exception { + runTest("compiler/testData/codegen/box/forLoop/forInIterator.kt"); + } + } + @Nested @TestMetadata("compiler/testData/codegen/box/fullJdk") @TestDataPath("$PROJECT_ROOT") diff --git a/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/loops/ForLoopsLowering.kt b/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/loops/ForLoopsLowering.kt index 4636053793b..0835820ca88 100644 --- a/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/loops/ForLoopsLowering.kt +++ b/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/loops/ForLoopsLowering.kt @@ -10,17 +10,20 @@ import org.jetbrains.kotlin.backend.common.CommonBackendContext import org.jetbrains.kotlin.backend.common.IrElementTransformerVoidWithContext import org.jetbrains.kotlin.backend.common.lower.createIrBuilder import org.jetbrains.kotlin.backend.common.phaser.makeIrFilePhase +import org.jetbrains.kotlin.descriptors.Modality import org.jetbrains.kotlin.ir.IrElement import org.jetbrains.kotlin.ir.IrStatement -import org.jetbrains.kotlin.ir.declarations.IrDeclaration -import org.jetbrains.kotlin.ir.declarations.IrDeclarationOrigin -import org.jetbrains.kotlin.ir.declarations.IrSymbolOwner -import org.jetbrains.kotlin.ir.declarations.IrVariable +import org.jetbrains.kotlin.ir.builders.irCall +import org.jetbrains.kotlin.ir.declarations.* import org.jetbrains.kotlin.ir.expressions.* import org.jetbrains.kotlin.ir.expressions.impl.IrCompositeImpl -import org.jetbrains.kotlin.ir.util.dump -import org.jetbrains.kotlin.ir.util.render +import org.jetbrains.kotlin.ir.types.classFqName +import org.jetbrains.kotlin.ir.types.classifierOrFail +import org.jetbrains.kotlin.ir.util.* import org.jetbrains.kotlin.ir.visitors.* +import org.jetbrains.kotlin.name.FqName +import org.jetbrains.kotlin.utils.addToStdlib.cast +import org.jetbrains.kotlin.utils.addToStdlib.safeAs val forLoopsPhase = makeIrFilePhase( ::ForLoopsLowering, @@ -175,7 +178,7 @@ private class RangeLoopTransformer( } val loopHeader = headerProcessor.extractHeader(iteratorVariable) - ?: return super.visitBlock(expression) // The iterable in the header is not supported. + ?: return super.visitBlock(expression.apply { specializeIteratorIfPossible(this) }) // The iterable in the header is not supported. val loweredHeader = lowerHeader(iteratorVariable, loopHeader) val (newLoop, loopReplacementExpression) = lowerWhileLoop(oldLoop, loopHeader) @@ -288,6 +291,70 @@ private class RangeLoopTransformer( return loopHeader.buildLoop(context.createIrBuilder(getScopeOwnerSymbol(), loop.startOffset, loop.endOffset), loop, newBody) } + /** + * This optimization is for the stdlib extension function in package `kotlin.collections`: + * ``` + * @kotlin.internal.InlineOnly + * public inline operator fun Iterator.iterator(): Iterator = this + * ``` + * Let's say we have an instance of `MyIterator`, which directly implements [kotlin.collections.Iterator], + * when it is used in a for-loop like: + * + * ``` + * val iterator = MyIterator() + * for (x in iterator) + * println(x) + * ``` + * Without this optimization, receiver type of call of `next` would be Iterator instead of MyIterator, which + * is an unnecessary boxing can be optimized out. + */ + private fun specializeIteratorIfPossible(irForLoopBlock: IrContainerExpression) { + val statements = irForLoopBlock.statements + + val iterator = statements[0] as IrVariable + val loop = statements[1] as IrWhileLoop + + // do optimization iff the initializer is call of `kotlin.collections.CollectionsKt.iterator` + val initializer = iterator.initializer.safeAs().takeIf { + it != null && it.symbol.owner.fqNameWhenAvailable == FqName("kotlin.collections.CollectionsKt.iterator") + } ?: return + + // for now, we only optimize for the case where extension receiver is a direct subclass of `kotlin.collections.Iterator` + // we can expand the scope later, if it considered to be necessary + val receiverType = initializer.extensionReceiver!!.type.takeIf { + it.superTypes().any { superType -> + superType.classFqName == FqName("kotlin.collections.Iterator") + } + } ?: return + + fun IrClass.findNextFunction(): IrFunction? = + this.functions.find { + it.nameForIrSerialization.asString() == "next" + && it.valueParameters.isEmpty() + && it.modality != Modality.ABSTRACT + } + + val next = receiverType.classifierOrFail.owner.cast().findNextFunction() ?: return + + val loopVariable = (loop.body as IrBlock).statements[0] as IrVariable + val loopCondition = loop.condition as IrCall + + iterator.apply { + type = receiverType + this.initializer = initializer.extensionReceiver + } + + loopCondition.apply { + dispatchReceiver = initializer.extensionReceiver!!.shallowCopy() + } + + loopVariable.initializer = with(next) { + context.createIrBuilder(symbol, startOffset, endOffset).irCall(this).apply { + dispatchReceiver = initializer.extensionReceiver!!.shallowCopy() + } + } + } + private data class LoopVariableInfo( val mainLoopVariable: IrVariable, val mainLoopVariableIndex: Int, diff --git a/compiler/testData/codegen/box/forLoop/forInIterator.kt b/compiler/testData/codegen/box/forLoop/forInIterator.kt new file mode 100644 index 00000000000..75f1f6bfe98 --- /dev/null +++ b/compiler/testData/codegen/box/forLoop/forInIterator.kt @@ -0,0 +1,28 @@ +// TARGET_BACKEND: JVM +// WITH_STDLIB +// CHECK_BYTECODE_LISTING + +fun myFun(iterator: MyUIntIterator) { + for (x in iterator) { + assert(x == 42u) + } +} + +class MyUIntIterator : Iterator { + private var count = 1 + override fun hasNext(): Boolean = count-- > 0 + override fun next(): UInt = 42u +} + +fun box(): String { + myFun(MyUIntIterator()) + return "OK" +} + +// CHECK_BYTECODE_TEXT +// JVM_IR_TEMPLATES +// 0 java/util/Iterator.hasNext \(\)Z +// 0 java/util/Iterator.next \(\)Ljava/lang/Object; +// 1 MyUIntIterator.hasNext \(\)Z +// 1 public synthetic bridge next\(\)Ljava/lang/Object; +// 2 MyUIntIterator.next-pVg5ArA \(\)I \ No newline at end of file diff --git a/compiler/testData/codegen/box/forLoop/forInIterator.txt b/compiler/testData/codegen/box/forLoop/forInIterator.txt new file mode 100644 index 00000000000..c1d5c03ebd2 --- /dev/null +++ b/compiler/testData/codegen/box/forLoop/forInIterator.txt @@ -0,0 +1,17 @@ +@kotlin.Metadata +public final class ForInIteratorKt { + // source: 'forInIterator.kt' + public final static @org.jetbrains.annotations.NotNull method box(): java.lang.String + public final static method myFun(@org.jetbrains.annotations.NotNull p0: MyUIntIterator): void +} + +@kotlin.Metadata +public final class MyUIntIterator { + // source: 'forInIterator.kt' + private field count: int + public method (): void + public method hasNext(): boolean + public synthetic bridge method next(): java.lang.Object + public method next-pVg5ArA(): int + public method remove(): void +} diff --git a/compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/codegen/BlackBoxCodegenTestGenerated.java b/compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/codegen/BlackBoxCodegenTestGenerated.java index 935f325c68d..1d4bd0a2397 100644 --- a/compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/codegen/BlackBoxCodegenTestGenerated.java +++ b/compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/codegen/BlackBoxCodegenTestGenerated.java @@ -16846,6 +16846,22 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest { } } + @Nested + @TestMetadata("compiler/testData/codegen/box/forLoop") + @TestDataPath("$PROJECT_ROOT") + public class ForLoop { + @Test + public void testAllFilesPresentInForLoop() throws Exception { + KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/testData/codegen/box/forLoop"), Pattern.compile("^(.+)\\.kt$"), null, TargetBackend.JVM, true); + } + + @Test + @TestMetadata("forInIterator.kt") + public void testForInIterator() throws Exception { + runTest("compiler/testData/codegen/box/forLoop/forInIterator.kt"); + } + } + @Nested @TestMetadata("compiler/testData/codegen/box/fullJdk") @TestDataPath("$PROJECT_ROOT") diff --git a/compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/codegen/IrBlackBoxCodegenTestGenerated.java b/compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/codegen/IrBlackBoxCodegenTestGenerated.java index 92c7579198d..bfc96faad76 100644 --- a/compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/codegen/IrBlackBoxCodegenTestGenerated.java +++ b/compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/codegen/IrBlackBoxCodegenTestGenerated.java @@ -17266,6 +17266,22 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes } } + @Nested + @TestMetadata("compiler/testData/codegen/box/forLoop") + @TestDataPath("$PROJECT_ROOT") + public class ForLoop { + @Test + public void testAllFilesPresentInForLoop() throws Exception { + KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/testData/codegen/box/forLoop"), Pattern.compile("^(.+)\\.kt$"), null, TargetBackend.JVM_IR, true); + } + + @Test + @TestMetadata("forInIterator.kt") + public void testForInIterator() throws Exception { + runTest("compiler/testData/codegen/box/forLoop/forInIterator.kt"); + } + } + @Nested @TestMetadata("compiler/testData/codegen/box/fullJdk") @TestDataPath("$PROJECT_ROOT") diff --git a/compiler/tests-gen/org/jetbrains/kotlin/codegen/LightAnalysisModeTestGenerated.java b/compiler/tests-gen/org/jetbrains/kotlin/codegen/LightAnalysisModeTestGenerated.java index 4265592878f..afde920f74e 100644 --- a/compiler/tests-gen/org/jetbrains/kotlin/codegen/LightAnalysisModeTestGenerated.java +++ b/compiler/tests-gen/org/jetbrains/kotlin/codegen/LightAnalysisModeTestGenerated.java @@ -13929,6 +13929,24 @@ public class LightAnalysisModeTestGenerated extends AbstractLightAnalysisModeTes } } + @TestMetadata("compiler/testData/codegen/box/forLoop") + @TestDataPath("$PROJECT_ROOT") + @RunWith(JUnit3RunnerWithInners.class) + public static class ForLoop extends AbstractLightAnalysisModeTest { + private void runTest(String testDataFilePath) throws Exception { + KotlinTestUtils.runTest(this::doTest, TargetBackend.JVM, testDataFilePath); + } + + public void testAllFilesPresentInForLoop() throws Exception { + KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/testData/codegen/box/forLoop"), Pattern.compile("^(.+)\\.kt$"), null, TargetBackend.JVM, true); + } + + @TestMetadata("forInIterator.kt") + public void testForInIterator() throws Exception { + runTest("compiler/testData/codegen/box/forLoop/forInIterator.kt"); + } + } + @TestMetadata("compiler/testData/codegen/box/fullJdk") @TestDataPath("$PROJECT_ROOT") @RunWith(JUnit3RunnerWithInners.class) diff --git a/js/js.tests/tests-gen/org/jetbrains/kotlin/js/test/JsCodegenBoxTestGenerated.java b/js/js.tests/tests-gen/org/jetbrains/kotlin/js/test/JsCodegenBoxTestGenerated.java index 839b74d522e..2cd8c1de76b 100644 --- a/js/js.tests/tests-gen/org/jetbrains/kotlin/js/test/JsCodegenBoxTestGenerated.java +++ b/js/js.tests/tests-gen/org/jetbrains/kotlin/js/test/JsCodegenBoxTestGenerated.java @@ -13014,6 +13014,16 @@ public class JsCodegenBoxTestGenerated extends AbstractJsCodegenBoxTest { } } + @Nested + @TestMetadata("compiler/testData/codegen/box/forLoop") + @TestDataPath("$PROJECT_ROOT") + public class ForLoop { + @Test + public void testAllFilesPresentInForLoop() throws Exception { + KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/testData/codegen/box/forLoop"), Pattern.compile("^(.+)\\.kt$"), null, TargetBackend.JS, true); + } + } + @Nested @TestMetadata("compiler/testData/codegen/box/fullJdk") @TestDataPath("$PROJECT_ROOT") diff --git a/js/js.tests/tests-gen/org/jetbrains/kotlin/js/test/ir/IrJsCodegenBoxTestGenerated.java b/js/js.tests/tests-gen/org/jetbrains/kotlin/js/test/ir/IrJsCodegenBoxTestGenerated.java index f956aaca8be..1ac57c08867 100644 --- a/js/js.tests/tests-gen/org/jetbrains/kotlin/js/test/ir/IrJsCodegenBoxTestGenerated.java +++ b/js/js.tests/tests-gen/org/jetbrains/kotlin/js/test/ir/IrJsCodegenBoxTestGenerated.java @@ -13056,6 +13056,16 @@ public class IrJsCodegenBoxTestGenerated extends AbstractIrJsCodegenBoxTest { } } + @Nested + @TestMetadata("compiler/testData/codegen/box/forLoop") + @TestDataPath("$PROJECT_ROOT") + public class ForLoop { + @Test + public void testAllFilesPresentInForLoop() throws Exception { + KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/testData/codegen/box/forLoop"), Pattern.compile("^(.+)\\.kt$"), null, TargetBackend.JS_IR, true); + } + } + @Nested @TestMetadata("compiler/testData/codegen/box/fullJdk") @TestDataPath("$PROJECT_ROOT") diff --git a/js/js.tests/tests-gen/org/jetbrains/kotlin/js/testOld/wasm/semantics/IrCodegenBoxWasmTestGenerated.java b/js/js.tests/tests-gen/org/jetbrains/kotlin/js/testOld/wasm/semantics/IrCodegenBoxWasmTestGenerated.java index 61b4a6f3374..f691beba5f2 100644 --- a/js/js.tests/tests-gen/org/jetbrains/kotlin/js/testOld/wasm/semantics/IrCodegenBoxWasmTestGenerated.java +++ b/js/js.tests/tests-gen/org/jetbrains/kotlin/js/testOld/wasm/semantics/IrCodegenBoxWasmTestGenerated.java @@ -10891,6 +10891,19 @@ public class IrCodegenBoxWasmTestGenerated extends AbstractIrCodegenBoxWasmTest } } + @TestMetadata("compiler/testData/codegen/box/forLoop") + @TestDataPath("$PROJECT_ROOT") + @RunWith(JUnit3RunnerWithInners.class) + public static class ForLoop extends AbstractIrCodegenBoxWasmTest { + private void runTest(String testDataFilePath) throws Exception { + KotlinTestUtils.runTest0(this::doTest, TargetBackend.WASM, testDataFilePath); + } + + public void testAllFilesPresentInForLoop() throws Exception { + KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/testData/codegen/box/forLoop"), Pattern.compile("^([^_](.+))\\.kt$"), null, TargetBackend.WASM, true); + } + } + @TestMetadata("compiler/testData/codegen/box/fullJdk") @TestDataPath("$PROJECT_ROOT") @RunWith(JUnit3RunnerWithInners.class) diff --git a/native/native.tests/tests-gen/org/jetbrains/kotlin/konan/blackboxtest/NativeCodegenBoxTestGenerated.java b/native/native.tests/tests-gen/org/jetbrains/kotlin/konan/blackboxtest/NativeCodegenBoxTestGenerated.java index 75b99345a57..a87266e7f7a 100644 --- a/native/native.tests/tests-gen/org/jetbrains/kotlin/konan/blackboxtest/NativeCodegenBoxTestGenerated.java +++ b/native/native.tests/tests-gen/org/jetbrains/kotlin/konan/blackboxtest/NativeCodegenBoxTestGenerated.java @@ -14030,6 +14030,18 @@ public class NativeCodegenBoxTestGenerated extends AbstractNativeCodegenBoxTest } } + @Nested + @TestMetadata("compiler/testData/codegen/box/forLoop") + @TestDataPath("$PROJECT_ROOT") + @Tag("codegen") + @UseExtTestCaseGroupProvider() + public class ForLoop { + @Test + public void testAllFilesPresentInForLoop() throws Exception { + KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/testData/codegen/box/forLoop"), Pattern.compile("^(.+)\\.kt$"), null, TargetBackend.NATIVE, true); + } + } + @Nested @TestMetadata("compiler/testData/codegen/box/fullJdk") @TestDataPath("$PROJECT_ROOT")