[K/N] Bounds check elimination in basic for loops forms

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

@@ -97,11 +97,11 @@ val forLoopsPhase = makeIrFilePhase(
  *   }
  * ```
  */
-class ForLoopsLowering(val context: CommonBackendContext) : BodyLoweringPass {
+class ForLoopsLowering(val context: CommonBackendContext, val loopBodyTransformer: ForLoopBodyTransformer? = null) : BodyLoweringPass {
 
     override fun lower(irBody: IrBody, container: IrDeclaration) {
         val oldLoopToNewLoop = mutableMapOf<IrLoop, IrLoop>()
-        val transformer = RangeLoopTransformer(context, container as IrSymbolOwner, oldLoopToNewLoop)
+        val transformer = RangeLoopTransformer(context, container as IrSymbolOwner, oldLoopToNewLoop, loopBodyTransformer)
         irBody.transformChildrenVoid(transformer)
 
         // Update references in break/continue.
@@ -114,10 +114,46 @@ class ForLoopsLowering(val context: CommonBackendContext) : BodyLoweringPass {
     }
 }
 
+/**
+ * Abstract class for additional for-loop bodies transformations.
+ */
+abstract class ForLoopBodyTransformer : IrElementTransformerVoid() {
+    protected lateinit var mainLoopVariable: IrVariable
+    protected lateinit var loopHeader: ForLoopHeader
+    protected lateinit var loopVariableComponents: Map<Int, IrVariable>
+    protected lateinit var context: CommonBackendContext
+
+    open fun initialize(
+        context: CommonBackendContext,
+        loopVariable: IrVariable,
+        forLoopHeader: ForLoopHeader,
+        loopComponents: Map<Int, IrVariable>
+    ) {
+        this.context = context
+        mainLoopVariable = loopVariable
+        loopHeader = forLoopHeader
+        loopVariableComponents = loopComponents
+    }
+
+    fun transform(
+        context: CommonBackendContext,
+        irExpression: IrExpression,
+        loopVariable: IrVariable,
+        forLoopHeader: ForLoopHeader,
+        loopComponents: Map<Int, IrVariable>
+    ) {
+        initialize(context, loopVariable, forLoopHeader, loopComponents)
+        irExpression.transformChildrenVoid(this)
+    }
+
+    open fun shouldTransform(context: CommonBackendContext) = true
+}
+
 private class RangeLoopTransformer(
     val context: CommonBackendContext,
     val container: IrSymbolOwner,
-    val oldLoopToNewLoop: MutableMap<IrLoop, IrLoop>
+    val oldLoopToNewLoop: MutableMap<IrLoop, IrLoop>,
+    val loopBodyTransformer: ForLoopBodyTransformer? = null
 ) : IrElementTransformerVoidWithContext() {
 
     private val headerInfoBuilder = DefaultHeaderInfoBuilder(context, this::getScopeOwnerSymbol)
@@ -260,6 +296,9 @@ private class RangeLoopTransformer(
                 it
             }
         }
+        if (newBody != null && loopBodyTransformer != null && loopBodyTransformer.shouldTransform(context)) {
+            loopBodyTransformer.transform(context, newBody, mainLoopVariable, loopHeader, loopVariableComponents)
+        }
 
         return loopHeader.buildLoop(context.createIrBuilder(getScopeOwnerSymbol(), loop.startOffset, loop.endOffset), loop, newBody)
     }

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

@@ -19,7 +19,7 @@ import org.jetbrains.kotlin.ir.expressions.IrExpression
 import org.jetbrains.kotlin.ir.symbols.IrSymbol
 import org.jetbrains.kotlin.ir.visitors.IrElementVisitor
 
-internal enum class ProgressionDirection {
+enum class ProgressionDirection {
     DECREASING {
         override fun asReversed() = INCREASING
     },
@@ -34,7 +34,7 @@ internal enum class ProgressionDirection {
 }
 
 /** Information about a loop that is required by [HeaderProcessor] to build a [ForLoopHeader]. */
-internal sealed class HeaderInfo {
+sealed class HeaderInfo {
     /**
      * Returns a copy of this [HeaderInfo] with the values reversed.
      * I.e., first and last are swapped, step is negated.
@@ -59,7 +59,7 @@ internal class ComparableRangeInfo(
     override fun asReversed(): HeaderInfo? = null
 }
 
-internal sealed class NumericHeaderInfo(
+sealed class NumericHeaderInfo(
     val progressionType: ProgressionType,
     val first: IrExpression,
     val last: IrExpression,
@@ -71,7 +71,7 @@ internal sealed class NumericHeaderInfo(
 ) : HeaderInfo()
 
 /** Information about a for-loop over a progression. */
-internal class ProgressionHeaderInfo(
+class ProgressionHeaderInfo(
     progressionType: ProgressionType,
     first: IrExpression,
     last: IrExpression,
@@ -169,7 +169,7 @@ internal class ProgressionHeaderInfo(
  * Information about a for-loop over an object with an indexed get method (such as arrays or character sequences).
  * The internal induction variable used is an Int.
  */
-internal class IndexedGetHeaderInfo(
+class IndexedGetHeaderInfo(
     symbols: Symbols<CommonBackendContext>,
     first: IrExpression,
     last: IrExpression,
@@ -203,7 +203,7 @@ internal class IndexedGetHeaderInfo(
 /**
  * Information about a for-loop over an iterable returned by `withIndex()`.
  */
-internal class WithIndexHeaderInfo(val nestedInfo: HeaderInfo) : HeaderInfo() {
+class WithIndexHeaderInfo(val nestedInfo: HeaderInfo) : HeaderInfo() {
     // We cannot easily reverse `withIndex()` so we do not attempt to handle it. We would have to start from the last value of the index,
     // which is not easily calculable (or even impossible) in most cases.
     override fun asReversed(): HeaderInfo? = null

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

@@ -29,12 +29,12 @@ import org.jetbrains.kotlin.util.OperatorNameConventions
  * @param replacementExpression The expression to use in place of the old loop. It is either `newLoop`, or a container
  * that contains `newLoop`.
  */
-internal data class LoopReplacement(
+data class LoopReplacement(
     val newLoop: IrLoop,
     val replacementExpression: IrExpression
 )
 
-internal interface ForLoopHeader {
+interface ForLoopHeader {
     /** Statements used to initialize the entire loop (e.g., declare induction variable). */
     val loopInitStatements: List<IrStatement>
 
@@ -56,8 +56,8 @@ internal interface ForLoopHeader {
     fun buildLoop(builder: DeclarationIrBuilder, oldLoop: IrLoop, newBody: IrExpression?): LoopReplacement
 }
 
-internal abstract class NumericForLoopHeader<T : NumericHeaderInfo>(
-    protected val headerInfo: T,
+abstract class NumericForLoopHeader<T : NumericHeaderInfo>(
+    val headerInfo: T,
     builder: DeclarationIrBuilder,
     context: CommonBackendContext
 ) : ForLoopHeader {
@@ -246,7 +246,7 @@ internal abstract class NumericForLoopHeader<T : NumericHeaderInfo>(
     }
 }
 
-internal class ProgressionLoopHeader(
+class ProgressionLoopHeader(
     headerInfo: ProgressionHeaderInfo,
     builder: DeclarationIrBuilder,
     context: CommonBackendContext
@@ -367,7 +367,7 @@ private class InitializerCallReplacer(val replacementCall: IrCall) : IrElementTr
     }
 }
 
-internal class IndexedGetLoopHeader(
+class IndexedGetLoopHeader(
     headerInfo: IndexedGetHeaderInfo,
     builder: DeclarationIrBuilder,
     context: CommonBackendContext
@@ -418,14 +418,14 @@ internal class IndexedGetLoopHeader(
     }
 }
 
-internal class WithIndexLoopHeader(
+class WithIndexLoopHeader(
     headerInfo: WithIndexHeaderInfo,
     builder: DeclarationIrBuilder,
     context: CommonBackendContext
 ) : ForLoopHeader {
 
-    private val nestedLoopHeader: ForLoopHeader
-    private val indexVariable: IrVariable
+    val nestedLoopHeader: ForLoopHeader
+    val indexVariable: IrVariable
     private val ownsIndexVariable: Boolean
     private val incrementIndexStatement: IrStatement?
 

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

@@ -24,7 +24,7 @@ import org.jetbrains.kotlin.ir.types.isSubtypeOfClass
 import org.jetbrains.kotlin.ir.util.defaultType
 
 /** Represents a progression type in the Kotlin stdlib. */
-internal sealed class ProgressionType(
+sealed class ProgressionType(
     val elementClass: IrClass,
     val stepClass: IrClass,
     val minValueAsLong: Long,

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

@@ -28,7 +28,7 @@ import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.util.OperatorNameConventions
 
 /** Builds a [HeaderInfo] for iteration over iterables using the `get / []` operator and an index. */
-internal abstract class IndexedGetIterationHandler(
+abstract class IndexedGetIterationHandler(
     protected val context: CommonBackendContext,
     private val canCacheLast: Boolean
 ) : ExpressionHandler {

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

@@ -15,6 +15,7 @@ import org.jetbrains.kotlin.backend.konan.lower.*
 import org.jetbrains.kotlin.backend.konan.lower.FinallyBlocksLowering
 import org.jetbrains.kotlin.backend.konan.lower.InitializersLowering
 import org.jetbrains.kotlin.backend.konan.lower.StringConcatenationLowering
+import org.jetbrains.kotlin.backend.konan.optimizations.KonanBCEForLoopBodyTransformer
 import org.jetbrains.kotlin.ir.declarations.IrFile
 import org.jetbrains.kotlin.ir.declarations.IrModuleFragment
 
@@ -257,8 +258,10 @@ internal val rangeContainsLoweringPhase = makeKonanFileLoweringPhase(
         description = "Optimizes calls to contains() for ClosedRanges"
 )
 
-internal val forLoopsPhase = makeKonanFileLoweringPhase(
-        ::ForLoopsLowering,
+internal val forLoopsPhase = makeKonanFileOpPhase(
+        { context, irFile ->
+            ForLoopsLowering(context, KonanBCEForLoopBodyTransformer()).lower(irFile)
+        },
         name = "ForLoops",
         description = "For loops lowering"
 )

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

@@ -0,0 +1,244 @@
+/*
+ * Copyright 2010-2021 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the LICENSE file.
+ */
+
+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
+import org.jetbrains.kotlin.ir.expressions.*
+import org.jetbrains.kotlin.ir.expressions.impl.IrCallImpl
+import org.jetbrains.kotlin.ir.symbols.IrValueSymbol
+import org.jetbrains.kotlin.ir.types.IrType
+import org.jetbrains.kotlin.ir.types.getClass
+import org.jetbrains.kotlin.ir.types.isArray
+import org.jetbrains.kotlin.ir.util.IdSignature
+import org.jetbrains.kotlin.ir.util.functions
+import org.jetbrains.kotlin.ir.util.getPropertyGetter
+import org.jetbrains.kotlin.ir.util.isPrimitiveArray
+import org.jetbrains.kotlin.util.OperatorNameConventions
+
+// Class contains information about analyzed loop.
+internal data class BoundsCheckAnalysisResult(val boundsAreSafe: Boolean, val arrayInLoop: IrValueSymbol?)
+// TODO: support `forEachIndexed`. Function is inlined and index is separate variable which isn't connected with loop induction variable.
+/**
+ * Transformer for for loops bodies replacing get/set operators on analogs without bounds check where it's possible.
+ */
+class KonanBCEForLoopBodyTransformer : ForLoopBodyTransformer() {
+    private var analysisResult: BoundsCheckAnalysisResult = BoundsCheckAnalysisResult(false, null)
+
+    override fun shouldTransform(context: CommonBackendContext): Boolean {
+        return context is Context && context.shouldOptimize()
+    }
+
+    override fun initialize(context: CommonBackendContext, loopVariable: IrVariable,
+                            forLoopHeader: ForLoopHeader, loopComponents: Map<Int, IrVariable>) {
+        super.initialize(context, loopVariable, forLoopHeader, loopComponents)
+        analysisResult = analyzeLoopHeader(loopHeader)
+    }
+
+    private fun IrExpression.compareIntegerNumericConst(compare: (Long) -> Boolean): Boolean {
+        @Suppress("UNCHECKED_CAST")
+        return this is IrConst<*> && value is Number && compare((value as Number).toLong())
+    }
+
+    private fun IrExpression.compareFloatNumericConst(compare: (Double) -> Boolean): Boolean {
+        @Suppress("UNCHECKED_CAST")
+        return this is IrConst<*> && value is Number && compare((value as Number).toDouble())
+    }
+
+    private fun IrType.isBasicArray() = isPrimitiveArray() || isArray()
+
+    private fun IrCall.isGetSizeCall() = dispatchReceiver?.type?.isBasicArray() == true &&
+            symbol.owner == dispatchReceiver!!.type.getClass()!!.getPropertyGetter("size")!!.owner
+
+    private fun IrCall.dispatchReceiverIsGetSizeCall() = (dispatchReceiver as? IrCall)?.let { it.isGetSizeCall() } ?: false
+
+    private fun lessThanSize(functionCall: IrCall): BoundsCheckAnalysisResult {
+        var result = false
+        when (functionCall.symbol.owner.name) {
+            OperatorNameConventions.DEC ->
+                if (functionCall.dispatchReceiverIsGetSizeCall()) {
+                    result = true
+                }
+            OperatorNameConventions.MINUS -> {
+                val value = functionCall.getValueArgument(0)
+                result = functionCall.dispatchReceiverIsGetSizeCall() &&
+                        value?.compareIntegerNumericConst { it > 0 } == true
+            }
+            OperatorNameConventions.DIV-> {
+                val value = functionCall.getValueArgument(0)
+                result = functionCall.dispatchReceiverIsGetSizeCall() &&
+                        value?.compareFloatNumericConst { it > 1 } == true
+            }
+        }
+        val array = ((functionCall.dispatchReceiver as? IrCall)?.dispatchReceiver as? IrGetValue)?.symbol
+        return BoundsCheckAnalysisResult(result, array)
+    }
+
+    private fun checkLastElement(last: IrExpression, loopHeader: ProgressionLoopHeader): BoundsCheckAnalysisResult {
+        var result = BoundsCheckAnalysisResult(false, null)
+        if (last is IrCall) {
+            result = lessThanSize(last)
+            if (last.isGetSizeCall() && !loopHeader.headerInfo.isLastInclusive) {
+                result = BoundsCheckAnalysisResult(true, (last.dispatchReceiver as? IrGetValue)?.symbol)
+            }
+        }
+        return result
+    }
+
+    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 &&
+                        dispatchReceiver?.type?.getClass()?.symbol in context.ir.symbols.progressionClasses
+
+    private fun analyzeLoopHeader(loopHeader: ForLoopHeader): BoundsCheckAnalysisResult {
+        analysisResult = BoundsCheckAnalysisResult(false, null)
+        when(loopHeader) {
+            is ProgressionLoopHeader ->
+                when (loopHeader.headerInfo.direction) {
+                    ProgressionDirection.INCREASING -> {
+                        // Analyze first element of progression.
+                        if (!loopHeader.headerInfo.first.compareIntegerNumericConst { it >= 0 }) {
+                            return analysisResult
+                        }
+                        // TODO: variable set to const value. Add constant propagation?
+                        // Analyze last element of progression.
+                        if (loopHeader.headerInfo.last is IrCall) {
+                            val functionCall = (loopHeader.headerInfo.last as IrCall)
+                            // Case of range with step - `for (i in 0..array.size - 1 step n)`.
+                            if (loopHeader.headerInfo.progressionType.getProgressionLastElementFunction == functionCall.symbol) {
+                                val nestedLastVariable = functionCall.getValueArgument(1)
+                                if (nestedLastVariable is IrGetValue && nestedLastVariable.symbol.owner is IrVariable) {
+                                    val nestedLast = (nestedLastVariable.symbol.owner as IrVariable).initializer
+                                    analysisResult = checkLastElement(nestedLast!!, loopHeader)
+                                }
+                            } else {
+                                // Simple progression.
+                                analysisResult = checkLastElement(functionCall, loopHeader)
+                            }
+                        }
+                    }
+                    ProgressionDirection.DECREASING -> {
+                        val valueToCompare = if (loopHeader.headerInfo.isLastInclusive) 0 else -1
+                        var boundsAreSafe = false
+                        if (loopHeader.headerInfo.last is IrCall) {
+                            val functionCall = (loopHeader.headerInfo.last as IrCall)
+                            // Case of range with step - for (i in array.size - 1 downTo 0 step n).
+                            if (loopHeader.headerInfo.progressionType.getProgressionLastElementFunction == functionCall.symbol) {
+                                if (functionCall.getValueArgument(1)?.compareIntegerNumericConst { it >= valueToCompare } == true) {
+                                    boundsAreSafe = true
+                                }
+                            }
+                        } else if (loopHeader.headerInfo.last.compareIntegerNumericConst { it >= valueToCompare }) {
+                            boundsAreSafe = true
+                        }
+                        if (!boundsAreSafe)
+                            return analysisResult
+                        when (loopHeader.headerInfo.first) {
+                            is IrCall -> {
+                                val functionCall = (loopHeader.headerInfo.first as IrCall)
+                                analysisResult = lessThanSize(functionCall)
+                            }
+                            is IrGetValue -> {
+                                (((loopHeader.headerInfo.first as IrGetValue).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)
+                        if (loopHeader.headerInfo.first.isProgressionPropertyGetter("first") &&
+                                loopHeader.headerInfo.last.isProgressionPropertyGetter("last")) {
+                            val firstReceiver = (loopHeader.headerInfo.first as IrCall).dispatchReceiver as? IrGetValue
+                            val lastReceiver = (loopHeader.headerInfo.last as IrCall).dispatchReceiver as? IrGetValue
+                            if (firstReceiver?.symbol?.owner == lastReceiver?.symbol?.owner) {
+                                val untilFunction =
+                                        ((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 = lessThanSize(last)
+                                        // `isLastInclusive` for current case is set to true.
+                                        // This case isn't fully optimized in ForLoopsLowering.
+                                        if (last.isGetSizeCall()) {
+                                            analysisResult = BoundsCheckAnalysisResult(true, (last.dispatchReceiver as? IrGetValue)?.symbol)
+                                        }
+                                    }
+                                }
+                            }
+                        }
+                }
+
+            is WithIndexLoopHeader ->
+                when(loopHeader.nestedLoopHeader) {
+                    is IndexedGetLoopHeader -> {
+                        analysisResult = BoundsCheckAnalysisResult(true,
+                                ((loopHeader.loopInitStatements[0] as? IrVariable)?.initializer as? IrGetValue)?.symbol)
+                    }
+                    is ProgressionLoopHeader -> analysisResult = analyzeLoopHeader(loopHeader.nestedLoopHeader)
+                }
+        }
+        return analysisResult
+    }
+
+    private fun replaceOperators(expression: IrCall, index: IrExpression, safeIndexVariables: List<IrVariable>): IrExpression {
+        if (index is IrGetValue && index.symbol.owner in safeIndexVariables) {
+            val operatorWithoutBC = expression.dispatchReceiver!!.type.getClass()!!.functions.singleOrNull {
+                if (expression.symbol.owner.name == OperatorNameConventions.SET)
+                    it.name == KonanNameConventions.setWithoutBC
+                else
+                    it.name == KonanNameConventions.getWithoutBC
+            } ?: return expression
+            return IrCallImpl(
+                    expression.startOffset, expression.endOffset, expression.type, operatorWithoutBC.symbol,
+                    typeArgumentsCount = expression.typeArgumentsCount,
+                    valueArgumentsCount = expression.valueArgumentsCount).apply {
+                dispatchReceiver = expression.dispatchReceiver
+                for (argIndex in 0 until expression.valueArgumentsCount) {
+                    putValueArgument(argIndex, expression.getValueArgument(argIndex))
+                }
+            }
+        }
+        return expression
+    }
+
+    override fun visitCall(expression: IrCall): IrExpression {
+        if (!analysisResult.boundsAreSafe || analysisResult.arrayInLoop == null)
+            return expression
+        if (expression.symbol.owner.name != OperatorNameConventions.SET && expression.symbol.owner.name != OperatorNameConventions.GET) {
+            expression.transformChildrenVoid()
+            return expression
+        }
+        if (expression.dispatchReceiver?.type?.isBasicArray() != true ||
+                (expression.dispatchReceiver as? IrGetValue)?.symbol != analysisResult.arrayInLoop)
+            return expression
+        // Analyze arguments of set/get operator.
+        val index = expression.getValueArgument(0)
+        return when(loopHeader) {
+            is ProgressionLoopHeader -> with(loopHeader as ProgressionLoopHeader) {
+                replaceOperators(expression, index!!, listOf(mainLoopVariable, inductionVariable))
+            }
+
+            is WithIndexLoopHeader -> with(loopHeader as WithIndexLoopHeader) {
+                when(nestedLoopHeader) {
+                    is IndexedGetLoopHeader ->
+                        replaceOperators(expression, index!!, listOfNotNull(indexVariable, loopVariableComponents[1]))
+                    is ProgressionLoopHeader ->
+                        replaceOperators(expression, index!!,
+                                listOfNotNull(indexVariable, loopVariableComponents[1], loopVariableComponents[2])
+                        )
+                    else -> expression
+                }
+            }
+
+            else -> expression
+        }
+    }
+}