FIR: expect nullable type for elvis LHS

compiler/fir/fir2ir/tests-gen/org/jetbrains/kotlin/test/runners/codegen/FirBlackBoxCodegenTestGenerated.java

@@ -13644,6 +13644,12 @@ public class FirBlackBoxCodegenTestGenerated extends AbstractFirBlackBoxCodegenT
             runTest("compiler/testData/codegen/box/elvis/nullNullOk.kt");
         }
 
+        @Test
+        @TestMetadata("ofNonNullableResultType.kt")
+        public void testOfNonNullableResultType() throws Exception {
+            runTest("compiler/testData/codegen/box/elvis/ofNonNullableResultType.kt");
+        }
+
         @Test
         @TestMetadata("primitive.kt")
         public void testPrimitive() throws Exception {

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/ResolutionMode.kt

@@ -26,6 +26,11 @@ sealed class ResolutionMode {
 fun withExpectedType(expectedTypeRef: FirTypeRef?): ResolutionMode =
     expectedTypeRef?.let { ResolutionMode.WithExpectedType(it) } ?: ResolutionMode.ContextDependent
 
+@JvmName("withExpectedTypeNullable")
+fun withExpectedType(coneType: ConeKotlinType?): ResolutionMode {
+    return coneType?.let { withExpectedType(it) } ?: ResolutionMode.ContextDependent
+}
+
 fun withExpectedType(coneType: ConeKotlinType): ResolutionMode {
     val typeRef = buildResolvedTypeRef {
         type = coneType

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

@@ -5,7 +5,6 @@
 
 package org.jetbrains.kotlin.fir.resolve.transformers.body.resolve
 
-import org.jetbrains.kotlin.fir.FirFakeSourceElementKind
 import org.jetbrains.kotlin.fir.FirTargetElement
 import org.jetbrains.kotlin.fir.diagnostics.ConeSimpleDiagnostic
 import org.jetbrains.kotlin.fir.diagnostics.DiagnosticKind
@@ -18,7 +17,7 @@ import org.jetbrains.kotlin.fir.resolve.transformers.FirSyntheticCallGenerator
 import org.jetbrains.kotlin.fir.resolve.transformers.FirWhenExhaustivenessTransformer
 import org.jetbrains.kotlin.fir.resolve.withExpectedType
 import org.jetbrains.kotlin.fir.resolvedTypeFromPrototype
-import org.jetbrains.kotlin.fir.types.FirImplicitTypeRef
+import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.fir.types.builder.buildErrorTypeRef
 import org.jetbrains.kotlin.fir.visitors.CompositeTransformResult
 import org.jetbrains.kotlin.fir.visitors.compose
@@ -228,12 +227,14 @@ class FirControlFlowStatementsResolveTransformer(transformer: FirBodyResolveTran
     ): CompositeTransformResult<FirStatement> {
         if (elvisExpression.calleeReference is FirResolvedNamedReference) return elvisExpression.compose()
         elvisExpression.transformAnnotations(transformer, data)
-        val expectedArgumentType =
+
-            if (data is ResolutionMode.WithExpectedType && data.expectedType !is FirImplicitTypeRef) data
+        val expectedType = data.expectedType?.coneTypeSafe<ConeKotlinType>()
-            else ResolutionMode.ContextDependent
+        val resolutionModeForLhs = withExpectedType(expectedType?.withNullability(ConeNullability.NULLABLE))
-        elvisExpression.transformLhs(transformer, expectedArgumentType)
+        elvisExpression.transformLhs(transformer, resolutionModeForLhs)
         dataFlowAnalyzer.exitElvisLhs(elvisExpression)
-        elvisExpression.transformRhs(transformer, expectedArgumentType)
+
+        val resolutionModeForRhs = withExpectedType(expectedType)
+        elvisExpression.transformRhs(transformer, resolutionModeForRhs)
 
         val result = syntheticCallGenerator.generateCalleeForElvisExpression(elvisExpression, resolutionContext)?.let {
             callCompleter.completeCall(it, data.expectedType).result

compiler/testData/codegen/box/elvis/ofNonNullableResultType.kt

@@ -0,0 +1,7 @@
+fun <T> f(): T? = "OK" as? T
+
+fun g(): Nothing = throw RuntimeException("fail")
+
+fun <T : Any> h(): T = run { f() } ?: run { g() }
+
+fun box(): String = h<String>()

compiler/testData/ir/irText/expressions/kt30796.fir.kt.txt

@@ -50,17 +50,17 @@ fun <T : Any?> test(value: T, value2: T) {
     }
   }
   val x5: Any = { // BLOCK
-    val <elvis>: Any = magic<Any>()
+    val <elvis>: Any? = magic<Any?>()
     when {
       EQEQ(arg0 = <elvis>, arg1 = null) -> 42
       else -> <elvis>
     }
   }
   val x6: Any = { // BLOCK
-    val <elvis>: Any = { // BLOCK
+    val <elvis>: Any? = { // BLOCK
       val <elvis>: T = value
       when {
-        EQEQ(arg0 = <elvis>, arg1 = null) -> magic<Any>()
+        EQEQ(arg0 = <elvis>, arg1 = null) -> magic<Any?>()
         else -> <elvis>
       }
     }
@@ -70,8 +70,8 @@ fun <T : Any?> test(value: T, value2: T) {
     }
   }
   val x7: Any = { // BLOCK
-    val <elvis>: Any = { // BLOCK
+    val <elvis>: Any? = { // BLOCK
-      val <elvis>: Any = magic<Any>()
+      val <elvis>: Any? = magic<Any?>()
       when {
         EQEQ(arg0 = <elvis>, arg1 = null) -> value
         else -> <elvis>
@@ -83,3 +83,4 @@ fun <T : Any?> test(value: T, value2: T) {
     }
   }
 }
+

compiler/testData/ir/irText/expressions/kt30796.fir.txt

@@ -97,65 +97,65 @@ FILE fqName:<root> fileName:/kt30796.kt
               then: GET_VAR 'val tmp_5: T of <root>.test [val] declared in <root>.test' type=T of <root>.test origin=null
       VAR name:x5 type:kotlin.Any [val]
         BLOCK type=kotlin.Any origin=ELVIS
-          VAR IR_TEMPORARY_VARIABLE name:tmp_7 type:kotlin.Any [val]
+          VAR IR_TEMPORARY_VARIABLE name:tmp_7 type:kotlin.Any? [val]
-            CALL 'public final fun magic <T> (): T of <root>.magic declared in <root>' type=kotlin.Any origin=null
+            CALL 'public final fun magic <T> (): T of <root>.magic declared in <root>' type=kotlin.Any? origin=null
-              <T>: kotlin.Any
+              <T>: kotlin.Any?
           WHEN type=kotlin.Any origin=ELVIS
             BRANCH
               if: CALL 'public final fun EQEQ (arg0: kotlin.Any?, arg1: kotlin.Any?): kotlin.Boolean declared in kotlin.internal.ir' type=kotlin.Boolean origin=EQEQ
-                arg0: GET_VAR 'val tmp_7: kotlin.Any [val] declared in <root>.test' type=kotlin.Any origin=null
+                arg0: GET_VAR 'val tmp_7: kotlin.Any? [val] declared in <root>.test' type=kotlin.Any? origin=null
                 arg1: CONST Null type=kotlin.Nothing? value=null
               then: CONST Int type=kotlin.Int value=42
             BRANCH
               if: CONST Boolean type=kotlin.Boolean value=true
-              then: GET_VAR 'val tmp_7: kotlin.Any [val] declared in <root>.test' type=kotlin.Any origin=null
+              then: GET_VAR 'val tmp_7: kotlin.Any? [val] declared in <root>.test' type=kotlin.Any? origin=null
       VAR name:x6 type:kotlin.Any [val]
         BLOCK type=kotlin.Any origin=ELVIS
-          VAR IR_TEMPORARY_VARIABLE name:tmp_8 type:kotlin.Any [val]
+          VAR IR_TEMPORARY_VARIABLE name:tmp_8 type:kotlin.Any? [val]
-            BLOCK type=kotlin.Any origin=ELVIS
+            BLOCK type=kotlin.Any? origin=ELVIS
               VAR IR_TEMPORARY_VARIABLE name:tmp_9 type:T of <root>.test [val]
                 GET_VAR 'value: T of <root>.test declared in <root>.test' type=T of <root>.test origin=null
-              WHEN type=kotlin.Any origin=ELVIS
+              WHEN type=kotlin.Any? origin=ELVIS
                 BRANCH
                   if: CALL 'public final fun EQEQ (arg0: kotlin.Any?, arg1: kotlin.Any?): kotlin.Boolean declared in kotlin.internal.ir' type=kotlin.Boolean origin=EQEQ
                     arg0: GET_VAR 'val tmp_9: T of <root>.test [val] declared in <root>.test' type=T of <root>.test origin=null
                     arg1: CONST Null type=kotlin.Nothing? value=null
-                  then: CALL 'public final fun magic <T> (): T of <root>.magic declared in <root>' type=kotlin.Any origin=null
+                  then: CALL 'public final fun magic <T> (): T of <root>.magic declared in <root>' type=kotlin.Any? origin=null
-                    <T>: kotlin.Any
+                    <T>: kotlin.Any?
                 BRANCH
                   if: CONST Boolean type=kotlin.Boolean value=true
                   then: GET_VAR 'val tmp_9: T of <root>.test [val] declared in <root>.test' type=T of <root>.test origin=null
           WHEN type=kotlin.Any origin=ELVIS
             BRANCH
               if: CALL 'public final fun EQEQ (arg0: kotlin.Any?, arg1: kotlin.Any?): kotlin.Boolean declared in kotlin.internal.ir' type=kotlin.Boolean origin=EQEQ
-                arg0: GET_VAR 'val tmp_8: kotlin.Any [val] declared in <root>.test' type=kotlin.Any origin=null
+                arg0: GET_VAR 'val tmp_8: kotlin.Any? [val] declared in <root>.test' type=kotlin.Any? origin=null
                 arg1: CONST Null type=kotlin.Nothing? value=null
               then: CONST Int type=kotlin.Int value=42
             BRANCH
               if: CONST Boolean type=kotlin.Boolean value=true
-              then: GET_VAR 'val tmp_8: kotlin.Any [val] declared in <root>.test' type=kotlin.Any origin=null
+              then: GET_VAR 'val tmp_8: kotlin.Any? [val] declared in <root>.test' type=kotlin.Any? origin=null
       VAR name:x7 type:kotlin.Any [val]
         BLOCK type=kotlin.Any origin=ELVIS
-          VAR IR_TEMPORARY_VARIABLE name:tmp_10 type:kotlin.Any [val]
+          VAR IR_TEMPORARY_VARIABLE name:tmp_10 type:kotlin.Any? [val]
-            BLOCK type=kotlin.Any origin=ELVIS
+            BLOCK type=kotlin.Any? origin=ELVIS
-              VAR IR_TEMPORARY_VARIABLE name:tmp_11 type:kotlin.Any [val]
+              VAR IR_TEMPORARY_VARIABLE name:tmp_11 type:kotlin.Any? [val]
-                CALL 'public final fun magic <T> (): T of <root>.magic declared in <root>' type=kotlin.Any origin=null
+                CALL 'public final fun magic <T> (): T of <root>.magic declared in <root>' type=kotlin.Any? origin=null
-                  <T>: kotlin.Any
+                  <T>: kotlin.Any?
-              WHEN type=kotlin.Any origin=ELVIS
+              WHEN type=kotlin.Any? origin=ELVIS
                 BRANCH
                   if: CALL 'public final fun EQEQ (arg0: kotlin.Any?, arg1: kotlin.Any?): kotlin.Boolean declared in kotlin.internal.ir' type=kotlin.Boolean origin=EQEQ
-                    arg0: GET_VAR 'val tmp_11: kotlin.Any [val] declared in <root>.test' type=kotlin.Any origin=null
+                    arg0: GET_VAR 'val tmp_11: kotlin.Any? [val] declared in <root>.test' type=kotlin.Any? origin=null
                     arg1: CONST Null type=kotlin.Nothing? value=null
                   then: GET_VAR 'value: T of <root>.test declared in <root>.test' type=T of <root>.test origin=null
                 BRANCH
                   if: CONST Boolean type=kotlin.Boolean value=true
-                  then: GET_VAR 'val tmp_11: kotlin.Any [val] declared in <root>.test' type=kotlin.Any origin=null
+                  then: GET_VAR 'val tmp_11: kotlin.Any? [val] declared in <root>.test' type=kotlin.Any? origin=null
           WHEN type=kotlin.Any origin=ELVIS
             BRANCH
               if: CALL 'public final fun EQEQ (arg0: kotlin.Any?, arg1: kotlin.Any?): kotlin.Boolean declared in kotlin.internal.ir' type=kotlin.Boolean origin=EQEQ
-                arg0: GET_VAR 'val tmp_10: kotlin.Any [val] declared in <root>.test' type=kotlin.Any origin=null
+                arg0: GET_VAR 'val tmp_10: kotlin.Any? [val] declared in <root>.test' type=kotlin.Any? origin=null
                 arg1: CONST Null type=kotlin.Nothing? value=null
               then: CONST Int type=kotlin.Int value=42
             BRANCH
               if: CONST Boolean type=kotlin.Boolean value=true
-              then: GET_VAR 'val tmp_10: kotlin.Any [val] declared in <root>.test' type=kotlin.Any origin=null
+              then: GET_VAR 'val tmp_10: kotlin.Any? [val] declared in <root>.test' type=kotlin.Any? origin=null

compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/codegen/BlackBoxCodegenTestGenerated.java

@@ -13644,6 +13644,12 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
             runTest("compiler/testData/codegen/box/elvis/nullNullOk.kt");
         }
 
+        @Test
+        @TestMetadata("ofNonNullableResultType.kt")
+        public void testOfNonNullableResultType() throws Exception {
+            runTest("compiler/testData/codegen/box/elvis/ofNonNullableResultType.kt");
+        }
+
         @Test
         @TestMetadata("primitive.kt")
         public void testPrimitive() throws Exception {

compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/codegen/IrBlackBoxCodegenTestGenerated.java

@@ -13644,6 +13644,12 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
             runTest("compiler/testData/codegen/box/elvis/nullNullOk.kt");
         }
 
+        @Test
+        @TestMetadata("ofNonNullableResultType.kt")
+        public void testOfNonNullableResultType() throws Exception {
+            runTest("compiler/testData/codegen/box/elvis/ofNonNullableResultType.kt");
+        }
+
         @Test
         @TestMetadata("primitive.kt")
         public void testPrimitive() throws Exception {

compiler/tests-gen/org/jetbrains/kotlin/codegen/LightAnalysisModeTestGenerated.java

@@ -11220,6 +11220,11 @@ public class LightAnalysisModeTestGenerated extends AbstractLightAnalysisModeTes
             runTest("compiler/testData/codegen/box/elvis/nullNullOk.kt");
         }
 
+        @TestMetadata("ofNonNullableResultType.kt")
+        public void testOfNonNullableResultType() throws Exception {
+            runTest("compiler/testData/codegen/box/elvis/ofNonNullableResultType.kt");
+        }
+
         @TestMetadata("primitive.kt")
         public void testPrimitive() throws Exception {
             runTest("compiler/testData/codegen/box/elvis/primitive.kt");

js/js.tests/tests-gen/org/jetbrains/kotlin/js/test/es6/semantics/IrJsCodegenBoxES6TestGenerated.java

@@ -10015,6 +10015,11 @@ public class IrJsCodegenBoxES6TestGenerated extends AbstractIrJsCodegenBoxES6Tes
             runTest("compiler/testData/codegen/box/elvis/nullNullOk.kt");
         }
 
+        @TestMetadata("ofNonNullableResultType.kt")
+        public void testOfNonNullableResultType() throws Exception {
+            runTest("compiler/testData/codegen/box/elvis/ofNonNullableResultType.kt");
+        }
+
         @TestMetadata("primitive.kt")
         public void testPrimitive() throws Exception {
             runTest("compiler/testData/codegen/box/elvis/primitive.kt");

js/js.tests/tests-gen/org/jetbrains/kotlin/js/test/ir/semantics/IrJsCodegenBoxTestGenerated.java

@@ -9500,6 +9500,11 @@ public class IrJsCodegenBoxTestGenerated extends AbstractIrJsCodegenBoxTest {
             runTest("compiler/testData/codegen/box/elvis/nullNullOk.kt");
         }
 
+        @TestMetadata("ofNonNullableResultType.kt")
+        public void testOfNonNullableResultType() throws Exception {
+            runTest("compiler/testData/codegen/box/elvis/ofNonNullableResultType.kt");
+        }
+
         @TestMetadata("primitive.kt")
         public void testPrimitive() throws Exception {
             runTest("compiler/testData/codegen/box/elvis/primitive.kt");

js/js.tests/tests-gen/org/jetbrains/kotlin/js/test/semantics/JsCodegenBoxTestGenerated.java

@@ -9500,6 +9500,11 @@ public class JsCodegenBoxTestGenerated extends AbstractJsCodegenBoxTest {
             runTest("compiler/testData/codegen/box/elvis/nullNullOk.kt");
         }
 
+        @TestMetadata("ofNonNullableResultType.kt")
+        public void testOfNonNullableResultType() throws Exception {
+            runTest("compiler/testData/codegen/box/elvis/ofNonNullableResultType.kt");
+        }
+
         @TestMetadata("primitive.kt")
         public void testPrimitive() throws Exception {
             runTest("compiler/testData/codegen/box/elvis/primitive.kt");

js/js.tests/tests-gen/org/jetbrains/kotlin/js/test/wasm/semantics/IrCodegenBoxWasmTestGenerated.java

@@ -4549,6 +4549,11 @@ public class IrCodegenBoxWasmTestGenerated extends AbstractIrCodegenBoxWasmTest
             runTest("compiler/testData/codegen/box/elvis/nullNullOk.kt");
         }
 
+        @TestMetadata("ofNonNullableResultType.kt")
+        public void testOfNonNullableResultType() throws Exception {
+            runTest("compiler/testData/codegen/box/elvis/ofNonNullableResultType.kt");
+        }
+
         @TestMetadata("primitive.kt")
         public void testPrimitive() throws Exception {
             runTest("compiler/testData/codegen/box/elvis/primitive.kt");