FIR: Minor. Extract common parts in FirCallCompletionResultsWriterTransformer

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

@@ -53,11 +53,7 @@ class FirCallCompletionResultsWriterTransformer(
             FirSimpleDiagnostic("Callee reference to candidate without return type: ${candidateFir.render()}")
         )
 
-        val initialType = calleeReference.candidate.substitutor.substituteOrNull(typeRef.type)
+        qualifiedAccessExpression.replaceTypeRefWithSubstituted(calleeReference, typeRef)
-        val finalType = finalSubstitutor.substituteOrNull(initialType)
-
-        val resultType = typeRef.withReplacedConeType(finalType)
-        qualifiedAccessExpression.replaceTypeRef(resultType)
 
         return qualifiedAccessExpression.transformCalleeReference(
             StoreCalleeReference,
@@ -69,6 +65,24 @@ class FirCallCompletionResultsWriterTransformer(
         ).compose()
     }
 
+    private fun <D : FirExpression> D.replaceTypeRefWithSubstituted(
+        calleeReference: FirNamedReferenceWithCandidate,
+        typeRef: FirResolvedTypeRef
+    ): D {
+        val resultTypeRef = typeRef.substituteTypeRef(calleeReference.candidate)
+        replaceTypeRef(resultTypeRef)
+        return this
+    }
+
+    private fun FirResolvedTypeRef.substituteTypeRef(
+        candidate: Candidate
+    ): FirResolvedTypeRef {
+        val initialType = candidate.substitutor.substituteOrNull(type)
+        val finalType = finalSubstitutor.substituteOrNull(initialType)
+
+        return withReplacedConeType(finalType)
+    }
+
     override fun transformCallableReferenceAccess(
         callableReferenceAccess: FirCallableReferenceAccess,
         data: Nothing?
@@ -149,10 +163,7 @@ class FirCallCompletionResultsWriterTransformer(
             }
         }
 
-        val initialType = subCandidate.substitutor.substituteOrNull(typeRef.type)
+        val resultType = typeRef.substituteTypeRef(subCandidate)
-        val finalType = finalSubstitutor.substituteOrNull(initialType)
-
-        val resultType = typeRef.withReplacedConeType(finalType)
 
         return functionCall.copy(
             resultType = resultType,
@@ -216,10 +227,7 @@ class FirCallCompletionResultsWriterTransformer(
 
         val typeRef = typeCalculator.tryCalculateReturnType(declaration)
 
-        val initialType = subCandidate.substitutor.substituteOrNull(typeRef.type)
+        whenExpression.resultType = typeRef.substituteTypeRef(subCandidate)
-        val finalType = finalSubstitutor.substituteOrNull(initialType)
-
-        whenExpression.resultType = typeRef.withReplacedConeType(finalType)
 
         return whenExpression.transformCalleeReference(
             StoreCalleeReference,
@@ -242,10 +250,7 @@ class FirCallCompletionResultsWriterTransformer(
 
         val typeRef = typeCalculator.tryCalculateReturnType(declaration)
 
-        val initialType = subCandidate.substitutor.substituteOrNull(typeRef.type)
+        tryExpression.resultType = typeRef.substituteTypeRef(subCandidate)
-        val finalType = finalSubstitutor.substituteOrNull(initialType)
-
-        tryExpression.resultType = typeRef.withReplacedConeType(finalType)
         return tryExpression.transformCalleeReference(
             StoreCalleeReference,
             FirResolvedNamedReferenceImpl(