FIR: pass elvis expected type to lhs/rhs

compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/Fir2IrVisitor.kt

@@ -587,8 +587,7 @@ class Fir2IrVisitor(
             fun irGetLhsValue(): IrGetValue =
                 IrGetValueImpl(startOffset, endOffset, irLhsVariable.type, irLhsVariable.symbol)
 
-            // TODO: replace with .coneType
-            val originalType = firLhsVariable.returnTypeRef.coneTypeUnsafe<ConeKotlinType>()
+            val originalType = firLhsVariable.returnTypeRef.coneType
             val notNullType = originalType.withNullability(ConeNullability.NOT_NULL)
             val irBranches = listOf(
                 IrBranchImpl(

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

@@ -216,9 +216,12 @@ class FirControlFlowStatementsResolveTransformer(transformer: FirBodyResolveTran
     ): CompositeTransformResult<FirStatement> {
         if (elvisExpression.calleeReference is FirResolvedNamedReference) return elvisExpression.compose()
         elvisExpression.transformAnnotations(transformer, data)
-        elvisExpression.transformLhs(transformer, ResolutionMode.ContextDependent)
+        val expectedArgumentType =
+            if (data is ResolutionMode.WithExpectedType && data.expectedType !is FirImplicitTypeRef) data
+            else ResolutionMode.ContextDependent
+        elvisExpression.transformLhs(transformer, expectedArgumentType)
         dataFlowAnalyzer.exitElvisLhs(elvisExpression)
-        elvisExpression.transformRhs(transformer, ResolutionMode.ContextDependent)
+        elvisExpression.transformRhs(transformer, expectedArgumentType)
 
         val result = syntheticCallGenerator.generateCalleeForElvisExpression(elvisExpression, resolutionContext)?.let {
             callCompleter.completeCall(it, data.expectedType).result