From ce6084302fbb20fca2651fdfdd2463fac2b090ba Mon Sep 17 00:00:00 2001
From: Ivan Kylchik <ivan.kylchik@jetbrains.com>
Date: Mon, 27 Dec 2021 14:58:40 +0300
Subject: [PATCH] Add `FunctionInlining` lowering to `JvmLower`

---
 .../common/lower/ReturnableBlockLowering.kt   | 26 ++++++++++++++-----
 .../common/lower/inline/FunctionInlining.kt   |  3 +++
 .../jetbrains/kotlin/backend/jvm/JvmLower.kt  | 26 +++++++++++++++++++
 3 files changed, 48 insertions(+), 7 deletions(-)

diff --git a/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/ReturnableBlockLowering.kt b/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/ReturnableBlockLowering.kt
index 740d0568d3a..5f4f5f2f489 100644
--- a/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/ReturnableBlockLowering.kt
+++ b/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/ReturnableBlockLowering.kt
@@ -8,12 +8,12 @@ package org.jetbrains.kotlin.backend.common.lower
 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.ir.IrStatement
 import org.jetbrains.kotlin.ir.builders.*
 import org.jetbrains.kotlin.ir.declarations.IrDeclaration
 import org.jetbrains.kotlin.ir.declarations.IrSymbolOwner
 import org.jetbrains.kotlin.ir.expressions.*
 import org.jetbrains.kotlin.ir.expressions.impl.IrBlockImpl
-import org.jetbrains.kotlin.ir.expressions.impl.IrCompositeImpl
 import org.jetbrains.kotlin.ir.expressions.impl.IrDoWhileLoopImpl
 import org.jetbrains.kotlin.ir.symbols.IrReturnableBlockSymbol
 import org.jetbrains.kotlin.ir.symbols.IrSymbol
@@ -113,14 +113,26 @@ class ReturnableBlockTransformer(val context: CommonBackendContext, val containe
             }
         }
 
-        val newStatements = expression.statements.mapIndexed { i, s ->
-            if (i == expression.statements.lastIndex && s is IrReturn && s.returnTargetSymbol == expression.symbol) {
-                s.transformChildrenVoid()
-                if (!hasReturned) s.value else {
-                    builder.irSet(variable.symbol, s.value)
+
+        fun transformSingleStatement(statement: IrStatement, isLastInList: Boolean): IrStatement {
+            return if (isLastInList && statement is IrReturn && statement.returnTargetSymbol == expression.symbol) {
+                statement.transformChildrenVoid()
+                if (!hasReturned) statement.value else {
+                    builder.irSet(variable.symbol, statement.value)
                 }
             } else {
-                s.transformStatement(this)
+                statement.transformStatement(this)
+            }
+        }
+
+        val newStatements = expression.statements.mapIndexed { i, s ->
+            if (expression.statements.size == 1 && s is IrInlinedFunctionBlock) {
+                val transformed = s.statements.mapIndexed { j, it -> transformSingleStatement(it, j == s.statements.lastIndex) }
+                s.statements.clear()
+                s.statements.addAll(transformed)
+                s
+            } else {
+                transformSingleStatement(s, i == expression.statements.lastIndex)
             }
         }
 
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 cf2ee367c5e..8d909c75cf3 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
@@ -128,6 +128,9 @@ class FunctionInlining(
             return expression
 
         val actualCallee = inlineFunctionResolver.getFunctionDeclaration(callee.symbol)
+        if (actualCallee.body == null) {
+            return expression
+        }
 
         val parent = allScopes.map { it.irElement }.filterIsInstance<IrDeclarationParent>().lastOrNull()
             ?: allScopes.map { it.irElement }.filterIsInstance<IrDeclaration>().lastOrNull()?.parent
diff --git a/compiler/ir/backend.jvm/lower/src/org/jetbrains/kotlin/backend/jvm/JvmLower.kt b/compiler/ir/backend.jvm/lower/src/org/jetbrains/kotlin/backend/jvm/JvmLower.kt
index 8acdd7f05b9..740c36006e5 100644
--- a/compiler/ir/backend.jvm/lower/src/org/jetbrains/kotlin/backend/jvm/JvmLower.kt
+++ b/compiler/ir/backend.jvm/lower/src/org/jetbrains/kotlin/backend/jvm/JvmLower.kt
@@ -9,6 +9,7 @@ import org.jetbrains.kotlin.backend.common.CommonBackendContext
 import org.jetbrains.kotlin.backend.common.FileLoweringPass
 import org.jetbrains.kotlin.backend.common.checkDeclarationParents
 import org.jetbrains.kotlin.backend.common.lower.*
+import org.jetbrains.kotlin.backend.common.lower.inline.*
 import org.jetbrains.kotlin.backend.common.lower.loops.forLoopsPhase
 import org.jetbrains.kotlin.backend.common.phaser.*
 import org.jetbrains.kotlin.backend.jvm.ir.constantValue
@@ -19,10 +20,12 @@ import org.jetbrains.kotlin.descriptors.DescriptorVisibilities
 import org.jetbrains.kotlin.descriptors.DescriptorVisibility
 import org.jetbrains.kotlin.ir.IrElement
 import org.jetbrains.kotlin.ir.declarations.*
+import org.jetbrains.kotlin.ir.symbols.IrFunctionSymbol
 import org.jetbrains.kotlin.ir.symbols.IrValueSymbol
 import org.jetbrains.kotlin.ir.util.PatchDeclarationParentsVisitor
 import org.jetbrains.kotlin.ir.util.isAnonymousObject
 import org.jetbrains.kotlin.ir.util.parentAsClass
+import org.jetbrains.kotlin.ir.util.resolveFakeOverride
 import org.jetbrains.kotlin.ir.visitors.IrElementVisitorVoid
 import org.jetbrains.kotlin.ir.visitors.acceptChildrenVoid
 import org.jetbrains.kotlin.ir.visitors.acceptVoid
@@ -277,6 +280,26 @@ private val kotlinNothingValueExceptionPhase = makeIrFilePhase<CommonBackendCont
     description = "Throw proper exception for calls returning value of type 'kotlin.Nothing'"
 )
 
+private val functionInliningPhase = makeIrModulePhase<JvmBackendContext>(
+    { context ->
+        class JvmInlineFunctionResolver : InlineFunctionResolver {
+            override fun getFunctionDeclaration(symbol: IrFunctionSymbol): IrFunction {
+                return (symbol.owner as? IrSimpleFunction)?.resolveFakeOverride() ?: symbol.owner
+            }
+
+            override fun getFunctionSymbol(irFunction: IrFunction): IrFunctionSymbol {
+                return irFunction.symbol
+            }
+        }
+        FunctionInlining(context, JvmInlineFunctionResolver(), context.innerClassesSupport)
+    },
+    name = "FunctionInliningPhase",
+    description = "Perform function inlining",
+    prerequisite = setOf(
+        expectDeclarationsRemovingPhase,
+    )
+)
+
 private val jvmFilePhases = listOf(
     typeAliasAnnotationMethodsPhase,
     provisionalFunctionExpressionPhase,
@@ -412,6 +435,9 @@ private fun buildJvmLoweringPhases(
                 fileClassPhase then
                 jvmStaticInObjectPhase then
                 repeatedAnnotationPhase then
+
+                functionInliningPhase then
+
                 buildLoweringsPhase(phases) then
                 generateMultifileFacadesPhase then
                 resolveInlineCallsPhase then