diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirControlFlowStatementsResolveTransformer.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirControlFlowStatementsResolveTransformer.kt
index f0d13a75632..26c08d8ac28 100644
--- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirControlFlowStatementsResolveTransformer.kt
+++ b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirControlFlowStatementsResolveTransformer.kt
@@ -174,10 +174,16 @@ class FirControlFlowStatementsResolveTransformer(transformer: FirBodyResolveTran
     // ------------------------------- Jumps -------------------------------
 
     override fun <E : FirTargetElement> transformJump(jump: FirJump<E>, data: ResolutionMode): CompositeTransformResult<FirStatement> {
-        var result = transformer.transformExpression(jump, data).single
+        val expectedTypeRef = (jump as? FirReturnExpression)?.target?.labeledElement?.returnTypeRef
+
+        val mode = if (expectedTypeRef != null) {
+            ResolutionMode.WithExpectedType(expectedTypeRef)
+        } else {
+            ResolutionMode.ContextIndependent
+        }
+        var result = transformer.transformExpression(jump, mode).single
         if (result is FirReturnExpression) {
-            val expectedType = result.target.labeledElement.returnTypeRef.coneTypeSafe<ConeKotlinType>()
-            result = result.transformResult(integerLiteralTypeApproximator, expectedType)
+            result = result.transformResult(integerLiteralTypeApproximator, expectedTypeRef!!.coneTypeSafe())
         }
         dataFlowAnalyzer.exitJump(jump)
         return result.compose()
diff --git a/compiler/fir/resolve/testData/resolve/arguments/untouchedReturnInIf.kt b/compiler/fir/resolve/testData/resolve/arguments/untouchedReturnInIf.kt
new file mode 100644
index 00000000000..ae7b8c7d106
--- /dev/null
+++ b/compiler/fir/resolve/testData/resolve/arguments/untouchedReturnInIf.kt
@@ -0,0 +1,11 @@
+interface Box<T>
+
+public fun <T> foo(nextFunction: (T) -> T): Box<T> = null!!
+
+fun leaves(value: String, forward: Boolean): Box<String> {
+    if (forward) {
+        return foo { "" }
+    } else {
+        return foo { "" }
+    }
+}
diff --git a/compiler/fir/resolve/testData/resolve/arguments/untouchedReturnInIf.txt b/compiler/fir/resolve/testData/resolve/arguments/untouchedReturnInIf.txt
new file mode 100644
index 00000000000..57ae4330819
--- /dev/null
+++ b/compiler/fir/resolve/testData/resolve/arguments/untouchedReturnInIf.txt
@@ -0,0 +1,23 @@
+FILE: untouchedReturnInIf.kt
+    public abstract interface Box<T> : R|kotlin/Any| {
+    }
+    public final fun <T> foo(nextFunction: R|(T) -> T|): R|Box<T>| {
+        ^foo Null(null)!!
+    }
+    public final fun leaves(value: R|kotlin/String|, forward: R|kotlin/Boolean|): R|Box<kotlin/String>| {
+        when () {
+            R|<local>/forward| ->  {
+                ^leaves R|/foo|<R|kotlin/String|>(<L> = foo@fun <anonymous>(it: R|kotlin/String|): R|kotlin/String| {
+                    ^ String()
+                }
+                )
+            }
+            else ->  {
+                ^leaves R|/foo|<R|kotlin/String|>(<L> = foo@fun <anonymous>(it: R|kotlin/String|): R|kotlin/String| {
+                    ^ String()
+                }
+                )
+            }
+        }
+
+    }
diff --git a/compiler/fir/resolve/testData/resolve/problems/typesInLocalFunctions.kt b/compiler/fir/resolve/testData/resolve/typesInLocalFunctions.kt
similarity index 100%
rename from compiler/fir/resolve/testData/resolve/problems/typesInLocalFunctions.kt
rename to compiler/fir/resolve/testData/resolve/typesInLocalFunctions.kt
diff --git a/compiler/fir/resolve/testData/resolve/problems/typesInLocalFunctions.txt b/compiler/fir/resolve/testData/resolve/typesInLocalFunctions.txt
similarity index 71%
rename from compiler/fir/resolve/testData/resolve/problems/typesInLocalFunctions.txt
rename to compiler/fir/resolve/testData/resolve/typesInLocalFunctions.txt
index 1f324e1353e..52dba0bbc73 100644
--- a/compiler/fir/resolve/testData/resolve/problems/typesInLocalFunctions.txt
+++ b/compiler/fir/resolve/testData/resolve/typesInLocalFunctions.txt
@@ -9,9 +9,9 @@ FILE: typesInLocalFunctions.kt
         lval s: R|Some| = R|/Some.Some|()
         when () {
             Boolean(true) ->  {
-                ^foo fun <implicit>.<anonymous>(): <implicit> {
-                    when () {
-                        (s# is Some) ->  {
+                ^foo fun <anonymous>(): R|kotlin/Boolean| {
+                    ^ when () {
+                        (R|<local>/s| is R|Some|) ->  {
                             Boolean(true)
                         }
                         else ->  {
@@ -23,8 +23,8 @@ FILE: typesInLocalFunctions.kt
 
             }
             else ->  {
-                ^foo fun <implicit>.<anonymous>(): <implicit> {
-                    Boolean(true)
+                ^foo fun <anonymous>(): R|kotlin/Boolean| {
+                    ^ Boolean(true)
                 }
 
             }
diff --git a/compiler/fir/resolve/tests/org/jetbrains/kotlin/fir/FirDiagnosticsTestGenerated.java b/compiler/fir/resolve/tests/org/jetbrains/kotlin/fir/FirDiagnosticsTestGenerated.java
index 3888118758b..d11e8f8b599 100644
--- a/compiler/fir/resolve/tests/org/jetbrains/kotlin/fir/FirDiagnosticsTestGenerated.java
+++ b/compiler/fir/resolve/tests/org/jetbrains/kotlin/fir/FirDiagnosticsTestGenerated.java
@@ -338,6 +338,11 @@ public class FirDiagnosticsTestGenerated extends AbstractFirDiagnosticsTest {
         runTest("compiler/fir/resolve/testData/resolve/typeParameterVsNested.kt");
     }
 
+    @TestMetadata("typesInLocalFunctions.kt")
+    public void testTypesInLocalFunctions() throws Exception {
+        runTest("compiler/fir/resolve/testData/resolve/typesInLocalFunctions.kt");
+    }
+
     @TestMetadata("varargInPrimaryConstructor.kt")
     public void testVarargInPrimaryConstructor() throws Exception {
         runTest("compiler/fir/resolve/testData/resolve/varargInPrimaryConstructor.kt");
@@ -460,6 +465,11 @@ public class FirDiagnosticsTestGenerated extends AbstractFirDiagnosticsTest {
             runTest("compiler/fir/resolve/testData/resolve/arguments/tryInLambda.kt");
         }
 
+        @TestMetadata("untouchedReturnInIf.kt")
+        public void testUntouchedReturnInIf() throws Exception {
+            runTest("compiler/fir/resolve/testData/resolve/arguments/untouchedReturnInIf.kt");
+        }
+
         @TestMetadata("vararg.kt")
         public void testVararg() throws Exception {
             runTest("compiler/fir/resolve/testData/resolve/arguments/vararg.kt");
@@ -1361,11 +1371,6 @@ public class FirDiagnosticsTestGenerated extends AbstractFirDiagnosticsTest {
         public void testJavaAccessorConversion() throws Exception {
             runTest("compiler/fir/resolve/testData/resolve/problems/javaAccessorConversion.kt");
         }
-
-        @TestMetadata("typesInLocalFunctions.kt")
-        public void testTypesInLocalFunctions() throws Exception {
-            runTest("compiler/fir/resolve/testData/resolve/problems/typesInLocalFunctions.kt");
-        }
     }
 
     @TestMetadata("compiler/fir/resolve/testData/resolve/references")
diff --git a/compiler/fir/resolve/tests/org/jetbrains/kotlin/fir/FirDiagnosticsWithLightTreeTestGenerated.java b/compiler/fir/resolve/tests/org/jetbrains/kotlin/fir/FirDiagnosticsWithLightTreeTestGenerated.java
index 7f8f7179479..6126fb37f5a 100644
--- a/compiler/fir/resolve/tests/org/jetbrains/kotlin/fir/FirDiagnosticsWithLightTreeTestGenerated.java
+++ b/compiler/fir/resolve/tests/org/jetbrains/kotlin/fir/FirDiagnosticsWithLightTreeTestGenerated.java
@@ -338,6 +338,11 @@ public class FirDiagnosticsWithLightTreeTestGenerated extends AbstractFirDiagnos
         runTest("compiler/fir/resolve/testData/resolve/typeParameterVsNested.kt");
     }
 
+    @TestMetadata("typesInLocalFunctions.kt")
+    public void testTypesInLocalFunctions() throws Exception {
+        runTest("compiler/fir/resolve/testData/resolve/typesInLocalFunctions.kt");
+    }
+
     @TestMetadata("varargInPrimaryConstructor.kt")
     public void testVarargInPrimaryConstructor() throws Exception {
         runTest("compiler/fir/resolve/testData/resolve/varargInPrimaryConstructor.kt");
@@ -460,6 +465,11 @@ public class FirDiagnosticsWithLightTreeTestGenerated extends AbstractFirDiagnos
             runTest("compiler/fir/resolve/testData/resolve/arguments/tryInLambda.kt");
         }
 
+        @TestMetadata("untouchedReturnInIf.kt")
+        public void testUntouchedReturnInIf() throws Exception {
+            runTest("compiler/fir/resolve/testData/resolve/arguments/untouchedReturnInIf.kt");
+        }
+
         @TestMetadata("vararg.kt")
         public void testVararg() throws Exception {
             runTest("compiler/fir/resolve/testData/resolve/arguments/vararg.kt");
@@ -1361,11 +1371,6 @@ public class FirDiagnosticsWithLightTreeTestGenerated extends AbstractFirDiagnos
         public void testJavaAccessorConversion() throws Exception {
             runTest("compiler/fir/resolve/testData/resolve/problems/javaAccessorConversion.kt");
         }
-
-        @TestMetadata("typesInLocalFunctions.kt")
-        public void testTypesInLocalFunctions() throws Exception {
-            runTest("compiler/fir/resolve/testData/resolve/problems/typesInLocalFunctions.kt");
-        }
     }
 
     @TestMetadata("compiler/fir/resolve/testData/resolve/references")
diff --git a/compiler/testData/diagnostics/tests/callableReference/lambdaResult.fir.kt b/compiler/testData/diagnostics/tests/callableReference/lambdaResult.fir.kt
index 92f1ec9ba32..00b7794fa25 100644
--- a/compiler/testData/diagnostics/tests/callableReference/lambdaResult.fir.kt
+++ b/compiler/testData/diagnostics/tests/callableReference/lambdaResult.fir.kt
@@ -36,7 +36,7 @@ fun main(x: Inv<Int>) {
     }
 
     foo1 {
-        if (x.hashCode() == 0) return@foo1 <!UNRESOLVED_REFERENCE!>::bar<!>
+        if (x.hashCode() == 0) return@foo1 ::bar
 
         ::bar
     }