Cherry-pick "Fix asBodyAndResultVar call in visitInlinedLambdaInComposableScope"

plugins/compose/compiler-hosted/integration-tests/testResources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testNonLocalReturnFromWhen.txt

@@ -0,0 +1,105 @@
+//
+// Source
+// ------------------------------------------
+
+import androidx.compose.runtime.Composable
+
+val failed = FakeResult(RuntimeException("error"))
+
+@Composable
+fun Test() {
+    Wrapper {
+        val value: String = failed.fold(
+            onSuccess = { "Ok" },
+            onFailure = {
+                Wrapper {}
+                return@Wrapper
+            }
+        )
+    }
+}
+
+//
+// Transformed IR
+// ------------------------------------------
+
+val failed: FakeResult = FakeResult(RuntimeException("error"))
+@Composable
+@FunctionKeyMeta(key = -1794342280, startOffset = 157, endOffset = 494)
+fun Test(%composer: Composer?, %changed: Int) {
+  %composer = %composer.startRestartGroup(<>)
+  sourceInformation(%composer, "C(Test)<Wrappe...>:Test.kt")
+  if (%composer.shouldExecute(%changed != 0, %changed and 0b0001)) {
+    if (isTraceInProgress()) {
+      traceEventStart(<>, %changed, -1, <>)
+    }
+    Wrapper(ComposableSingletons%TestKt.lambda%682646544, %composer, 0b0110)
+    if (isTraceInProgress()) {
+      traceEventEnd()
+    }
+  } else {
+    %composer.skipToGroupEnd()
+  }
+  %composer.endRestartGroup()?.updateScope { %composer: Composer?, %force: Int ->
+    Test(%composer, updateChangedFlags(%changed or 0b0001))
+  }
+}
+internal object ComposableSingletons%TestKt {
+  val lambda%-222111391: Function2<Composer, Int, Unit>?
+    get() {
+      if (field == null) {
+        field = composableLambdaInstance(<>, false) { %composer: Composer?, %changed: Int ->
+          sourceInformation(%composer, "C:Test.kt")
+          if (%composer.shouldExecute(%changed and 0b0011 != 0b0010, %changed and 0b0001)) {
+            if (isTraceInProgress()) {
+              traceEventStart(<>, %changed, -1, <>)
+            }
+            Unit
+            if (isTraceInProgress()) {
+              traceEventEnd()
+            }
+          } else {
+            %composer.skipToGroupEnd()
+          }
+        }
+      }
+      return field
+    }
+  val lambda%682646544: Function2<Composer, Int, Unit>?
+    get() {
+      if (field == null) {
+        field = composableLambdaInstance(<>, false) { %composer: Composer?, %changed: Int ->
+          sourceInformation(%composer, "C:Test.kt")
+          if (%composer.shouldExecute(%changed and 0b0011 != 0b0010, %changed and 0b0001)) {
+            if (isTraceInProgress()) {
+              traceEventStart(<>, %changed, -1, <>)
+            }
+            val value = <block>{
+              val tmp0_group = failed.fold({
+                val tmp1_return = "Ok"
+                tmp1_return
+              }
+              ) { it: Exception ->
+                %composer.startReplaceGroup(<>)
+                sourceInformation(%composer, "invalid source info at 1: 'CN(it)*11@361L10:Test.kt'")
+                Wrapper(ComposableSingletons%TestKt.lambda%-222111391, %composer, 0b0110)
+                %composer.endReplaceGroup()
+                if (isTraceInProgress()) {
+                  traceEventEnd()
+                }
+                return@composableLambdaInstance
+                %composer.endReplaceGroup()
+              }
+              tmp0_group
+            }
+            if (isTraceInProgress()) {
+              traceEventEnd()
+            }
+          } else {
+            %composer.skipToGroupEnd()
+          }
+        }
+      }
+      return field
+    }
+}

plugins/compose/compiler-hosted/integration-tests/tests/androidx/compose/compiler/plugins/kotlin/ControlFlowTransformTests.kt

@@ -2771,4 +2771,51 @@ class ControlFlowTransformTests : AbstractControlFlowTransformTests() {
             }
         """
     )
+
+    /**
+     * This is a regression test against a bug that, in certain cases, prevented execution of
+     * non-local `return` statements inside lambdas called from `when` expressions.
+     * For more details, see https://issuetracker.google.com/issues/549552317.
+     */
+    @Test
+    fun testNonLocalReturnFromWhen() = verifyGoldenComposeIrTransform(
+        source = """
+            import androidx.compose.runtime.Composable
+
+            val failed = FakeResult(RuntimeException("error"))
+
+            @Composable
+            fun Test() {
+                Wrapper {
+                    val value: String = failed.fold(
+                        onSuccess = { "Ok" },
+                        onFailure = {
+                            Wrapper {}
+                            return@Wrapper
+                        }
+                    )
+                }
+            }
+        """,
+        extra = """
+            import androidx.compose.runtime.Composable
+
+            @Composable
+            fun Wrapper(content: @Composable () -> Unit) {
+                content()
+            }
+
+            class FakeResult(val value: Exception?) {
+                inline fun fold(
+                    onSuccess: () -> String,
+                    onFailure: (exception: Exception) -> String,
+                ): String {
+                    return when (value) {
+                        null -> onSuccess()
+                        else -> onFailure(value)
+                    }
+                }
+            }
+        """
+    )
 }

plugins/compose/compiler-hosted/runtime-tests/src/commonTest/kotlin/androidx/compose/compiler/test/CompositionTests.kt

@@ -568,6 +568,39 @@ class CompositionTests {
             movableContent.content
         }
     }
+
+    /**
+     * This is a regression test against a bug that, in certain cases, prevented execution of
+     * non-local `return` statements inside lambdas called from `when` expressions.
+     * For more details, see https://issuetracker.google.com/issues/549552317.
+     */
+    @Test
+    @OptIn(InternalComposeApi::class)
+    fun testNonLocalReturnFromWhen() = compositionTest {
+        val failed = FakeResult(RuntimeException("error"))
+        var reachedCodeAfterGuard = false
+        var leaked: String? = "guard-not-reached"
+
+        compose {
+            Wrapper {
+                val value: String = failed.fold(
+                    onSuccess = { "Ok" },
+                    onFailure = {
+                        stringResource()
+                        return@Wrapper
+                    }
+                )
+                leaked = value
+                reachedCodeAfterGuard = true
+            }
+
+            if (reachedCodeAfterGuard) {
+                error(
+                    "return@Wrapper had no effect and the String-typed `val value` held: $leaked"
+                )
+            }
+        }
+    }
 }
 
 @Composable
@@ -759,3 +792,20 @@ internal fun ConsumeChildState(state: ChildState, result: MutableState<Boolean>)
         }
     }
 }
+
+@Composable
+fun Wrapper(content: @Composable () -> Unit) {
+    content()
+}
+
+class FakeResult(val value: Exception?) {
+    inline fun fold(
+        onSuccess: () -> String,
+        onFailure: (exception: Exception) -> String,
+    ): String {
+        return when (value) {
+            null -> onSuccess()
+            else -> onFailure(value)
+        }
+    }
+}

plugins/compose/compiler-hosted/src/main/java/androidx/compose/compiler/plugins/kotlin/lower/ComposableFunctionBodyTransformer.kt

@@ -1272,7 +1272,7 @@ class ComposableFunctionBodyTransformer(
         }
 
         val originalBody = declaration.body ?: return super.visitFunction(declaration)
-        val [body, returnVar] = originalBody.asBodyAndResultVar()
+        val [body, returnVar] = originalBody.asBodyAndResultVar(expectedTarget = declaration)
         body.transformChildrenVoid()
 
         // Avoid transforming functions that are not referencing anything composable, as they cannot use slots (read-only is fine).