From a7b31ba42b65912458197359485b7adbb7b70cb7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Steven=20Sch=C3=A4fer?= <stevenschaefer@google.com>
Date: Tue, 6 Aug 2019 11:38:00 +0200
Subject: [PATCH] Add a utility function to compute type substitutions for call

---
 .../backend/jvm/codegen/ExpressionCodegen.kt    |  8 +-------
 .../jvm/lower/CallableReferenceLowering.kt      |  8 +-------
 .../src/org/jetbrains/kotlin/ir/util/IrUtils.kt | 17 +++++++++++++++++
 3 files changed, 19 insertions(+), 14 deletions(-)

diff --git a/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/codegen/ExpressionCodegen.kt b/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/codegen/ExpressionCodegen.kt
index 96903a6e9d8..06d1a02d302 100644
--- a/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/codegen/ExpressionCodegen.kt
+++ b/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/codegen/ExpressionCodegen.kt
@@ -378,12 +378,6 @@ class ExpressionCodegen(
         }
 
         callGenerator.beforeValueParametersStart()
-        val typeParameters = if (callee is IrConstructor)
-            callee.parentAsClass.typeParameters + callee.typeParameters
-        else
-            callee.typeParameters
-        val typeArguments = (0 until typeParameters.size).map { expression.getTypeArgument(it)!! }
-        val typeSubstitutionMap = typeParameters.map { it.symbol }.zip(typeArguments).toMap()
         expression.symbol.owner.valueParameters.forEachIndexed { i, irParameter ->
             val arg = expression.getValueArgument(i)
             val parameterType = callable.valueParameterTypes[i]
@@ -429,7 +423,7 @@ class ExpressionCodegen(
         }
 
         return when {
-            returnType.substitute(typeSubstitutionMap).isNothing() -> {
+            returnType.substitute(expression.typeSubstitutionMap).isNothing() -> {
                 mv.aconst(null)
                 mv.athrow()
                 immaterialUnitValue
diff --git a/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/lower/CallableReferenceLowering.kt b/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/lower/CallableReferenceLowering.kt
index b5661f7672c..b24b772e376 100644
--- a/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/lower/CallableReferenceLowering.kt
+++ b/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/lower/CallableReferenceLowering.kt
@@ -131,13 +131,7 @@ internal class CallableReferenceLowering(private val context: JvmBackendContext)
         private val parameterTypes = (irFunctionReference.type as IrSimpleType).arguments.map { (it as IrTypeProjection).type }
         private val argumentTypes = parameterTypes.dropLast(1)
 
-        private val typeParameters = if (callee is IrConstructor)
-            callee.parentAsClass.typeParameters + callee.typeParameters
-        else
-            callee.typeParameters
-        private val typeArgumentsMap = typeParameters.associate { typeParam ->
-            typeParam.symbol to irFunctionReference.getTypeArgument(typeParam.index)!!
-        }
+        private val typeArgumentsMap = irFunctionReference.typeSubstitutionMap
 
         private val functionReferenceClass = buildClass {
             setSourceRange(irFunctionReference)
diff --git a/compiler/ir/ir.tree/src/org/jetbrains/kotlin/ir/util/IrUtils.kt b/compiler/ir/ir.tree/src/org/jetbrains/kotlin/ir/util/IrUtils.kt
index 9370a42c10f..2f235c7e575 100644
--- a/compiler/ir/ir.tree/src/org/jetbrains/kotlin/ir/util/IrUtils.kt
+++ b/compiler/ir/ir.tree/src/org/jetbrains/kotlin/ir/util/IrUtils.kt
@@ -590,3 +590,20 @@ val IrDeclaration.file: IrFile
             else -> TODO("Unexpected declaration parent")
         }
     }
+
+val IrFunction.allTypeParameters: List<IrTypeParameter>
+    get() = if (this is IrConstructor)
+        parentAsClass.typeParameters + typeParameters
+    else
+        typeParameters
+
+fun IrMemberAccessExpression.getTypeSubstitutionMap(irFunction: IrFunction): Map<IrTypeParameterSymbol, IrType> =
+    irFunction.allTypeParameters.withIndex().associate {
+        it.value.symbol to getTypeArgument(it.index)!!
+    }
+
+val IrFunctionReference.typeSubstitutionMap: Map<IrTypeParameterSymbol, IrType>
+    get() = getTypeSubstitutionMap(symbol.owner)
+
+val IrFunctionAccessExpression.typeSubstitutionMap: Map<IrTypeParameterSymbol, IrType>
+    get() = getTypeSubstitutionMap(symbol.owner)