K2: drop nullable results for generated synthetic callees

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

@@ -57,11 +57,16 @@ class FirSyntheticCallGenerator(
     private val checkNotNullFunction: FirSimpleFunction = generateSyntheticCheckNotNullFunction()
     private val elvisFunction: FirSimpleFunction = generateSyntheticElvisFunction()
 
-    fun generateCalleeForWhenExpression(whenExpression: FirWhenExpression, context: ResolutionContext): FirWhenExpression? {
-        val stubReference = whenExpression.calleeReference
-        // TODO: Investigate: assertion failed in ModularizedTest
-        // assert(stubReference is FirStubReference)
-        if (stubReference !is FirStubReference) return null
+    private fun assertSyntheticResolvableReferenceIsNotResolved(resolvable: FirResolvable) {
+        // All synthetic calls (FirWhenExpression, FirTryExpression, FirElvisExpression, FirCheckNotNullCall)
+        // contains FirStubReference on creation.
+        // generateCallee... functions below replace these references with resolved references.
+        // This check ensures that we don't enter their resolve twice.
+        assert(resolvable.calleeReference is FirStubReference)
+    }
+
+    fun generateCalleeForWhenExpression(whenExpression: FirWhenExpression, context: ResolutionContext): FirWhenExpression {
+        assertSyntheticResolvableReferenceIsNotResolved(whenExpression)
 
         val argumentList = buildArgumentList {
             arguments += whenExpression.branches.map { it.result }
@@ -78,8 +83,7 @@ class FirSyntheticCallGenerator(
     }
 
     fun generateCalleeForTryExpression(tryExpression: FirTryExpression, context: ResolutionContext): FirTryExpression {
-        val stubReference = tryExpression.calleeReference
-        assert(stubReference is FirStubReference)
+        assertSyntheticResolvableReferenceIsNotResolved(tryExpression)
 
         val argumentList = buildArgumentList {
             with(tryExpression) {
@@ -101,9 +105,8 @@ class FirSyntheticCallGenerator(
         return tryExpression.transformCalleeReference(UpdateReference, reference)
     }
 
-    fun generateCalleeForCheckNotNullCall(checkNotNullCall: FirCheckNotNullCall, context: ResolutionContext): FirCheckNotNullCall? {
-        val stubReference = checkNotNullCall.calleeReference
-        if (stubReference !is FirStubReference) return null
+    fun generateCalleeForCheckNotNullCall(checkNotNullCall: FirCheckNotNullCall, context: ResolutionContext): FirCheckNotNullCall {
+        assertSyntheticResolvableReferenceIsNotResolved(checkNotNullCall)
 
         val reference = generateCalleeReferenceWithCandidate(
             checkNotNullCall,
@@ -116,8 +119,8 @@ class FirSyntheticCallGenerator(
         return checkNotNullCall.transformCalleeReference(UpdateReference, reference)
     }
 
-    fun generateCalleeForElvisExpression(elvisExpression: FirElvisExpression, context: ResolutionContext): FirElvisExpression? {
-        if (elvisExpression.calleeReference !is FirStubReference) return null
+    fun generateCalleeForElvisExpression(elvisExpression: FirElvisExpression, context: ResolutionContext): FirElvisExpression {
+        assertSyntheticResolvableReferenceIsNotResolved(elvisExpression)
 
         val argumentList = buildArgumentList {
             arguments += elvisExpression.lhs

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

@@ -6,8 +6,6 @@
 package org.jetbrains.kotlin.fir.resolve.transformers.body.resolve
 
 import org.jetbrains.kotlin.fir.FirTargetElement
-import org.jetbrains.kotlin.fir.diagnostics.ConeSimpleDiagnostic
-import org.jetbrains.kotlin.fir.diagnostics.DiagnosticKind
 import org.jetbrains.kotlin.fir.expressions.*
 import org.jetbrains.kotlin.fir.expressions.impl.FirElseIfTrueCondition
 import org.jetbrains.kotlin.fir.expressions.impl.FirEmptyExpressionBlock
@@ -18,7 +16,6 @@ import org.jetbrains.kotlin.fir.resolve.transformers.FirSyntheticCallGenerator
 import org.jetbrains.kotlin.fir.resolve.transformers.FirWhenExhaustivenessTransformer
 import org.jetbrains.kotlin.fir.resolvedTypeFromPrototype
 import org.jetbrains.kotlin.fir.types.*
-import org.jetbrains.kotlin.fir.types.builder.buildErrorTypeRef
 import org.jetbrains.kotlin.fir.visitors.transformSingle
 
 class FirControlFlowStatementsResolveTransformer(transformer: FirAbstractBodyResolveTransformerDispatcher) :
@@ -85,15 +82,7 @@ class FirControlFlowStatementsResolveTransformer(transformer: FirAbstractBodyRes
                             resolutionModeForBranches,
                         )
 
-                        whenExpression = syntheticCallGenerator.generateCalleeForWhenExpression(whenExpression, resolutionContext) ?: run {
-                            whenExpression = whenExpression.transformSingle(whenExhaustivenessTransformer, null)
-                            dataFlowAnalyzer.exitWhenExpression(whenExpression, callCompleted = false)
-                            whenExpression.resultType = buildErrorTypeRef {
-                                diagnostic = ConeSimpleDiagnostic("Can't resolve when expression", DiagnosticKind.InferenceError)
-                            }
-                            return@withWhenSubjectType whenExpression
-                        }
-
+                        whenExpression = syntheticCallGenerator.generateCalleeForWhenExpression(whenExpression, resolutionContext)
                         completionNeeded = true
                     }
                 }
@@ -268,16 +257,9 @@ class FirControlFlowStatementsResolveTransformer(transformer: FirAbstractBodyRes
         )
         elvisExpression.transformRhs(transformer, resolutionModeForRhs)
 
-        var callCompleted = false
-        val result = syntheticCallGenerator.generateCalleeForElvisExpression(elvisExpression, resolutionContext)?.let {
-            val completed = callCompleter.completeCall(it, data)
-            callCompleted = completed.callCompleted
-            completed.result
-        } ?: elvisExpression.also {
-            it.resultType = buildErrorTypeRef {
-                diagnostic = ConeSimpleDiagnostic("Can't resolve ?: operator call", DiagnosticKind.InferenceError)
-            }
-        }
+        val (result, callCompleted) = callCompleter.completeCall(
+            syntheticCallGenerator.generateCalleeForElvisExpression(elvisExpression, resolutionContext), data
+        )
 
         var isLhsNotNull = false
         if (result.rhs.typeRef.coneTypeSafe<ConeKotlinType>()?.isNothing == true) {

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

@@ -954,19 +954,10 @@ open class FirExpressionsResolveTransformer(transformer: FirAbstractBodyResolveT
             .transformAnnotations(transformer, ResolutionMode.ContextIndependent)
             .replaceArgumentList(checkNotNullCall.argumentList.transform(transformer, ResolutionMode.ContextDependent))
 
-        var callCompleted = false
-        val result = components.syntheticCallGenerator.generateCalleeForCheckNotNullCall(checkNotNullCall, resolutionContext)?.let {
-            val completionResult = callCompleter.completeCall(it, data)
-            callCompleted = completionResult.callCompleted
-            completionResult.result
-        } ?: run {
-            checkNotNullCall.resultType =
-                buildErrorTypeRef {
-                    diagnostic = ConeSimpleDiagnostic("Can't resolve !! operator call", DiagnosticKind.InferenceError)
-                }
-            callCompleted = true
-            checkNotNullCall
-        }
+        val (result, callCompleted) = callCompleter.completeCall(
+            components.syntheticCallGenerator.generateCalleeForCheckNotNullCall(checkNotNullCall, resolutionContext), data
+        )
+
         dataFlowAnalyzer.exitCheckNotNullCall(result, callCompleted)
         return result
     }