Provide optimized code generation for for-in-withIndex for CharSequences

#KT-5177 In Progress

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

@@ -143,22 +143,26 @@ fun isArrayOrPrimitiveArrayWithIndex(descriptor: CallableDescriptor) =
         KotlinBuiltIns.isArray(it) || KotlinBuiltIns.isPrimitiveArray(it)
     }
 
-fun isIterableWithIndex(descriptor: CallableDescriptor) =
-    descriptor.isTopLevelExtensionOnType("withIndex", "kotlin.collections") {
-        val typeDescriptor = it.constructor.declarationDescriptor ?: return false
-        isTopLevelInPackage(typeDescriptor, "Iterable", "kotlin.collections")
-    }
-
 fun isCollectionIndices(descriptor: CallableDescriptor) =
     descriptor.isTopLevelExtensionOnType("indices", "kotlin.collections") {
         KotlinBuiltIns.isCollectionOrNullableCollection(it)
     }
 
+fun isIterableWithIndex(descriptor: CallableDescriptor) =
+    descriptor.isTopLevelExtensionOnType("withIndex", "kotlin.collections") {
+        KotlinBuiltIns.isIterableOrNullableIterable(it)
+    }
+
 fun isCharSequenceIndices(descriptor: CallableDescriptor) =
     descriptor.isTopLevelExtensionOnType("indices", "kotlin.text") {
         KotlinBuiltIns.isCharSequenceOrNullableCharSequence(it)
     }
 
+fun isCharSequenceWithIndex(descriptor: CallableDescriptor) =
+    descriptor.isTopLevelExtensionOnType("withIndex", "kotlin.text") {
+        KotlinBuiltIns.isCharSequenceOrNullableCharSequence(it)
+    }
+
 fun isComparableRangeTo(descriptor: CallableDescriptor) =
     descriptor.isTopLevelExtensionOnType("rangeTo", "kotlin.ranges") {
         val extensionReceiverTypeDescriptor = it.constructor.declarationDescriptor as? TypeParameterDescriptor ?: return false

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

@@ -0,0 +1,35 @@
+/*
+ * Copyright 2000-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.codegen.range
+
+import org.jetbrains.kotlin.builtins.KotlinBuiltIns
+import org.jetbrains.kotlin.codegen.ExpressionCodegen
+import org.jetbrains.kotlin.codegen.range.forLoop.CharSequenceWithIndexForLoopGenerator
+import org.jetbrains.kotlin.codegen.range.forLoop.ForLoopGenerator
+import org.jetbrains.kotlin.descriptors.CallableDescriptor
+import org.jetbrains.kotlin.psi.KtDestructuringDeclaration
+import org.jetbrains.kotlin.psi.KtForExpression
+import org.jetbrains.kotlin.resolve.calls.model.ResolvedCall
+
+class CharSequenceWithIndexRangeValue(rangeCall: ResolvedCall<out CallableDescriptor>) : AbstractDestructuredPairRangeValue(rangeCall) {
+
+    override fun createDestructuredPairForLoopGenerator(
+        codegen: ExpressionCodegen,
+        forExpression: KtForExpression,
+        loopParameter: KtDestructuringDeclaration,
+        rangeCall: ResolvedCall<out CallableDescriptor>
+    ): ForLoopGenerator =
+        CharSequenceWithIndexForLoopGenerator(
+            codegen, forExpression, loopParameter, rangeCall,
+            canCacheLength = canCacheLength(rangeCall)
+        )
+
+    private fun canCacheLength(rangeCall: ResolvedCall<out CallableDescriptor>): Boolean {
+        val receiverType = rangeCall.extensionReceiver?.type ?: return false
+        return KotlinBuiltIns.isString(receiverType)
+    }
+
+}

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

@@ -130,6 +130,8 @@ private fun ExpressionCodegen.createIntrinsifiedRangeValueOrNull(rangeCall: Reso
             ArrayWithIndexRangeValue(rangeCall)
         isIterableWithIndex(rangeCallee) ->
             IterableWithIndexRangeValue(rangeCall)
+        isCharSequenceWithIndex(rangeCallee) ->
+            CharSequenceWithIndexRangeValue(rangeCall)
         isComparableRangeTo(rangeCallee) ->
             ComparableRangeLiteralRangeValue(this, rangeCall)
         isPrimitiveProgressionReverse(rangeCallee) ->

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

@@ -0,0 +1,80 @@
+/*
+ * Copyright 2000-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.codegen.range.forLoop
+
+import org.jetbrains.kotlin.codegen.ExpressionCodegen
+import org.jetbrains.kotlin.codegen.StackValue
+import org.jetbrains.kotlin.codegen.generateCallReceiver
+import org.jetbrains.kotlin.descriptors.CallableDescriptor
+import org.jetbrains.kotlin.psi.KtDestructuringDeclaration
+import org.jetbrains.kotlin.psi.KtForExpression
+import org.jetbrains.kotlin.resolve.calls.model.ResolvedCall
+import org.jetbrains.org.objectweb.asm.Label
+import org.jetbrains.org.objectweb.asm.Type
+
+class CharSequenceWithIndexForLoopGenerator(
+    codegen: ExpressionCodegen,
+    forExpression: KtForExpression,
+    loopParameter: KtDestructuringDeclaration,
+    rangeCall: ResolvedCall<out CallableDescriptor>,
+    private val canCacheLength: Boolean
+) : AbstractWithIndexForLoopGenerator(codegen, forExpression, loopParameter, rangeCall) {
+
+    private val charSeqType = codegen.asmType(ExpressionCodegen.getExpectedReceiverType(rangeCall))
+    private var charSeqVar = -1
+    private var lengthVar = -1
+
+    private var indexVar = -1
+    private var indexType = Type.INT_TYPE
+
+    override fun beforeLoop() {
+        charSeqVar = createLoopTempVariable(charSeqType)
+        val charSeqValue = StackValue.local(charSeqVar, charSeqType)
+        charSeqValue.store(codegen.generateCallReceiver(rangeCall), v)
+
+        if (canCacheLength) {
+            lengthVar = createLoopTempVariable(Type.INT_TYPE)
+            evalCharSeqLengthOnStack()
+            v.store(lengthVar, Type.INT_TYPE)
+        }
+
+        indexVar = indexLoopComponent?.parameterVar ?: createLoopTempVariable(Type.INT_TYPE)
+        indexType = indexLoopComponent?.parameterType ?: Type.INT_TYPE
+
+        StackValue.local(indexVar, indexType)
+            .store(StackValue.constant(0, Type.INT_TYPE), v)
+    }
+
+    private fun evalCharSeqLengthOnStack() {
+        v.load(charSeqVar, charSeqType)
+        v.invokeinterface("java/lang/CharSequence", "length", "()I")
+    }
+
+    override fun checkPreCondition(loopExit: Label) {
+        v.load(indexVar, Type.INT_TYPE)
+        if (canCacheLength) {
+            v.load(lengthVar, Type.INT_TYPE)
+        } else {
+            evalCharSeqLengthOnStack()
+        }
+        v.ificmpge(loopExit)
+    }
+
+    override fun assignLoopParametersNextValues() {
+        if (elementLoopComponent != null) {
+            v.load(charSeqVar, charSeqType)
+            v.load(indexVar, Type.INT_TYPE)
+            v.invokeinterface("java/lang/CharSequence", "charAt", "(I)C")
+            StackValue.local(elementLoopComponent.parameterVar, elementLoopComponent.parameterType)
+                .store(StackValue.onStack(Type.CHAR_TYPE), v)
+        }
+    }
+
+    override fun incrementAndCheckPostCondition(loopExit: Label) {
+        v.iinc(indexVar, 1)
+    }
+
+}