From 59522a61305564b84b428ab96246acd3277d930a Mon Sep 17 00:00:00 2001
From: Ivan Kylchik <ivan.kylchik@jetbrains.com>
Date: Wed, 9 Feb 2022 18:54:51 +0300
Subject: [PATCH] Improve IR inliner to be able to inline property references

---
 .../common/lower/inline/FunctionInlining.kt   | 49 ++++++++++++++++---
 1 file changed, 41 insertions(+), 8 deletions(-)

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 aa40d0a2e2a..129f8f05c38 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
@@ -265,6 +265,8 @@ class FunctionInlining(
                     functionArgument is IrFunctionReference ->
                         inlineFunctionReference(expression, functionArgument, functionArgument.symbol.owner)
 
+                    functionArgument is IrPropertyReference -> inlinePropertyReference(expression, functionArgument)
+
                     functionArgument.isAdaptedFunctionReference() ->
                         inlineAdaptedFunctionReference(expression, functionArgument as IrBlock)
 
@@ -287,6 +289,30 @@ class FunctionInlining(
                 return newExpression.transform(this, null)
             }
 
+            private fun inlinePropertyReference(expression: IrCall, propertyReference: IrPropertyReference): IrExpression {
+                val getterCall = IrCallImpl.fromSymbolOwner(
+                    expression.startOffset, expression.endOffset, expression.type, propertyReference.getter!!
+                )
+
+                fun tryToGetArg(i: Int): IrExpression? {
+                    if (i >= expression.valueArgumentsCount) return null
+                    return expression.getValueArgument(i)?.transform(this, null)
+                }
+
+                val receiverFromField = propertyReference.dispatchReceiver ?: propertyReference.extensionReceiver
+                getterCall.dispatchReceiver = getterCall.symbol.owner.dispatchReceiverParameter?.let {
+                    receiverFromField ?: tryToGetArg(0)
+                }
+                getterCall.extensionReceiver = getterCall.symbol.owner.extensionReceiverParameter?.let {
+                    when (getterCall.symbol.owner.dispatchReceiverParameter) {
+                        null -> receiverFromField ?: tryToGetArg(0)
+                        else -> tryToGetArg(if (receiverFromField != null) 0 else 1)
+                    }
+                }
+
+                return getterCall
+            }
+
             fun inlineAdaptedFunctionReference(irCall: IrCall, irBlock: IrBlock): IrExpression {
                 val irFunction = irBlock.statements[0].let {
                     it.transformChildrenVoid(this)
@@ -436,6 +462,9 @@ class FunctionInlining(
                                 || argumentExpression is IrFunctionExpression
                                 || argumentExpression.isAdaptedFunctionReference())
 
+            val isInlinablePropertyReference: Boolean
+                get() = parameter.isInlineParameter() && argumentExpression is IrPropertyReference
+
             val isImmutableVariableLoad: Boolean
                 get() = argumentExpression.let { argument ->
                     argument is IrGetValue && !argument.symbol.owner.let { it is IrVariable && it.isVar }
@@ -552,11 +581,15 @@ class FunctionInlining(
 
         //-------------------------------------------------------------------------//
 
-        private fun evaluateArguments(functionReference: IrFunctionReference): List<IrStatement> {
-            val arguments = functionReference.getArgumentsWithIr().map { ParameterToArgument(it.first, it.second) }
+        private fun evaluateArguments(reference: IrCallableReference<*>): List<IrStatement> {
+            val arguments = reference.getArgumentsWithIr().map { ParameterToArgument(it.first, it.second) }
             val evaluationStatements = mutableListOf<IrStatement>()
             val substitutor = ParameterSubstitutor()
-            val referenced = functionReference.symbol.owner
+            val referenced = when (reference) {
+                is IrFunctionReference -> reference.symbol.owner
+                is IrPropertyReference -> reference.getter!!.owner
+                else -> error(this)
+            }
             arguments.forEach {
                 val newArgument = if (it.isImmutableVariableLoad) {
                     it.argumentExpression.transform( // Arguments may reference the previous ones - substitute them.
@@ -579,9 +612,9 @@ class FunctionInlining(
                     IrGetValueWithoutLocation(newVariable.symbol)
                 }
                 when (it.parameter) {
-                    referenced.dispatchReceiverParameter -> functionReference.dispatchReceiver = newArgument
-                    referenced.extensionReceiverParameter -> functionReference.extensionReceiver = newArgument
-                    else -> functionReference.putValueArgument(it.parameter.index, newArgument)
+                    referenced.dispatchReceiverParameter -> reference.dispatchReceiver = newArgument
+                    referenced.extensionReceiverParameter -> reference.extensionReceiver = newArgument
+                    else -> reference.putValueArgument(it.parameter.index, newArgument)
                 }
             }
             return evaluationStatements
@@ -598,9 +631,9 @@ class FunctionInlining(
                  * For simplicity and to produce simpler IR we don't create temporaries for every immutable variable,
                  * not only for those referring to inlinable lambdas.
                  */
-                if (argument.isInlinableLambdaArgument) {
+                if (argument.isInlinableLambdaArgument || argument.isInlinablePropertyReference) {
                     substituteMap[parameter] = argument.argumentExpression
-                    (argument.argumentExpression as? IrFunctionReference)?.let { evaluationStatements += evaluateArguments(it) }
+                    (argument.argumentExpression as? IrCallableReference<*>)?.let { evaluationStatements += evaluateArguments(it) }
 
                     return@forEach
                 }