Optimize iteration over CharSequence's on JVM

KT-7931 Optimize iteration over strings/charsequences on JVM

compiler/backend/src/org/jetbrains/kotlin/codegen/forLoop/ForInCharSequenceLoopGenerator.kt

@@ -0,0 +1,77 @@
+/*
+ * Copyright 2010-2017 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.jetbrains.kotlin.codegen.forLoop
+
+import org.jetbrains.kotlin.codegen.ExpressionCodegen
+import org.jetbrains.kotlin.codegen.StackValue
+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
+
+class ForInCharSequenceLoopGenerator(codegen: ExpressionCodegen, forExpression: KtForExpression)
+    : AbstractForLoopGenerator(codegen, forExpression)
+{
+    private var indexVar: Int = 0
+    private var charSequenceVar: Int = 0
+
+    override fun beforeLoop() {
+        super.beforeLoop()
+
+        indexVar = createLoopTempVariable(Type.INT_TYPE)
+
+        val loopRange = forExpression.loopRange
+        val value = codegen.gen(loopRange)
+        val loopRangeType: KotlinType = bindingContext.getType(forExpression.loopRange!!)!!
+        val asmLoopRangeType = codegen.asmType(loopRangeType)
+
+        // 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)
+        value.put(asmLoopRangeType, v)
+        v.store(charSequenceVar, CHAR_SEQUENCE_TYPE)
+
+        v.iconst(0)
+        v.store(indexVar, Type.INT_TYPE)
+    }
+
+    override fun checkEmptyLoop(loopExit: Label) {}
+
+    override fun checkPreCondition(loopExit: Label) {
+        v.load(indexVar, Type.INT_TYPE)
+        v.load(charSequenceVar, CHAR_SEQUENCE_TYPE)
+        v.invokeinterface(CHAR_SEQUENCE_TYPE.internalName, "length", "()I")
+        v.ificmpge(loopExit)
+    }
+
+    override fun assignToLoopParameter() {
+        v.load(charSequenceVar, CHAR_SEQUENCE_TYPE)
+        v.load(indexVar, Type.INT_TYPE)
+        v.invokeinterface(CHAR_SEQUENCE_TYPE.internalName, "charAt", "(I)C")
+        StackValue.onStack(Type.CHAR_TYPE).put(asmElementType, codegen.v)
+        v.store(loopParameterVar, asmElementType)
+    }
+
+    override fun increment(loopExit: Label) {
+        v.iinc(indexVar, 1)
+    }
+
+    companion object {
+        val CHAR_SEQUENCE_TYPE: Type = Type.getObjectType("java/lang/CharSequence")
+    }
+}

compiler/backend/src/org/jetbrains/kotlin/codegen/forLoop/ForLoopGenerators.kt

@@ -16,13 +16,17 @@
 
 package org.jetbrains.kotlin.codegen.forLoop
 
+import org.jetbrains.kotlin.builtins.KotlinBuiltIns
 import org.jetbrains.kotlin.codegen.ExpressionCodegen
 import org.jetbrains.kotlin.codegen.RangeCodegenUtil
 import org.jetbrains.kotlin.descriptors.CallableDescriptor
+import org.jetbrains.kotlin.name.Name
 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.types.KotlinType
+import org.jetbrains.kotlin.types.checker.KotlinTypeChecker
 import org.jetbrains.org.objectweb.asm.Type
 
 fun ExpressionCodegen.getForLoopGenerator(forExpression: KtForExpression) : AbstractForLoopGenerator {
@@ -43,11 +47,16 @@ fun ExpressionCodegen.getForLoopGenerator(forExpression: KtForExpression) : Abst
             ForInRangeInstanceLoopGenerator(this, forExpression)
         RangeCodegenUtil.isProgression(loopRangeType) ->
             ForInProgressionExpressionLoopGenerator(this, forExpression)
+        isSubtypeOfCharSequence(loopRangeType, state.module.builtIns) ->
+            ForInCharSequenceLoopGenerator(this, forExpression)
         else ->
             IteratorForLoopGenerator(this, forExpression)
     }
 }
 
+private fun isSubtypeOfCharSequence(type: KotlinType, builtIns: KotlinBuiltIns) =
+        KotlinTypeChecker.DEFAULT.isSubtypeOf(type, builtIns.getBuiltInClassByName(Name.identifier("CharSequence")).defaultType)
+
 private fun ExpressionCodegen.createOptimizedForLoopGeneratorOrNull(
         forExpression: KtForExpression,
         loopRangeCall: ResolvedCall<out CallableDescriptor>

compiler/testData/codegen/box/controlStructures/forInCharSequence.kt

@@ -0,0 +1,9 @@
+// WITH_RUNTIME
+
+fun box(): String {
+    var s = ""
+    for (c in StringBuilder("OK")) {
+        s += c
+    }
+    return s
+}

compiler/testData/codegen/box/controlStructures/forInCharSequenceMut.kt

@@ -0,0 +1,19 @@
+// WITH_RUNTIME
+
+fun box(): String {
+    var clist = listOf('1', '2', '3', '4')
+    var res1 = ""
+    for (ch in clist) {
+        res1 += ch
+        clist = listOf()
+    }
+
+    var s = "1234"
+    var res2 = ""
+    for (ch in s) {
+        res2 += ch
+        s = ""
+    }
+
+    return if (res1 == res2) "OK" else "'$res1' != '$res2'"
+}

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

@@ -0,0 +1,18 @@
+// WITH_RUNTIME
+
+fun test() {
+    var s = ""
+    for (c in "testString") {
+        s += c
+    }
+    for (c in StringBuilder("testStringBuilder")) {
+        s += c
+    }
+    for (c in "testCharSequence".apply<CharSequence>{}) {
+        s += c
+    }
+}
+
+// 0 iterator
+// 0 hasNext
+// 0 nextChar

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

@@ -4250,6 +4250,18 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
             doTest(fileName);
         }
 
+        @TestMetadata("forInCharSequence.kt")
+        public void testForInCharSequence() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/controlStructures/forInCharSequence.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("forInCharSequenceMut.kt")
+        public void testForInCharSequenceMut() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/controlStructures/forInCharSequenceMut.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("forInSmartCastToArray.kt")
         public void testForInSmartCastToArray() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/controlStructures/forInSmartCastToArray.kt");

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

@@ -4250,6 +4250,18 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
             doTest(fileName);
         }
 
+        @TestMetadata("forInCharSequence.kt")
+        public void testForInCharSequence() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/controlStructures/forInCharSequence.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("forInCharSequenceMut.kt")
+        public void testForInCharSequenceMut() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/controlStructures/forInCharSequenceMut.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("forInSmartCastToArray.kt")
         public void testForInSmartCastToArray() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/controlStructures/forInSmartCastToArray.kt");

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

@@ -1127,6 +1127,12 @@ public class BytecodeTextTestGenerated extends AbstractBytecodeTextTest {
             KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("compiler/testData/codegen/bytecodeText/forLoop"), Pattern.compile("^(.+)\\.kt$"), TargetBackend.ANY, true);
         }
 
+        @TestMetadata("forInCharSequence.kt")
+        public void testForInCharSequence() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/bytecodeText/forLoop/forInCharSequence.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("forInRangeWithImplicitReceiver.kt")
         public void testForInRangeWithImplicitReceiver() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/bytecodeText/forLoop/forInRangeWithImplicitReceiver.kt");

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

@@ -4250,6 +4250,18 @@ public class LightAnalysisModeTestGenerated extends AbstractLightAnalysisModeTes
             doTest(fileName);
         }
 
+        @TestMetadata("forInCharSequence.kt")
+        public void testForInCharSequence() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/controlStructures/forInCharSequence.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("forInCharSequenceMut.kt")
+        public void testForInCharSequenceMut() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/controlStructures/forInCharSequenceMut.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("forInSmartCastToArray.kt")
         public void testForInSmartCastToArray() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/controlStructures/forInSmartCastToArray.kt");

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

@@ -4911,6 +4911,18 @@ public class JsCodegenBoxTestGenerated extends AbstractJsCodegenBoxTest {
             doTest(fileName);
         }
 
+        @TestMetadata("forInCharSequence.kt")
+        public void testForInCharSequence() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/controlStructures/forInCharSequence.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("forInCharSequenceMut.kt")
+        public void testForInCharSequenceMut() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/controlStructures/forInCharSequenceMut.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("forInSmartCastToArray.kt")
         public void testForInSmartCastToArray() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/controlStructures/forInSmartCastToArray.kt");