From 84bd580f92e5c347518c2caab7b135b5bebc2a6d Mon Sep 17 00:00:00 2001
From: Victor Petukhov <victor.petukhov@jetbrains.com>
Date: Fri, 21 May 2021 15:24:55 +0300
Subject: [PATCH] Complete block expressions to substitute stub type variables
 into last callable references there

---
 .../calls/tower/ResolvedAtomCompleter.kt      | 26 ++++++++++++++++---
 .../org/jetbrains/kotlin/types/TypeUtils.kt   |  3 ++-
 2 files changed, 25 insertions(+), 4 deletions(-)

diff --git a/compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/tower/ResolvedAtomCompleter.kt b/compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/tower/ResolvedAtomCompleter.kt
index bd7869d407d..6a91f40d894 100644
--- a/compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/tower/ResolvedAtomCompleter.kt
+++ b/compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/tower/ResolvedAtomCompleter.kt
@@ -29,6 +29,7 @@ import org.jetbrains.kotlin.resolve.calls.components.SuspendConversionStrategy
 import org.jetbrains.kotlin.resolve.calls.components.isVararg
 import org.jetbrains.kotlin.resolve.calls.context.BasicCallResolutionContext
 import org.jetbrains.kotlin.resolve.calls.context.ContextDependency
+import org.jetbrains.kotlin.resolve.calls.inference.BuilderInferenceSession
 import org.jetbrains.kotlin.resolve.calls.inference.components.NewTypeSubstitutor
 import org.jetbrains.kotlin.resolve.calls.inference.components.NewTypeSubstitutorByConstructorMap
 import org.jetbrains.kotlin.resolve.calls.model.*
@@ -46,9 +47,7 @@ import org.jetbrains.kotlin.types.expressions.CoercionStrategy
 import org.jetbrains.kotlin.types.expressions.DoubleColonExpressionResolver
 import org.jetbrains.kotlin.types.expressions.ExpressionTypingServices
 import org.jetbrains.kotlin.types.expressions.typeInfoFactory.createTypeInfo
-import org.jetbrains.kotlin.types.typeUtil.contains
-import org.jetbrains.kotlin.types.typeUtil.isUnit
-import org.jetbrains.kotlin.types.typeUtil.shouldBeUpdated
+import org.jetbrains.kotlin.types.typeUtil.*
 import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 
 class ResolvedAtomCompleter(
@@ -89,9 +88,22 @@ class ResolvedAtomCompleter(
             is ResolvedLambdaAtom -> completeLambda(resolvedAtom)
             is ResolvedCallAtom -> completeResolvedCall(resolvedAtom, emptyList())
             is ResolvedSubCallArgument -> completeSubCallArgument(resolvedAtom)
+            is ResolvedExpressionAtom -> completeExpression(resolvedAtom)
         }
     }
 
+    // We run completion on expressions only for last statements of block expression to substitute freshly inferred stub types variables
+    fun completeExpression(resolvedAtom: ResolvedExpressionAtom) {
+        val argumentExpression = resolvedAtom.atom.psiExpression
+        val inferenceSession = topLevelCallContext.inferenceSession
+
+        if (argumentExpression !is KtBlockExpression || inferenceSession !is BuilderInferenceSession) return
+
+        val callableReference = argumentExpression.statements.lastOrNull() as? KtCallableReferenceExpression ?: return
+
+        inferenceSession.completeDoubleColonExpression(callableReference, inferenceSession.getNotFixedToInferredTypesSubstitutor())
+    }
+
     fun completeAll(resolvedAtom: ResolvedAtom) {
         if (!resolvedAtom.analyzed)
             return
@@ -142,6 +154,14 @@ class ResolvedAtomCompleter(
             else -> atom.psiKotlinCall.psiCall
         }
 
+        val callElement = psiCallForResolutionContext.callElement
+        if (callElement is KtExpression) {
+            val recordedType = topLevelCallContext.trace.getType(callElement)
+            if ((recordedType.shouldBeUpdated()) && resolvedCall.resultingDescriptor.returnType != null) {
+                topLevelCallContext.trace.recordType(callElement, resolvedCall.resultingDescriptor.returnType)
+            }
+        }
+
         val resolutionContextForPartialCall =
             topLevelCallContext.trace[BindingContext.PARTIAL_CALL_RESOLUTION_CONTEXT, psiCallForResolutionContext]
 
diff --git a/core/descriptors/src/org/jetbrains/kotlin/types/TypeUtils.kt b/core/descriptors/src/org/jetbrains/kotlin/types/TypeUtils.kt
index 954ea780b61..037180ce0a5 100644
--- a/core/descriptors/src/org/jetbrains/kotlin/types/TypeUtils.kt
+++ b/core/descriptors/src/org/jetbrains/kotlin/types/TypeUtils.kt
@@ -375,4 +375,5 @@ private fun NewCapturedType.unCaptureTopLevelType(): UnwrappedType {
     return constructor.projection.type.unwrap()
 }
 
-fun KotlinType.shouldBeUpdated() = contains { it is StubTypeForBuilderInference || it.constructor is TypeVariableTypeConstructorMarker || it.isError }
+fun KotlinType?.shouldBeUpdated() =
+    this == null || contains { it is StubTypeForBuilderInference || it.constructor is TypeVariableTypeConstructorMarker || it.isError }