azalea/kotlin-fork
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Handle withIndex() on Iterables (including progressions) and Sequences

Browse Source 
in ForLoopsLowering.
This commit is contained in:
Mark Punzalan
2019-11-05 11:15:33 -08:00
committed by max-kammerer
parent a54d9482dd
commit 7adffe0007
49 changed files with 1357 additions and 42 deletions
Download Patch File Download Diff File Expand all files Collapse all files
compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/loops/ForLoopsLowering.kt
+1 -1
View File
@@ -118,7 +118,7 @@ private class RangeLoopTransformer(
) : IrElementTransformerVoidWithContext() {
private val symbols = context.ir.symbols
private val headerInfoBuilder = HeaderInfoBuilder(context, this::getScopeOwnerSymbol)
private val headerInfoBuilder = DefaultHeaderInfoBuilder(context, this::getScopeOwnerSymbol)
private val headerProcessor = HeaderProcessor(context, headerInfoBuilder, this::getScopeOwnerSymbol)
fun getScopeOwnerSymbol() = currentScope!!.scope.scopeOwnerSymbol
compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/loops/HeaderInfo.kt
+51 -14
View File
@@ -212,6 +212,13 @@ internal class WithIndexHeaderInfo(val nestedInfo: HeaderInfo) : HeaderInfo() {
override fun asReversed(): HeaderInfo? = null
}
/**
* Information about a for-loop over an Iterable or Sequence.
*/
internal class IterableHeaderInfo(val iteratorVariable: IrVariable) : HeaderInfo() {
override fun asReversed(): HeaderInfo? = null
}
/** Matches an iterable expression and builds a [HeaderInfo] from the expression. */
internal interface HeaderInfoHandler<E : IrExpression, D> {
/** Returns true if the handler can build a [HeaderInfo] from the iterable expression. */
@@ -249,7 +256,7 @@ internal interface HeaderInfoFromCallHandler<D> : HeaderInfoHandler<IrCall, D> {
internal typealias ProgressionHandler = HeaderInfoFromCallHandler<ProgressionType>
internal class HeaderInfoBuilder(context: CommonBackendContext, private val scopeOwnerSymbol: () -> IrSymbol) :
internal abstract class HeaderInfoBuilder(context: CommonBackendContext, private val scopeOwnerSymbol: () -> IrSymbol) :
IrElementVisitor<HeaderInfo?, IrCall?> {
private val symbols = context.ir.symbols
@@ -272,19 +279,8 @@ internal class HeaderInfoBuilder(context: CommonBackendContext, private val scop
StepHandler(context, this)
)
private val callHandlers = listOf(
ReversedHandler(context, this),
WithIndexHandler(context, this)
)
// NOTE: StringIterationHandler MUST come before CharSequenceIterationHandler.
// String is subtype of CharSequence and therefore its handler is more specialized.
private val expressionHandlers = listOf(
ArrayIterationHandler(context),
DefaultProgressionHandler(context),
StringIterationHandler(context),
CharSequenceIterationHandler(context)
)
protected abstract val callHandlers: List<HeaderInfoFromCallHandler<Nothing?>>
protected abstract val expressionHandlers: List<ExpressionHandler>
override fun visitElement(element: IrElement, data: IrCall?): HeaderInfo? = null
@@ -309,4 +305,45 @@ internal class HeaderInfoBuilder(context: CommonBackendContext, private val scop
return expressionHandlers.firstNotNullResult { it.handle(iterable, iteratorCall, null, scopeOwnerSymbol()) }
?: super.visitExpression(iterable, iteratorCall)
}
}
internal class DefaultHeaderInfoBuilder(context: CommonBackendContext, scopeOwnerSymbol: () -> IrSymbol) :
HeaderInfoBuilder(context, scopeOwnerSymbol) {
override val callHandlers = listOf(
ReversedHandler(context, this),
WithIndexHandler(context, NestedHeaderInfoBuilderForWithIndex(context, scopeOwnerSymbol))
)
// NOTE: StringIterationHandler MUST come before CharSequenceIterationHandler.
// String is subtype of CharSequence and therefore its handler is more specialized.
override val expressionHandlers = listOf(
ArrayIterationHandler(context),
DefaultProgressionHandler(context),
StringIterationHandler(context),
CharSequenceIterationHandler(context)
)
}
// WithIndexHandler attempts to retrieve the HeaderInfo from the underlying index, using NestedHeaderInfoBuilderForWithIndex instead of
// DefaultHeaderInfoBuilder. The differences between the two are that NestedHeaderInfoBuilderForWithIndex:
//
// - Has NO WithIndexHandler. We do not attempt to optimize `*.withIndex().withIndex()`.
// - Has DefaultIterableHandler. This allows us to optimize `Iterable<*>.withIndex()` and `Sequence<*>.withIndex()`.
internal class NestedHeaderInfoBuilderForWithIndex(context: CommonBackendContext, scopeOwnerSymbol: () -> IrSymbol) :
HeaderInfoBuilder(context, scopeOwnerSymbol) {
// NOTE: No WithIndexHandler; we cannot lower `iterable.withIndex().withIndex()`.
override val callHandlers = listOf(
ReversedHandler(context, this)
)
// NOTE: StringIterationHandler MUST come before CharSequenceIterationHandler.
// String is subtype of CharSequence and therefore its handler is more specialized.
// DefaultIterableHandler must come last as it is handles iterables not handled by more specialized handlers.
override val expressionHandlers = listOf(
ArrayIterationHandler(context),
DefaultProgressionHandler(context),
StringIterationHandler(context),
CharSequenceIterationHandler(context),
DefaultIterableHandler(context)
)
}
compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/loops/HeaderProcessor.kt
+77 -11
View File
@@ -22,6 +22,7 @@ import org.jetbrains.kotlin.ir.expressions.impl.IrDoWhileLoopImpl
import org.jetbrains.kotlin.ir.expressions.impl.IrWhileLoopImpl
import org.jetbrains.kotlin.ir.symbols.IrSymbol
import org.jetbrains.kotlin.ir.types.*
import org.jetbrains.kotlin.ir.util.coerceToUnitIfNeeded
import org.jetbrains.kotlin.ir.util.deepCopyWithSymbols
import org.jetbrains.kotlin.ir.util.functions
import org.jetbrains.kotlin.util.OperatorNameConventions
@@ -348,7 +349,7 @@ internal class IndexedGetLoopHeader(
}
internal class WithIndexLoopHeader(
private val headerInfo: WithIndexHeaderInfo,
headerInfo: WithIndexHeaderInfo,
builder: DeclarationIrBuilder
) : ForLoopHeader {
@@ -361,8 +362,9 @@ internal class WithIndexLoopHeader(
// To build the optimized/lowered `for` loop over a `withIndex()` call, we first need the header for the underlying iterable so
// so that we know how to build the loop for that iterable. More info in comments in initializeIteration().
nestedLoopHeader = when (val nestedInfo = headerInfo.nestedInfo) {
is IndexedGetHeaderInfo -> IndexedGetLoopHeader(nestedInfo, builder)
is ProgressionHeaderInfo -> ProgressionLoopHeader(nestedInfo, builder)
is IndexedGetHeaderInfo -> IndexedGetLoopHeader(nestedInfo, this@with)
is ProgressionHeaderInfo -> ProgressionLoopHeader(nestedInfo, this@with)
is IterableHeaderInfo -> IterableLoopHeader(nestedInfo)
is WithIndexHeaderInfo -> throw IllegalStateException("Nested WithIndexHeaderInfo not allowed for WithIndexLoopHeader")
}
@@ -424,9 +426,9 @@ internal class WithIndexLoopHeader(
// val step = 2
// if (inductionVar <= last) {
// do {
// val v = inductionVar
// inductionVar += step
// // Loop body
// val v = inductionVar
// inductionVar += step
// // Loop body
// } while (inductionVar <= last)
// }
//
@@ -438,14 +440,30 @@ internal class WithIndexLoopHeader(
// var index = 0 // ADDED
// if (inductionVar <= last) {
// do {
// val i = index // ADDED
// val v = inductionVar
// inductionVar += step
// // Loop body
// checkIndexOverflow(index++) // ADDED
// val i = index // ADDED
// val v = inductionVar
// inductionVar += step
// // Loop body
// checkIndexOverflow(index++) // ADDED
// } while (inductionVar <= last)
// }
//
// As another example, in a for-loop over a call to `Iterable<*>.withIndex()` or `Sequence<*>.withIndex()`, e.g.:
//
// for ((i, v) in listOf(2, 3, 5, 7, 11).withIndex()) { /* Loop body */ }
//
// For-loops over an Iterable are normally not optimized, but when getting the underlying iterable for `withIndex()` (and ONLY
// in this case), we use DefaultIterableHandler to match it and IterableLoopHeader to build the underlying loop. The optimized
// loop with `withIndex()` looks something like this:
//
// val iterator = listOf(2, 3, 5, 7, 11).iterator()
// var index = 0
// while (it.hasNext())
// val i = index
// val v = it.next()
// checkIndexOverflow(index++)
// }
//
// We "wire" the 1st destructured component to index, and the 2nd to the loop variable value from the underlying iterable.
loopVariableComponents[1]?.initializer = irGet(indexVariable)
listOfNotNull(loopVariableComponents[1]) + nestedLoopHeader.initializeIteration(
@@ -482,6 +500,53 @@ internal class WithIndexLoopHeader(
}
}
internal class IterableLoopHeader(
private val headerInfo: IterableHeaderInfo
) : ForLoopHeader {
override val loopInitStatements = listOf(headerInfo.iteratorVariable)
override val consumesLoopVariableComponents = false
override fun initializeIteration(
loopVariable: IrVariable?,
loopVariableComponents: Map<Int, IrVariable>,
symbols: Symbols<CommonBackendContext>,
builder: DeclarationIrBuilder
) =
with(builder) {
// loopVariable = iteratorVar.next()
val iteratorClass = headerInfo.iteratorVariable.type.getClass()!!
val next =
irCall(iteratorClass.functions.first { it.name == OperatorNameConventions.NEXT && it.valueParameters.isEmpty() }).apply {
dispatchReceiver = irGet(headerInfo.iteratorVariable)
}
loopVariable?.initializer = next
// Even if there is no loop variable, we always want to call `next()` for iterables and sequences.
listOf(loopVariable ?: next.coerceToUnitIfNeeded(next.type, context.irBuiltIns))
}
override fun buildLoop(builder: DeclarationIrBuilder, oldLoop: IrLoop, newBody: IrExpression?): LoopReplacement = with(builder) {
// Loop is lowered into something like:
//
// var iteratorVar = someIterable.iterator()
// while (iteratorVar.hasNext()) {
// val loopVar = iteratorVar.next()
// // Loop body
// }
val iteratorClass = headerInfo.iteratorVariable.type.getClass()!!
val hasNext =
irCall(iteratorClass.functions.first { it.name == OperatorNameConventions.HAS_NEXT && it.valueParameters.isEmpty() }).apply {
dispatchReceiver = irGet(headerInfo.iteratorVariable)
}
val newLoop = IrWhileLoopImpl(oldLoop.startOffset, oldLoop.endOffset, oldLoop.type, oldLoop.origin).apply {
label = oldLoop.label
condition = hasNext
body = newBody
}
LoopReplacement(newLoop, newLoop)
}
}
/**
* Given the for-loop iterator variable, extract information about the iterable subject
* and create a [ForLoopHeader] from it.
@@ -529,6 +594,7 @@ internal class HeaderProcessor(
is IndexedGetHeaderInfo -> IndexedGetLoopHeader(headerInfo, builder)
is ProgressionHeaderInfo -> ProgressionLoopHeader(headerInfo, builder)
is WithIndexHeaderInfo -> WithIndexLoopHeader(headerInfo, builder)
is IterableHeaderInfo -> IterableLoopHeader(headerInfo)
}
}
}
compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/loops/ProgressionHandlers.kt
+30 -5
View File
@@ -697,14 +697,14 @@ internal class StringIterationHandler(context: CommonBackendContext) : CharSeque
}
/** Builds a [HeaderInfo] for calls to `withIndex()`. */
internal class WithIndexHandler(context: CommonBackendContext, private val visitor: HeaderInfoBuilder) :
internal class WithIndexHandler(context: CommonBackendContext, private val visitor: NestedHeaderInfoBuilderForWithIndex) :
HeaderInfoFromCallHandler<Nothing?> {
// Use Quantifier.ANY so we can handle all `withIndex()` calls in the same manner.
override val matcher = createIrCallMatcher(Quantifier.ANY) {
callee {
fqName { it == FqName("kotlin.collections.withIndex") }
extensionReceiver { it != null && it.type.run { isArray() || isPrimitiveArray() } }
extensionReceiver { it != null && it.type.run { isArray() || isPrimitiveArray() || isIterable() } }
parameterCount { it == 0 }
}
callee {
@@ -712,8 +712,11 @@ internal class WithIndexHandler(context: CommonBackendContext, private val visit
extensionReceiver { it != null && it.type.isSubtypeOfClass(context.ir.symbols.charSequence) }
parameterCount { it == 0 }
}
// TODO: Handle Iterable.withIndex(), Sequence.withIndex()
callee {
fqName { it == FqName("kotlin.sequences.withIndex") }
extensionReceiver { it != null && it.type.run { isSequence() } }
parameterCount { it == 0 }
}
}
override fun build(expression: IrCall, data: Nothing?, scopeOwner: IrSymbol): HeaderInfo? {
@@ -721,8 +724,30 @@ internal class WithIndexHandler(context: CommonBackendContext, private val visit
val nestedInfo = expression.extensionReceiver!!.accept(visitor, null) ?: return null
// We cannot lower `iterable.withIndex().withIndex()`.
if (nestedInfo is WithIndexHeaderInfo) return null
// NestedHeaderInfoBuilderForWithIndex should not be yielding a WithIndexHeaderInfo, hence the assert.
assert(nestedInfo !is WithIndexHeaderInfo)
return WithIndexHeaderInfo(nestedInfo)
}
}
/** Builds a [HeaderInfo] for iterables not handled by more specialized handlers. */
internal class DefaultIterableHandler(private val context: CommonBackendContext) : ExpressionHandler {
override fun matchIterable(expression: IrExpression) = true
override fun build(expression: IrExpression, scopeOwner: IrSymbol): HeaderInfo? =
with(context.createIrBuilder(scopeOwner, expression.startOffset, expression.endOffset)) {
val iterableClass = expression.type.getClass()!!
val iterator =
irCall(iterableClass.functions.single {
it.name == OperatorNameConventions.ITERATOR &&
it.valueParameters.isEmpty()
}).apply {
dispatchReceiver = expression
}
IterableHeaderInfo(
scope.createTemporaryVariable(iterator, nameHint = "iterator")
)
}
}