Cache array length in for-in-array loop if possible

If the range expression is not a local variable (which can be updated in
the loop body affecting loop behavior, see KT-21354), we can cache the
array length, thus turning a for-in-array loop into a simple optimizable
counter loop.

 #KT-21321 In Progress

compiler/backend/src/org/jetbrains/kotlin/codegen/range/ArrayRangeValue.kt

@@ -23,9 +23,9 @@ import org.jetbrains.kotlin.codegen.range.inExpression.InExpressionGenerator
 import org.jetbrains.kotlin.psi.KtForExpression
 import org.jetbrains.kotlin.psi.KtSimpleNameExpression
 
-class ArrayRangeValue : RangeValue {
+class ArrayRangeValue(private val canCacheArrayLength: Boolean) : RangeValue {
     override fun createForLoopGenerator(codegen: ExpressionCodegen, forExpression: KtForExpression) =
-            ForInArrayLoopGenerator(codegen, forExpression)
+            ForInArrayLoopGenerator(codegen, forExpression, canCacheArrayLength)
 
     override fun createInExpressionGenerator(codegen: ExpressionCodegen, operatorReference: KtSimpleNameExpression): InExpressionGenerator =
             CallBasedInExpressionGenerator(codegen, operatorReference)

compiler/backend/src/org/jetbrains/kotlin/codegen/range/RangeValues.kt

@@ -19,6 +19,7 @@ package org.jetbrains.kotlin.codegen.range
 import org.jetbrains.kotlin.builtins.KotlinBuiltIns
 import org.jetbrains.kotlin.codegen.*
 import org.jetbrains.kotlin.descriptors.CallableDescriptor
+import org.jetbrains.kotlin.descriptors.impl.LocalVariableDescriptor
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.resolve.BindingContext
@@ -44,7 +45,7 @@ fun ExpressionCodegen.createRangeValueForExpression(rangeExpression: KtExpressio
 
     return when {
         asmRangeType.sort == Type.ARRAY ->
-            ArrayRangeValue()
+            ArrayRangeValue(!isLocalVarReference(rangeExpression, bindingContext))
         isPrimitiveRange(rangeType) ->
             PrimitiveRangeRangeValue()
         isPrimitiveProgression(rangeType) ->
@@ -56,6 +57,12 @@ fun ExpressionCodegen.createRangeValueForExpression(rangeExpression: KtExpressio
     }
 }
 
+fun isLocalVarReference(rangeExpression: KtExpression, bindingContext: BindingContext): Boolean {
+    if (rangeExpression !is KtSimpleNameExpression) return false
+    val resultingDescriptor = rangeExpression.getResolvedCall(bindingContext)?.resultingDescriptor ?: return false
+    return resultingDescriptor is LocalVariableDescriptor && resultingDescriptor.isVar
+}
+
 private fun isSubtypeOfCharSequence(type: KotlinType, builtIns: KotlinBuiltIns) =
         KotlinTypeChecker.DEFAULT.isSubtypeOf(type, builtIns.getBuiltInClassByName(Name.identifier("CharSequence")).defaultType)
 

compiler/backend/src/org/jetbrains/kotlin/codegen/range/forLoop/ForInArrayLoopGenerator.kt

@@ -17,21 +17,25 @@
 package org.jetbrains.kotlin.codegen.range.forLoop
 
 import org.jetbrains.kotlin.builtins.KotlinBuiltIns
+import org.jetbrains.kotlin.codegen.AsmUtil.boxType
+import org.jetbrains.kotlin.codegen.ExpressionCodegen
+import org.jetbrains.kotlin.codegen.StackValue
 import org.jetbrains.kotlin.psi.KtForExpression
+import org.jetbrains.kotlin.resolve.jvm.AsmTypes.OBJECT_TYPE
 import org.jetbrains.kotlin.types.KotlinType
 import org.jetbrains.org.objectweb.asm.Label
 import org.jetbrains.org.objectweb.asm.Type
 
-import org.jetbrains.kotlin.codegen.AsmUtil.boxType
+class ForInArrayLoopGenerator(
-import org.jetbrains.kotlin.codegen.ExpressionCodegen
+        codegen: ExpressionCodegen,
-import org.jetbrains.kotlin.codegen.StackValue
+        forExpression: KtForExpression,
-import org.jetbrains.kotlin.resolve.jvm.AsmTypes.OBJECT_TYPE
+        // We can cache array length if the corresponding range expression is not a local var.
-
+        // See https://youtrack.jetbrains.com/issue/KT-21354.
-class ForInArrayLoopGenerator(codegen: ExpressionCodegen, forExpression: KtForExpression)
+        private val canCacheArrayLength: Boolean
-    : AbstractForLoopGenerator(codegen, forExpression)
+) : AbstractForLoopGenerator(codegen, forExpression) {
-{
     private var indexVar: Int = 0
     private var arrayVar: Int = 0
+    private var arrayLengthVar: Int = 0
     private val loopRangeType: KotlinType = bindingContext.getType(forExpression.loopRange!!)!!
 
     override fun beforeLoop() {
@@ -51,6 +55,13 @@ class ForInArrayLoopGenerator(codegen: ExpressionCodegen, forExpression: KtForEx
             v.store(arrayVar, OBJECT_TYPE)
         }
 
+        if (canCacheArrayLength) {
+            arrayLengthVar = createLoopTempVariable(Type.INT_TYPE)
+            v.load(arrayVar, OBJECT_TYPE)
+            v.arraylength()
+            v.store(arrayLengthVar, Type.INT_TYPE)
+        }
+
         v.iconst(0)
         v.store(indexVar, Type.INT_TYPE)
     }
@@ -59,8 +70,13 @@ class ForInArrayLoopGenerator(codegen: ExpressionCodegen, forExpression: KtForEx
 
     override fun checkPreCondition(loopExit: Label) {
         v.load(indexVar, Type.INT_TYPE)
-        v.load(arrayVar, OBJECT_TYPE)
+        if (canCacheArrayLength) {
-        v.arraylength()
+            v.load(arrayLengthVar, Type.INT_TYPE)
+        }
+        else {
+            v.load(arrayVar, OBJECT_TYPE)
+            v.arraylength()
+        }
         v.ificmpge(loopExit)
     }
 

compiler/testData/checkLocalVariablesTable/underscoreNames.kt

@@ -21,7 +21,7 @@ fun box() {
 }
 
 // METHOD : UnderscoreNamesKt$box$1.invoke(LA;Ljava/lang/String;I)Ljava/lang/String;
-// VARIABLE : NAME=q TYPE=Ljava/lang/String; INDEX=15
+// VARIABLE : NAME=q TYPE=Ljava/lang/String; INDEX=16
 // VARIABLE : NAME=d TYPE=C INDEX=11
 // VARIABLE : NAME=_ TYPE=Ljava/lang/String; INDEX=10
 // VARIABLE : NAME=c TYPE=C INDEX=9

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

@@ -0,0 +1,12 @@
+// WITH_RUNTIME
+
+var xs = intArrayOf(1, 2, 3)
+
+fun box(): String {
+    var sum = 0
+    for (x in xs) {
+        sum = sum * 10 + x
+        xs = IntArray(0)
+    }
+    return if (sum == 123) "OK" else "Fail: $sum"
+}

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

@@ -0,0 +1,23 @@
+// WITH_RUNTIME
+// IGNORE_BACKEND: JS
+
+// In Kotlin 1.0, in a for-in-array loop, range expression is cached in
+// a local variable unless it is already a local variable.
+// This caused the following quirky behavior:
+// if an array variable is updated in the loop body, it affects the loop
+// execution (see https://youtrack.jetbrains.com/issue/KT-21354).
+// Most likely it is a bug, however, fixing it right now is a breaking
+// change requiring a proper deprecation loop.
+// When the design decision is made, it might be required to update this
+// test (e.g., by adding a proper LANGUAGE_VERSION directive).
+// Note that JS back-end handles this case "correctly".
+
+fun box(): String {
+    var xs = intArrayOf(1, 2, 3)
+    var sum = 0
+    for (x in xs) {
+        sum = sum * 10 + x
+        xs = intArrayOf(4, 5)
+    }
+    return if (sum == 15) "OK" else "Fail: $sum"
+}

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

@@ -14537,6 +14537,18 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
             doTest(fileName);
         }
 
+        @TestMetadata("forInArrayWithArrayPropertyUpdatedInLoopBody.kt")
+        public void testForInArrayWithArrayPropertyUpdatedInLoopBody() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/ranges/forInArrayWithArrayPropertyUpdatedInLoopBody.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("forInArrayWithArrayVarUpdatedInLoopBody.kt")
+        public void testForInArrayWithArrayVarUpdatedInLoopBody() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/ranges/forInArrayWithArrayVarUpdatedInLoopBody.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("forInRangeLiteralWithMixedTypeBounds.kt")
         public void testForInRangeLiteralWithMixedTypeBounds() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/ranges/forInRangeLiteralWithMixedTypeBounds.kt");

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

@@ -14537,6 +14537,18 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
             doTest(fileName);
         }
 
+        @TestMetadata("forInArrayWithArrayPropertyUpdatedInLoopBody.kt")
+        public void testForInArrayWithArrayPropertyUpdatedInLoopBody() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/ranges/forInArrayWithArrayPropertyUpdatedInLoopBody.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("forInArrayWithArrayVarUpdatedInLoopBody.kt")
+        public void testForInArrayWithArrayVarUpdatedInLoopBody() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/ranges/forInArrayWithArrayVarUpdatedInLoopBody.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("forInRangeLiteralWithMixedTypeBounds.kt")
         public void testForInRangeLiteralWithMixedTypeBounds() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/ranges/forInRangeLiteralWithMixedTypeBounds.kt");

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

@@ -14537,6 +14537,18 @@ public class LightAnalysisModeTestGenerated extends AbstractLightAnalysisModeTes
             doTest(fileName);
         }
 
+        @TestMetadata("forInArrayWithArrayPropertyUpdatedInLoopBody.kt")
+        public void testForInArrayWithArrayPropertyUpdatedInLoopBody() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/ranges/forInArrayWithArrayPropertyUpdatedInLoopBody.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("forInArrayWithArrayVarUpdatedInLoopBody.kt")
+        public void testForInArrayWithArrayVarUpdatedInLoopBody() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/ranges/forInArrayWithArrayVarUpdatedInLoopBody.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("forInRangeLiteralWithMixedTypeBounds.kt")
         public void testForInRangeLiteralWithMixedTypeBounds() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/ranges/forInRangeLiteralWithMixedTypeBounds.kt");

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

@@ -26,7 +26,7 @@ import org.junit.runner.RunWith;
 import java.io.File;
 import java.util.regex.Pattern;
 
-/** This class is generated by {@link org.jetbrains.kotlin.generators.tests.TestsPackage}. DO NOT MODIFY MANUALLY */
+/** This class is generated by {@link org.jetbrains.kotlin.generators.tests.TestsPackage}. DO NOT MODIFY MANUALLY test*/
 @SuppressWarnings("all")
 @TestMetadata("compiler/testData/codegen/box")
 @TestDataPath("$PROJECT_ROOT")
@@ -15803,6 +15803,24 @@ public class JsCodegenBoxTestGenerated extends AbstractJsCodegenBoxTest {
             doTest(fileName);
         }
 
+        @TestMetadata("forInArrayWithArrayPropertyUpdatedInLoopBody.kt")
+        public void testForInArrayWithArrayPropertyUpdatedInLoopBody() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/ranges/forInArrayWithArrayPropertyUpdatedInLoopBody.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("forInArrayWithArrayVarUpdatedInLoopBody.kt")
+        public void testForInArrayWithArrayVarUpdatedInLoopBody() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/ranges/forInArrayWithArrayVarUpdatedInLoopBody.kt");
+            try {
+                doTest(fileName);
+            }
+            catch (Throwable ignore) {
+                return;
+            }
+            throw new AssertionError("Looks like this test can be unmuted. Remove IGNORE_BACKEND directive for that.");
+        }
+
         @TestMetadata("forInRangeLiteralWithMixedTypeBounds.kt")
         public void testForInRangeLiteralWithMixedTypeBounds() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/ranges/forInRangeLiteralWithMixedTypeBounds.kt");