Set original type for expression inlined as arguments

This is mainly needed to fix cases like in "kt44429".

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

@@ -84,7 +84,8 @@ class FunctionInlining(
     val inlineFunctionResolver: InlineFunctionResolver,
     val innerClassesSupport: InnerClassesSupport? = null,
     val insertAdditionalImplicitCasts: Boolean = false,
-    val alwaysCreateTemporaryVariablesForArguments: Boolean = false
+    val alwaysCreateTemporaryVariablesForArguments: Boolean = false,
+    val inlineArgumentsWithTheirOriginalType: Boolean = false,
 ) : IrElementTransformerVoidWithContext(), BodyLoweringPass {
 
     constructor(context: CommonBackendContext) : this(context, DefaultInlineFunctionResolver(context), null)
@@ -604,18 +605,15 @@ class FunctionInlining(
                 else -> error(this)
             }
             arguments.forEach {
+                // Arguments may reference the previous ones - substitute them.
+                val irExpression = it.argumentExpression.transform(substitutor, data = null)
                 val newArgument = if (it.isImmutableVariableLoad) {
-                    it.argumentExpression.transform( // Arguments may reference the previous ones - substitute them.
-                        substitutor,
-                        data = null
-                    )
+                    IrGetValueWithoutLocation((irExpression as IrGetValue).symbol)
                 } else {
                     val newVariable =
                         currentScope.scope.createTemporaryVariable(
-                            irExpression = it.argumentExpression.transform( // Arguments may reference the previous ones - substitute them.
-                                substitutor,
-                                data = null
-                            ),
+                            irExpression = irExpression,
+                            irType = if (inlineArgumentsWithTheirOriginalType) it.parameter.getOriginalType() else irExpression.type,
                             nameHint = callee.symbol.owner.name.asStringStripSpecialMarkers(),
                             isMutable = false
                         )
@@ -633,6 +631,30 @@ class FunctionInlining(
             return evaluationStatements
         }
 
+        // In short this is needed for `kt44429` test. We need to get original generic type to trick type system on JVM backend.
+        private fun IrValueParameter.getOriginalType(): IrType {
+            if (this.parent !is IrFunction) return type
+            val callee = this.parent as IrFunction
+            val original = callee.originalFunction
+
+            fun IrValueParameter?.getTypeIfFromTypeParameter(): IrType? {
+                val typeClassifier = this?.type?.classifierOrNull?.owner as? IrTypeParameter ?: return null
+                if (typeClassifier.parent != this.parent) return null
+                return callee.typeParameters[typeClassifier.index].defaultType
+            }
+
+            return when (this) {
+                callee.dispatchReceiverParameter -> original.dispatchReceiverParameter?.getTypeIfFromTypeParameter()
+                    ?: callee.dispatchReceiverParameter!!.type
+                callee.extensionReceiverParameter -> original.extensionReceiverParameter?.getTypeIfFromTypeParameter()
+                    ?: callee.extensionReceiverParameter!!.type
+                else -> callee.valueParameters.first { it == this }.let { valueParameter ->
+                    original.valueParameters[valueParameter.index].getTypeIfFromTypeParameter()
+                        ?: valueParameter.type
+                }
+            }
+        }
+
         private fun evaluateArguments(callSite: IrFunctionAccessExpression, callee: IrFunction): List<IrStatement> {
             val arguments = buildParameterToArgument(callSite, callee)
             val evaluationStatements = mutableListOf<IrVariable>()
@@ -695,7 +717,7 @@ class FunctionInlining(
                 irExpression = IrBlockImpl(
                     variableInitializer.startOffset,
                     variableInitializer.endOffset,
-                    variableInitializer.type,
+                    if (inlineArgumentsWithTheirOriginalType) 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

@@ -291,7 +291,10 @@ private val functionInliningPhase = makeIrModulePhase<JvmBackendContext>(
                 return irFunction.symbol
             }
         }
-        FunctionInlining(context, JvmInlineFunctionResolver(), context.innerClassesSupport)
+        FunctionInlining(
+            context, JvmInlineFunctionResolver(), context.innerClassesSupport,
+            inlineArgumentsWithTheirOriginalType = true
+        )
     },
     name = "FunctionInliningPhase",
     description = "Perform function inlining",

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

@@ -45,7 +45,10 @@ private class ReplaceKFunctionInvokeWithFunctionInvoke : FileLoweringPass, IrEle
         if (callee.name != OperatorNameConventions.INVOKE) return super.visitCall(expression)
 
         val parentClass = callee.parent as? IrClass ?: return super.visitCall(expression)
-        if (!parentClass.defaultType.isKFunction() && !parentClass.defaultType.isKSuspendFunction()) return super.visitCall(expression)
+        if (!parentClass.defaultType.isKFunction() && !parentClass.defaultType.isKSuspendFunction()) {
+            implicitCastKFunctionReceiverIntoFunctionIfNeeded(expression, parentClass)
+            return super.visitCall(expression)
+        }
 
         // The single overridden function of KFunction{n}.invoke must be Function{n}.invoke.
         val newCallee = callee.overriddenSymbols.single()
@@ -63,4 +66,16 @@ private class ReplaceKFunctionInvokeWithFunctionInvoke : FileLoweringPass, IrEle
             }
         }
     }
+
+    // This method suppose to cover case when we have `Function{n}.invoke` but receiver has type of `KFunction`
+    private fun implicitCastKFunctionReceiverIntoFunctionIfNeeded(expression: IrCall, parentClass: IrClass) {
+        val receiver = expression.dispatchReceiver
+        if (receiver != null && (receiver.type.isKFunction() || receiver.type.isKSuspendFunction())) {
+            val newType = parentClass.defaultType
+
+            expression.dispatchReceiver = IrTypeOperatorCallImpl(
+                expression.startOffset, expression.endOffset, newType, IrTypeOperator.IMPLICIT_CAST, newType, receiver.transform(this, null)
+            )
+        }
+    }
 }