Optimize for-in-string loops

For-in-string loop can be generated using specialized 'length' and
'charAt' method calls, and with cached string length.
Note that update of the string variable in loop body doesn't affect
loop execution semantics.

 #KT-21322 Fixed Target versions 1.2.20

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

@@ -22,10 +22,11 @@ import org.jetbrains.kotlin.codegen.range.inExpression.CallBasedInExpressionGene
 import org.jetbrains.kotlin.codegen.range.inExpression.InExpressionGenerator
 import org.jetbrains.kotlin.psi.KtForExpression
 import org.jetbrains.kotlin.psi.KtSimpleNameExpression
+import org.jetbrains.org.objectweb.asm.Type
 
-class CharSequenceRangeValue : RangeValue {
+class CharSequenceRangeValue(private val canCacheLength: Boolean, private val charSequenceClassType: Type?) : RangeValue {
     override fun createForLoopGenerator(codegen: ExpressionCodegen, forExpression: KtForExpression) =
-            ForInCharSequenceLoopGenerator(codegen, forExpression)
+            ForInCharSequenceLoopGenerator(codegen, forExpression, canCacheLength, charSequenceClassType)
 
     override fun createInExpressionGenerator(codegen: ExpressionCodegen, operatorReference: KtSimpleNameExpression): InExpressionGenerator =
             CallBasedInExpressionGenerator(codegen, operatorReference)

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

@@ -25,6 +25,7 @@ import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.calls.callUtil.getResolvedCall
 import org.jetbrains.kotlin.resolve.calls.model.ResolvedCall
+import org.jetbrains.kotlin.resolve.jvm.AsmTypes
 import org.jetbrains.kotlin.types.KotlinType
 import org.jetbrains.kotlin.types.checker.KotlinTypeChecker
 import org.jetbrains.org.objectweb.asm.Type
@@ -43,6 +44,7 @@ fun ExpressionCodegen.createRangeValueForExpression(rangeExpression: KtExpressio
     val rangeType = bindingContext.getType(rangeExpression)!!
     val asmRangeType = asmType(rangeType)
 
+    val builtIns = state.module.builtIns
     return when {
         asmRangeType.sort == Type.ARRAY ->
             ArrayRangeValue(!isLocalVarReference(rangeExpression, bindingContext))
@@ -50,8 +52,10 @@ fun ExpressionCodegen.createRangeValueForExpression(rangeExpression: KtExpressio
             PrimitiveRangeRangeValue()
         isPrimitiveProgression(rangeType) ->
             PrimitiveProgressionRangeValue()
-        isSubtypeOfCharSequence(rangeType, state.module.builtIns) ->
+        isSubtypeOfString(rangeType, builtIns) ->
-            CharSequenceRangeValue()
+            CharSequenceRangeValue(true, AsmTypes.JAVA_STRING_TYPE)
+        isSubtypeOfCharSequence(rangeType, builtIns) ->
+            CharSequenceRangeValue(false, null)
         else ->
             IterableRangeValue()
     }
@@ -63,6 +67,9 @@ fun isLocalVarReference(rangeExpression: KtExpression, bindingContext: BindingCo
     return resultingDescriptor is LocalVariableDescriptor && resultingDescriptor.isVar
 }
 
+private fun isSubtypeOfString(type: KotlinType, builtIns: KotlinBuiltIns) =
+        KotlinTypeChecker.DEFAULT.isSubtypeOf(type, builtIns.stringType)
+
 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/ForInCharSequenceLoopGenerator.kt

@@ -22,12 +22,19 @@ import org.jetbrains.kotlin.psi.KtForExpression
 import org.jetbrains.kotlin.types.KotlinType
 import org.jetbrains.org.objectweb.asm.Label
 import org.jetbrains.org.objectweb.asm.Type
+import org.jetbrains.org.objectweb.asm.commons.InstructionAdapter
 
-class ForInCharSequenceLoopGenerator(codegen: ExpressionCodegen, forExpression: KtForExpression)
+class ForInCharSequenceLoopGenerator(
-    : AbstractForLoopGenerator(codegen, forExpression)
+        codegen: ExpressionCodegen,
-{
+        forExpression: KtForExpression,
+        private val canCacheLength: Boolean,
+        private val charSequenceClassType: Type?
+) : AbstractForLoopGenerator(codegen, forExpression) {
     private var indexVar: Int = 0
     private var charSequenceVar: Int = 0
+    private var charSequenceLengthVar: Int = 0
+
+    private val charSequenceType = charSequenceClassType ?: CHAR_SEQUENCE_TYPE
 
     override fun beforeLoop() {
         super.beforeLoop()
@@ -42,9 +49,16 @@ class ForInCharSequenceLoopGenerator(codegen: ExpressionCodegen, forExpression:
         // NB even if we already have a loop range stored in local variable, that variable might be modified in the loop body
         // (see controlStructures/forInCharSequenceMut.kt).
         // We should always store the corresponding CharSequence to a local variable to preserve the Iterator-based behavior.
-        charSequenceVar = createLoopTempVariable(CHAR_SEQUENCE_TYPE)
+        charSequenceVar = createLoopTempVariable(charSequenceType)
         value.put(asmLoopRangeType, v)
-        v.store(charSequenceVar, CHAR_SEQUENCE_TYPE)
+        v.store(charSequenceVar, charSequenceType)
+
+        if (canCacheLength) {
+            charSequenceLengthVar = createLoopTempVariable(Type.INT_TYPE)
+            v.load(charSequenceVar, charSequenceType)
+            v.invokeCharSequenceMethod("length", "()I")
+            v.store(charSequenceLengthVar, Type.INT_TYPE)
+        }
 
         v.iconst(0)
         v.store(indexVar, Type.INT_TYPE)
@@ -54,15 +68,20 @@ class ForInCharSequenceLoopGenerator(codegen: ExpressionCodegen, forExpression:
 
     override fun checkPreCondition(loopExit: Label) {
         v.load(indexVar, Type.INT_TYPE)
-        v.load(charSequenceVar, CHAR_SEQUENCE_TYPE)
+        if (canCacheLength) {
-        v.invokeinterface(CHAR_SEQUENCE_TYPE.internalName, "length", "()I")
+            v.load(charSequenceLengthVar, Type.INT_TYPE)
+        }
+        else {
+            v.load(charSequenceVar, charSequenceType)
+            v.invokeCharSequenceMethod("length", "()I")
+        }
         v.ificmpge(loopExit)
     }
 
     override fun assignToLoopParameter() {
-        v.load(charSequenceVar, CHAR_SEQUENCE_TYPE)
+        v.load(charSequenceVar, charSequenceType)
         v.load(indexVar, Type.INT_TYPE)
-        v.invokeinterface(CHAR_SEQUENCE_TYPE.internalName, "charAt", "(I)C")
+        v.invokeCharSequenceMethod("charAt", "(I)C")
         StackValue.onStack(Type.CHAR_TYPE).put(asmElementType, codegen.v)
         v.store(loopParameterVar, asmElementType)
     }
@@ -71,6 +90,16 @@ class ForInCharSequenceLoopGenerator(codegen: ExpressionCodegen, forExpression:
         v.iinc(indexVar, 1)
     }
 
+    private fun InstructionAdapter.invokeCharSequenceMethod(name: String, desc: String) {
+        val charSequenceClassType = charSequenceClassType
+        if (charSequenceClassType != null) {
+            invokevirtual(charSequenceClassType.internalName, name, desc, false)
+        }
+        else {
+            invokeinterface(CHAR_SEQUENCE_TYPE.internalName, name, desc)
+        }
+    }
+
     companion object {
         val CHAR_SEQUENCE_TYPE: Type = Type.getObjectType("java/lang/CharSequence")
     }

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

@@ -0,0 +1,11 @@
+// WITH_RUNTIME
+
+fun box(): String {
+    var str = "OK"
+    var r = ""
+    for (ch in str) {
+        r += ch
+        str = "zzz"
+    }
+    return r
+}

compiler/testData/codegen/bytecodeText/forLoop/forInStringSpecialized.kt

@@ -0,0 +1,15 @@
+// WITH_RUNTIME
+
+fun test() {
+    var s = ""
+    for (c in "testString") {
+        s += c
+    }
+}
+
+// 0 iterator
+// 0 hasNext
+// 0 nextChar
+// 0 INVOKEINTERFACE
+// 1 ISTORE 4
+// 1 ILOAD 4

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

@@ -14567,6 +14567,12 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
             doTest(fileName);
         }
 
+        @TestMetadata("forInStringVarUpdatedInLoopBody.kt")
+        public void testForInStringVarUpdatedInLoopBody() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/ranges/forInStringVarUpdatedInLoopBody.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("forIntRange.kt")
         public void testForIntRange() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/ranges/forIntRange.kt");

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

@@ -14567,6 +14567,12 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
             doTest(fileName);
         }
 
+        @TestMetadata("forInStringVarUpdatedInLoopBody.kt")
+        public void testForInStringVarUpdatedInLoopBody() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/ranges/forInStringVarUpdatedInLoopBody.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("forIntRange.kt")
         public void testForIntRange() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/ranges/forIntRange.kt");

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

@@ -1301,6 +1301,12 @@ public class BytecodeTextTestGenerated extends AbstractBytecodeTextTest {
             doTest(fileName);
         }
 
+        @TestMetadata("forInStringSpecialized.kt")
+        public void testForInStringSpecialized() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/bytecodeText/forLoop/forInStringSpecialized.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("forInUntil.kt")
         public void testForInUntil() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/bytecodeText/forLoop/forInUntil.kt");

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

@@ -14567,6 +14567,12 @@ public class LightAnalysisModeTestGenerated extends AbstractLightAnalysisModeTes
             doTest(fileName);
         }
 
+        @TestMetadata("forInStringVarUpdatedInLoopBody.kt")
+        public void testForInStringVarUpdatedInLoopBody() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/ranges/forInStringVarUpdatedInLoopBody.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("forIntRange.kt")
         public void testForIntRange() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/ranges/forIntRange.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 test*/
+/** This class is generated by {@link org.jetbrains.kotlin.generators.tests.TestsPackage}. DO NOT MODIFY MANUALLY */
 @SuppressWarnings("all")
 @TestMetadata("compiler/testData/codegen/box")
 @TestDataPath("$PROJECT_ROOT")
@@ -15839,6 +15839,12 @@ public class JsCodegenBoxTestGenerated extends AbstractJsCodegenBoxTest {
             doTest(fileName);
         }
 
+        @TestMetadata("forInStringVarUpdatedInLoopBody.kt")
+        public void testForInStringVarUpdatedInLoopBody() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/ranges/forInStringVarUpdatedInLoopBody.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("forIntRange.kt")
         public void testForIntRange() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/ranges/forIntRange.kt");