Re-enable early exit detection

Re-enable the early exit code generation for inline functions
not no does not add early exit for inlined composable lambdas
inlined into composable non-inline functions.

Fixes: 255722247
Test: ./gradlew :compose:compiler:c-h:i-t:tDUT :compose:r:r:tDUT

This reverts commit 2d2693582b04069b68149616b1c3ef539543ac0b.

Change-Id: Ic37b9d02cc424d920e6cab15077b7404bc66baa1
diff --git a/compose/compiler/compiler-hosted/integration-tests/src/test/java/androidx/compose/compiler/plugins/kotlin/ControlFlowTransformTests.kt b/compose/compiler/compiler-hosted/integration-tests/src/test/java/androidx/compose/compiler/plugins/kotlin/ControlFlowTransformTests.kt
index db43f57..6843730 100644
--- a/compose/compiler/compiler-hosted/integration-tests/src/test/java/androidx/compose/compiler/plugins/kotlin/ControlFlowTransformTests.kt
+++ b/compose/compiler/compiler-hosted/integration-tests/src/test/java/androidx/compose/compiler/plugins/kotlin/ControlFlowTransformTests.kt
@@ -262,7 +262,12 @@
                     }
                 }
             }
-        """.trimIndent()
+
+            inline fun <T> Identity(block: () -> T): T = block()
+
+            @Composable
+            fun Stack(content: @Composable () -> Unit) = content()
+        """
     )
 
     @Test
@@ -287,6 +292,7 @@
             fun Test(condition: Boolean, %composer: Composer?, %changed: Int) {
               %composer = %composer.startRestartGroup(<>)
               sourceInformation(%composer, "C(Test)<A()>,<M3>,<A()>:Test.kt")
+              val tmp0_marker = %composer.currentMarker
               val %dirty = %changed
               if (%changed and 0b1110 === 0) {
                 %dirty = %dirty or if (%composer.changed(condition)) 0b0100 else 0b0010
@@ -302,6 +308,7 @@
                   if (%changed and 0b1011 !== 0b0010 || !%composer.skipping) {
                     A(%composer, 0)
                     if (condition) {
+                      %composer.endToMarker(tmp0_marker)
                       if (isTraceInProgress()) {
                         traceEventEnd()
                       }
@@ -359,6 +366,7 @@
             fun Test(a: Boolean, b: Boolean, %composer: Composer?, %changed: Int) {
               %composer = %composer.startRestartGroup(<>)
               sourceInformation(%composer, "C(Test)<A()>,<M3>,<M3>,<A()>:Test.kt")
+              val tmp0_marker = %composer.currentMarker
               val %dirty = %changed
               if (%changed and 0b1110 === 0) {
                 %dirty = %dirty or if (%composer.changed(a)) 0b0100 else 0b0010
@@ -377,6 +385,7 @@
                   if (%changed and 0b1011 !== 0b0010 || !%composer.skipping) {
                     A(%composer, 0)
                     if (a) {
+                      %composer.endToMarker(tmp0_marker)
                       if (isTraceInProgress()) {
                         traceEventEnd()
                       }
@@ -397,6 +406,7 @@
                   if (%changed and 0b1011 !== 0b0010 || !%composer.skipping) {
                     A(%composer, 0)
                     if (b) {
+                      %composer.endToMarker(tmp0_marker)
                       if (isTraceInProgress()) {
                         traceEventEnd()
                       }
@@ -456,12 +466,14 @@
                   traceEventStart(<>, %changed, -1, <>)
                 }
                 A(%composer, 0)
+                val tmp0_marker = %composer.currentMarker
                 M3({ %composer: Composer?, %changed: Int ->
                   %composer.startReplaceableGroup(<>)
                   sourceInformation(%composer, "C<A()>,<A()>:Test.kt")
                   if (%changed and 0b1011 !== 0b0010 || !%composer.skipping) {
                     A(%composer, 0)
                     if (condition) {
+                      %composer.endToMarker(tmp0_marker)
                     }
                     A(%composer, 0)
                   } else {
@@ -501,6 +513,7 @@
               T {
                 %this%T.compose(composableLambdaInstance(<>, true) { %composer: Composer?, %changed: Int ->
                   sourceInformation(%composer, "C<M1>:Test.kt")
+                  val tmp0_marker = %composer.currentMarker
                   if (%changed and 0b1011 !== 0b0010 || !%composer.skipping) {
                     if (isTraceInProgress()) {
                       traceEventStart(<>, %changed, -1, <>)
@@ -510,6 +523,7 @@
                       sourceInformation(%composer, "C:Test.kt")
                       if (%changed and 0b1011 !== 0b0010 || !%composer.skipping) {
                         if (condition) {
+                          %composer.endToMarker(tmp0_marker)
                           if (isTraceInProgress()) {
                             traceEventEnd()
                           }
@@ -584,11 +598,13 @@
                   sourceInformation(%composer, "C<A()>,<M1>,<A()>:Test.kt")
                   if (%changed and 0b1011 !== 0b0010 || !%composer.skipping) {
                     A(%composer, 0)
+                    val tmp0_marker = %composer.currentMarker
                     M1({ %composer: Composer?, %changed: Int ->
                       %composer.startReplaceableGroup(<>)
                       sourceInformation(%composer, "C:Test.kt")
                       if (%changed and 0b1011 !== 0b0010 || !%composer.skipping) {
                         if (condition) {
+                          %composer.endToMarker(tmp0_marker)
                         }
                       } else {
                         %composer.skipToGroupEnd()
@@ -647,6 +663,7 @@
                   traceEventStart(<>, %changed, -1, <>)
                 }
                 A(%composer, 0)
+                val tmp0_marker = %composer.currentMarker
                 M3({ %composer: Composer?, %changed: Int ->
                   %composer.startReplaceableGroup(<>)
                   sourceInformation(%composer, "C<A()>,<M1>,<A()>:Test.kt")
@@ -657,6 +674,7 @@
                       sourceInformation(%composer, "C:Test.kt")
                       if (%changed and 0b1011 !== 0b0010 || !%composer.skipping) {
                         if (condition) {
+                          %composer.endToMarker(tmp0_marker)
                         }
                       } else {
                         %composer.skipToGroupEnd()
@@ -709,6 +727,7 @@
             fun testInline_M1_W_Return_Func(condition: Boolean, %composer: Composer?, %changed: Int) {
               %composer = %composer.startRestartGroup(<>)
               sourceInformation(%composer, "C(testInline_M1_W_Return_Func)<A()>,<M1>,<A()>:Test.kt")
+              val tmp0_marker = %composer.currentMarker
               val %dirty = %changed
               if (%changed and 0b1110 === 0) {
                 %dirty = %dirty or if (%composer.changed(condition)) 0b0100 else 0b0010
@@ -728,8 +747,7 @@
                     while (true) {
                       A(%composer, 0)
                       if (condition) {
-                        %composer.endReplaceableGroup()
-                        %composer.endReplaceableGroup()
+                        %composer.endToMarker(tmp0_marker)
                         if (isTraceInProgress()) {
                           traceEventEnd()
                         }
@@ -799,12 +817,14 @@
                   traceEventStart(<>, %changed, -1, <>)
                 }
                 A(%composer, 0)
+                val tmp0_marker = %composer.currentMarker
                 M3({ %composer: Composer?, %changed: Int ->
                   %composer.startReplaceableGroup(<>)
                   sourceInformation(%composer, "C<A()>,<A()>:Test.kt")
                   if (%changed and 0b1011 !== 0b0010 || !%composer.skipping) {
                     A(%composer, 0)
                     if (condition) {
+                      %composer.endToMarker(tmp0_marker)
                     }
                     A(%composer, 0)
                   } else {
@@ -812,12 +832,14 @@
                   }
                   %composer.endReplaceableGroup()
                 }, %composer, 0)
+                val tmp1_marker = %composer.currentMarker
                 M3({ %composer: Composer?, %changed: Int ->
                   %composer.startReplaceableGroup(<>)
                   sourceInformation(%composer, "C<A()>,<A()>:Test.kt")
                   if (%changed and 0b1011 !== 0b0010 || !%composer.skipping) {
                     A(%composer, 0)
                     if (condition) {
+                      %composer.endToMarker(tmp1_marker)
                     }
                     A(%composer, 0)
                   } else {
@@ -865,6 +887,7 @@
             fun test_CM1_CCM1_RetFun(condition: Boolean, %composer: Composer?, %changed: Int) {
               %composer = %composer.startRestartGroup(<>)
               sourceInformation(%composer, "C(test_CM1_CCM1_RetFun)<Text("...>,<M1>,<Text("...>:Test.kt")
+              val tmp0_marker = %composer.currentMarker
               val %dirty = %changed
               if (%changed and 0b1110 === 0) {
                 %dirty = %dirty or if (%composer.changed(condition)) 0b0100 else 0b0010
@@ -888,8 +911,7 @@
                         sourceInformation(%composer, "C<Text("...>:Test.kt")
                         if (%changed and 0b1011 !== 0b0010 || !%composer.skipping) {
                           Text("In CCM1", %composer, 0b0110)
-                          %composer.endReplaceableGroup()
-                          %composer.endReplaceableGroup()
+                          %composer.endToMarker(tmp0_marker)
                           if (isTraceInProgress()) {
                             traceEventEnd()
                           }
@@ -951,11 +973,13 @@
               if (isTraceInProgress()) {
                 traceEventStart(<>, %changed, -1, <>)
               }
+              val tmp0_marker = %composer.currentMarker
               FakeBox({ %composer: Composer?, %changed: Int ->
                 %composer.startReplaceableGroup(<>)
                 sourceInformation(%composer, "C<A()>:Test.kt")
                 if (%changed and 0b1011 !== 0b0010 || !%composer.skipping) {
                   if (condition) {
+                    %composer.endToMarker(tmp0_marker)
                   }
                   A(%composer, 0)
                 } else {
@@ -1127,11 +1151,13 @@
                 if (isTraceInProgress()) {
                   traceEventStart(<>, %changed, -1, <>)
                 }
+                val tmp0_marker = %composer.currentMarker
                 IW({ %composer: Composer?, %changed: Int ->
                   %composer.startReplaceableGroup(<>)
                   sourceInformation(%composer, "C<A()>:Test.kt")
                   if (%changed and 0b1011 !== 0b0010 || !%composer.skipping) {
                     if (condition) {
+                      %composer.endToMarker(tmp0_marker)
                     }
                     A(%composer, 0)
                   } else {
@@ -1153,6 +1179,178 @@
     )
 
     @Test
+    fun testVerifyEarlyExitFromNonComposable() = verifyInlineReturn(
+        source = """
+            @Composable
+            fun Test(condition: Boolean) {
+                Text("Some text")
+                Identity {
+                    if (condition) return@Test
+                }
+                Text("Some more text")
+            }
+        """,
+        expectedTransformed = """
+            @Composable
+            fun Test(condition: Boolean, %composer: Composer?, %changed: Int) {
+              %composer = %composer.startRestartGroup(<>)
+              sourceInformation(%composer, "C(Test)<Text("...>,<Text("...>:Test.kt")
+              val %dirty = %changed
+              if (%changed and 0b1110 === 0) {
+                %dirty = %dirty or if (%composer.changed(condition)) 0b0100 else 0b0010
+              }
+              if (%dirty and 0b1011 !== 0b0010 || !%composer.skipping) {
+                if (isTraceInProgress()) {
+                  traceEventStart(<>, %changed, -1, <>)
+                }
+                Text("Some text", %composer, 0b0110)
+                Identity {
+                  if (condition) {
+                    if (isTraceInProgress()) {
+                      traceEventEnd()
+                    }
+                    %composer.endRestartGroup()?.updateScope { %composer: Composer?, %force: Int ->
+                      Test(condition, %composer, updateChangedFlags(%changed or 0b0001))
+                    }
+                    return
+                  }
+                }
+                Text("Some more text", %composer, 0b0110)
+                if (isTraceInProgress()) {
+                  traceEventEnd()
+                }
+              } else {
+                %composer.skipToGroupEnd()
+              }
+              %composer.endRestartGroup()?.updateScope { %composer: Composer?, %force: Int ->
+                Test(condition, %composer, updateChangedFlags(%changed or 0b0001))
+              }
+            }
+        """
+    )
+
+    @Test
+    fun testVerifyEarlyExitFromNonComposable_M1() = verifyInlineReturn(
+        source = """
+            @Composable
+            fun Test(condition: Boolean) {
+                Text("Some text")
+                M1 {
+                    Identity {
+                        if (condition) return@Test
+                    }
+                }
+                Text("Some more text")
+            }
+        """,
+        expectedTransformed = """
+            @Composable
+            fun Test(condition: Boolean, %composer: Composer?, %changed: Int) {
+              %composer = %composer.startRestartGroup(<>)
+              sourceInformation(%composer, "C(Test)<Text("...>,<M1>,<Text("...>:Test.kt")
+              val tmp0_marker = %composer.currentMarker
+              val %dirty = %changed
+              if (%changed and 0b1110 === 0) {
+                %dirty = %dirty or if (%composer.changed(condition)) 0b0100 else 0b0010
+              }
+              if (%dirty and 0b1011 !== 0b0010 || !%composer.skipping) {
+                if (isTraceInProgress()) {
+                  traceEventStart(<>, %changed, -1, <>)
+                }
+                Text("Some text", %composer, 0b0110)
+                M1({ %composer: Composer?, %changed: Int ->
+                  %composer.startReplaceableGroup(<>)
+                  sourceInformation(%composer, "C:Test.kt")
+                  if (%changed and 0b1011 !== 0b0010 || !%composer.skipping) {
+                    Identity {
+                      if (condition) {
+                        %composer.endToMarker(tmp0_marker)
+                        if (isTraceInProgress()) {
+                          traceEventEnd()
+                        }
+                        %composer.endRestartGroup()?.updateScope { %composer: Composer?, %force: Int ->
+                          Test(condition, %composer, updateChangedFlags(%changed or 0b0001))
+                        }
+                        return
+                      }
+                    }
+                  } else {
+                    %composer.skipToGroupEnd()
+                  }
+                  %composer.endReplaceableGroup()
+                }, %composer, 0)
+                Text("Some more text", %composer, 0b0110)
+                if (isTraceInProgress()) {
+                  traceEventEnd()
+                }
+              } else {
+                %composer.skipToGroupEnd()
+              }
+              %composer.endRestartGroup()?.updateScope { %composer: Composer?, %force: Int ->
+                Test(condition, %composer, updateChangedFlags(%changed or 0b0001))
+              }
+            }
+        """
+    )
+
+    @Test
+    fun testVerifyEarlyExitFromNonComposable_M1_RM1() = verifyInlineReturn(
+        source = """
+            @Composable
+            fun Test(condition: Boolean) {
+                Text("Some text")
+                M1 {
+                    Identity {
+                        if (condition) return@M1
+                    }
+                }
+                Text("Some more text")
+            }
+        """,
+        expectedTransformed = """
+            @Composable
+            fun Test(condition: Boolean, %composer: Composer?, %changed: Int) {
+              %composer = %composer.startRestartGroup(<>)
+              sourceInformation(%composer, "C(Test)<Text("...>,<M1>,<Text("...>:Test.kt")
+              val %dirty = %changed
+              if (%changed and 0b1110 === 0) {
+                %dirty = %dirty or if (%composer.changed(condition)) 0b0100 else 0b0010
+              }
+              if (%dirty and 0b1011 !== 0b0010 || !%composer.skipping) {
+                if (isTraceInProgress()) {
+                  traceEventStart(<>, %changed, -1, <>)
+                }
+                Text("Some text", %composer, 0b0110)
+                val tmp0_marker = %composer.currentMarker
+                M1({ %composer: Composer?, %changed: Int ->
+                  %composer.startReplaceableGroup(<>)
+                  sourceInformation(%composer, "C:Test.kt")
+                  if (%changed and 0b1011 !== 0b0010 || !%composer.skipping) {
+                    Identity {
+                      if (condition) {
+                        %composer.endToMarker(tmp0_marker)
+                      }
+                    }
+                  } else {
+                    %composer.skipToGroupEnd()
+                  }
+                  %composer.endReplaceableGroup()
+                }, %composer, 0)
+                Text("Some more text", %composer, 0b0110)
+                if (isTraceInProgress()) {
+                  traceEventEnd()
+                }
+              } else {
+                %composer.skipToGroupEnd()
+              }
+              %composer.endRestartGroup()?.updateScope { %composer: Composer?, %force: Int ->
+                Test(condition, %composer, updateChangedFlags(%changed or 0b0001))
+              }
+            }
+        """
+    )
+
+    @Test
     fun testEnsureRuntimeTestWillCompile_CL() = ensureSetup {
         classLoader(
             """
@@ -1203,6 +1401,7 @@
             fun test_CM1_RetFun(condition: Boolean, %composer: Composer?, %changed: Int) {
               %composer = %composer.startRestartGroup(<>)
               sourceInformation(%composer, "C(test_CM1_RetFun)<Text("...>,<M1>,<Text("...>:Test.kt")
+              val tmp0_marker = %composer.currentMarker
               val %dirty = %changed
               if (%changed and 0b1110 === 0) {
                 %dirty = %dirty or if (%composer.changed(condition)) 0b0100 else 0b0010
@@ -1218,6 +1417,7 @@
                   if (%changed and 0b1011 !== 0b0010 || !%composer.skipping) {
                     Text("M1 - before", %composer, 0b0110)
                     if (condition) {
+                      %composer.endToMarker(tmp0_marker)
                       if (isTraceInProgress()) {
                         traceEventEnd()
                       }
@@ -5828,4 +6028,4 @@
             }
         """
     )
-}
+}
\ No newline at end of file
diff --git a/compose/compiler/compiler-hosted/integration-tests/src/test/java/androidx/compose/compiler/plugins/kotlin/TraceInformationTest.kt b/compose/compiler/compiler-hosted/integration-tests/src/test/java/androidx/compose/compiler/plugins/kotlin/TraceInformationTest.kt
index e8ffd3f..d085612 100644
--- a/compose/compiler/compiler-hosted/integration-tests/src/test/java/androidx/compose/compiler/plugins/kotlin/TraceInformationTest.kt
+++ b/compose/compiler/compiler-hosted/integration-tests/src/test/java/androidx/compose/compiler/plugins/kotlin/TraceInformationTest.kt
@@ -118,6 +118,7 @@
             fun Test(condition: Boolean, %composer: Composer?, %changed: Int) {
               %composer = %composer.startRestartGroup(<>)
               sourceInformation(%composer, "C(Test)<A()>,<Wrappe...>,<A()>:Test.kt")
+              val tmp0_marker = %composer.currentMarker
               val %dirty = %changed
               if (%changed and 0b1110 === 0) {
                 %dirty = %dirty or if (%composer.changed(condition)) 0b0100 else 0b0010
@@ -133,6 +134,7 @@
                   if (%changed and 0b1011 !== 0b0010 || !%composer.skipping) {
                     A(%composer, 0)
                     if (!condition) {
+                      %composer.endToMarker(tmp0_marker)
                       if (isTraceInProgress()) {
                         traceEventEnd()
                       }
diff --git a/compose/compiler/compiler-hosted/src/main/java/androidx/compose/compiler/plugins/kotlin/lower/ComposableFunctionBodyTransformer.kt b/compose/compiler/compiler-hosted/src/main/java/androidx/compose/compiler/plugins/kotlin/lower/ComposableFunctionBodyTransformer.kt
index 79b1f93c..6e9b83d 100644
--- a/compose/compiler/compiler-hosted/src/main/java/androidx/compose/compiler/plugins/kotlin/lower/ComposableFunctionBodyTransformer.kt
+++ b/compose/compiler/compiler-hosted/src/main/java/androidx/compose/compiler/plugins/kotlin/lower/ComposableFunctionBodyTransformer.kt
@@ -616,9 +616,8 @@
             }
     }
 
-    private val rollbackGroupMarkerEnabled get() = false
-        // Temporarily disabled for b/255722247
-        // currentMarkerProperty != null && endToMarkerFunction != null
+    private val rollbackGroupMarkerEnabled get() =
+        currentMarkerProperty != null && endToMarkerFunction != null
 
     private val endRestartGroupFunction by guardedLazy {
         composerIrClass
@@ -2633,7 +2632,8 @@
 
                         break@loop
                     }
-                    if (scope.isInlinedLambda) leavingInlinedLambda = true
+                    if (scope.isInlinedLambda && scope.inComposableCall)
+                        leavingInlinedLambda = true
                 }
                 is Scope.BlockScope -> {
                     blockScopeMarks.add(scope)
diff --git a/compose/runtime/runtime/src/commonTest/kotlin/androidx/compose/runtime/CompositionTests.kt b/compose/runtime/runtime/src/commonTest/kotlin/androidx/compose/runtime/CompositionTests.kt
index 5c32e29..285838a 100644
--- a/compose/runtime/runtime/src/commonTest/kotlin/androidx/compose/runtime/CompositionTests.kt
+++ b/compose/runtime/runtime/src/commonTest/kotlin/androidx/compose/runtime/CompositionTests.kt
@@ -54,7 +54,6 @@
 import kotlinx.coroutines.coroutineScope
 import kotlinx.coroutines.launch
 import kotlinx.coroutines.test.runTest
-import org.junit.Ignore
 
 @Composable
 fun Container(content: @Composable () -> Unit) = content()
@@ -3235,7 +3234,6 @@
     }
 
     @Test
-    @Ignore("b/255722247")
     fun testNonLocalReturn_CM1_RetFunc_FalseTrue() = compositionTest {
         var condition by mutableStateOf(false)
 
@@ -3255,7 +3253,6 @@
     }
 
     @Test
-    @Ignore("b/255722247")
     fun testNonLocalReturn_CM1_RetFunc_TrueFalse() = compositionTest {
         var condition by mutableStateOf(true)
 
@@ -3275,7 +3272,6 @@
     }
 
     @Test
-    @Ignore("b/255722247")
     fun test_CM1_CCM1_RetFun_FalseTrue() = compositionTest {
         var condition by mutableStateOf(false)
 
@@ -3295,7 +3291,6 @@
     }
 
     @Test
-    @Ignore("b/255722247")
     fun test_CM1_CCM1_RetFun_TrueFalse() = compositionTest {
         var condition by mutableStateOf(true)