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

#KT-5177 In Progress

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

@@ -143,6 +143,12 @@ 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)

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

@@ -0,0 +1,26 @@
+/*
+ * 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.codegen.ExpressionCodegen
+import org.jetbrains.kotlin.codegen.range.forLoop.ForLoopGenerator
+import org.jetbrains.kotlin.codegen.range.forLoop.IterableWithIndexForLoopGenerator
+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 IterableWithIndexRangeValue(rangeCall: ResolvedCall<out CallableDescriptor>) : AbstractDestructuredPairRangeValue(rangeCall) {
+
+    override fun createDestructuredPairForLoopGenerator(
+        codegen: ExpressionCodegen,
+        forExpression: KtForExpression,
+        loopParameter: KtDestructuringDeclaration,
+        rangeCall: ResolvedCall<out CallableDescriptor>
+    ): ForLoopGenerator =
+        IterableWithIndexForLoopGenerator(codegen, forExpression, loopParameter, rangeCall)
+
+}

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

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

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

@@ -0,0 +1,65 @@
+/*
+ * 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.kotlin.resolve.jvm.AsmTypes
+import org.jetbrains.org.objectweb.asm.Label
+import org.jetbrains.org.objectweb.asm.Type
+
+class IterableWithIndexForLoopGenerator(
+    codegen: ExpressionCodegen,
+    forExpression: KtForExpression,
+    loopParameter: KtDestructuringDeclaration,
+    rangeCall: ResolvedCall<out CallableDescriptor>
+) : AbstractWithIndexForLoopGenerator(codegen, forExpression, loopParameter, rangeCall) {
+    private var iteratorVar = -1
+
+    private var indexVar = -1
+    private var indexType = Type.INT_TYPE
+
+    private val asmTypeForIterator = Type.getType(Iterator::class.java)
+
+    override fun beforeLoop() {
+        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)
+
+        iteratorVar = createLoopTempVariable(asmTypeForIterator)
+        codegen.generateCallReceiver(rangeCall).put(Type.getType(Iterable::class.java), v)
+        v.invokeinterface("java/lang/Iterable", "iterator", "()Ljava/util/Iterator;")
+        StackValue.local(iteratorVar, asmTypeForIterator).store(StackValue.onStack(asmTypeForIterator), v)
+    }
+
+    override fun checkPreCondition(loopExit: Label) {
+        v.load(iteratorVar, asmTypeForIterator)
+        v.invokeinterface("java/util/Iterator", "hasNext", "()Z")
+        v.ifeq(loopExit)
+    }
+
+    override fun assignLoopParametersNextValues() {
+        v.load(iteratorVar, asmTypeForIterator)
+        v.invokeinterface("java/util/Iterator", "next", "()Ljava/lang/Object;")
+        if (elementLoopComponent != null) {
+            StackValue.local(elementLoopComponent.parameterVar, elementLoopComponent.parameterType)
+                .store(StackValue.onStack(AsmTypes.OBJECT_TYPE), v)
+        } else {
+            v.pop()
+        }
+    }
+
+    override fun incrementAndCheckPostCondition(loopExit: Label) {
+        v.iinc(indexVar, 1)
+    }
+}