JVM_IR generate range-based loop closer to Java counter loop

KT-48435 KT-48507
This commit is contained in:
Dmitry Petrov
2021-09-01 15:31:19 +03:00
parent 9ed08438d5
commit b669de1663
137 changed files with 1518 additions and 573 deletions
@@ -11,9 +11,11 @@ import org.jetbrains.kotlin.ir.IrElement
import org.jetbrains.kotlin.ir.builders.IrBuilderWithScope
import org.jetbrains.kotlin.ir.builders.irCall
import org.jetbrains.kotlin.ir.builders.irString
import org.jetbrains.kotlin.ir.declarations.IrDeclarationOrigin
import org.jetbrains.kotlin.ir.declarations.IrFile
import org.jetbrains.kotlin.ir.expressions.IrCall
import org.jetbrains.kotlin.ir.expressions.IrExpression
import org.jetbrains.kotlin.ir.expressions.IrStatementOrigin
import org.jetbrains.kotlin.ir.symbols.IrClassSymbol
import org.jetbrains.kotlin.ir.symbols.IrFileSymbol
import org.jetbrains.kotlin.ir.symbols.IrSimpleFunctionSymbol
@@ -53,4 +55,13 @@ interface CommonBackendContext : BackendContext, LoggingContext {
val preferJavaLikeCounterLoop: Boolean
get() = false
val reuseLoopVariableAsInductionVariable: Boolean
get() = false
val doWhileCounterLoopOrigin: IrStatementOrigin?
get() = null
val inductionVariableOrigin: IrDeclarationOrigin
get() = IrDeclarationOrigin.IR_TEMPORARY_VARIABLE
}
@@ -47,7 +47,7 @@ abstract class AbstractVariableRemapper : IrElementTransformerVoid() {
override fun visitSetValue(expression: IrSetValue): IrExpression {
expression.transformChildrenVoid()
return remapVariable(expression.symbol.owner)?.let {
IrSetValueImpl(expression.startOffset, expression.endOffset, it.type, it.symbol, expression.value, expression.origin)
IrSetValueImpl(expression.startOffset, expression.endOffset, expression.type, it.symbol, expression.value, expression.origin)
} ?: expression
}
}
@@ -8,16 +8,18 @@ package org.jetbrains.kotlin.backend.common.lower.loops
import org.jetbrains.kotlin.backend.common.BodyLoweringPass
import org.jetbrains.kotlin.backend.common.CommonBackendContext
import org.jetbrains.kotlin.backend.common.IrElementTransformerVoidWithContext
import org.jetbrains.kotlin.backend.common.lower.AbstractVariableRemapper
import org.jetbrains.kotlin.backend.common.lower.createIrBuilder
import org.jetbrains.kotlin.backend.common.phaser.makeIrFilePhase
import org.jetbrains.kotlin.ir.IrElement
import org.jetbrains.kotlin.ir.IrStatement
import org.jetbrains.kotlin.ir.declarations.IrDeclaration
import org.jetbrains.kotlin.ir.declarations.IrDeclarationOrigin
import org.jetbrains.kotlin.ir.declarations.IrSymbolOwner
import org.jetbrains.kotlin.ir.declarations.IrVariable
import org.jetbrains.kotlin.ir.declarations.*
import org.jetbrains.kotlin.ir.declarations.impl.IrVariableImpl
import org.jetbrains.kotlin.ir.expressions.*
import org.jetbrains.kotlin.ir.expressions.impl.IrCompositeImpl
import org.jetbrains.kotlin.ir.symbols.impl.IrVariableSymbolImpl
import org.jetbrains.kotlin.ir.types.classifierOrNull
import org.jetbrains.kotlin.ir.types.isNullable
import org.jetbrains.kotlin.ir.util.dump
import org.jetbrains.kotlin.ir.util.render
import org.jetbrains.kotlin.ir.visitors.*
@@ -97,7 +99,10 @@ val forLoopsPhase = makeIrFilePhase(
* }
* ```
*/
class ForLoopsLowering(val context: CommonBackendContext, val loopBodyTransformer: ForLoopBodyTransformer? = null) : BodyLoweringPass {
class ForLoopsLowering(
val context: CommonBackendContext,
private val loopBodyTransformer: ForLoopBodyTransformer? = null
) : BodyLoweringPass {
override fun lower(irBody: IrBody, container: IrDeclaration) {
val oldLoopToNewLoop = mutableMapOf<IrLoop, IrLoop>()
@@ -160,30 +165,127 @@ private class RangeLoopTransformer(
return super.visitBlock(expression) // Not a for-loop block.
}
with(expression.statements) {
assert(size == 2) { "Expected 2 statements in for-loop block, was:\n${expression.dump()}" }
val iteratorVariable = get(0) as IrVariable
assert(iteratorVariable.origin == IrDeclarationOrigin.FOR_LOOP_ITERATOR) { "Expected FOR_LOOP_ITERATOR origin for iterator variable, was:\n${iteratorVariable.dump()}" }
val loopHeader = headerProcessor.extractHeader(iteratorVariable)
?: return super.visitBlock(expression) // The iterable in the header is not supported.
val loweredHeader = lowerHeader(iteratorVariable, loopHeader)
val statements = expression.statements
assert(statements.size == 2) { "Expected 2 statements in for-loop block, was:\n${expression.dump()}" }
val iteratorVariable = statements[0] as IrVariable
assert(iteratorVariable.origin == IrDeclarationOrigin.FOR_LOOP_ITERATOR) {
"Expected FOR_LOOP_ITERATOR origin for iterator variable, was:\n${iteratorVariable.dump()}"
}
val oldLoop = statements[1] as IrWhileLoop
assert(oldLoop.origin == IrStatementOrigin.FOR_LOOP_INNER_WHILE) {
"Expected FOR_LOOP_INNER_WHILE origin for while loop, was:\n${oldLoop.dump()}"
}
val oldLoop = get(1) as IrWhileLoop
assert(oldLoop.origin == IrStatementOrigin.FOR_LOOP_INNER_WHILE) { "Expected FOR_LOOP_INNER_WHILE origin for while loop, was:\n${oldLoop.dump()}" }
val (newLoop, loopReplacementExpression) = lowerWhileLoop(oldLoop, loopHeader)
?: return super.visitBlock(expression) // Cannot lower the loop.
val loopHeader = headerProcessor.extractHeader(iteratorVariable)
?: return super.visitBlock(expression) // The iterable in the header is not supported.
val loweredHeader = lowerHeader(iteratorVariable, loopHeader)
// We can lower both the header and while loop.
// Update mapping from old to new loop so we can later update references in break/continue.
oldLoopToNewLoop[oldLoop] = newLoop
val (newLoop, loopReplacementExpression) = lowerWhileLoop(oldLoop, loopHeader)
?: return super.visitBlock(expression) // Cannot lower the loop.
set(0, loweredHeader)
set(1, loopReplacementExpression)
// We can lower both the header and while loop.
// Update mapping from old to new loop so we can later update references in break/continue.
oldLoopToNewLoop[oldLoop] = newLoop
statements[0] = loweredHeader
statements[1] = loopReplacementExpression
if (context.reuseLoopVariableAsInductionVariable && loopHeader.canReuseLoopVariableAsInductionVariable) {
reuseLoopVariableAsInductionVariable(expression)
}
return super.visitBlock(expression)
}
private fun reuseLoopVariableAsInductionVariable(irBlock: IrBlock) {
// Given a loop in the form:
// {
// var inductionVariable = <start>
// }
// do {
// if (!(<whileLoopCondition>)) break
// val loopVariable = inductionVariable
// <originalLoopBody>
// } while ( { inductionVariable += <step>; true } )
// replace it with:
// {
// var loopVariable' = <start>
// }
// do {
// if (!(<whileLoopCondition'>)) break
// <originalLoopBody'>
// } while ( { loopVariable' += <step>; true } )
// where whenLoopCondition' and originalLoopBody' are corresponding statements
// with inductionVariable and loopVariable remapped to loopVariable'.
//
// NB we can do so only with a do-while counter loop as described above,
// otherwise it changes semantics of 'continue' inside the loop.
val header = irBlock.statements[0] as? IrStatementContainer ?: return
val inductionVariableIndex = header.statements.indexOfFirst { it.isInductionVariable(context) }
if (inductionVariableIndex < 0) return
val inductionVariable = header.statements[inductionVariableIndex] as IrVariable
val innerLoop = findInnerDoWhileLoop(irBlock.statements[1]) ?: return
if (innerLoop.origin != context.doWhileCounterLoopOrigin) return
val loopVariableContainerAndIndex = findLoopVariable(innerLoop) ?: return
val (loopVariableContainer, loopVariableIndex) = loopVariableContainerAndIndex
val loopVariable = loopVariableContainer.statements[loopVariableIndex] as IrVariable
val inductionVariableType = inductionVariable.type
val loopVariableType = loopVariable.type
if (loopVariableType.isNullable()) return
if (loopVariableType.classifierOrNull != inductionVariableType.classifierOrNull) return
val newLoopVariable = IrVariableImpl(
loopVariable.startOffset, loopVariable.endOffset, loopVariable.origin,
IrVariableSymbolImpl(),
loopVariable.name, loopVariableType,
isVar = true, // NB original loop variable is 'val'
isConst = false, isLateinit = false
)
newLoopVariable.initializer = inductionVariable.initializer
header.statements[inductionVariableIndex] = newLoopVariable
loopVariableContainer.statements.removeAt(loopVariableIndex)
val remapper = object : AbstractVariableRemapper() {
override fun remapVariable(value: IrValueDeclaration): IrValueDeclaration? =
if (value == inductionVariable || value == loopVariable) newLoopVariable else null
}
irBlock.statements[1].transformChildren(remapper, null)
}
private fun findInnerDoWhileLoop(statement: IrStatement): IrDoWhileLoop? {
if (statement is IrDoWhileLoop) {
return statement
}
if (statement is IrWhen) {
val branch0Result = statement.branches[0].result
if (branch0Result is IrDoWhileLoop)
return branch0Result
}
return null
}
private fun findLoopVariable(doWhileLoop: IrDoWhileLoop): Pair<IrContainerExpression, Int>? {
val loopBody = doWhileLoop.body as? IrContainerExpression ?: return null
for ((index, statement) in loopBody.statements.withIndex()) {
if (statement.isLoopVariable())
return Pair(loopBody, index)
else if (statement is IrContainerExpression && statement.origin == IrStatementOrigin.FOR_LOOP_NEXT) {
val loopVarIndex = statement.statements.indexOfFirst { it.isLoopVariable() }
if (loopVarIndex < 0) return null
return Pair(statement, loopVarIndex)
}
}
return null
}
private fun IrStatement.isLoopVariable() =
this is IrVariable && origin == IrDeclarationOrigin.FOR_LOOP_VARIABLE
/**
* Lowers the "header" statement that stores the iterator into the loop variable
* (e.g., `val it = someIterable.iterator()`) and gather information for building the for-loop
@@ -204,9 +306,8 @@ private class RangeLoopTransformer(
private fun lowerWhileLoop(loop: IrWhileLoop, loopHeader: ForLoopHeader): LoopReplacement? {
val loopBodyStatements = (loop.body as? IrContainerExpression)?.statements ?: return null
val (mainLoopVariable, mainLoopVariableIndex, loopVariableComponents, loopVariableComponentIndices) = gatherLoopVariableInfo(
loopBodyStatements
)
val (mainLoopVariable, mainLoopVariableIndex, loopVariableComponents, loopVariableComponentIndices) =
gatherLoopVariableInfo(loopBodyStatements)
if (loopHeader.consumesLoopVariableComponents && mainLoopVariable.origin != IrDeclarationOrigin.IR_TEMPORARY_VARIABLE) {
// We determine if there is a destructuring declaration by checking if the main loop variable is temporary.
@@ -44,6 +44,10 @@ interface ForLoopHeader {
*/
val consumesLoopVariableComponents: Boolean
/** `true` if it's possible to use loop variable as induction variable in this kind of loop */
val canReuseLoopVariableAsInductionVariable: Boolean
get() = false
/** Statements used to initialize an iteration of the loop (e.g., assign loop variable). */
fun initializeIteration(
loopVariable: IrVariable?,
@@ -56,6 +60,13 @@ interface ForLoopHeader {
fun buildLoop(builder: DeclarationIrBuilder, oldLoop: IrLoop, newBody: IrExpression?): LoopReplacement
}
internal const val inductionVariableName = "inductionVariable"
internal fun IrStatement.isInductionVariable(context: CommonBackendContext) =
this is IrVariable &&
origin == context.inductionVariableOrigin &&
name.asString() == inductionVariableName
abstract class NumericForLoopHeader<T : NumericHeaderInfo>(
val headerInfo: T,
builder: DeclarationIrBuilder,
@@ -64,6 +75,8 @@ abstract class NumericForLoopHeader<T : NumericHeaderInfo>(
override val consumesLoopVariableComponents = false
override val canReuseLoopVariableAsInductionVariable get() = true
val inductionVariable: IrVariable
protected val stepVariable: IrVariable?
@@ -92,8 +105,9 @@ abstract class NumericForLoopHeader<T : NumericHeaderInfo>(
inductionVariable =
scope.createTmpVariable(
headerInfo.first.asElementType(),
nameHint = "inductionVariable",
nameHint = inductionVariableName,
isMutable = true,
origin = this@NumericForLoopHeader.context.inductionVariableOrigin,
irType = elementClass.defaultType
)
@@ -317,7 +331,11 @@ class ProgressionLoopHeader(
// // Loop body
// } while (loopVar != last)
// }
val newLoop = IrDoWhileLoopImpl(oldLoop.startOffset, oldLoop.endOffset, oldLoop.type, oldLoop.origin).apply {
val newLoopOrigin = if (preferJavaLikeCounterLoop)
this@ProgressionLoopHeader.context.doWhileCounterLoopOrigin
else
oldLoop.origin
val newLoop = IrDoWhileLoopImpl(oldLoop.startOffset, oldLoop.endOffset, oldLoop.type, newLoopOrigin).apply {
val loopVariableExpression = irGet(loopVariable!!).let {
headerInfo.progressionType.run {
if (this is UnsignedProgressionType) {
@@ -331,6 +349,10 @@ class ProgressionLoopHeader(
body = newBody
}
if (preferJavaLikeCounterLoop) {
moveInductionVariableUpdateToLoopCondition(newLoop)
}
val loopCondition = buildLoopCondition(this@with)
LoopReplacement(newLoop, irIfThen(loopCondition, newLoop))
} else if (preferJavaLikeCounterLoop && !headerInfo.isLastInclusive) {
@@ -338,25 +360,17 @@ class ProgressionLoopHeader(
// (`for (int i = first; i < lastExclusive; ++i) { ... }`).
// Otherwise loop-related optimizations will not kick in, resulting in significant performance degradation.
//
// If possible, use a do-while loop:
// Use a do-while loop:
// do {
// if ( !( inductionVariable < last ) ) break
// val loopVariable = inductionVariable
// <body>
// } while ( { inductionVariable += step; true } )
// This loop form is equivalent to the Java counter loop shown above.
//
// Otherwise, use a simple while loop:
// while (inductionVar < last) {
// val loopVar = inductionVar
// inductionVar += step
// // Loop body
// }
val newLoopCondition = buildLoopCondition(this@with)
buildJavaLikeDoWhileCounterLoop(oldLoop, newLoopCondition, newBody)
?: buildJavaLikeWhileCounterLoop(oldLoop, newLoopCondition, newBody)
} else {
// Use an if-guarded do-while loop (note the difference in loop condition):
//
@@ -383,11 +397,52 @@ class ProgressionLoopHeader(
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? {
): LoopReplacement {
// Transform loop:
// while (<newLoopCondition>) {
// {
@@ -402,31 +457,19 @@ class ProgressionLoopHeader(
// val forLoopVariable = inductionVariable
// <originalLoopBody>
// } while ( { inductionVariable += step; true } )
val bodyBlock = newBody as? IrContainerExpression ?: return null
val forLoopNextBlock = bodyBlock.statements[0] as? IrContainerExpression ?: return null
if (forLoopNextBlock.origin != IrStatementOrigin.FOR_LOOP_NEXT) return null
val loopStep = forLoopNextBlock.statements.last() as? IrSetValue ?: return null
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 loopStep = forLoopNextBlock.statements.last() as? IrSetValue
?: throw AssertionError("forLoopNextBlock.last: ${forLoopNextBlock.statements.last().dump()}")
val doWhileLoop = IrDoWhileLoopImpl(oldLoop.startOffset, oldLoop.endOffset, oldLoop.type, oldLoop.origin)
val doWhileLoop = IrDoWhileLoopImpl(oldLoop.startOffset, oldLoop.endOffset, oldLoop.type, context.doWhileCounterLoopOrigin)
doWhileLoop.label = oldLoop.label
val conditionStartOffset = newLoopCondition.startOffset
val conditionEndOffset = newLoopCondition.endOffset
val negatedCondition =
IrCallImpl.fromSymbolOwner(conditionStartOffset, conditionEndOffset, booleanNot.symbol).apply {
dispatchReceiver = newLoopCondition
}
val negatedConditionCheck =
IrWhenImpl(
conditionStartOffset, conditionEndOffset, context.irBuiltIns.unitType, null,
listOf(
IrBranchImpl(
negatedCondition,
IrBreakImpl(conditionStartOffset, conditionEndOffset, context.irBuiltIns.nothingType, doWhileLoop)
)
)
)
val negatedConditionCheck = createNegatedConditionCheck(newLoopCondition, doWhileLoop)
bodyBlock.statements[0] = negatedConditionCheck
val loopVarAssignments =
@@ -455,23 +498,25 @@ class ProgressionLoopHeader(
return LoopReplacement(doWhileLoop, doWhileLoop)
}
private fun buildJavaLikeWhileCounterLoop(
oldLoop: IrLoop,
newLoopCondition: IrExpression,
newBody: IrExpression?
): LoopReplacement {
// while (inductionVar < last) {
// val loopVar = inductionVar
// inductionVar += step
// // Loop body
// }
val newLoop = IrWhileLoopImpl(oldLoop.startOffset, oldLoop.endOffset, oldLoop.type, oldLoop.origin).apply {
label = oldLoop.label
condition = newLoopCondition
body = newBody
}
return LoopReplacement(newLoop, newLoop)
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)
)
)
)
}
}
private class InitializerCallReplacer(val replacementCall: IrCall) : IrElementTransformerVoid() {
@@ -30,6 +30,7 @@ import org.jetbrains.kotlin.ir.declarations.*
import org.jetbrains.kotlin.ir.declarations.impl.IrFactoryImpl
import org.jetbrains.kotlin.ir.expressions.IrExpression
import org.jetbrains.kotlin.ir.expressions.IrFunctionReference
import org.jetbrains.kotlin.ir.expressions.IrStatementOrigin
import org.jetbrains.kotlin.ir.symbols.*
import org.jetbrains.kotlin.ir.types.IrType
import org.jetbrains.kotlin.ir.types.IrTypeSystemContext
@@ -217,6 +218,12 @@ class JvmBackendContext(
override val preferJavaLikeCounterLoop: Boolean
get() = true
override val reuseLoopVariableAsInductionVariable: Boolean
get() = true
override val doWhileCounterLoopOrigin: IrStatementOrigin
get() = JvmLoweredStatementOrigin.DO_WHILE_COUNTER_LOOP
inner class JvmIr(
irModuleFragment: IrModuleFragment,
symbolTable: SymbolTable
@@ -9,4 +9,5 @@ import org.jetbrains.kotlin.ir.expressions.IrStatementOriginImpl
interface JvmLoweredStatementOrigin {
object DEFAULT_STUB_CALL_TO_IMPLEMENTATION : IrStatementOriginImpl("DEFAULT_STUB_CALL_TO_IMPLEMENTATION")
object DO_WHILE_COUNTER_LOOP: IrStatementOriginImpl("DO_WHILE_COUNTER_LOOP")
}