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

2021-06-18 17:32:57 +03:00
parent 7975311ca2
commit 0e04c21625
7 changed files with 308 additions and 22 deletions
@@ -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)
    }
@@ -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
@@ -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?

@@ -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,
@@ -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 {