[FIR] Add lValueTypeRef to FirVariableAssignment, fix tests with ASSIGNMENT_TYPE_MISMATCH

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/inference/FirCallCompleter.kt

@@ -12,6 +12,7 @@ import org.jetbrains.kotlin.fir.declarations.builder.buildValueParameter
 import org.jetbrains.kotlin.fir.expressions.FirExpression
 import org.jetbrains.kotlin.fir.expressions.FirResolvable
 import org.jetbrains.kotlin.fir.expressions.FirStatement
+import org.jetbrains.kotlin.fir.expressions.FirVariableAssignment
 import org.jetbrains.kotlin.fir.resolve.ResolutionMode
 import org.jetbrains.kotlin.fir.resolve.calls.Candidate
 import org.jetbrains.kotlin.fir.resolve.calls.FirNamedReferenceWithCandidate
@@ -56,7 +57,12 @@ class FirCallCompleter(
             where T : FirResolvable, T : FirStatement {
         val typeRef = components.typeFromCallee(call)
 
+        if (call is FirVariableAssignment) {
+            call.replaceLValueTypeRef(typeRef)
+        }
+
         val reference = call.calleeReference as? FirNamedReferenceWithCandidate ?: return CompletionResult(call, true)
+
         val candidate = reference.candidate
         val initialType = components.initialTypeOfCandidate(candidate, call)
 

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/FirCallCompletionResultsWriterTransformer.kt

@@ -66,7 +66,7 @@ class FirCallCompletionResultsWriterTransformer(
         }
     }
 
-    private fun <T : FirQualifiedAccessExpression> prepareQualifiedTransform(
+    private fun <T : FirQualifiedAccess> prepareQualifiedTransform(
         qualifiedAccessExpression: T, calleeReference: FirNamedReferenceWithCandidate
     ): T {
         val subCandidate = calleeReference.candidate
@@ -111,7 +111,13 @@ class FirCallCompletionResultsWriterTransformer(
             )
             .transformDispatchReceiver(StoreReceiver, subCandidate.dispatchReceiverExpression())
             .transformExtensionReceiver(StoreReceiver, subCandidate.extensionReceiverExpression()) as T
-        result.replaceTypeRef(typeRef)
+
+        if (result is FirQualifiedAccessExpression) {
+            result.replaceTypeRef(typeRef)
+        } else if (result is FirVariableAssignment) {
+            result.replaceLValueTypeRef(typeRef)
+        }
+
         if (declaration !is FirErrorFunction) {
             result.replaceTypeArguments(typeArguments)
         }
@@ -354,13 +360,19 @@ class FirCallCompletionResultsWriterTransformer(
     ): FirStatement {
         val calleeReference = variableAssignment.calleeReference as? FirNamedReferenceWithCandidate
             ?: return variableAssignment
-        val typeArguments = computeTypeArguments(variableAssignment, calleeReference.candidate)
+
+        // Initialize lValueTypeRef
+        val qualifiedTransform = prepareQualifiedTransform(variableAssignment, calleeReference)
+        val lValueTypeRef = qualifiedTransform.lValueTypeRef as FirResolvedTypeRef
+        val resultLValueType = lValueTypeRef.substituteTypeRef(calleeReference.candidate)
+        resultLValueType.ensureResolvedTypeDeclaration(session)
+        variableAssignment.replaceLValueTypeRef(resultLValueType)
+        session.lookupTracker?.recordTypeResolveAsLookup(resultLValueType, variableAssignment.lValue.source, null)
+
         return variableAssignment.transformCalleeReference(
             StoreCalleeReference,
             calleeReference.toResolvedReference(),
-        ).apply {
-            replaceTypeArguments(typeArguments)
-        }
+        )
     }
 
     private inner class TypeUpdaterForDelegateArguments : FirTransformer<Any?>() {

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirExpressionsResolveTransformer.kt

@@ -677,8 +677,7 @@ open class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransform
         val resolvedAssignment = callResolver.resolveVariableAccessAndSelectCandidate(variableAssignment)
         val result = if (resolvedAssignment is FirVariableAssignment) {
             val completeAssignment = callCompleter.completeCall(resolvedAssignment, noExpectedType).result // TODO: check
-            val expectedType = components.typeFromCallee(completeAssignment)
-            completeAssignment.transformRValue(transformer, withExpectedType(expectedType))
+            completeAssignment.transformRValue(transformer, withExpectedType(variableAssignment.lValueTypeRef))
         } else {
             // This can happen in erroneous code only
             resolvedAssignment