[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.

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<IrDeclaration>) {
+    declarations.forEach { this.addChild(it) }
+}
+
+fun <T : IrElement> T.setDeclarationsParent(parent: IrDeclarationParent): T {
+    accept(SetDeclarationsParentVisitor, parent)
+    return this
 }
 
 object SetDeclarationsParentVisitor : IrElementVisitor<Unit, IrDeclarationParent> {

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<IrValueSymbol>, val capturedTypeParameters: List<IrTypeParameter>)
 
-class ClosureAnnotator(body: IrBody, declaration: IrDeclaration) {
+class ClosureAnnotator(irElement: IrElement, declaration: IrDeclaration) {
     private val closureBuilders = mutableMapOf<IrDeclaration, ClosureBuilder>()
 
     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)

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<IrClass>) {
+        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<IrClass>?
+    ) {
         val localFunctions: MutableMap<IrFunction, LocalFunctionContext> = LinkedHashMap()
         val localClasses: MutableMap<IrClass, LocalClassContext> = LinkedHashMap()
         val localClassConstructors: MutableMap<IrConstructor, LocalClassConstructorContext> = 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

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 <crossinline> 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)
 

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

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<IrClass>()
+                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<IrClass>()
+        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<IrClass>()
+
+    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)
+    }
+}