KT-5075 Optimize array/collection indices usage in 'for' loop

Use specialized 'for' loop code generation strategy for loops over array indices and collection indices.

compiler/backend/src/org/jetbrains/kotlin/codegen/ExpressionCodegen.java

@@ -584,6 +584,19 @@ public class ExpressionCodegen extends KtVisitor<StackValue, StackValue> impleme
             }
         }
 
+        ResolvedCall<? extends CallableDescriptor> loopRangeCall = RangeCodegenUtil.getLoopRangeResolvedCall(forExpression, bindingContext);
+        if (loopRangeCall != null) {
+            CallableDescriptor loopRangeCallee = loopRangeCall.getResultingDescriptor();
+            if (RangeCodegenUtil.isArrayOrPrimitiveArrayIndices(loopRangeCallee)) {
+                generateForLoop(createForInArrayIndicesRangeLoopGenerator(forExpression, loopRangeCall));
+                return StackValue.none();
+            }
+            else if (RangeCodegenUtil.isCollectionIndices(loopRangeCallee)) {
+                generateForLoop(createForInCollectionIndicesRangeLoopGenerator(forExpression, loopRangeCall));
+                return StackValue.none();
+            }
+        }
+
         KtExpression loopRange = forExpression.getLoopRange();
         assert loopRange != null;
         KotlinType loopRangeType = bindingContext.getType(loopRange);
@@ -608,6 +621,24 @@ public class ExpressionCodegen extends KtVisitor<StackValue, StackValue> impleme
         return StackValue.none();
     }
 
+    private AbstractForLoopGenerator createForInCollectionIndicesRangeLoopGenerator(
+            @NotNull KtForExpression forExpression,
+            @NotNull ResolvedCall<? extends CallableDescriptor> loopRangeCall
+    ) {
+        ReceiverValue extensionReceiver = loopRangeCall.getExtensionReceiver();
+        assert extensionReceiver != null : "Extension receiver should be non-null for optimizable 'indices' call";
+        return new ForInCollectionIndicesRangeLoopGenerator(forExpression, extensionReceiver);
+    }
+
+    private AbstractForLoopGenerator createForInArrayIndicesRangeLoopGenerator(
+            @NotNull KtForExpression forExpression,
+            @NotNull ResolvedCall<? extends CallableDescriptor> loopRangeCall
+    ) {
+        ReceiverValue extensionReceiver = loopRangeCall.getExtensionReceiver();
+        assert extensionReceiver != null : "Extension receiver should be non-null for optimizable 'indices' call";
+        return new ForInArrayIndicesRangeLoopGenerator(forExpression, extensionReceiver);
+    }
+
     private OwnerKind contextKind() {
         return context.getContextKind();
     }
@@ -1083,6 +1114,54 @@ public class ExpressionCodegen extends KtVisitor<StackValue, StackValue> impleme
         }
     }
 
+    private abstract class ForInOptimizedIndicesLoopGenerator extends AbstractForInRangeLoopGenerator {
+        protected final ReceiverValue receiverValue;
+
+        private ForInOptimizedIndicesLoopGenerator(@NotNull KtForExpression forExpression, @NotNull ReceiverValue receiverValue) {
+            super(forExpression);
+            this.receiverValue = receiverValue;
+        }
+
+        @Override
+        protected void storeRangeStartAndEnd() {
+            StackValue.local(loopParameterVar, loopParameterType).store(StackValue.constant(0, asmElementType), v);
+
+            StackValue receiver = generateReceiverValue(receiverValue, false);
+            receiver.put(receiver.type, v);
+            getReceiverSizeAsInt();
+            v.iconst(1);
+            v.sub(Type.INT_TYPE);
+            StackValue.local(endVar, asmElementType).store(StackValue.onStack(Type.INT_TYPE), v);
+        }
+
+        /**
+         * <code>(receiver -> size:I)</code>
+         */
+        protected abstract void getReceiverSizeAsInt();
+    }
+
+    private class ForInCollectionIndicesRangeLoopGenerator extends ForInOptimizedIndicesLoopGenerator {
+        private ForInCollectionIndicesRangeLoopGenerator(@NotNull KtForExpression forExpression, @NotNull ReceiverValue receiverValue) {
+            super(forExpression, receiverValue);
+        }
+
+        @Override
+        protected void getReceiverSizeAsInt() {
+            v.invokeinterface("java/util/Collection", "size", "()I");
+        }
+    }
+
+    private class ForInArrayIndicesRangeLoopGenerator extends ForInOptimizedIndicesLoopGenerator {
+        private ForInArrayIndicesRangeLoopGenerator(@NotNull KtForExpression forExpression, @NotNull ReceiverValue receiverValue) {
+            super(forExpression, receiverValue);
+        }
+
+        @Override
+        protected void getReceiverSizeAsInt() {
+            v.arraylength();
+        }
+    }
+
     private class ForInProgressionExpressionLoopGenerator extends AbstractForInProgressionOrRangeLoopGenerator {
         private int incrementVar;
         private Type incrementType;
@@ -2873,7 +2952,7 @@ public class ExpressionCodegen extends KtVisitor<StackValue, StackValue> impleme
         if (hasSpread) {
             boolean arrayOfReferences = KotlinBuiltIns.isArray(outType);
             if (size == 1) {
-                // Arrays.copyOf(array, newLength)
+                // Arrays.copyOf(receiverValue, newLength)
                 ValueArgument argument = arguments.get(0);
                 Type arrayType = arrayOfReferences ? Type.getType("[Ljava/lang/Object;")
                                                    : Type.getType("[" + elementType.getDescriptor());

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

@@ -19,14 +19,17 @@ package org.jetbrains.kotlin.codegen;
 import com.google.common.collect.ImmutableMap;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
+import org.jetbrains.kotlin.builtins.KotlinBuiltIns;
 import org.jetbrains.kotlin.builtins.PrimitiveType;
-import org.jetbrains.kotlin.descriptors.CallableDescriptor;
-import org.jetbrains.kotlin.descriptors.ClassifierDescriptor;
+import org.jetbrains.kotlin.descriptors.*;
 import org.jetbrains.kotlin.name.FqName;
 import org.jetbrains.kotlin.name.FqNameUnsafe;
 import org.jetbrains.kotlin.name.Name;
 import org.jetbrains.kotlin.psi.*;
+import org.jetbrains.kotlin.resolve.BindingContext;
 import org.jetbrains.kotlin.resolve.DescriptorUtils;
+import org.jetbrains.kotlin.resolve.calls.callUtil.CallUtilKt;
+import org.jetbrains.kotlin.resolve.calls.model.ResolvedCall;
 import org.jetbrains.kotlin.types.KotlinType;
 
 import java.util.Arrays;
@@ -97,6 +100,24 @@ public class RangeCodegenUtil {
         return null;
     }
 
+    @Nullable
+    public static ResolvedCall<? extends CallableDescriptor> getLoopRangeResolvedCall(@NotNull KtForExpression forExpression, @NotNull BindingContext bindingContext) {
+        KtExpression loopRange = KtPsiUtil.deparenthesize(forExpression.getLoopRange());
+
+        if (loopRange instanceof KtQualifiedExpression) {
+            KtQualifiedExpression qualifiedExpression = (KtQualifiedExpression) loopRange;
+            KtExpression selector = qualifiedExpression.getSelectorExpression();
+            if (selector instanceof KtCallExpression || selector instanceof KtSimpleNameExpression) {
+                return CallUtilKt.getResolvedCall(selector, bindingContext);
+            }
+        }
+        else if (loopRange instanceof KtSimpleNameExpression) {
+            return CallUtilKt.getResolvedCall(loopRange, bindingContext);
+        }
+
+        return null;
+    }
+
     @Nullable
     private static PrimitiveType getPrimitiveRangeElementType(KotlinType rangeType) {
         return getPrimitiveRangeOrProgressionElementType(rangeType, RANGE_TO_ELEMENT_TYPE);
@@ -138,6 +159,39 @@ public class RangeCodegenUtil {
         return false;
     }
 
+    public static boolean isArrayOrPrimitiveArrayIndices(@NotNull CallableDescriptor descriptor) {
+        if (!isTopLevelInPackage(descriptor, "indices", "kotlin.collections")) return false;
+
+        ReceiverParameterDescriptor extensionReceiver = descriptor.getExtensionReceiverParameter();
+        if (extensionReceiver == null) return false;
+        KotlinType extensionReceiverType = extensionReceiver.getType();
+        if (!KotlinBuiltIns.isArray(extensionReceiverType) && !KotlinBuiltIns.isPrimitiveArray(extensionReceiverType)) return false;
+
+        return true;
+    }
+
+    public static boolean isCollectionIndices(@NotNull CallableDescriptor descriptor) {
+        if (!isTopLevelInPackage(descriptor, "indices", "kotlin.collections")) return false;
+
+        ReceiverParameterDescriptor extensionReceiver = descriptor.getExtensionReceiverParameter();
+        if (extensionReceiver == null) return false;
+        KotlinType extensionReceiverType = extensionReceiver.getType();
+        if (!KotlinBuiltIns.isCollectionOrNullableCollection(extensionReceiverType)) return false;
+
+        return true;
+    }
+
+    private static boolean isTopLevelInPackage(@NotNull CallableDescriptor descriptor, @NotNull String name, @NotNull String packageName) {
+        if (!name.equals(descriptor.getName().asString())) return false;
+
+        DeclarationDescriptor containingDeclaration = descriptor.getContainingDeclaration();
+        if (!(containingDeclaration instanceof PackageFragmentDescriptor)) return false;
+        String packageFqName = ((PackageFragmentDescriptor) containingDeclaration).getFqName().asString();
+        if (!packageName.equals(packageFqName)) return false;
+
+        return true;
+    }
+
     public static class BinaryCall {
         public final KtExpression left;
         public final KtExpression op;

compiler/testData/codegen/box/ranges/forInIndices/forInCollectionImplicitReceiverIndices.kt

@@ -0,0 +1,19 @@
+// WITH_RUNTIME
+
+import kotlin.test.assertEquals
+
+fun Collection<Int>.sumIndices(): Int {
+    var sum = 0
+    for (i in indices) {
+        sum += i
+    }
+    return sum
+}
+
+fun box(): String {
+    val list = listOf(0, 0, 0, 0)
+    val sum = list.sumIndices()
+    assertEquals(6, sum)
+
+    return "OK"
+}

compiler/testData/codegen/box/ranges/forInIndices/forInCollectionIndices.kt

@@ -0,0 +1,13 @@
+// WITH_RUNTIME
+
+import kotlin.test.assertEquals
+
+fun box(): String {
+    var sum = 0
+    for (i in listOf(0, 0, 0, 0).indices) {
+        sum += i
+    }
+    assertEquals(6, sum)
+
+    return "OK"
+}

compiler/testData/codegen/box/ranges/forInIndices/forInNonOptimizedIndices.kt

@@ -0,0 +1,18 @@
+// WITH_RUNTIME
+
+import kotlin.test.assertEquals
+
+fun sumIndices(coll: Collection<*>?): Int {
+    var sum = 0
+    for (i in coll?.indices ?: return 0) {
+        sum += i
+    }
+    return sum
+}
+
+fun box(): String {
+    assertEquals(6, sumIndices(listOf(0, 0, 0, 0)))
+    assertEquals(0, sumIndices(null))
+
+    return "OK"
+}

compiler/testData/codegen/box/ranges/forInIndices/forInObjectArrayIndices.kt

@@ -0,0 +1,13 @@
+// WITH_RUNTIME
+
+import kotlin.test.assertEquals
+
+fun box(): String {
+    var sum = 0
+    for (i in arrayOf("", "", "", "").indices) {
+        sum += i
+    }
+    assertEquals(6, sum)
+
+    return "OK"
+}

compiler/testData/codegen/box/ranges/forInIndices/forInPrimitiveArrayIndices.kt

@@ -0,0 +1,13 @@
+// WITH_RUNTIME
+
+import kotlin.test.assertEquals
+
+fun box(): String {
+    var sum = 0
+    for (i in intArrayOf(0, 0, 0, 0).indices) {
+        sum += i
+    }
+    assertEquals(6, sum)
+
+    return "OK"
+}

compiler/testData/codegen/box/ranges/forInIndices/forNullableIntInArrayIndices.kt

@@ -0,0 +1,16 @@
+// WITH_RUNTIME
+
+import kotlin.test.assertEquals
+
+fun suppressBoxingOptimization(ni: Int?) {}
+
+fun box(): String {
+    var sum = 0
+    for (i: Int? in arrayOf("", "", "", "").indices) {
+        suppressBoxingOptimization(i)
+        sum += i ?: 0
+    }
+    assertEquals(6, sum)
+
+    return "OK"
+}

compiler/testData/codegen/box/ranges/forInIndices/forNullableIntInCollectionIndices.kt

@@ -0,0 +1,16 @@
+// WITH_RUNTIME
+
+import kotlin.test.assertEquals
+
+fun suppressBoxingOptimization(ni: Int?) {}
+
+fun box(): String {
+    var sum = 0
+    for (i: Int? in listOf("", "", "", "").indices) {
+        suppressBoxingOptimization(i)
+        sum += i ?: 0
+    }
+    assertEquals(6, sum)
+
+    return "OK"
+}

compiler/testData/codegen/bytecodeText/forLoop/forInIndices/forInCollectionImplicitReceiverIndices.kt

@@ -0,0 +1,15 @@
+// WITH_RUNTIME
+
+fun Collection<Int>.sumIndices(): Int {
+    var sum = 0
+    for (i in indices) {
+        sum += i
+    }
+    return sum
+}
+
+// 0 iterator
+// 0 getStart
+// 0 getEnd
+// 0 getFirst
+// 0 getLast

compiler/testData/codegen/bytecodeText/forLoop/forInIndices/forInCollectionIndices.kt

@@ -0,0 +1,14 @@
+// WITH_RUNTIME
+
+fun test() {
+    var sum = 0
+    for (i in listOf(0, 0, 0, 0).indices) {
+        sum += i
+    }
+}
+
+// 0 iterator
+// 0 getStart
+// 0 getEnd
+// 0 getFirst
+// 0 getLast

compiler/testData/codegen/bytecodeText/forLoop/forInIndices/forInNonOptimizedIndices.kt

@@ -0,0 +1,18 @@
+// WITH_RUNTIME
+
+import kotlin.test.assertEquals
+
+fun test(coll: Collection<*>?): Int {
+    var sum = 0
+    for (i in coll?.indices ?: return 0) {
+        sum += i
+    }
+    return sum
+}
+
+// 0 iterator
+// 0 getStart
+// 0 getEnd
+// 1 getIndices
+// 1 getFirst
+// 1 getLast

compiler/testData/codegen/bytecodeText/forLoop/forInIndices/forInObjectArrayIndices.kt

@@ -0,0 +1,14 @@
+// WITH_RUNTIME
+
+fun test() {
+    var sum = 0
+    for (i in arrayOf("", "", "", "").indices) {
+        sum += i
+    }
+}
+
+// 0 iterator
+// 0 getStart
+// 0 getEnd
+// 0 getFirst
+// 0 getLast

compiler/testData/codegen/bytecodeText/forLoop/forInIndices/forInPrimitiveArrayIndices.kt

@@ -0,0 +1,14 @@
+// WITH_RUNTIME
+
+fun test() {
+    var sum = 0
+    for (i in intArrayOf(0, 0, 0, 0).indices) {
+        sum += i
+    }
+}
+
+// 0 iterator
+// 0 getStart
+// 0 getEnd
+// 0 getFirst
+// 0 getLast

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

@@ -10687,6 +10687,57 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
             }
         }
 
+        @TestMetadata("compiler/testData/codegen/box/ranges/forInIndices")
+        @TestDataPath("$PROJECT_ROOT")
+        @RunWith(JUnit3RunnerWithInners.class)
+        public static class ForInIndices extends AbstractBlackBoxCodegenTest {
+            public void testAllFilesPresentInForInIndices() throws Exception {
+                KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("compiler/testData/codegen/box/ranges/forInIndices"), Pattern.compile("^(.+)\\.kt$"), true);
+            }
+
+            @TestMetadata("forInCollectionImplicitReceiverIndices.kt")
+            public void testForInCollectionImplicitReceiverIndices() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/ranges/forInIndices/forInCollectionImplicitReceiverIndices.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("forInCollectionIndices.kt")
+            public void testForInCollectionIndices() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/ranges/forInIndices/forInCollectionIndices.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("forInNonOptimizedIndices.kt")
+            public void testForInNonOptimizedIndices() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/ranges/forInIndices/forInNonOptimizedIndices.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("forInObjectArrayIndices.kt")
+            public void testForInObjectArrayIndices() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/ranges/forInIndices/forInObjectArrayIndices.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("forInPrimitiveArrayIndices.kt")
+            public void testForInPrimitiveArrayIndices() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/ranges/forInIndices/forInPrimitiveArrayIndices.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("forNullableIntInArrayIndices.kt")
+            public void testForNullableIntInArrayIndices() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/ranges/forInIndices/forNullableIntInArrayIndices.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("forNullableIntInCollectionIndices.kt")
+            public void testForNullableIntInCollectionIndices() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/ranges/forInIndices/forNullableIntInCollectionIndices.kt");
+                doTest(fileName);
+            }
+        }
+
         @TestMetadata("compiler/testData/codegen/box/ranges/literal")
         @TestDataPath("$PROJECT_ROOT")
         @RunWith(JUnit3RunnerWithInners.class)

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

@@ -882,6 +882,45 @@ public class BytecodeTextTestGenerated extends AbstractBytecodeTextTest {
             String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/bytecodeText/forLoop/primitiveRange.kt");
             doTest(fileName);
         }
+
+        @TestMetadata("compiler/testData/codegen/bytecodeText/forLoop/forInIndices")
+        @TestDataPath("$PROJECT_ROOT")
+        @RunWith(JUnit3RunnerWithInners.class)
+        public static class ForInIndices extends AbstractBytecodeTextTest {
+            public void testAllFilesPresentInForInIndices() throws Exception {
+                KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("compiler/testData/codegen/bytecodeText/forLoop/forInIndices"), Pattern.compile("^(.+)\\.kt$"), true);
+            }
+
+            @TestMetadata("forInCollectionImplicitReceiverIndices.kt")
+            public void testForInCollectionImplicitReceiverIndices() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/bytecodeText/forLoop/forInIndices/forInCollectionImplicitReceiverIndices.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("forInCollectionIndices.kt")
+            public void testForInCollectionIndices() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/bytecodeText/forLoop/forInIndices/forInCollectionIndices.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("forInNonOptimizedIndices.kt")
+            public void testForInNonOptimizedIndices() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/bytecodeText/forLoop/forInIndices/forInNonOptimizedIndices.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("forInObjectArrayIndices.kt")
+            public void testForInObjectArrayIndices() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/bytecodeText/forLoop/forInIndices/forInObjectArrayIndices.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("forInPrimitiveArrayIndices.kt")
+            public void testForInPrimitiveArrayIndices() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/bytecodeText/forLoop/forInIndices/forInPrimitiveArrayIndices.kt");
+                doTest(fileName);
+            }
+        }
     }
 
     @TestMetadata("compiler/testData/codegen/bytecodeText/inline")