Support inlining of property reference with non-nullable field

compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/inline/FunctionInlining.kt

@@ -17,9 +17,11 @@ import org.jetbrains.kotlin.backend.common.lower.at
 import org.jetbrains.kotlin.backend.common.lower.createIrBuilder
 import org.jetbrains.kotlin.builtins.StandardNames
 import org.jetbrains.kotlin.contracts.parsing.ContractsDslNames
+import org.jetbrains.kotlin.descriptors.DescriptorVisibilities
 import org.jetbrains.kotlin.ir.IrElement
 import org.jetbrains.kotlin.ir.IrStatement
 import org.jetbrains.kotlin.ir.UNDEFINED_OFFSET
+import org.jetbrains.kotlin.ir.builders.declarations.buildFun
 import org.jetbrains.kotlin.ir.builders.irReturn
 import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.ir.expressions.*
@@ -31,6 +33,7 @@ import org.jetbrains.kotlin.ir.symbols.impl.IrReturnableBlockSymbolImpl
 import org.jetbrains.kotlin.ir.types.*
 import org.jetbrains.kotlin.ir.util.*
 import org.jetbrains.kotlin.ir.visitors.*
+import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.util.OperatorNameConventions
 
 fun IrValueParameter.isInlineParameter(type: IrType = this.type) =
@@ -86,7 +89,7 @@ class FunctionInlining(
     val alwaysCreateTemporaryVariablesForArguments: Boolean = false,
     val inlinePureArguments: Boolean = true,
     val regenerateInlinedAnonymousObjects: Boolean = false,
-    val inlineArgumentsWithTheirOriginalType: Boolean = false,
+    val inlineArgumentsWithTheirOriginalTypeAndOffset: Boolean = false,
 ) : IrElementTransformerVoidWithContext(), BodyLoweringPass {
     constructor(context: CommonBackendContext) : this(context, DefaultInlineFunctionResolver(context), null)
     constructor(context: CommonBackendContext, innerClassesSupport: InnerClassesSupport) : this(
@@ -303,6 +306,8 @@ class FunctionInlining(
                     functionArgument is IrFunctionReference ->
                         inlineFunctionReference(expression, functionArgument, functionArgument.symbol.owner)
 
+                    functionArgument is IrPropertyReference && functionArgument.field != null -> inlineField(expression, functionArgument)
+
                     functionArgument is IrPropertyReference -> inlinePropertyReference(expression, functionArgument)
 
                     functionArgument.isAdaptedFunctionReference() ->
@@ -325,9 +330,14 @@ class FunctionInlining(
                 return newExpression.transform(this, null)
             }
 
+            private fun inlineField(invokeCall: IrCall, propertyReference: IrPropertyReference): IrExpression {
+                return wrapInStubFunction(invokeCall, invokeCall, propertyReference)
+            }
+
             private fun inlinePropertyReference(expression: IrCall, propertyReference: IrPropertyReference): IrExpression {
                 val getterCall = IrCallImpl.fromSymbolOwner(
-                    expression.startOffset, expression.endOffset, expression.type, propertyReference.getter!!
+                    expression.startOffset, expression.endOffset, expression.type, propertyReference.getter!!,
+                    origin = INLINED_FUNCTION_REFERENCE
                 )
 
                 fun tryToGetArg(i: Int): IrExpression? {
@@ -346,7 +356,36 @@ class FunctionInlining(
                     }
                 }
 
-                return getterCall
+                return wrapInStubFunction(super.visitExpression(getterCall), expression, propertyReference)
+            }
+
+            private fun wrapInStubFunction(
+                inlinedCall: IrExpression, invokeCall: IrFunctionAccessExpression, reference: IrCallableReference<*>
+            ): IrReturnableBlock {
+                // Note: This function is not exist in tree. It is appeared only in `IrInlinedFunctionBlock` as intermediate callee.
+                val stubForInline = context.irFactory.buildFun {
+                    startOffset = inlinedCall.startOffset
+                    endOffset = inlinedCall.endOffset
+                    name = Name.identifier("stub_for_ir_inlining")
+                    visibility = DescriptorVisibilities.LOCAL
+                    returnType = inlinedCall.type
+                    isSuspend = reference.symbol.isSuspend
+                }.apply {
+                    body = context.irFactory.createBlockBody(UNDEFINED_OFFSET, UNDEFINED_OFFSET).apply {
+                        val statement = if (reference is IrPropertyReference && reference.field != null) {
+                            val field = reference.field!!.owner
+                            val boundReceiver = reference.dispatchReceiver ?: reference.extensionReceiver
+                            val fieldReceiver = if (field.isStatic) null else boundReceiver
+                            IrGetFieldImpl(UNDEFINED_OFFSET, UNDEFINED_OFFSET, field.symbol, field.type, fieldReceiver)
+                        } else {
+                            IrReturnImpl(UNDEFINED_OFFSET, UNDEFINED_OFFSET, context.irBuiltIns.nothingType, symbol, inlinedCall)
+                        }
+                        statements += statement
+                    }
+                    parent = callee.parent
+                }
+
+                return inlineFunction(invokeCall, stubForInline, reference, false)
             }
 
             fun inlineAdaptedFunctionReference(irCall: IrCall, irBlock: IrBlock): IrExpression {
@@ -385,32 +424,29 @@ class FunctionInlining(
                     typeParam.symbol to superType.arguments[typeParam.index].typeOrNull!!
                 }
 
-                val immediateCall = with(irCall) {
+                val immediateCall = when (inlinedFunction) {
-                    when (inlinedFunction) {
+                    is IrConstructor -> {
-                        is IrConstructor -> {
+                        val classTypeParametersCount = inlinedFunction.parentAsClass.typeParameters.size
-                            val classTypeParametersCount = inlinedFunction.parentAsClass.typeParameters.size
+                        IrConstructorCallImpl.fromSymbolOwner(
-                            IrConstructorCallImpl.fromSymbolOwner(
+                            if (inlineArgumentsWithTheirOriginalTypeAndOffset) irFunctionReference.startOffset else irCall.startOffset,
-                                startOffset,
+                            if (inlineArgumentsWithTheirOriginalTypeAndOffset) irFunctionReference.endOffset else irCall.endOffset,
-                                endOffset,
+                            inlinedFunction.returnType,
-                                inlinedFunction.returnType,
+                            inlinedFunction.symbol,
-                                inlinedFunction.symbol,
+                            classTypeParametersCount,
-                                classTypeParametersCount,
+                            INLINED_FUNCTION_REFERENCE
-                                INLINED_FUNCTION_REFERENCE
+                        )
-                            )
-                        }
-                        is IrSimpleFunction ->
-                            IrCallImpl(
-                                startOffset,
-                                endOffset,
-                                inlinedFunction.returnType,
-                                inlinedFunction.symbol,
-                                inlinedFunction.typeParameters.size,
-                                inlinedFunction.valueParameters.size,
-                                INLINED_FUNCTION_REFERENCE
-                            )
-                        else ->
-                            error("Unknown function kind : ${inlinedFunction.render()}")
                     }
+                    is IrSimpleFunction ->
+                        IrCallImpl(
+                            if (inlineArgumentsWithTheirOriginalTypeAndOffset) irFunctionReference.startOffset else irCall.startOffset,
+                            if (inlineArgumentsWithTheirOriginalTypeAndOffset) irFunctionReference.endOffset else irCall.endOffset,
+                            inlinedFunction.returnType,
+                            inlinedFunction.symbol,
+                            inlinedFunction.typeParameters.size,
+                            inlinedFunction.valueParameters.size,
+                            INLINED_FUNCTION_REFERENCE
+                        )
+                    else -> error("Unknown function kind : ${inlinedFunction.render()}")
                 }.apply {
                     for (parameter in functionParameters) {
                         val argument =
@@ -462,13 +498,27 @@ class FunctionInlining(
                     assert(unboundIndex == valueParameters.size) { "Not all arguments of the callee are used" }
                     for (index in 0 until irFunctionReference.typeArgumentsCount)
                         putTypeArgument(index, irFunctionReference.getTypeArgument(index))
+                }
+
+                return if (inlinedFunction.needsInlining && inlinedFunction.body != null) {
+                    inlineFunction(immediateCall, inlinedFunction, irFunctionReference, performRecursiveInline = true)
+                } else {
+                    val transformedExpression = super.visitExpression(immediateCall).transform(this@FunctionInlining, null)
+                    wrapInStubFunction(transformedExpression, irCall, irFunctionReference)
                 }.implicitCastIfNeededTo(irCall.type)
-                return super.visitExpression(immediateCall).transform(this@FunctionInlining, null)
             }
 
             override fun visitElement(element: IrElement) = element.accept(this, null)
         }
 
+        private fun IrExpression.implicitCastIfNeededTo(type: IrType): IrExpression {
+            // No need to cast expressions of type nothing
+            return if (type == this.type || !insertAdditionalImplicitCasts || this.type == context.irBuiltIns.nothingType)
+                this
+            else
+                IrTypeOperatorCallImpl(startOffset, endOffset, type, IrTypeOperator.IMPLICIT_CAST, type, this)
+        }
+
         // With `insertAdditionalImplicitCasts` flag we sometimes insert
         // casts to inline lambda parameters before calling `invoke` on them.
         // Unwrapping these casts helps us satisfy inline lambda call detection logic.
@@ -640,8 +690,8 @@ class FunctionInlining(
                     val newVariable =
                         currentScope.scope.createTemporaryVariable(
                             irExpression = irExpression,
-                            irType = if (inlineArgumentsWithTheirOriginalType) it.parameter.getOriginalType() else irExpression.type,
+                            irType = if (inlineArgumentsWithTheirOriginalTypeAndOffset) it.parameter.getOriginalType() else irExpression.type,
-                            nameHint = callee.symbol.owner.name.asStringStripSpecialMarkers(),
+                            nameHint = callee.symbol.owner.name.asStringStripSpecialMarkers() + "_" + it.parameter.name.asStringStripSpecialMarkers(),
                             isMutable = false
                         )
 
@@ -658,6 +708,28 @@ class FunctionInlining(
             return evaluationStatements
         }
 
+        private fun evaluateReceiverForPropertyWithField(reference: IrPropertyReference): IrVariable? {
+            val argument = reference.dispatchReceiver ?: reference.extensionReceiver ?: return null
+            // Arguments may reference the previous ones - substitute them.
+            val irExpression = argument.transform(ParameterSubstitutor(), data = null)
+
+            val newVariable = currentScope.scope.createTemporaryVariable(
+                startOffset = UNDEFINED_OFFSET,
+                endOffset = UNDEFINED_OFFSET,
+                irExpression = irExpression,
+                nameHint = callee.symbol.owner.name.asStringStripSpecialMarkers() + "_this",
+                isMutable = false
+            )
+
+            val newArgument = IrGetValueWithoutLocation(newVariable.symbol)
+            when {
+                reference.dispatchReceiver != null -> reference.dispatchReceiver = newArgument
+                reference.extensionReceiver != null -> reference.extensionReceiver = newArgument
+            }
+
+            return newVariable
+        }
+
         private fun IrValueParameter.getOriginalParameter(): IrValueParameter {
             if (this.parent !is IrFunction) return this
             val original = (this.parent as IrFunction).originalFunction
@@ -702,7 +774,15 @@ class FunctionInlining(
                  */
                 if (argument.isInlinableLambdaArgument || argument.isInlinablePropertyReference) {
                     substituteMap[parameter] = argument.argumentExpression
-                    (argument.argumentExpression as? IrCallableReference<*>)?.let { evaluationStatements += evaluateArguments(it) }
+                    val arg = argument.argumentExpression
+                    when {
+                        // This first branch is required to avoid assertion in `getArgumentsWithIr`
+                        arg is IrPropertyReference && arg.field != null -> evaluateReceiverForPropertyWithField(arg)?.let { evaluationStatements += it }
+                        arg is IrCallableReference<*> -> evaluationStatements += evaluateArguments(arg)
+                        arg is IrBlock -> if (arg.origin == IrStatementOrigin.ADAPTED_FUNCTION_REFERENCE) {
+                            evaluationStatements += evaluateArguments(arg.statements.last() as IrFunctionReference)
+                        }
+                    }
 
                     return@forEach
                 }
@@ -713,14 +793,14 @@ class FunctionInlining(
                     (alwaysCreateTemporaryVariablesForArguments && !parameter.isInlineParameter()) ||
                             argument.shouldBeSubstitutedViaTemporaryVariable()
 
-                if (shouldCreateTemporaryVariable) {
+                if (!shouldCreateTemporaryVariable) {
-                    val newVariable = createTemporaryVariable(parameter, variableInitializer, callee)
-
-                    if (argument.isDefaultArg) evaluationStatementsFromDefault.add(newVariable) else evaluationStatements.add(newVariable)
-                    substituteMap[parameter] = IrGetValueWithoutLocation(newVariable.symbol)
-                } else {
                     substituteMap[parameter] = variableInitializer
+                    return@forEach
                 }
+
+                val newVariable = createTemporaryVariable(parameter, variableInitializer, callee)
+                if (argument.isDefaultArg) evaluationStatementsFromDefault.add(newVariable) else evaluationStatements.add(newVariable)
+                substituteMap[parameter] = IrGetValueWithoutLocation(newVariable.symbol)
             }
 
 
@@ -751,7 +831,7 @@ class FunctionInlining(
                 irExpression = IrBlockImpl(
                     variableInitializer.startOffset,
                     variableInitializer.endOffset,
-                    if (inlineArgumentsWithTheirOriginalType) parameter.getOriginalType() else variableInitializer.type,
+                    if (inlineArgumentsWithTheirOriginalTypeAndOffset) parameter.getOriginalType() else variableInitializer.type,
                     InlinerExpressionLocationHint((currentScope.irElement as IrSymbolOwner).symbol)
                 ).apply {
                     statements.add(variableInitializer)

compiler/ir/backend.jvm/lower/src/org/jetbrains/kotlin/backend/jvm/JvmLower.kt

@@ -297,7 +297,7 @@ internal val functionInliningPhase = makeIrModulePhase<JvmBackendContext>(
             context, JvmInlineFunctionResolver(), context.innerClassesSupport,
             inlinePureArguments = false,
             regenerateInlinedAnonymousObjects = true,
-            inlineArgumentsWithTheirOriginalType = true
+            inlineArgumentsWithTheirOriginalTypeAndOffset = true
         )
     },
     name = "FunctionInliningPhase",