JVM_IR KT-50078 fix for-in-array loop bytecode shape

compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/loops/IndexedGetLoopHeader.kt

@@ -21,6 +21,9 @@ class IndexedGetLoopHeader(
     context: CommonBackendContext
 ) : NumericForLoopHeader<IndexedGetHeaderInfo>(headerInfo, builder, context) {
 
+    private val preferJavaLikeCounterLoop = context.preferJavaLikeCounterLoop
+    private val javaLikeCounterLoopBuilder = JavaLikeCounterLoopBuilder(context)
+
     override val loopInitStatements =
         listOfNotNull(headerInfo.objectVariable, inductionVariable, lastVariableIfCanCacheLast, stepVariable)
 
@@ -48,20 +51,25 @@ class IndexedGetLoopHeader(
         }
 
     override fun buildLoop(builder: DeclarationIrBuilder, oldLoop: IrLoop, newBody: IrExpression?): LoopReplacement = with(builder) {
-        // Loop is lowered into something like:
+        val newLoopCondition = buildLoopCondition(this@with)
-        //
+        if (preferJavaLikeCounterLoop) {
-        //   var inductionVar = 0
+            javaLikeCounterLoopBuilder.buildJavaLikeDoWhileCounterLoop(oldLoop, newLoopCondition, newBody, loopOrigin = null)
-        //   var last = objectVariable.size
+        } else {
-        //   while (inductionVar < last) {
+            // Loop is lowered into something like:
-        //       val loopVar = objectVariable.get(inductionVar)
+            //
-        //       inductionVar++
+            //   var inductionVar = 0
-        //       // Loop body
+            //   var last = objectVariable.size
-        //   }
+            //   while (inductionVar < last) {
-        val newLoop = IrWhileLoopImpl(oldLoop.startOffset, oldLoop.endOffset, oldLoop.type, oldLoop.origin).apply {
+            //       val loopVar = objectVariable.get(inductionVar)
-            label = oldLoop.label
+            //       inductionVar++
-            condition = buildLoopCondition(this@with)
+            //       // Loop body
-            body = newBody
+            //   }
+            val newLoop = IrWhileLoopImpl(oldLoop.startOffset, oldLoop.endOffset, oldLoop.type, oldLoop.origin).apply {
+                label = oldLoop.label
+                condition = newLoopCondition
+                body = newBody
+            }
+            LoopReplacement(newLoop, newLoop)
         }
-        LoopReplacement(newLoop, newLoop)
     }
 }

compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/loops/JavaLikeCounterLoopBuilder.kt

@@ -0,0 +1,152 @@
+/*
+ * Copyright 2010-2021 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * 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.backend.common.lower.loops
+
+import org.jetbrains.kotlin.backend.common.CommonBackendContext
+import org.jetbrains.kotlin.ir.declarations.IrSimpleFunction
+import org.jetbrains.kotlin.ir.expressions.*
+import org.jetbrains.kotlin.ir.expressions.impl.*
+import org.jetbrains.kotlin.ir.util.dump
+import org.jetbrains.kotlin.ir.util.findDeclaration
+import org.jetbrains.kotlin.ir.util.render
+import org.jetbrains.kotlin.util.OperatorNameConventions
+
+class JavaLikeCounterLoopBuilder(private val context: CommonBackendContext) {
+    private val booleanNot =
+        context.irBuiltIns.booleanClass.owner.findDeclaration<IrSimpleFunction> {
+            it.name == OperatorNameConventions.NOT
+        } ?: error("No '${OperatorNameConventions.NOT}' in ${context.irBuiltIns.booleanClass.owner.render()}")
+
+    fun buildJavaLikeDoWhileCounterLoop(
+        oldLoop: IrLoop,
+        newLoopCondition: IrExpression,
+        newBody: IrExpression?,
+        loopOrigin: IrStatementOrigin?
+    ): LoopReplacement {
+        // Transform loop:
+        //      while (<newLoopCondition>) {
+        //          { // FOR_LOOP_NEXT
+        //              <initializeLoopIteration>
+        //              <inductionVariableUpdate>
+        //          }
+        //          <originalLoopBody>
+        //      }
+        // to:
+        //      do {
+        //          { // FOR_LOOP_NEXT
+        //              if (!(<newLoopCondition>)) break
+        //              <initializeLoopIteration>
+        //          }
+        //          <originalLoopBody>
+        //      } while (
+        //          {
+        //              <inductionVariableUpdate>
+        //              true
+        //          }
+        //      )
+        val bodyBlock = newBody as? IrContainerExpression
+            ?: throw AssertionError("newBody: ${newBody?.dump()}")
+        val forLoopNextBlock = bodyBlock.statements[0] as? IrContainerExpression
+            ?: throw AssertionError("bodyBlock[0]: ${bodyBlock.statements[0].dump()}")
+        if (forLoopNextBlock.origin != IrStatementOrigin.FOR_LOOP_NEXT)
+            throw AssertionError("FOR_LOOP_NEXT expected: ${forLoopNextBlock.dump()}")
+        val inductionVariableUpdate = forLoopNextBlock.statements.last() as? IrSetValue
+            ?: throw AssertionError("forLoopNextBlock.last: ${forLoopNextBlock.statements.last().dump()}")
+
+        val doWhileLoop = IrDoWhileLoopImpl(oldLoop.startOffset, oldLoop.endOffset, oldLoop.type, loopOrigin)
+        doWhileLoop.label = oldLoop.label
+
+        bodyBlock.statements[0] = IrCompositeImpl(
+            forLoopNextBlock.startOffset, forLoopNextBlock.endOffset,
+            forLoopNextBlock.type,
+            forLoopNextBlock.origin,
+        ).apply {
+            statements.add(createNegatedConditionCheck(newLoopCondition, doWhileLoop))
+            if (forLoopNextBlock.statements.size >= 2)
+                statements.addAll(forLoopNextBlock.statements.subList(0, forLoopNextBlock.statements.lastIndex))
+        }
+
+        doWhileLoop.body = bodyBlock
+
+        val stepStartOffset = inductionVariableUpdate.startOffset
+        val stepEndOffset = inductionVariableUpdate.endOffset
+        val doWhileCondition =
+            IrCompositeImpl(
+                stepStartOffset, stepEndOffset, context.irBuiltIns.booleanType, null,
+                listOf(
+                    inductionVariableUpdate,
+                    IrConstImpl.boolean(stepStartOffset, stepEndOffset, context.irBuiltIns.booleanType, true)
+                )
+            )
+        doWhileLoop.condition = doWhileCondition
+
+        return LoopReplacement(doWhileLoop, doWhileLoop)
+    }
+
+
+    private fun createNegatedConditionCheck(
+        newLoopCondition: IrExpression,
+        doWhileLoop: IrDoWhileLoop
+    ): IrWhenImpl {
+        val conditionStartOffset = newLoopCondition.startOffset
+        val conditionEndOffset = newLoopCondition.endOffset
+        val negatedCondition =
+            IrCallImpl.fromSymbolOwner(conditionStartOffset, conditionEndOffset, booleanNot.symbol).apply {
+                dispatchReceiver = newLoopCondition
+            }
+
+        return IrWhenImpl(
+            conditionStartOffset, conditionEndOffset, context.irBuiltIns.unitType, null,
+            listOf(
+                IrBranchImpl(
+                    negatedCondition,
+                    IrBreakImpl(conditionStartOffset, conditionEndOffset, context.irBuiltIns.nothingType, doWhileLoop)
+                )
+            )
+        )
+    }
+
+    fun moveInductionVariableUpdateToLoopCondition(doWhileLoop: IrDoWhileLoop) {
+        // On JVM, it's important that induction variable update happens in the end of the loop
+        // (otherwise HotSpot will not treat it as a counter loop).
+        // Moving induction variable update to loop condition (instead of just placing it in the end of loop body)
+        // also allows reusing loop variable as induction variable later.
+        //
+        // Transform a loop in the form:
+        //      do {
+        //          { <next> }
+        //          <body>
+        //      } while (<condition>)
+        // to
+        //      do {
+        //          { <next'> }
+        //          <body>
+        //      } while ( { if (!<condition>) break; <updateInductionVar>; true } )
+        val doWhileBody = doWhileLoop.body as? IrContainerExpression ?: return
+        if (doWhileBody.origin != IrStatementOrigin.FOR_LOOP_INNER_WHILE) return
+        val doWhileLoopNext = doWhileBody.statements[0] as? IrContainerExpression ?: return
+        if (doWhileLoopNext.origin != IrStatementOrigin.FOR_LOOP_NEXT) return
+
+        val updateInductionVarIndex = doWhileLoopNext.statements
+            .indexOfFirst { it is IrSetValue && it.symbol.owner.isInductionVariable(context) }
+        if (updateInductionVarIndex < 0) return
+        val updateInductionVar = doWhileLoopNext.statements[updateInductionVarIndex]
+        doWhileLoopNext.statements.removeAt(updateInductionVarIndex)
+
+        val loopCondition = doWhileLoop.condition
+        val loopConditionStartOffset = loopCondition.startOffset
+        val loopConditionEndOffset = loopCondition.endOffset
+        doWhileLoop.condition = IrCompositeImpl(
+            loopConditionStartOffset, loopConditionEndOffset, loopCondition.type,
+            origin = null,
+            statements = listOf(
+                createNegatedConditionCheck(doWhileLoop.condition, doWhileLoop),
+                updateInductionVar,
+                IrConstImpl.boolean(loopConditionStartOffset, loopConditionEndOffset, context.irBuiltIns.booleanType, true)
+            )
+        )
+    }
+}

compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/loops/ProgressionLoopHeader.kt

@@ -12,14 +12,10 @@ import org.jetbrains.kotlin.ir.IrStatement
 import org.jetbrains.kotlin.ir.builders.createTmpVariable
 import org.jetbrains.kotlin.ir.builders.irGet
 import org.jetbrains.kotlin.ir.builders.irNotEquals
-import org.jetbrains.kotlin.ir.declarations.IrSimpleFunction
 import org.jetbrains.kotlin.ir.declarations.IrVariable
-import org.jetbrains.kotlin.ir.expressions.*
+import org.jetbrains.kotlin.ir.expressions.IrExpression
-import org.jetbrains.kotlin.ir.expressions.impl.*
+import org.jetbrains.kotlin.ir.expressions.IrLoop
-import org.jetbrains.kotlin.ir.util.dump
+import org.jetbrains.kotlin.ir.expressions.impl.IrDoWhileLoopImpl
-import org.jetbrains.kotlin.ir.util.findDeclaration
-import org.jetbrains.kotlin.ir.util.render
-import org.jetbrains.kotlin.util.OperatorNameConventions
 
 class ProgressionLoopHeader(
     headerInfo: ProgressionHeaderInfo,
@@ -28,6 +24,7 @@ class ProgressionLoopHeader(
 ) : NumericForLoopHeader<ProgressionHeaderInfo>(headerInfo, builder, context) {
 
     private val preferJavaLikeCounterLoop = context.preferJavaLikeCounterLoop
+    private val javaLikeCounterLoopBuilder = JavaLikeCounterLoopBuilder(context)
 
     // For this loop:
     //
@@ -119,7 +116,7 @@ class ProgressionLoopHeader(
                 }
 
                 if (preferJavaLikeCounterLoop) {
-                    moveInductionVariableUpdateToLoopCondition(newLoop)
+                    javaLikeCounterLoopBuilder.moveInductionVariableUpdateToLoopCondition(newLoop)
                 }
 
                 val loopCondition = buildLoopCondition(this@with)
@@ -139,7 +136,10 @@ class ProgressionLoopHeader(
 
                 val newLoopCondition = buildLoopCondition(this@with)
 
-                buildJavaLikeDoWhileCounterLoop(oldLoop, newLoopCondition, newBody)
+                javaLikeCounterLoopBuilder.buildJavaLikeDoWhileCounterLoop(
+                    oldLoop, newLoopCondition, newBody,
+                    this@ProgressionLoopHeader.context.doWhileCounterLoopOrigin
+                )
             } else {
                 // Use an if-guarded do-while loop (note the difference in loop condition):
                 //
@@ -161,134 +161,5 @@ class ProgressionLoopHeader(
             }
         }
 
-    private val booleanNot =
+}
-        context.irBuiltIns.booleanClass.owner.findDeclaration<IrSimpleFunction> {
-            it.name == OperatorNameConventions.NOT
-        } ?: error("No '${OperatorNameConventions.NOT}' in ${context.irBuiltIns.booleanClass.owner.render()}")
 
-    private fun moveInductionVariableUpdateToLoopCondition(doWhileLoop: IrDoWhileLoop) {
-        // On JVM, it's important that induction variable update happens in the end of the loop
-        // (otherwise HotSpot will not treat it as a counter loop).
-        // Moving induction variable update to loop condition (instead of just placing it in the end of loop body)
-        // also allows reusing loop variable as induction variable later.
-        //
-        // Transform a loop in the form:
-        //      do {
-        //          { <next> }
-        //          <body>
-        //      } while (<condition>)
-        // to
-        //      do {
-        //          { <next'> }
-        //          <body>
-        //      } while ( { if (!<condition>) break; <updateInductionVar>; true } )
-        val doWhileBody = doWhileLoop.body as? IrContainerExpression ?: return
-        if (doWhileBody.origin != IrStatementOrigin.FOR_LOOP_INNER_WHILE) return
-        val doWhileLoopNext = doWhileBody.statements[0] as? IrContainerExpression ?: return
-        if (doWhileLoopNext.origin != IrStatementOrigin.FOR_LOOP_NEXT) return
-
-        val updateInductionVarIndex = doWhileLoopNext.statements
-            .indexOfFirst { it is IrSetValue && it.symbol.owner.isInductionVariable(context) }
-        if (updateInductionVarIndex < 0) return
-        val updateInductionVar = doWhileLoopNext.statements[updateInductionVarIndex]
-        doWhileLoopNext.statements.removeAt(updateInductionVarIndex)
-
-        val loopCondition = doWhileLoop.condition
-        val loopConditionStartOffset = loopCondition.startOffset
-        val loopConditionEndOffset = loopCondition.endOffset
-        doWhileLoop.condition = IrCompositeImpl(
-            loopConditionStartOffset, loopConditionEndOffset, loopCondition.type,
-            origin = null,
-            statements = listOf(
-                createNegatedConditionCheck(doWhileLoop.condition, doWhileLoop),
-                updateInductionVar,
-                IrConstImpl.boolean(loopConditionStartOffset, loopConditionEndOffset, context.irBuiltIns.booleanType, true)
-            )
-        )
-    }
-
-    private fun buildJavaLikeDoWhileCounterLoop(
-        oldLoop: IrLoop,
-        newLoopCondition: IrExpression,
-        newBody: IrExpression?
-    ): LoopReplacement {
-        // Transform loop:
-        //      while (<newLoopCondition>) {
-        //          { // FOR_LOOP_NEXT
-        //              <initializeLoopIteration>
-        //              <inductionVariableUpdate>
-        //          }
-        //          <originalLoopBody>
-        //      }
-        // to:
-        //      do {
-        //          { // FOR_LOOP_NEXT
-        //              if (!(<newLoopCondition>)) break
-        //              <initializeLoopIteration>
-        //          }
-        //          <originalLoopBody>
-        //      } while (
-        //          {
-        //              <inductionVariableUpdate>
-        //              true
-        //          }
-        //      )
-        val bodyBlock = newBody as? IrContainerExpression
-            ?: throw AssertionError("newBody: ${newBody?.dump()}")
-        val forLoopNextBlock = bodyBlock.statements[0] as? IrContainerExpression
-            ?: throw AssertionError("bodyBlock[0]: ${bodyBlock.statements[0].dump()}")
-        if (forLoopNextBlock.origin != IrStatementOrigin.FOR_LOOP_NEXT)
-            throw AssertionError("FOR_LOOP_NEXT expected: ${forLoopNextBlock.dump()}")
-        val inductionVariableUpdate = forLoopNextBlock.statements.last() as? IrSetValue
-            ?: throw AssertionError("forLoopNextBlock.last: ${forLoopNextBlock.statements.last().dump()}")
-
-        val doWhileLoop = IrDoWhileLoopImpl(oldLoop.startOffset, oldLoop.endOffset, oldLoop.type, context.doWhileCounterLoopOrigin)
-        doWhileLoop.label = oldLoop.label
-
-        bodyBlock.statements[0] = IrCompositeImpl(
-            forLoopNextBlock.startOffset, forLoopNextBlock.endOffset,
-            forLoopNextBlock.type,
-            forLoopNextBlock.origin,
-        ).apply {
-            statements.add(createNegatedConditionCheck(newLoopCondition, doWhileLoop))
-            if (forLoopNextBlock.statements.size >= 2)
-                statements.addAll(forLoopNextBlock.statements.subList(0, forLoopNextBlock.statements.lastIndex))
-        }
-
-        doWhileLoop.body = bodyBlock
-
-        val stepStartOffset = inductionVariableUpdate.startOffset
-        val stepEndOffset = inductionVariableUpdate.endOffset
-        val doWhileCondition =
-            IrCompositeImpl(
-                stepStartOffset, stepEndOffset, context.irBuiltIns.booleanType, null,
-                listOf(
-                    inductionVariableUpdate,
-                    IrConstImpl.boolean(stepStartOffset, stepEndOffset, context.irBuiltIns.booleanType, true)
-                )
-            )
-        doWhileLoop.condition = doWhileCondition
-
-        return LoopReplacement(doWhileLoop, doWhileLoop)
-    }
-
-    private fun createNegatedConditionCheck(newLoopCondition: IrExpression, doWhileLoop: IrDoWhileLoop): IrWhenImpl {
-        val conditionStartOffset = newLoopCondition.startOffset
-        val conditionEndOffset = newLoopCondition.endOffset
-        val negatedCondition =
-            IrCallImpl.fromSymbolOwner(conditionStartOffset, conditionEndOffset, booleanNot.symbol).apply {
-                dispatchReceiver = newLoopCondition
-            }
-
-        return IrWhenImpl(
-            conditionStartOffset, conditionEndOffset, context.irBuiltIns.unitType, null,
-            listOf(
-                IrBranchImpl(
-                    negatedCondition,
-                    IrBreakImpl(conditionStartOffset, conditionEndOffset, context.irBuiltIns.nothingType, doWhileLoop)
-                )
-            )
-        )
-    }
-
-}

compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/loops/WithIndexLoopHeader.kt

@@ -23,8 +23,8 @@ class WithIndexLoopHeader(
     context: CommonBackendContext
 ) : ForLoopHeader {
 
-    private val nestedLoopHeader: ForLoopHeader
+    val nestedLoopHeader: ForLoopHeader
-    private val indexVariable: IrVariable
+    val indexVariable: IrVariable
     private val ownsIndexVariable: Boolean
     private val incrementIndexStatement: IrStatement?
 

compiler/ir/backend.jvm/lower/src/org/jetbrains/kotlin/backend/jvm/lower/JvmOptimizationLowering.kt

@@ -315,6 +315,9 @@ class JvmOptimizationLowering(val context: JvmBackendContext) : FileLoweringPass
             val loopVariablePosition = findLoopVariablePosition(irForLoopBlock.statements[1]) ?: return
             val (loopVariableContainer, loopVariableIndex) = loopVariablePosition
             val loopVariable = loopVariableContainer.statements[loopVariableIndex] as? IrVariable ?: return
+            val loopVariableInitializer = loopVariable.initializer ?: return
+            if (loopVariableInitializer !is IrGetValue) return
+            if (loopVariableInitializer.symbol != inductionVariable.symbol) return
 
             val inductionVariableType = inductionVariable.type
             val loopVariableType = loopVariable.type

compiler/testData/codegen/box/controlStructures/forInArray/forIntArray.kt

@@ -1,14 +1,13 @@
-fun box() : String {
+fun box(): String {
     val a = arrayOfNulls<Int>(5)
-    var i = 0
     var sum = 0
-    for(el in 0..4) {
+    for (i in 0..4) {
-       a[i] = i++
+        a[i] = i + 1
     }
     for (el in (a as Array<Int>)) {
         sum = sum + el
     }
-    if(sum != 10) return "a failed"
+    if (sum != 15) return "failed: sum=$sum"
 
     return "OK"
 }

compiler/testData/codegen/box/controlStructures/forInArray/forPrimitiveIntArray.kt

@@ -1,14 +1,15 @@
-fun box() : String {
+// WITH_RUNTIME
-    val a = IntArray (5)
+
-    var i = 0
+fun box(): String {
+    val a = IntArray(5)
     var sum = 0
-    for(el in 0..4) {
+    for (i in 0..4) {
-       a[i] = i++
+        a[i] = i + 1
     }
     for (el in a) {
         sum = sum + el
     }
-    if(sum != 10) return "a failed"
+    if (sum != 15) return "failed: sum=$sum"
 
     return "OK"
 }

compiler/testData/codegen/bytecodeText/forLoop/forInArrayWithIndex/forInEmptyArrayWithIndex.kt

@@ -33,4 +33,4 @@ fun box(): String {
 // 4 ISTORE
 // 0 IADD
 // 0 ISUB
-// 1 IINC
+// 0 IINC

compiler/testData/codegen/bytecodeText/forLoop/forInCharSequenceWithIndex/forInEmptyStringWithIndex.kt

@@ -31,4 +31,4 @@ fun box(): String {
 // 4 ISTORE
 // 0 IADD
 // 0 ISUB
-// 1 IINC
+// 0 IINC

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

@@ -21,6 +21,6 @@ fun f() {
 // 1 LOCALVARIABLE c C L3 L\d+ 0
 
 // JVM_IR_TEMPLATES
-// 1 ISTORE 2\s+L3
+// 1 ISTORE 2\s+L4
 // 1 ILOAD 2\s+INVOKEVIRTUAL java/io/PrintStream.print \(C\)V
-// 1 LOCALVARIABLE c C L3 L\d+ 2
+// 1 LOCALVARIABLE c C L4 L\d+ 2