Substitute stub type variables in local variable descriptors

compiler/frontend/src/org/jetbrains/kotlin/resolve/LocalVariableResolver.kt

@@ -32,6 +32,7 @@ import org.jetbrains.kotlin.psi.KtPsiUtil
 import org.jetbrains.kotlin.psi.KtVariableDeclaration
 import org.jetbrains.kotlin.resolve.calls.components.InferenceSession
 import org.jetbrains.kotlin.resolve.calls.context.ContextDependency
+import org.jetbrains.kotlin.resolve.calls.inference.BuilderInferenceSession
 import org.jetbrains.kotlin.resolve.calls.smartcasts.DataFlowInfo
 import org.jetbrains.kotlin.resolve.calls.smartcasts.DataFlowValueFactory
 import org.jetbrains.kotlin.resolve.calls.util.isSingleUnderscore
@@ -189,6 +190,9 @@ class LocalVariableResolver(
             initializeWithDefaultGetterSetter(propertyDescriptor)
             trace.record(BindingContext.VARIABLE, variable, propertyDescriptor)
             result = propertyDescriptor
+            if (inferenceSession is BuilderInferenceSession) {
+                inferenceSession.addLocalVariable(variable)
+            }
         } else {
             val variableDescriptor = resolveLocalVariableDescriptorWithType(scope, variable, null, trace)
             // For a local variable the type must not be deferred
@@ -196,6 +200,9 @@ class LocalVariableResolver(
                 variableDescriptor, scope, variable, dataFlowInfo, inferenceSession, trace, local = true
             )
             variableDescriptor.setOutType(type)
+            if (inferenceSession is BuilderInferenceSession) {
+                inferenceSession.addLocalVariable(variable)
+            }
             result = variableDescriptor
         }
         variableTypeAndInitializerResolver

compiler/frontend/src/org/jetbrains/kotlin/resolve/VariableTypeAndInitializerResolver.kt

@@ -15,7 +15,6 @@ import org.jetbrains.kotlin.psi.KtProperty
 import org.jetbrains.kotlin.psi.KtVariableDeclaration
 import org.jetbrains.kotlin.resolve.DescriptorResolver.transformAnonymousTypeIfNeeded
 import org.jetbrains.kotlin.resolve.calls.components.InferenceSession
-import org.jetbrains.kotlin.resolve.calls.inference.BuilderInferenceSession
 import org.jetbrains.kotlin.resolve.calls.smartcasts.DataFlowInfo
 import org.jetbrains.kotlin.resolve.constants.evaluate.ConstantExpressionEvaluator
 import org.jetbrains.kotlin.resolve.scopes.LexicalScope
@@ -176,13 +175,8 @@ class VariableTypeAndInitializerResolver(
         val inferredType = expressionTypingServices.safeGetType(
             scope, initializer, TypeUtils.NO_EXPECTED_TYPE, dataFlowInfo, inferenceSession, trace
         )
-        val preparedType = approximateType(
-            if (inferenceSession is BuilderInferenceSession) {
-                inferenceSession.getNotFixedToInferredTypesSubstitutor().safeSubstitute(inferredType.unwrap())
-            } else inferredType,
-            local
-        )
-        return declarationReturnTypeSanitizer.sanitizeReturnType(preparedType, wrappedTypeFactory, trace, languageVersionSettings)
+        val approximatedType = approximateType(inferredType, local)
+        return declarationReturnTypeSanitizer.sanitizeReturnType(approximatedType, wrappedTypeFactory, trace, languageVersionSettings)
     }
 
     private fun approximateType(type: KotlinType, local: Boolean): UnwrappedType =

compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/inference/BuilderInferenceSession.kt

@@ -35,6 +35,7 @@ import org.jetbrains.kotlin.types.model.freshTypeConstructor
 import org.jetbrains.kotlin.types.model.safeSubstitute
 import org.jetbrains.kotlin.types.typeUtil.asTypeProjection
 import org.jetbrains.kotlin.types.typeUtil.contains
+import org.jetbrains.kotlin.types.typeUtil.shouldBeUpdated
 import org.jetbrains.kotlin.utils.addToStdlib.cast
 
 class BuilderInferenceSession(
@@ -72,6 +73,8 @@ class BuilderInferenceSession(
 
     private val commonCalls = arrayListOf<PSICompletedCallInfo>()
 
+    private val localVariables = arrayListOf<KtVariableDeclaration>()
+
     // These calls come from the old type inference
     private val oldDoubleColonExpressionCalls = arrayListOf<KtExpression>()
 
@@ -149,6 +152,10 @@ class BuilderInferenceSession(
         }
     }
 
+    fun addLocalVariable(variable: KtVariableDeclaration) {
+        localVariables.add(variable)
+    }
+
     private fun anyReceiverContainStubType(descriptor: CallableDescriptor): Boolean {
         return descriptor.dispatchReceiverParameter?.type?.contains { it is StubTypeForBuilderInference } == true ||
                 descriptor.extensionReceiverParameter?.type?.contains { it is StubTypeForBuilderInference } == true
@@ -426,6 +433,13 @@ class BuilderInferenceSession(
         )
     }
 
+    private fun updateLocalVariable(localVariable: KtVariableDeclaration, substitutor: NewTypeSubstitutor) {
+        val descriptor = trace[BindingContext.VARIABLE, localVariable] as? LocalVariableDescriptor
+        if (descriptor != null && descriptor.type.shouldBeUpdated()) {
+            descriptor.setOutType(substitutor.safeSubstitute(descriptor.type.unwrap()))
+        }
+    }
+
     private fun updateCall(
         completedCall: PSICompletedCallInfo,
         nonFixedTypesToResultSubstitutor: NewTypeSubstitutor,
@@ -544,6 +558,10 @@ class BuilderInferenceSession(
                 topLevelCallContext.replaceBindingTrace(findTopLevelTrace()).replaceInferenceSession(this)
             )
 
+            for (localVariable in localVariables) {
+                updateLocalVariable(localVariable, nonFixedTypesToResultSubstitutor)
+            }
+
             for (completedCall in commonCalls) {
                 updateCall(completedCall, nonFixedTypesToResultSubstitutor, nonFixedTypesToResult)
                 reportErrors(completedCall, completedCall.resolvedCall, errors)