JS: optimize loop over array.withIndex()

See KT-20932

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

@@ -2494,6 +2494,18 @@ public class BoxJsTestGenerated extends AbstractBoxJsTest {
                 doTest(fileName);
             }
 
+            @TestMetadata("overArrayWithIndex.kt")
+            public void testOverArrayWithIndex() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("js/js.translator/testData/box/expression/for/overArrayWithIndex.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("overCollectionWithIndex.kt")
+            public void testOverCollectionWithIndex() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("js/js.translator/testData/box/expression/for/overCollectionWithIndex.kt");
+                doTest(fileName);
+            }
+
             @TestMetadata("rangeOptimization.kt")
             public void testRangeOptimization() throws Exception {
                 String fileName = KotlinTestUtils.navigationMetadata("js/js.translator/testData/box/expression/for/rangeOptimization.kt");

js/js.translator/src/org/jetbrains/kotlin/js/translate/expression/LoopTranslator.kt

@@ -19,8 +19,7 @@
 package org.jetbrains.kotlin.js.translate.expression
 
 import org.jetbrains.kotlin.builtins.KotlinBuiltIns
-import org.jetbrains.kotlin.descriptors.CallableDescriptor
+import org.jetbrains.kotlin.descriptors.*
-import org.jetbrains.kotlin.descriptors.FunctionDescriptor
 import org.jetbrains.kotlin.js.backend.ast.*
 import org.jetbrains.kotlin.js.translate.callTranslator.CallTranslator
 import org.jetbrains.kotlin.js.translate.context.TranslationContext
@@ -29,14 +28,22 @@ import org.jetbrains.kotlin.js.translate.intrinsic.functions.factories.ArrayFIF
 import org.jetbrains.kotlin.js.translate.utils.BindingUtils.*
 import org.jetbrains.kotlin.js.translate.utils.JsAstUtils
 import org.jetbrains.kotlin.js.translate.utils.JsAstUtils.*
+import org.jetbrains.kotlin.js.translate.utils.JsDescriptorUtils.getReceiverParameterForReceiver
 import org.jetbrains.kotlin.js.translate.utils.PsiUtils.getLoopRange
 import org.jetbrains.kotlin.js.translate.utils.TranslationUtils
+import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.resolve.BindingContext
+import org.jetbrains.kotlin.resolve.DescriptorUtils.getFunctionByName
+import org.jetbrains.kotlin.resolve.DescriptorUtils.getPropertyByName
 import org.jetbrains.kotlin.resolve.calls.callUtil.getResolvedCall
 import org.jetbrains.kotlin.resolve.calls.model.ResolvedCall
 import org.jetbrains.kotlin.resolve.descriptorUtil.fqNameSafe
 import org.jetbrains.kotlin.resolve.scopes.receivers.ExpressionReceiver
+import org.jetbrains.kotlin.resolve.scopes.receivers.ImplicitReceiver
+import org.jetbrains.kotlin.types.KotlinType
 
 fun createWhile(doWhile: Boolean, expression: KtWhileExpressionBase, context: TranslationContext): JsNode {
     val conditionExpression = expression.condition ?:
@@ -85,6 +92,12 @@ private val stepFunctionName = FqName("kotlin.ranges.step")
 private val intRangeName = FqName("kotlin.ranges.IntRange")
 private val intProgressionName = FqName("kotlin.ranges.IntProgression")
 
+private val withIndexFqName = FqName("kotlin.collections.withIndex")
+private val sequenceWithIndexFqName = FqName("kotlin.sequences.withIndex")
+private val indicesFqName = FqName("kotlin.collections.indices")
+
+private val sequenceFqName = FqName("kotlin.sequences.Sequence")
+
 fun translateForExpression(expression: KtForExpression, context: TranslationContext): JsStatement {
     val loopRange = KtPsiUtil.deparenthesize(getLoopRange(expression))!!
     val rangeType = getTypeForExpression(context.bindingContext(), loopRange)
@@ -134,6 +147,45 @@ fun translateForExpression(expression: KtForExpression, context: TranslationCont
         JsScope.declareTemporary()
     }
 
+    fun KtDeclaration.extractDescriptor() = context.bindingContext()[BindingContext.VARIABLE, this]?.takeUnless { it.name.isSpecial }
+
+    fun extractWithIndexCall(): WithIndexInfo? {
+        val resolvedCall = loopRange.getResolvedCall(context.bindingContext()) ?: return null
+        val fqName = resolvedCall.resultingDescriptor.fqNameSafe
+        val (indexDescriptor, elementDescriptor) = when (fqName) {
+            withIndexFqName, sequenceWithIndexFqName -> {
+                if (destructuringParameter == null) return null
+                destructuringParameter.entries.let { Pair(it[0].extractDescriptor(), it[1].extractDescriptor()) }
+            }
+            indicesFqName -> {
+                if (destructuringParameter != null) return null
+                val varDescriptor = context.bindingContext()[BindingContext.DECLARATION_TO_DESCRIPTOR, loopParameter] as?
+                                            VariableDescriptor ?: return null
+                Pair(varDescriptor, null)
+            }
+            else -> return null
+        }
+
+        val receiverClass = resolvedCall.resultingDescriptor.extensionReceiverParameter?.type?.constructor?.declarationDescriptor as?
+                                    ClassDescriptor ?: return null
+        val receiverType = when {
+            KotlinBuiltIns.isArrayOrPrimitiveArray(receiverClass) -> WithIndexReceiverType.ARRAY
+            KotlinBuiltIns.isCollectionOrNullableCollection(receiverClass.defaultType) -> WithIndexReceiverType.COLLECTION
+            KotlinBuiltIns.isIterableOrNullableIterable(receiverClass.defaultType) -> WithIndexReceiverType.ITERABLE
+            receiverClass.fqNameSafe == sequenceFqName -> WithIndexReceiverType.SEQUENCE
+            else -> return null
+        }
+
+        val receiver = resolvedCall.extensionReceiver ?: return null
+        val arrayExpr = when (receiver) {
+            is ExpressionReceiver -> Translation.translateAsExpression(receiver.expression, context)
+            is ImplicitReceiver -> context.getDispatchReceiver(getReceiverParameterForReceiver(receiver))
+            else -> return null
+        }
+
+        return WithIndexInfo(receiverType, indexDescriptor, elementDescriptor, arrayExpr)
+    }
+
     fun translateBody(itemValue: JsExpression?): JsStatement? {
         val realBody = expression.body?.let { Translation.translateAsStatementAndMergeInBlockIfNeeded(it, context) }
         if (itemValue == null && destructuringParameter == null) {
@@ -242,6 +294,87 @@ fun translateForExpression(expression: KtForExpression, context: TranslationCont
         return JsFor(initExpression, conditionExpression, incrementExpression, body)
     }
 
+    fun translateForOverArrayWithIndex(info: WithIndexInfo): JsStatement {
+        val range = context.cacheExpressionIfNeeded(info.range)
+        val indexVar = info.index?.let { context.getNameForDescriptor(it) } ?: JsScope.declareTemporary()
+        val valueVar = info.value?.let { context.getNameForDescriptor(it) }
+
+        val initExpression = newVar(indexVar, JsIntLiteral(0)).apply { source = expression }
+        val conditionExpression = inequality(indexVar.makeRef(), JsNameRef("length", range)).source(expression)
+        val incrementExpression = JsPrefixOperation(JsUnaryOperator.INC, indexVar.makeRef()).source(expression)
+
+        val body = JsBlock()
+        if (valueVar != null) {
+            body.statements += newVar(valueVar, JsArrayAccess(range, indexVar.makeRef())).apply { source = expression }
+        }
+        expression.body?.let { body.statements += Translation.translateAsStatement(it, context.innerBlock(body)) }
+
+        return JsFor(initExpression, conditionExpression, incrementExpression, body)
+    }
+
+    fun findCollection() =
+            context.currentModule.findClassAcrossModuleDependencies(ClassId.topLevel(KotlinBuiltIns.FQ_NAMES.collection))!!
+
+    fun translateForOverCollectionIndices(info: WithIndexInfo): JsStatement {
+        val range = context.cacheExpressionIfNeeded(info.range)
+        val indexVar = info.index?.let { context.getNameForDescriptor(it) } ?: JsScope.declareTemporary()
+
+        val initExpression = newVar(indexVar, JsIntLiteral(0)).apply { source = expression }
+
+        val sizeDescriptor = getPropertyByName(findCollection().unsubstitutedMemberScope, Name.identifier("size"))
+        val sizeName = context.getNameForDescriptor(sizeDescriptor)
+        val conditionExpression = inequality(indexVar.makeRef(), JsNameRef(sizeName, range)).source(expression)
+
+        val incrementExpression = JsPrefixOperation(JsUnaryOperator.INC, indexVar.makeRef()).source(expression)
+
+        val body = JsBlock()
+        expression.body?.let { body.statements += Translation.translateAsStatement(it, context.innerBlock(body)) }
+
+        return JsFor(initExpression, conditionExpression, incrementExpression, body)
+    }
+
+    fun findIterable() =
+        context.currentModule.findClassAcrossModuleDependencies(ClassId.topLevel(KotlinBuiltIns.FQ_NAMES.iterable))!!
+
+    fun findSequence() =
+            context.currentModule.findClassAcrossModuleDependencies(ClassId.topLevel(sequenceFqName))!!
+
+    fun translateForOverCollectionWithIndex(info: WithIndexInfo): JsStatement {
+        val range = context.cacheExpressionIfNeeded(info.range)
+        val indexVar = info.index?.let { context.getNameForDescriptor(it) }
+        val valueVar = info.value?.let { context.getNameForDescriptor(it) }
+
+        indexVar?.let { context.addStatementToCurrentBlock(newVar(it, JsIntLiteral(0)).apply { source = expression }) }
+
+        val iteratorVar = JsScope.declareTemporary()
+        val rangeOwner = if (info.receiverType == WithIndexReceiverType.SEQUENCE) findSequence() else findIterable()
+        val iteratorDescriptor =  getFunctionByName(rangeOwner.unsubstitutedMemberScope, Name.identifier("iterator"))
+        val iteratorName = context.getNameForDescriptor(iteratorDescriptor)
+        val initExpression = newVar(iteratorVar, JsInvocation(pureFqn(iteratorName, range))).apply { source = expression }
+
+        val iteratorClassDescriptor = iteratorDescriptor.returnType!!.constructor.declarationDescriptor as ClassDescriptor
+
+        val hasNextDescriptor = getFunctionByName(iteratorClassDescriptor.unsubstitutedMemberScope, Name.identifier("hasNext"))
+        val hasNextName = context.getNameForDescriptor(hasNextDescriptor)
+        val hasNextInvocation = JsInvocation(pureFqn(hasNextName, iteratorVar.makeRef())).source(expression)
+
+        val nextDescriptor = getFunctionByName(iteratorClassDescriptor.unsubstitutedMemberScope, Name.identifier("next"))
+        val nextName = context.getNameForDescriptor(nextDescriptor)
+        val nextInvocation = JsInvocation(pureFqn(nextName, iteratorVar.makeRef())).source(expression)
+
+        val incrementExpression = indexVar?.let { JsPrefixOperation(JsUnaryOperator.INC, it.makeRef()).source(expression) }
+
+        val body = JsBlock()
+        body.statements += if (valueVar != null) {
+            newVar(valueVar, nextInvocation).apply { source = expression }
+        }
+        else {
+            asSyntheticStatement(nextInvocation)
+        }
+        expression.body?.let { body.statements += Translation.translateAsStatement(it, context.innerBlock(body)) }
+        return JsFor(initExpression, hasNextInvocation, incrementExpression).also { it.body = body }
+    }
+
     fun translateForOverIterator(): JsStatement {
 
         fun translateMethodInvocation(
@@ -284,11 +417,27 @@ fun translateForExpression(expression: KtForExpression, context: TranslationCont
     }
 
     val rangeLiteral = extractForOverRangeLiteral()
+    val withIndexCall = extractWithIndexCall()
 
     val result = when {
         rangeLiteral != null ->
             translateForOverLiteralRange(rangeLiteral)
 
+        withIndexCall != null ->
+            when (withIndexCall.receiverType) {
+                WithIndexReceiverType.ARRAY -> translateForOverArrayWithIndex(withIndexCall)
+                WithIndexReceiverType.ITERABLE,
+                WithIndexReceiverType.SEQUENCE,
+                WithIndexReceiverType.COLLECTION -> {
+                    if (withIndexCall.value == null && withIndexCall.receiverType == WithIndexReceiverType.COLLECTION) {
+                        translateForOverCollectionIndices(withIndexCall)
+                    }
+                    else {
+                        translateForOverCollectionWithIndex(withIndexCall)
+                    }
+                }
+            }
+
         isForOverRange() ->
             translateForOverRange()
 
@@ -309,3 +458,16 @@ private enum class RangeType {
 }
 
 private class RangeLiteral(val type: RangeType, val first: KtExpression, val second: KtExpression, var step: KtExpression?)
+
+private class WithIndexInfo(
+        val receiverType: WithIndexReceiverType,
+        val index: VariableDescriptor?, val value: VariableDescriptor?,
+        val range: JsExpression
+)
+
+private enum class WithIndexReceiverType {
+    ARRAY,
+    COLLECTION,
+    ITERABLE,
+    SEQUENCE
+}

js/js.translator/testData/box/expression/for/overArrayWithIndex.kt

@@ -0,0 +1,55 @@
+// EXPECTED_REACHABLE_NODES: 1128
+// CHECK_CONTAINS_NO_CALLS: test1 except=toString
+// CHECK_CONTAINS_NO_CALLS: test2 except=toString
+// CHECK_CONTAINS_NO_CALLS: test3 except=toString
+// CHECK_CONTAINS_NO_CALLS: test4 except=toString
+
+fun test1(a: Array<String>): String {
+    var s = ""
+    for ((i, x) in a.withIndex()) {
+        s += "$i:$x;"
+    }
+    return s
+}
+
+fun test2(a: Array<String>): String {
+    var s = ""
+    for ((_, x) in a.withIndex()) {
+        s += "$x;"
+    }
+    return s
+}
+
+fun test3(a: Array<String>): String {
+    var s = ""
+    for ((i, _) in a.withIndex()) {
+        s += "$i;"
+    }
+    return s
+}
+
+fun test4(a: Array<String>): String {
+    var s = ""
+    for (i in a.indices) {
+        s += "$i;"
+    }
+    return s
+}
+
+fun box(): String {
+    val array = arrayOf("foo", "bar", "baz")
+
+    var r = test1(array)
+    if (r != "0:foo;1:bar;2:baz;") return "fail1: $r"
+
+    r = test2(array)
+    if (r != "foo;bar;baz;") return "fail2: $r"
+
+    r = test3(array)
+    if (r != "0;1;2;") return "fail3: $r"
+
+    r = test4(array)
+    if (r != "0;1;2;") return "fail4: $r"
+
+    return "OK"
+}

js/js.translator/testData/box/expression/for/overCollectionWithIndex.kt

@@ -0,0 +1,55 @@
+// EXPECTED_REACHABLE_NODES: 1461
+// CHECK_NOT_CALLED_IN_SCOPE: function=component1 scope=test1
+// CHECK_NOT_CALLED_IN_SCOPE: function=component1 scope=test2
+// CHECK_NOT_CALLED_IN_SCOPE: function=component1 scope=test3
+// CHECK_CONTAINS_NO_CALLS: test4 except=toString
+
+fun test1(a: Sequence<String>): String {
+    var s = ""
+    for ((i, x) in a.withIndex()) {
+        s += "$i:$x;"
+    }
+    return s
+}
+
+fun test2(a: Collection<String>): String {
+    var s = ""
+    for ((_, x) in a.withIndex()) {
+        s += "$x;"
+    }
+    return s
+}
+
+fun test3(a: List<String>): String {
+    var s = ""
+    for ((i, _) in a.withIndex()) {
+        s += "$i;"
+    }
+    return s
+}
+
+fun test4(a: Set<String>): String {
+    var s = ""
+    for (i in a.indices) {
+        s += "$i;"
+    }
+    return s
+}
+
+fun box(): String {
+    val list = listOf("foo", "bar", "baz")
+
+    var r = test1(list.asSequence())
+    if (r != "0:foo;1:bar;2:baz;") return "fail1: $r"
+
+    r = test2(list)
+    if (r != "foo;bar;baz;") return "fail2: $r"
+
+    r = test3(list)
+    if (r != "0;1;2;") return "fail3: $r"
+
+    r = test4(list.toSet())
+    if (r != "0;1;2;") return "fail4: $r"
+
+    return "OK"
+}

js/js.translator/testData/lineNumbers/for.kt

@@ -17,4 +17,4 @@ fun box() {
     }
 }
 
-// LINES: 18 2 2 10 15 2 2 2 2 2 2 3 3 6 6 6 6 7 7 10 10 10 10 10 10 11 11 14 14 15 15 15 15 15 15 15 15 15 15 15 15 16 16
+// LINES: 18 2 2 10 2 2 2 2 2 2 3 3 6 6 6 6 7 7 10 10 10 10 10 10 11 11 14 14 15 15 15 15 16 16