JVM_IR: Specialize iterator if possible

fixes KT-47171
2022-02-22 17:51:58 +08:00
parent 8bcdc70adc
commit 1b776bd5b6
11 changed files with 230 additions and 7 deletions
@@ -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")
@@ -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 <T> Iterator<T>.iterator(): Iterator<T> = 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<T> 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<IrCall>().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<IrClass>().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,
@@ -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<UInt> {
+    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
@@ -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 <init>(): void
+    public method hasNext(): boolean
+    public synthetic bridge method next(): java.lang.Object
+    public method next-pVg5ArA(): int
+    public method remove(): void
+}
@@ -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")
@@ -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")
@@ -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)
@@ -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")
@@ -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")
@@ -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)
@@ -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")