From 643e5814484db14471e22601ad9ef8cd38164689 Mon Sep 17 00:00:00 2001 From: Igor Chevdar Date: Wed, 25 Mar 2020 16:31:35 +0300 Subject: [PATCH] [IR] Supported extraction of local classes from inline bodies The mental model for inline functions/lambdas is as following: inline functions have the `inline` keyword before them so it makes sense for their bodies to be copied at call sites although the compiler might optimize that if there is no usage of reified type parameters or crossinline value parameters. But since there is no `inline` keyword around an inline lambda, then it is logical to think that there will be no copy of it and that it will be just embedded at the corresponding call site as is. --- .../kotlin/backend/common/ir/IrUtils.kt | 11 +- .../backend/common/lower/ClosureAnnotator.kt | 4 +- .../common/lower/LocalDeclarationsLowering.kt | 31 +++- .../common/lower/SharedVariablesLowering.kt | 38 ++-- .../common/lower/inline/FunctionInlining.kt | 10 +- .../common/lower/inline/LocalClasses.kt | 173 ++++++++++++++++++ 6 files changed, 236 insertions(+), 31 deletions(-) create mode 100644 compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/inline/LocalClasses.kt diff --git a/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/ir/IrUtils.kt b/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/ir/IrUtils.kt index 32e5d80fd81..26883a9516a 100644 --- a/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/ir/IrUtils.kt +++ b/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/ir/IrUtils.kt @@ -357,7 +357,16 @@ fun IrDeclarationContainer.addChild(declaration: IrDeclaration) { stageController.unrestrictDeclarationListsAccess { this.declarations += declaration } - declaration.accept(SetDeclarationsParentVisitor, this) + declaration.setDeclarationsParent(this) +} + +fun IrDeclarationContainer.addChildren(declarations: List) { + declarations.forEach { this.addChild(it) } +} + +fun T.setDeclarationsParent(parent: IrDeclarationParent): T { + accept(SetDeclarationsParentVisitor, parent) + return this } object SetDeclarationsParentVisitor : IrElementVisitor { diff --git a/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/ClosureAnnotator.kt b/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/ClosureAnnotator.kt index 625805d9c74..1f55972a9dd 100644 --- a/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/ClosureAnnotator.kt +++ b/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/ClosureAnnotator.kt @@ -32,12 +32,12 @@ import org.jetbrains.kotlin.ir.visitors.IrElementVisitor data class Closure(val capturedValues: List, val capturedTypeParameters: List) -class ClosureAnnotator(body: IrBody, declaration: IrDeclaration) { +class ClosureAnnotator(irElement: IrElement, declaration: IrDeclaration) { private val closureBuilders = mutableMapOf() init { // Collect all closures for classes and functions. Collect call graph - body.accept(ClosureCollectorVisitor(), declaration.closureBuilderOrNull ?: declaration.parentClosureBuilder) + irElement.accept(ClosureCollectorVisitor(), declaration.closureBuilderOrNull ?: declaration.parentClosureBuilder) } fun getFunctionClosure(declaration: IrFunction) = getClosure(declaration) diff --git a/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/LocalDeclarationsLowering.kt b/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/LocalDeclarationsLowering.kt index d30c289414c..8a3df356063 100644 --- a/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/LocalDeclarationsLowering.kt +++ b/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/LocalDeclarationsLowering.kt @@ -39,10 +39,7 @@ import org.jetbrains.kotlin.ir.types.impl.IrSimpleTypeImpl import org.jetbrains.kotlin.ir.types.impl.IrTypeAbbreviationImpl import org.jetbrains.kotlin.ir.types.impl.IrUninitializedType import org.jetbrains.kotlin.ir.types.impl.makeTypeProjection -import org.jetbrains.kotlin.ir.util.constructedClass -import org.jetbrains.kotlin.ir.util.dump -import org.jetbrains.kotlin.ir.util.file -import org.jetbrains.kotlin.ir.util.parentAsClass +import org.jetbrains.kotlin.ir.util.* import org.jetbrains.kotlin.ir.visitors.IrElementTransformerVoid import org.jetbrains.kotlin.ir.visitors.acceptChildrenVoid import org.jetbrains.kotlin.ir.visitors.acceptVoid @@ -113,7 +110,11 @@ class LocalDeclarationsLowering( IrStatementOriginImpl("INITIALIZER_OF_FIELD_FOR_CAPTURED_VALUE") override fun lower(irBody: IrBody, container: IrDeclaration) { - LocalDeclarationsTransformer(irBody, container).lowerLocalDeclarations() + LocalDeclarationsTransformer(irBody, container, null).lowerLocalDeclarations() + } + + fun lower(irElement: IrElement, container: IrDeclaration, classesToLower: Set) { + LocalDeclarationsTransformer(irElement, container, classesToLower).lowerLocalDeclarations() } private class ScopeWithCounter(scope: Scope, irElement: IrElement) : ScopeWithIr(scope, irElement) { @@ -222,7 +223,9 @@ class LocalDeclarationsLowering( abbreviation.annotations ) - private inner class LocalDeclarationsTransformer(val irBody: IrBody, val container: IrDeclaration) { + private inner class LocalDeclarationsTransformer( + val irElement: IrElement, val container: IrDeclaration, val classesToLower: Set? + ) { val localFunctions: MutableMap = LinkedHashMap() val localClasses: MutableMap = LinkedHashMap() val localClassConstructors: MutableMap = LinkedHashMap() @@ -503,7 +506,7 @@ class LocalDeclarationsLowering( rewriteClassMembers(it.declaration, it) } - rewriteFunctionBody(container, null) + rewriteFunctionBody(irElement, null) } private fun createNewCall(oldCall: IrCall, newCallee: IrFunction) = @@ -817,7 +820,7 @@ class LocalDeclarationsLowering( private fun collectClosureForLocalDeclarations() { //TODO: maybe use for granular declarations - val annotator = ClosureAnnotator(irBody, container) + val annotator = ClosureAnnotator(irElement, container) localFunctions.forEach { (declaration, context) -> context.closure = annotator.getFunctionClosure(declaration) @@ -840,7 +843,7 @@ class LocalDeclarationsLowering( currentParent as? IrClass }?.scopeWithCounter - irBody.acceptVoid(object : IrElementVisitorVoidWithContext() { + irElement.acceptVoid(object : IrElementVisitorVoidWithContext() { override fun visitElement(element: IrElement) { element.acceptChildrenVoid(this) @@ -850,6 +853,15 @@ class LocalDeclarationsLowering( return ScopeWithCounter(Scope(declaration.symbol), declaration) // Don't cache local declarations } + override fun visitFunctionExpression(expression: IrFunctionExpression) { + // TODO: For now IrFunctionExpression can only be encountered here if this was called from the inliner, + // then all IrFunctionExpression will be replaced by IrFunctionReferenceExpression. + // Don't forget to fix this when that replacement has been dropped. + // Also, a note: even if a lambda is not an inline one, there still cannot be a reference to it + // from an outside declaration, so it is safe to skip them here and correctly handle later, after the above conversion. + expression.function.acceptChildrenVoid(this) + } + override fun visitSimpleFunction(declaration: IrSimpleFunction) { super.visitSimpleFunction(declaration) @@ -878,6 +890,7 @@ class LocalDeclarationsLowering( } override fun visitClassNew(declaration: IrClass) { + if (classesToLower?.contains(declaration) == false) return super.visitClassNew(declaration) if (!declaration.isLocalNotInner()) return diff --git a/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/SharedVariablesLowering.kt b/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/SharedVariablesLowering.kt index a8a06d5628c..a6d06ff7d2c 100644 --- a/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/SharedVariablesLowering.kt +++ b/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/SharedVariablesLowering.kt @@ -18,6 +18,7 @@ package org.jetbrains.kotlin.backend.common.lower import org.jetbrains.kotlin.backend.common.BackendContext import org.jetbrains.kotlin.backend.common.BodyLoweringPass +import org.jetbrains.kotlin.backend.common.lower.inline.isInlineParameter import org.jetbrains.kotlin.backend.common.phaser.makeIrFilePhase import org.jetbrains.kotlin.ir.IrElement import org.jetbrains.kotlin.ir.IrStatement @@ -26,6 +27,7 @@ import org.jetbrains.kotlin.ir.expressions.* import org.jetbrains.kotlin.ir.symbols.IrValueSymbol import org.jetbrains.kotlin.ir.symbols.IrVariableSymbol import org.jetbrains.kotlin.ir.util.dump +import org.jetbrains.kotlin.ir.util.getArgumentsWithIr import org.jetbrains.kotlin.ir.visitors.IrElementTransformerVoid import org.jetbrains.kotlin.ir.visitors.IrElementVisitor import org.jetbrains.kotlin.ir.visitors.transformChildrenVoid @@ -36,8 +38,6 @@ val sharedVariablesPhase = makeIrFilePhase( description = "Transform shared variables" ) -object CoroutineIntrinsicLambdaOrigin : IrStatementOriginImpl("Coroutine intrinsic lambda") - class SharedVariablesLowering(val context: BackendContext) : BodyLoweringPass { override fun lower(irBody: IrBody, container: IrDeclaration) { @@ -66,20 +66,32 @@ class SharedVariablesLowering(val context: BackendContext) : BodyLoweringPass { element.acceptChildren(this, data) } + override fun visitCall(expression: IrCall, data: IrDeclarationParent?) { + val callee = expression.symbol.owner + if (!callee.isInline) { + super.visitCall(expression, data) + return + } + expression.dispatchReceiver?.accept(this, data) + expression.extensionReceiver?.accept(this, data) + for (param in callee.valueParameters) { + val arg = expression.getValueArgument(param.index) ?: continue + if (param.isInlineParameter() + // This is somewhat conservative but simple. + // If a user put redundant modifier on a parameter, + // may be it's their fault? + && !param.isCrossinline + && arg is IrFunctionExpression + ) + arg.function.acceptChildren(this, data) + else + arg.accept(this, data) + } + } + override fun visitDeclaration(declaration: IrDeclaration, data: IrDeclarationParent?) = super.visitDeclaration(declaration, declaration as? IrDeclarationParent ?: data) - override fun visitContainerExpression(expression: IrContainerExpression, data: IrDeclarationParent?) = - super.visitContainerExpression( - expression, - if (expression is IrReturnableBlock - && expression.origin == CoroutineIntrinsicLambdaOrigin - ) - null - else - data - ) - override fun visitVariable(declaration: IrVariable, data: IrDeclarationParent?) { declaration.acceptChildren(this, data) diff --git a/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/inline/FunctionInlining.kt b/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/inline/FunctionInlining.kt index 952fb7cdff2..c20e4ea4b39 100644 --- a/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/inline/FunctionInlining.kt +++ b/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/inline/FunctionInlining.kt @@ -9,10 +9,8 @@ package org.jetbrains.kotlin.backend.common.lower.inline import org.jetbrains.kotlin.backend.common.* import org.jetbrains.kotlin.backend.common.ir.Symbols import org.jetbrains.kotlin.backend.common.ir.createTemporaryVariableWithWrappedDescriptor -import org.jetbrains.kotlin.backend.common.lower.CoroutineIntrinsicLambdaOrigin import org.jetbrains.kotlin.backend.common.lower.createIrBuilder import org.jetbrains.kotlin.config.languageVersionSettings -import org.jetbrains.kotlin.descriptors.ValueDescriptor import org.jetbrains.kotlin.ir.IrElement import org.jetbrains.kotlin.ir.IrStatement import org.jetbrains.kotlin.ir.UNDEFINED_OFFSET @@ -31,6 +29,9 @@ import org.jetbrains.kotlin.ir.visitors.IrElementVisitor import org.jetbrains.kotlin.ir.visitors.transformChildrenVoid import org.jetbrains.kotlin.util.OperatorNameConventions +fun IrValueParameter.isInlineParameter(type: IrType = this.type) = + index >= 0 && !isNoinline && !type.isNullable() && (type.isFunction() || type.isSuspendFunction()) + class FunctionInlining(val context: CommonBackendContext) : IrElementTransformerVoidWithContext(), BodyLoweringPass { private var containerScope: ScopeWithIr? = null @@ -158,7 +159,7 @@ class FunctionInlining(val context: CommonBackendContext) : IrElementTransformer endOffset = callSite.endOffset, type = callSite.type, symbol = irReturnableBlockSymbol, - origin = if (isCoroutineIntrinsicCall) CoroutineIntrinsicLambdaOrigin else null, + origin = null, statements = statements, inlineFunctionSymbol = callee.symbol ).apply { @@ -320,9 +321,6 @@ class FunctionInlining(val context: CommonBackendContext) : IrElementTransformer //-------------------------------------------------------------------------// - private fun IrValueParameter.isInlineParameter() = - !isNoinline && !type.isNullable() && (type.isFunction() || type.isSuspendFunction()) - private inner class ParameterToArgument( val parameter: IrValueParameter, val argumentExpression: IrExpression diff --git a/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/inline/LocalClasses.kt b/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/inline/LocalClasses.kt new file mode 100644 index 00000000000..938c9a2a823 --- /dev/null +++ b/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/inline/LocalClasses.kt @@ -0,0 +1,173 @@ +/* + * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors. + * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file. + */ + +package org.jetbrains.kotlin.backend.common.lower.inline + +import org.jetbrains.kotlin.backend.common.* +import org.jetbrains.kotlin.backend.common.ir.setDeclarationsParent +import org.jetbrains.kotlin.backend.common.lower.LocalClassPopupLowering +import org.jetbrains.kotlin.backend.common.lower.LocalDeclarationsLowering +import org.jetbrains.kotlin.ir.IrElement +import org.jetbrains.kotlin.ir.IrStatement +import org.jetbrains.kotlin.ir.declarations.* +import org.jetbrains.kotlin.ir.expressions.IrBody +import org.jetbrains.kotlin.ir.expressions.IrCall +import org.jetbrains.kotlin.ir.expressions.IrExpression +import org.jetbrains.kotlin.ir.expressions.IrGetValue +import org.jetbrains.kotlin.ir.expressions.impl.IrBlockImpl +import org.jetbrains.kotlin.ir.expressions.impl.IrCompositeImpl +import org.jetbrains.kotlin.ir.visitors.* + +/* + Here we're extracting some local classes from inline bodies. + The mental model of inlining is as following: + - for inline lambdas, since we don't see the keyword `inline` at a callsite, + it is logical to think that the lambda won't be copied but will be embedded as is at the callsite, + so all local classes declared in those inline lambdas are NEVER COPIED. + - as for the bodies of inline functions, then it is the opposite - we see the `inline` keyword, + so it is only logical to think that this is a macro substitution, so the bodies of inline functions + are copied. But the compiler could optimize the usage of some local classes and not copy them. + So in this case all local classes MIGHT BE COPIED. + */ + +class LocalClassesInInlineLambdasLowering(val context: CommonBackendContext) : BodyLoweringPass { + override fun lower(irFile: IrFile) { + runOnFilePostfix(irFile, allowDeclarationModification = true) + } + + override fun lower(irBody: IrBody, container: IrDeclaration) { + irBody.transformChildrenVoid(object : IrElementTransformerVoidWithContext() { + override fun visitCall(expression: IrCall): IrExpression { + if (!expression.symbol.owner.isInline) + return super.visitCall(expression) + + val localClasses = mutableSetOf() + expression.acceptChildrenVoid(object : IrElementVisitorVoid { + override fun visitElement(element: IrElement) { + element.acceptChildrenVoid(this) + } + + override fun visitClass(declaration: IrClass) { + localClasses.add(declaration) + } + }) + if (localClasses.isEmpty()) + return expression + + LocalDeclarationsLowering(context).lower(expression, container, localClasses) + + expression.transformChildrenVoid(object : IrElementTransformerVoid() { + override fun visitClass(declaration: IrClass): IrStatement { + return IrCompositeImpl( + declaration.startOffset, declaration.endOffset, + context.irBuiltIns.unitType + ) + } + }) + localClasses.forEach { + it.setDeclarationsParent( + currentDeclarationParent + ?: (container as? IrDeclarationParent) + ?: container.parent + ) + } + return IrBlockImpl(expression.startOffset, expression.endOffset, expression.type).apply { + statements += localClasses + statements += expression + } + } + }) + } +} + +class LocalClassesInInlineFunctionsLowering(val context: CommonBackendContext) : BodyLoweringPass { + override fun lower(irFile: IrFile) { + runOnFilePostfix(irFile, allowDeclarationModification = true) + } + + override fun lower(irBody: IrBody, container: IrDeclaration) { + val function = container as? IrFunction ?: return + if (!function.isInline) return + // Conservatively assume that functions with reified type parameters must be copied. + if (function.typeParameters.any { it.isReified }) return + + val crossinlineParameters = function.valueParameters.filter { it.isCrossinline }.toSet() + val classesToExtract = mutableSetOf() + function.acceptChildrenVoid(object : IrElementVisitorVoid { + override fun visitElement(element: IrElement) { + element.acceptChildrenVoid(this) + } + + override fun visitClass(declaration: IrClass) { + var canExtract = true + if (crossinlineParameters.isNotEmpty()) { + declaration.acceptVoid(object : IrElementVisitorVoid { + override fun visitElement(element: IrElement) { + element.acceptChildrenVoid(this) + } + + override fun visitGetValue(expression: IrGetValue) { + if (expression.symbol.owner in crossinlineParameters) + canExtract = false + } + }) + } + if (canExtract) + classesToExtract.add(declaration) + } + }) + if (classesToExtract.isEmpty()) + return + + LocalDeclarationsLowering(context).lower(function, function, classesToExtract) + } +} + +class LocalClassesExtractionFromInlineFunctionsLowering(context: CommonBackendContext) : LocalClassPopupLowering(context) { + private val classesToExtract = mutableSetOf() + + override fun lower(irBody: IrBody, container: IrDeclaration) { + val function = container as? IrFunction ?: return + if (!function.isInline) return + // Conservatively assume that functions with reified type parameters must be copied. + if (function.typeParameters.any { it.isReified }) return + + val crossinlineParameters = function.valueParameters.filter { it.isCrossinline }.toSet() + + function.acceptChildrenVoid(object : IrElementVisitorVoid { + override fun visitElement(element: IrElement) { + element.acceptChildrenVoid(this) + } + + override fun visitClass(declaration: IrClass) { + var canExtract = true + if (crossinlineParameters.isNotEmpty()) { + declaration.acceptVoid(object : IrElementVisitorVoid { + override fun visitElement(element: IrElement) { + element.acceptChildrenVoid(this) + } + + override fun visitGetValue(expression: IrGetValue) { + if (expression.symbol.owner in crossinlineParameters) + canExtract = false + } + }) + } + if (canExtract) + classesToExtract.add(declaration) + } + }) + if (classesToExtract.isEmpty()) + return + + super.lower(irBody, container) + + classesToExtract.clear() + } + + override fun shouldPopUp(klass: IrClass, currentScope: ScopeWithIr?): Boolean { + return classesToExtract.contains(klass) + } +} \ No newline at end of file