[K/N] Analyze vals in bounds check elimination

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

@@ -121,7 +121,7 @@ abstract class ForLoopBodyTransformer : IrElementTransformerVoid() {
 
     abstract fun transform(
         context: CommonBackendContext,
-        irExpression: IrExpression,
+        loopBody: IrExpression,
         loopVariable: IrVariable,
         forLoopHeader: ForLoopHeader,
         loopComponents: Map<Int, IrVariable>

kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/optimizations/KonanBCEForLoopBodyTransformer.kt

@@ -7,7 +7,6 @@ package org.jetbrains.kotlin.backend.konan.optimizations
 
 import org.jetbrains.kotlin.backend.common.CommonBackendContext
 import org.jetbrains.kotlin.backend.common.lower.loops.*
-import org.jetbrains.kotlin.backend.konan.Context
 import org.jetbrains.kotlin.backend.konan.ir.KonanNameConventions
 import org.jetbrains.kotlin.ir.declarations.IrDeclarationOrigin
 import org.jetbrains.kotlin.ir.declarations.IrVariable
@@ -38,7 +37,7 @@ class KonanBCEForLoopBodyTransformer : ForLoopBodyTransformer() {
 
     private var analysisResult: BoundsCheckAnalysisResult = BoundsCheckAnalysisResult(false, null)
 
-    override fun transform(context: CommonBackendContext, irExpression: IrExpression, loopVariable: IrVariable,
+    override fun transform(context: CommonBackendContext, loopBody: IrExpression, loopVariable: IrVariable,
                            forLoopHeader: ForLoopHeader, loopComponents: Map<Int, IrVariable>) {
         this.context = context
         mainLoopVariable = loopVariable
@@ -46,17 +45,32 @@ class KonanBCEForLoopBodyTransformer : ForLoopBodyTransformer() {
         loopVariableComponents = loopComponents
         analysisResult = analyzeLoopHeader(loopHeader)
         if (analysisResult.boundsAreSafe && analysisResult.arrayInLoop != null)
-            irExpression.transformChildrenVoid(this)
+            loopBody.transformChildrenVoid(this)
+    }
+
+    private inline fun IrGetValue.compareConstValue(compare: (IrExpression) -> Boolean): Boolean {
+        val variable = symbol.owner
+        return if (variable is IrVariable && !variable.isVar && variable.initializer != null) {
+            compare(variable.initializer!!)
+        } else false
     }
 
     private fun IrExpression.compareIntegerNumericConst(compare: (Long) -> Boolean): Boolean {
         @Suppress("UNCHECKED_CAST")
-        return this is IrConst<*> && value is Number && compare((value as Number).toLong())
+        return when (this) {
+            is IrConst<*> -> value is Number && compare((value as Number).toLong())
+            is IrGetValue -> compareConstValue { it.compareIntegerNumericConst(compare) }
+            else -> false
+        }
     }
 
     private fun IrExpression.compareFloatNumericConst(compare: (Double) -> Boolean): Boolean {
         @Suppress("UNCHECKED_CAST")
-        return this is IrConst<*> && value is Number && compare((value as Number).toDouble())
+        return when (this) {
+            is IrConst<*> -> value is Number && compare((value as Number).toDouble())
+            is IrGetValue -> compareConstValue { it.compareFloatNumericConst(compare) }
+            else -> false
+        }
     }
 
     private fun IrType.isBasicArray() = isPrimitiveArray() || isArray()
@@ -86,18 +100,31 @@ class KonanBCEForLoopBodyTransformer : ForLoopBodyTransformer() {
         return BoundsCheckAnalysisResult(boundsAreSafe, array)
     }
 
-    private fun checkLastElement(last: IrExpression, loopHeader: ProgressionLoopHeader): BoundsCheckAnalysisResult {
+    private inline fun checkIrGetValue(value: IrGetValue, condition: (IrExpression) -> BoundsCheckAnalysisResult): BoundsCheckAnalysisResult {
-        return if (last is IrCall) {
+        val variable = value.symbol.owner
-            if (last.isGetSizeCall() && !loopHeader.headerInfo.isLastInclusive) {
+        return if (variable is IrVariable && !variable.isVar && variable.initializer != null) {
-                BoundsCheckAnalysisResult(true, (last.dispatchReceiver as? IrGetValue)?.symbol)
+            condition(variable.initializer!!)
-            } else {
-                lessThanSize(last)
-            }
         } else {
             BoundsCheckAnalysisResult(false, null)
         }
     }
 
+    private fun checkIrCallCondition(expression: IrExpression, condition: (IrCall) -> BoundsCheckAnalysisResult): BoundsCheckAnalysisResult =
+            when (expression) {
+                is IrCall -> condition(expression)
+                is IrGetValue -> checkIrGetValue(expression) { valueInitializer -> checkIrCallCondition(valueInitializer, condition) }
+                else -> BoundsCheckAnalysisResult(false, null)
+            }
+
+    private fun checkLastElement(last: IrExpression, loopHeader: ProgressionLoopHeader): BoundsCheckAnalysisResult =
+            checkIrCallCondition(last) { call ->
+                if (call.isGetSizeCall() && !loopHeader.headerInfo.isLastInclusive) {
+                    BoundsCheckAnalysisResult(true, (call.dispatchReceiver as? IrGetValue)?.symbol)
+                } else {
+                    lessThanSize(call)
+                }
+            }
+
     private fun IrExpression.isProgressionPropertyGetter(propertyName: String) =
             this is IrCall && symbol.owner.origin == IrDeclarationOrigin.DEFAULT_PROPERTY_ACCESSOR &&
                     (symbol.signature as? IdSignature.AccessorSignature)?.propertySignature?.asPublic()?.shortName == propertyName &&
@@ -113,7 +140,7 @@ class KonanBCEForLoopBodyTransformer : ForLoopBodyTransformer() {
                         if (!loopHeader.headerInfo.first.compareIntegerNumericConst { it >= 0 }) {
                             return analysisResult
                         }
-                        // TODO: variable set to const value. Add constant propagation?
+                        // TODO: variable set to const value and field getters. Add constant propagation?
                         // Analyze last element of progression.
                         if (loopHeader.headerInfo.last is IrCall) {
                             val functionCall = (loopHeader.headerInfo.last as IrCall)
@@ -130,6 +157,8 @@ class KonanBCEForLoopBodyTransformer : ForLoopBodyTransformer() {
                                 // Simple progression.
                                 analysisResult = checkLastElement(functionCall, loopHeader)
                             }
+                        } else {
+                            analysisResult = checkLastElement(loopHeader.headerInfo.last, loopHeader)
                         }
                     }
                     ProgressionDirection.DECREASING -> {
@@ -150,14 +179,8 @@ class KonanBCEForLoopBodyTransformer : ForLoopBodyTransformer() {
                         }
                         if (!boundsAreSafe)
                             return analysisResult
-                        when (val first = loopHeader.headerInfo.first) {
+
-                            is IrCall ->
+                        analysisResult = checkIrCallCondition(loopHeader.headerInfo.first, ::lessThanSize)
-                                analysisResult = lessThanSize(first)
-                            is IrGetValue ->
-                                ((first.symbol.owner as? IrVariable)?.initializer as? IrCall)?.let {
-                                    analysisResult = lessThanSize(it)
-                                }
-                        }
                     }
                     ProgressionDirection.UNKNOWN ->
                         // Case of progression - for (i in 0 until array.size step n)
@@ -170,13 +193,13 @@ class KonanBCEForLoopBodyTransformer : ForLoopBodyTransformer() {
                                         ((firstReceiver?.symbol?.owner as? IrVariable)?.initializer as? IrCall)?.extensionReceiver as? IrCall
                                 if (untilFunction?.symbol?.owner?.name?.asString() == "until" && untilFunction.extensionReceiver?.compareIntegerNumericConst { it >= 0 } == true) {
                                     val last = untilFunction.getValueArgument(0)!!
-                                    if (last is IrCall) {
+                                    analysisResult = checkIrCallCondition(last) { call ->
-                                        analysisResult = lessThanSize(last)
                                         // `isLastInclusive` for current case is set to true.
                                         // This case isn't fully optimized in ForLoopsLowering.
-                                        if (last.isGetSizeCall()) {
+                                        if (call.isGetSizeCall())
-                                            analysisResult = BoundsCheckAnalysisResult(true, (last.dispatchReceiver as? IrGetValue)?.symbol)
+                                            BoundsCheckAnalysisResult(true, (call.dispatchReceiver as? IrGetValue)?.symbol)
-                                        }
+                                        else
+                                            lessThanSize(call)
                                     }
                                 }
                             }

kotlin-native/backend.native/tests/codegen/bce/arraysForLoops.kt

@@ -96,6 +96,16 @@ import kotlin.test.*
         }
     }
 
+    var a = array.size - 1
+    val b = ++a
+    val c = b
+
+    assertFailsWith<ArrayIndexOutOfBoundsException> {
+        for (i in c downTo 0) {
+            array[i] = 6
+        }
+    }
+
     assertFailsWith<ArrayIndexOutOfBoundsException> {
         for (i in array.size + 1 downTo 0) {
             array[i] = 6
@@ -127,6 +137,14 @@ import kotlin.test.*
         }
     }
 
+    var length = array.size - 1
+    length = 2 * length
+    assertFailsWith<ArrayIndexOutOfBoundsException> {
+        for (i in 0..length) {
+            array[i] = 6
+        }
+    }
+
     assertFailsWith<ArrayIndexOutOfBoundsException> {
         for (i in 0..array.size - 1) {
             array[i + 1] = 6
@@ -373,3 +391,78 @@ import kotlin.test.*
         }
     }
 }
+
+fun foo(a: Int, b : Int): Int = a + b * 2
+
+@Test fun bceCases() {
+    val array = Array(10) { 100 }
+    val array1 = Array(3) { 0 }
+    var length = array.size  - 1
+    var sum = 0
+
+    array.forEach {
+        sum += it
+    }
+
+    for (i in array.indices) {
+        array[i] = 6
+    }
+
+    for (i in 0 until array.size) {
+        array[i] = 7
+    }
+
+    for (i in array.size - 1 downTo 1) {
+        array[i] = 7
+    }
+
+    for (it in array) {
+        sum += it
+    }
+
+    for (i in 0..array.size - 1 step 2) {
+        array[i] = 7
+    }
+
+    for (i in 0 until array.size step 2) {
+        array[i] = 7
+    }
+
+    for (i in array.indices step 2) {
+        array[i] = 6
+    }
+
+    for (i in array.size - 1 downTo 1 step 2) {
+        array[i] = 7
+    }
+
+    for ((index, value) in array.withIndex()) {
+        array[index] = 8
+    }
+
+    for ((i, v) in (0..array.size - 1 step 2).withIndex()) {
+        array[v] = 8
+        array[i] = 6
+    }
+    for (i in array.reversed()) {
+        sum += i
+    }
+
+    for (i in (0..array.size-1).reversed()) {
+        array [i] = 10
+    }
+
+    for (i in 0 until array.size) {
+        array[i] = 7
+        for (j in 0 until array1.size) {
+            array1[j] = array[i]
+        }
+    }
+
+    val size = array.size - 1
+    val size1 = size
+
+    for (i in 0..size1) {
+        foo(array[i], array[i])
+    }
+}