diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/kotlin/androidx/compose/compiler/plugins/kotlin/AbstractIrTransformTest.kt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/kotlin/androidx/compose/compiler/plugins/kotlin/AbstractIrTransformTest.kt
index 967b560cf1fae..629501fc0b552 100644
--- a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/kotlin/androidx/compose/compiler/plugins/kotlin/AbstractIrTransformTest.kt
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/kotlin/androidx/compose/compiler/plugins/kotlin/AbstractIrTransformTest.kt
@@ -33,9 +33,10 @@ abstract class AbstractIrTransformTest(useFir: Boolean) : AbstractCodegenTest(us
     override fun CompilerConfiguration.updateConfiguration() {
         put(ComposeConfiguration.SOURCE_INFORMATION_ENABLED_KEY, true)
         put(ComposeConfiguration.FEATURE_FLAGS, listOf(
-            FeatureFlag.StrongSkipping.featureName,
-            FeatureFlag.OptimizeNonSkippingGroups.featureName,
-        ))
+                FeatureFlag.StrongSkipping.featureName,
+                FeatureFlag.OptimizeNonSkippingGroups.featureName,
+            )
+        )
     }
 
     @JvmField
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/kotlin/androidx/compose/compiler/plugins/kotlin/FunctionBodySkippingTransformTests.kt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/kotlin/androidx/compose/compiler/plugins/kotlin/FunctionBodySkippingTransformTests.kt
index dd7865cd7a0dd..d0313835e8534 100644
--- a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/kotlin/androidx/compose/compiler/plugins/kotlin/FunctionBodySkippingTransformTests.kt
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/kotlin/androidx/compose/compiler/plugins/kotlin/FunctionBodySkippingTransformTests.kt
@@ -1306,6 +1306,7 @@ class FunctionBodySkippingTransformTests(
                 fun ReceiveValue(value: Int) { }
             """
         )
+
 }
 
 class FunctionBodySkippingTransformTestsNoSource(
@@ -1314,6 +1315,13 @@ class FunctionBodySkippingTransformTestsNoSource(
     override fun CompilerConfiguration.updateConfiguration() {
         put(ComposeConfiguration.SOURCE_INFORMATION_ENABLED_KEY, false)
         put(ComposeConfiguration.TRACE_MARKERS_ENABLED_KEY, false)
+        put(
+            ComposeConfiguration.FEATURE_FLAGS,
+            listOf(
+                FeatureFlag.StrongSkipping.featureName,
+                FeatureFlag.OptimizeNonSkippingGroups.featureName,
+            )
+        )
     }
 
     @Test
@@ -1457,4 +1465,32 @@ class FunctionBodySkippingTransformTestsNoSource(
             fun Text(value: String) {}
         """
     )
+
+    @Test
+    fun testIfStatementGroups() = verifyGoldenComposeIrTransform(
+        source = """
+            import androidx.compose.runtime.*
+
+            @Composable
+            fun Test(level: Int) {
+                Wrap {
+                    if (level > 0) {
+                        used(remember { "Before" })
+                        Wrap {
+                            used(remember { "Middle" })
+                        }
+                        used(remember { "End" })
+                    }
+                }
+            }
+        """,
+        """
+            import androidx.compose.runtime.*
+
+            @Composable
+            fun Wrap(content: @Composable () -> Unit) = content()
+
+            fun used(value: Any) { }
+        """
+    )
 }
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/ifInsideInlineComposableFunction[useFir = false].txt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/ifInsideInlineComposableFunction[useFir = false].txt
index 0c1d946262673..2465e8c777f40 100644
--- a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/ifInsideInlineComposableFunction[useFir = false].txt	
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/ifInsideInlineComposableFunction[useFir = false].txt	
@@ -35,12 +35,15 @@ fun Label(test: Boolean, %composer: Composer?, %changed: Int) {
     Layout({ %composer: Composer?, %changed: Int ->
       sourceInformationMarkerStart(%composer, <>, "C<Box()>:Test.kt")
       Box(%composer, 0)
-      %composer.startReplaceGroup(<>)
-      sourceInformation(%composer, "<Box()>")
       if (test) {
+        %composer.startReplaceGroup(<>)
+        sourceInformation(%composer, "<Box()>")
         Box(%composer, 0)
+        %composer.endReplaceGroup()
+      } else {
+        %composer.startReplaceGroup(<>)
+        %composer.endReplaceGroup()
       }
-      %composer.endReplaceGroup()
       sourceInformationMarkerEnd(%composer)
     }, %composer, 0)
     if (isTraceInProgress()) {
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/ifInsideInlineComposableFunction[useFir = true].txt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/ifInsideInlineComposableFunction[useFir = true].txt
index 0c1d946262673..2465e8c777f40 100644
--- a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/ifInsideInlineComposableFunction[useFir = true].txt	
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/ifInsideInlineComposableFunction[useFir = true].txt	
@@ -35,12 +35,15 @@ fun Label(test: Boolean, %composer: Composer?, %changed: Int) {
     Layout({ %composer: Composer?, %changed: Int ->
       sourceInformationMarkerStart(%composer, <>, "C<Box()>:Test.kt")
       Box(%composer, 0)
-      %composer.startReplaceGroup(<>)
-      sourceInformation(%composer, "<Box()>")
       if (test) {
+        %composer.startReplaceGroup(<>)
+        sourceInformation(%composer, "<Box()>")
         Box(%composer, 0)
+        %composer.endReplaceGroup()
+      } else {
+        %composer.startReplaceGroup(<>)
+        %composer.endReplaceGroup()
       }
-      %composer.endReplaceGroup()
       sourceInformationMarkerEnd(%composer)
     }, %composer, 0)
     if (isTraceInProgress()) {
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testAND[useFir = false].txt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testAND[useFir = false].txt
index 87b3736de1099..9040a49c5d7b8 100644
--- a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testAND[useFir = false].txt	
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testAND[useFir = false].txt	
@@ -26,8 +26,12 @@ fun Example(%composer: Composer?, %changed: Int) {
     traceEventStart(<>, %changed, -1, <>)
   }
   %composer.startReplaceGroup(<>)
-  sourceInformation(%composer, "<B()>,<B()>")
-  val tmp0_group = B(%composer, 0) && B(%composer, 0)
+  sourceInformation(%composer, "<B()>")
+  val tmp1_group = B(%composer, 0)
+  %composer.endReplaceGroup()
+  tmp1_group && %composer.startReplaceGroup(<>)
+  sourceInformation(%composer, "<B()>")
+  val tmp0_group = B(%composer, 0)
   %composer.endReplaceGroup()
   tmp0_group
   if (isTraceInProgress()) {
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testAND[useFir = true].txt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testAND[useFir = true].txt
index 87b3736de1099..9040a49c5d7b8 100644
--- a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testAND[useFir = true].txt	
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testAND[useFir = true].txt	
@@ -26,8 +26,12 @@ fun Example(%composer: Composer?, %changed: Int) {
     traceEventStart(<>, %changed, -1, <>)
   }
   %composer.startReplaceGroup(<>)
-  sourceInformation(%composer, "<B()>,<B()>")
-  val tmp0_group = B(%composer, 0) && B(%composer, 0)
+  sourceInformation(%composer, "<B()>")
+  val tmp1_group = B(%composer, 0)
+  %composer.endReplaceGroup()
+  tmp1_group && %composer.startReplaceGroup(<>)
+  sourceInformation(%composer, "<B()>")
+  val tmp0_group = B(%composer, 0)
   %composer.endReplaceGroup()
   tmp0_group
   if (isTraceInProgress()) {
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testConditionalReturnFromInline[useFir = false].txt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testConditionalReturnFromInline[useFir = false].txt
index 44394252e2c6c..0e8d86e3167cf 100644
--- a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testConditionalReturnFromInline[useFir = false].txt	
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testConditionalReturnFromInline[useFir = false].txt	
@@ -41,27 +41,31 @@ fun Test(test: Boolean, %composer: Composer?, %changed: Int) {
     Column({ %composer: Composer?, %changed: Int ->
       %composer.startReplaceGroup(<>)
       sourceInformation(%composer, "C<Text("...>:Test.kt")
-      %composer.startReplaceGroup(<>)
-      sourceInformation(%composer, "<Text("...>")
       if (!test) {
+        %composer.startReplaceGroup(<>)
+        sourceInformation(%composer, "<Text("...>")
         Text("Say", %composer, 0b0110)
         %composer.endReplaceGroup()
         %composer.endReplaceGroup()
         return@Column
+      } else {
+        %composer.startReplaceGroup(<>)
+        %composer.endReplaceGroup()
       }
-      %composer.endReplaceGroup()
       Text("Hello", %composer, 0b0110)
       %composer.endReplaceGroup()
     }, %composer, 0)
     NonComposable {
-      %composer.startReplaceGroup(<>)
-      sourceInformation(%composer, "<Text("...>")
       if (!test) {
+        %composer.startReplaceGroup(<>)
+        sourceInformation(%composer, "<Text("...>")
         Text("Say", %composer, 0b0110)
         %composer.endReplaceGroup()
         return@NonComposable
+      } else {
+        %composer.startReplaceGroup(<>)
+        %composer.endReplaceGroup()
       }
-      %composer.endReplaceGroup()
       Text("Hello", %composer, 0b0110)
     }
     if (isTraceInProgress()) {
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testConditionalReturnFromInline[useFir = true].txt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testConditionalReturnFromInline[useFir = true].txt
index 44394252e2c6c..0e8d86e3167cf 100644
--- a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testConditionalReturnFromInline[useFir = true].txt	
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testConditionalReturnFromInline[useFir = true].txt	
@@ -41,27 +41,31 @@ fun Test(test: Boolean, %composer: Composer?, %changed: Int) {
     Column({ %composer: Composer?, %changed: Int ->
       %composer.startReplaceGroup(<>)
       sourceInformation(%composer, "C<Text("...>:Test.kt")
-      %composer.startReplaceGroup(<>)
-      sourceInformation(%composer, "<Text("...>")
       if (!test) {
+        %composer.startReplaceGroup(<>)
+        sourceInformation(%composer, "<Text("...>")
         Text("Say", %composer, 0b0110)
         %composer.endReplaceGroup()
         %composer.endReplaceGroup()
         return@Column
+      } else {
+        %composer.startReplaceGroup(<>)
+        %composer.endReplaceGroup()
       }
-      %composer.endReplaceGroup()
       Text("Hello", %composer, 0b0110)
       %composer.endReplaceGroup()
     }, %composer, 0)
     NonComposable {
-      %composer.startReplaceGroup(<>)
-      sourceInformation(%composer, "<Text("...>")
       if (!test) {
+        %composer.startReplaceGroup(<>)
+        sourceInformation(%composer, "<Text("...>")
         Text("Say", %composer, 0b0110)
         %composer.endReplaceGroup()
         return@NonComposable
+      } else {
+        %composer.startReplaceGroup(<>)
+        %composer.endReplaceGroup()
       }
-      %composer.endReplaceGroup()
       Text("Hello", %composer, 0b0110)
     }
     if (isTraceInProgress()) {
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testDynamicWrappingGroupWithReturnValue[useFir = false].txt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testDynamicWrappingGroupWithReturnValue[useFir = false].txt
index 8f914995c363a..7d41355026344 100644
--- a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testDynamicWrappingGroupWithReturnValue[useFir = false].txt	
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testDynamicWrappingGroupWithReturnValue[useFir = false].txt	
@@ -28,11 +28,11 @@ fun Example(x: Int, %composer: Composer?, %changed: Int): Int {
   if (isTraceInProgress()) {
     traceEventStart(<>, %changed, -1, <>)
   }
-  val tmp0 = <block>{
+  val tmp0 = if (x > 0) {
     %composer.startReplaceGroup(<>)
     sourceInformation(%composer, "")
-    val tmp4_group = if (x > 0) {
-      val tmp3_group = if (%composer.startReplaceGroup(<>)
+    val tmp3_group = <block>{
+      if (%composer.startReplaceGroup(<>)
       sourceInformation(%composer, "<B()>")
       val tmp1_group = B(%composer, 0)
       %composer.endReplaceGroup()
@@ -40,13 +40,16 @@ fun Example(x: Int, %composer: Composer?, %changed: Int): Int {
       sourceInformation(%composer, "<B()>")
       val tmp2_group = B(%composer, 0)
       %composer.endReplaceGroup()
-      tmp2_group) 2 else 3
-      tmp3_group
-    } else {
-      4
+      tmp2_group) 2 else if (%composer.startReplaceGroup(<>)
+      %composer.endReplaceGroup()
+      true) 3
     }
     %composer.endReplaceGroup()
-    tmp4_group
+    tmp3_group
+  } else {
+    %composer.startReplaceGroup(<>)
+    %composer.endReplaceGroup()
+    4
   }
   if (isTraceInProgress()) {
     traceEventEnd()
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testDynamicWrappingGroupWithReturnValue[useFir = true].txt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testDynamicWrappingGroupWithReturnValue[useFir = true].txt
index 8f914995c363a..7d41355026344 100644
--- a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testDynamicWrappingGroupWithReturnValue[useFir = true].txt	
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testDynamicWrappingGroupWithReturnValue[useFir = true].txt	
@@ -28,11 +28,11 @@ fun Example(x: Int, %composer: Composer?, %changed: Int): Int {
   if (isTraceInProgress()) {
     traceEventStart(<>, %changed, -1, <>)
   }
-  val tmp0 = <block>{
+  val tmp0 = if (x > 0) {
     %composer.startReplaceGroup(<>)
     sourceInformation(%composer, "")
-    val tmp4_group = if (x > 0) {
-      val tmp3_group = if (%composer.startReplaceGroup(<>)
+    val tmp3_group = <block>{
+      if (%composer.startReplaceGroup(<>)
       sourceInformation(%composer, "<B()>")
       val tmp1_group = B(%composer, 0)
       %composer.endReplaceGroup()
@@ -40,13 +40,16 @@ fun Example(x: Int, %composer: Composer?, %changed: Int): Int {
       sourceInformation(%composer, "<B()>")
       val tmp2_group = B(%composer, 0)
       %composer.endReplaceGroup()
-      tmp2_group) 2 else 3
-      tmp3_group
-    } else {
-      4
+      tmp2_group) 2 else if (%composer.startReplaceGroup(<>)
+      %composer.endReplaceGroup()
+      true) 3
     }
     %composer.endReplaceGroup()
-    tmp4_group
+    tmp3_group
+  } else {
+    %composer.startReplaceGroup(<>)
+    %composer.endReplaceGroup()
+    4
   }
   if (isTraceInProgress()) {
     traceEventEnd()
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testEarlyReturnCallValue[useFir = false].txt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testEarlyReturnCallValue[useFir = false].txt
index 78cdf3e8cc14e..97cf15334342f 100644
--- a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testEarlyReturnCallValue[useFir = false].txt	
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testEarlyReturnCallValue[useFir = false].txt	
@@ -28,9 +28,9 @@ fun Example(x: Int, %composer: Composer?, %changed: Int): Int {
   if (isTraceInProgress()) {
     traceEventStart(<>, %changed, -1, <>)
   }
-  %composer.startReplaceGroup(<>)
-  sourceInformation(%composer, "<R()>")
   if (x > 0) {
+    %composer.startReplaceGroup(<>)
+    sourceInformation(%composer, "<R()>")
     val tmp1_return = R(%composer, 0)
     %composer.endReplaceGroup()
     if (isTraceInProgress()) {
@@ -38,8 +38,10 @@ fun Example(x: Int, %composer: Composer?, %changed: Int): Int {
     }
     %composer.endReplaceGroup()
     return tmp1_return
+  } else {
+    %composer.startReplaceGroup(<>)
+    %composer.endReplaceGroup()
   }
-  %composer.endReplaceGroup()
   val tmp0 = R(%composer, 0)
   if (isTraceInProgress()) {
     traceEventEnd()
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testEarlyReturnCallValue[useFir = true].txt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testEarlyReturnCallValue[useFir = true].txt
index 78cdf3e8cc14e..97cf15334342f 100644
--- a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testEarlyReturnCallValue[useFir = true].txt	
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testEarlyReturnCallValue[useFir = true].txt	
@@ -28,9 +28,9 @@ fun Example(x: Int, %composer: Composer?, %changed: Int): Int {
   if (isTraceInProgress()) {
     traceEventStart(<>, %changed, -1, <>)
   }
-  %composer.startReplaceGroup(<>)
-  sourceInformation(%composer, "<R()>")
   if (x > 0) {
+    %composer.startReplaceGroup(<>)
+    sourceInformation(%composer, "<R()>")
     val tmp1_return = R(%composer, 0)
     %composer.endReplaceGroup()
     if (isTraceInProgress()) {
@@ -38,8 +38,10 @@ fun Example(x: Int, %composer: Composer?, %changed: Int): Int {
     }
     %composer.endReplaceGroup()
     return tmp1_return
+  } else {
+    %composer.startReplaceGroup(<>)
+    %composer.endReplaceGroup()
   }
-  %composer.endReplaceGroup()
   val tmp0 = R(%composer, 0)
   if (isTraceInProgress()) {
     traceEventEnd()
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testEarlyReturnFromCrossInlinedLambda[useFir = false].txt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testEarlyReturnFromCrossInlinedLambda[useFir = false].txt
index 3f4bf08766990..248d1a240a0fc 100644
--- a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testEarlyReturnFromCrossInlinedLambda[useFir = false].txt	
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testEarlyReturnFromCrossInlinedLambda[useFir = false].txt	
@@ -30,13 +30,16 @@ private fun Test(param: String?, %composer: Composer?, %changed: Int) {
     Dialog({ %composer: Composer?, %changed: Int ->
       %composer.startReplaceGroup(<>)
       sourceInformation(%composer, "C:Test.kt")
-      %composer.startReplaceGroup(<>)
-      sourceInformation(%composer, "<Test(p...>")
       if (false) {
+        %composer.startReplaceGroup(<>)
+        sourceInformation(%composer, "<Test(p...>")
         Test(param, %composer, 0)
+        %composer.endReplaceGroup()
+      } else {
+        %composer.startReplaceGroup(<>)
+        %composer.endReplaceGroup()
       }
       %composer.endReplaceGroup()
-      %composer.endReplaceGroup()
     }, %composer, 0)
     if (isTraceInProgress()) {
       traceEventEnd()
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testEarlyReturnFromCrossInlinedLambda[useFir = true].txt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testEarlyReturnFromCrossInlinedLambda[useFir = true].txt
index 3f4bf08766990..248d1a240a0fc 100644
--- a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testEarlyReturnFromCrossInlinedLambda[useFir = true].txt	
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testEarlyReturnFromCrossInlinedLambda[useFir = true].txt	
@@ -30,13 +30,16 @@ private fun Test(param: String?, %composer: Composer?, %changed: Int) {
     Dialog({ %composer: Composer?, %changed: Int ->
       %composer.startReplaceGroup(<>)
       sourceInformation(%composer, "C:Test.kt")
-      %composer.startReplaceGroup(<>)
-      sourceInformation(%composer, "<Test(p...>")
       if (false) {
+        %composer.startReplaceGroup(<>)
+        sourceInformation(%composer, "<Test(p...>")
         Test(param, %composer, 0)
+        %composer.endReplaceGroup()
+      } else {
+        %composer.startReplaceGroup(<>)
+        %composer.endReplaceGroup()
       }
       %composer.endReplaceGroup()
-      %composer.endReplaceGroup()
     }, %composer, 0)
     if (isTraceInProgress()) {
       traceEventEnd()
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testEarlyReturnValue[useFir = false].txt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testEarlyReturnValue[useFir = false].txt
index 06474c2c6d264..e414324c3100c 100644
--- a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testEarlyReturnValue[useFir = false].txt	
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testEarlyReturnValue[useFir = false].txt	
@@ -25,18 +25,24 @@ fun Example(x: Int): Int {
 @Composable
 fun Example(x: Int, %composer: Composer?, %changed: Int): Int {
   %composer.startReplaceGroup(<>)
-  sourceInformation(%composer, "C(Example)<A()>:Test.kt")
+  sourceInformation(%composer, "C(Example):Test.kt")
   if (isTraceInProgress()) {
     traceEventStart(<>, %changed, -1, <>)
   }
   if (x > 0) {
+    %composer.startReplaceGroup(<>)
+    sourceInformation(%composer, "<A()>")
     A(%composer, 0)
     val tmp1_return = 1
+    %composer.endReplaceGroup()
     if (isTraceInProgress()) {
       traceEventEnd()
     }
     %composer.endReplaceGroup()
     return tmp1_return
+  } else {
+    %composer.startReplaceGroup(<>)
+    %composer.endReplaceGroup()
   }
   val tmp0 = 2
   if (isTraceInProgress()) {
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testEarlyReturnValue[useFir = true].txt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testEarlyReturnValue[useFir = true].txt
index 06474c2c6d264..e414324c3100c 100644
--- a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testEarlyReturnValue[useFir = true].txt	
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testEarlyReturnValue[useFir = true].txt	
@@ -25,18 +25,24 @@ fun Example(x: Int): Int {
 @Composable
 fun Example(x: Int, %composer: Composer?, %changed: Int): Int {
   %composer.startReplaceGroup(<>)
-  sourceInformation(%composer, "C(Example)<A()>:Test.kt")
+  sourceInformation(%composer, "C(Example):Test.kt")
   if (isTraceInProgress()) {
     traceEventStart(<>, %changed, -1, <>)
   }
   if (x > 0) {
+    %composer.startReplaceGroup(<>)
+    sourceInformation(%composer, "<A()>")
     A(%composer, 0)
     val tmp1_return = 1
+    %composer.endReplaceGroup()
     if (isTraceInProgress()) {
       traceEventEnd()
     }
     %composer.endReplaceGroup()
     return tmp1_return
+  } else {
+    %composer.startReplaceGroup(<>)
+    %composer.endReplaceGroup()
   }
   val tmp0 = 2
   if (isTraceInProgress()) {
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testEarlyReturnWithCallsBeforeButNotAfter[useFir = false].txt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testEarlyReturnWithCallsBeforeButNotAfter[useFir = false].txt
index 4d0b5f4ec27f5..031cc78e6c546 100644
--- a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testEarlyReturnWithCallsBeforeButNotAfter[useFir = false].txt	
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testEarlyReturnWithCallsBeforeButNotAfter[useFir = false].txt	
@@ -26,17 +26,23 @@ fun Example(x: Int) {
 @Composable
 fun Example(x: Int, %composer: Composer?, %changed: Int) {
   %composer.startReplaceGroup(<>)
-  sourceInformation(%composer, "C(Example)<A()>:Test.kt")
+  sourceInformation(%composer, "C(Example):Test.kt")
   if (isTraceInProgress()) {
     traceEventStart(<>, %changed, -1, <>)
   }
   if (x > 0) {
+    %composer.startReplaceGroup(<>)
+    sourceInformation(%composer, "<A()>")
     A(%composer, 0)
+    %composer.endReplaceGroup()
     if (isTraceInProgress()) {
       traceEventEnd()
     }
     %composer.endReplaceGroup()
     return
+  } else {
+    %composer.startReplaceGroup(<>)
+    %composer.endReplaceGroup()
   }
   print("hello")
   if (isTraceInProgress()) {
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testEarlyReturnWithCallsBeforeButNotAfter[useFir = true].txt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testEarlyReturnWithCallsBeforeButNotAfter[useFir = true].txt
index 4d0b5f4ec27f5..031cc78e6c546 100644
--- a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testEarlyReturnWithCallsBeforeButNotAfter[useFir = true].txt	
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testEarlyReturnWithCallsBeforeButNotAfter[useFir = true].txt	
@@ -26,17 +26,23 @@ fun Example(x: Int) {
 @Composable
 fun Example(x: Int, %composer: Composer?, %changed: Int) {
   %composer.startReplaceGroup(<>)
-  sourceInformation(%composer, "C(Example)<A()>:Test.kt")
+  sourceInformation(%composer, "C(Example):Test.kt")
   if (isTraceInProgress()) {
     traceEventStart(<>, %changed, -1, <>)
   }
   if (x > 0) {
+    %composer.startReplaceGroup(<>)
+    sourceInformation(%composer, "<A()>")
     A(%composer, 0)
+    %composer.endReplaceGroup()
     if (isTraceInProgress()) {
       traceEventEnd()
     }
     %composer.endReplaceGroup()
     return
+  } else {
+    %composer.startReplaceGroup(<>)
+    %composer.endReplaceGroup()
   }
   print("hello")
   if (isTraceInProgress()) {
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testElvis[useFir = false].txt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testElvis[useFir = false].txt
index da5c0d3dcc229..218498f1cb93e 100644
--- a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testElvis[useFir = false].txt	
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testElvis[useFir = false].txt	
@@ -29,18 +29,20 @@ fun Example(x: Int?, %composer: Composer?, %changed: Int) {
   }
   val y = <block>{
     val <elvis> = x
-    %composer.startReplaceGroup(<>)
-    sourceInformation(%composer, "<R()>")
-    val tmp0_group = when {
+    when {
       <elvis> == null -> {
-        R(%composer, 0)
+        %composer.startReplaceGroup(<>)
+        sourceInformation(%composer, "<R()>")
+        val tmp0_group = R(%composer, 0)
+        %composer.endReplaceGroup()
+        tmp0_group
       }
       else -> {
+        %composer.startReplaceGroup(<>)
+        %composer.endReplaceGroup()
         <elvis>
       }
     }
-    %composer.endReplaceGroup()
-    tmp0_group
   }
   if (isTraceInProgress()) {
     traceEventEnd()
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testElvis[useFir = true].txt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testElvis[useFir = true].txt
index ef8adf188bd16..c15a1d498fd18 100644
--- a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testElvis[useFir = true].txt	
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testElvis[useFir = true].txt	
@@ -29,18 +29,20 @@ fun Example(x: Int?, %composer: Composer?, %changed: Int) {
   }
   val y = <block>{
     val tmp0_elvis_lhs = x
-    %composer.startReplaceGroup(<>)
-    sourceInformation(%composer, "<R()>")
-    val tmp0_group = when {
+    when {
       tmp0_elvis_lhs == null -> {
-        R(%composer, 0)
+        %composer.startReplaceGroup(<>)
+        sourceInformation(%composer, "<R()>")
+        val tmp0_group = R(%composer, 0)
+        %composer.endReplaceGroup()
+        tmp0_group
       }
       else -> {
+        %composer.startReplaceGroup(<>)
+        %composer.endReplaceGroup()
         tmp0_elvis_lhs
       }
     }
-    %composer.endReplaceGroup()
-    tmp0_group
   }
   if (isTraceInProgress()) {
     traceEventEnd()
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testGroupAroundExtensionFunctions[useFir = false].txt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testGroupAroundExtensionFunctions[useFir = false].txt
index 57991856a5eb1..6002a0a9140b3 100644
--- a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testGroupAroundExtensionFunctions[useFir = false].txt	
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testGroupAroundExtensionFunctions[useFir = false].txt	
@@ -46,12 +46,15 @@ fun Test(start: Int, end: Int, %composer: Composer?, %changed: Int) {
     while (<iterator>.hasNext()) {
       val i = <iterator>.next()
       val b = a.get(bKey, %composer, 0b00110110)
-      %composer.startReplaceGroup(<>)
-      sourceInformation(%composer, "<get(cK...>")
       if (i == 0b0010) {
+        %composer.startReplaceGroup(<>)
+        sourceInformation(%composer, "<get(cK...>")
         a.get(cKey, %composer, 0b00110110)
+        %composer.endReplaceGroup()
+      } else {
+        %composer.startReplaceGroup(<>)
+        %composer.endReplaceGroup()
       }
-      %composer.endReplaceGroup()
     }
     if (isTraceInProgress()) {
       traceEventEnd()
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testGroupAroundExtensionFunctions[useFir = true].txt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testGroupAroundExtensionFunctions[useFir = true].txt
index 57991856a5eb1..6002a0a9140b3 100644
--- a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testGroupAroundExtensionFunctions[useFir = true].txt	
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testGroupAroundExtensionFunctions[useFir = true].txt	
@@ -46,12 +46,15 @@ fun Test(start: Int, end: Int, %composer: Composer?, %changed: Int) {
     while (<iterator>.hasNext()) {
       val i = <iterator>.next()
       val b = a.get(bKey, %composer, 0b00110110)
-      %composer.startReplaceGroup(<>)
-      sourceInformation(%composer, "<get(cK...>")
       if (i == 0b0010) {
+        %composer.startReplaceGroup(<>)
+        sourceInformation(%composer, "<get(cK...>")
         a.get(cKey, %composer, 0b00110110)
+        %composer.endReplaceGroup()
+      } else {
+        %composer.startReplaceGroup(<>)
+        %composer.endReplaceGroup()
       }
-      %composer.endReplaceGroup()
     }
     if (isTraceInProgress()) {
       traceEventEnd()
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testGroupsInLoops[useFir = false].txt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testGroupsInLoops[useFir = false].txt
index 2cc133876b53a..292f95b278f41 100644
--- a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testGroupsInLoops[useFir = false].txt	
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testGroupsInLoops[useFir = false].txt	
@@ -47,9 +47,9 @@ private fun KeyContent1(items: List<Int>, %composer: Composer?, %changed: Int) {
       traceEventStart(<>, %dirty, -1, <>)
     }
     items.forEach { item: Int ->
-      %composer.startReplaceGroup(<>)
-      sourceInformation(%composer, "")
       if (item > -1) {
+        %composer.startReplaceGroup(<>)
+        sourceInformation(%composer, "")
         %composer.startMovableGroup(<>, item)
         sourceInformation(%composer, "<rememb...>")
         val tmp0 = <block>{
@@ -62,8 +62,11 @@ private fun KeyContent1(items: List<Int>, %composer: Composer?, %changed: Int) {
         }
         %composer.endMovableGroup()
         tmp0
+        %composer.endReplaceGroup()
+      } else {
+        %composer.startReplaceGroup(<>)
+        %composer.endReplaceGroup()
       }
-      %composer.endReplaceGroup()
     }
     if (isTraceInProgress()) {
       traceEventEnd()
@@ -90,9 +93,9 @@ private fun KeyContent2(items: List<Int>, %composer: Composer?, %changed: Int) {
     val <iterator> = items.iterator()
     while (<iterator>.hasNext()) {
       val item = <iterator>.next()
-      %composer.startReplaceGroup(<>)
-      sourceInformation(%composer, "")
       if (item > -1) {
+        %composer.startReplaceGroup(<>)
+        sourceInformation(%composer, "")
         %composer.startMovableGroup(<>, item)
         sourceInformation(%composer, "<rememb...>")
         val tmp0 = <block>{
@@ -105,8 +108,11 @@ private fun KeyContent2(items: List<Int>, %composer: Composer?, %changed: Int) {
         }
         %composer.endMovableGroup()
         tmp0
+        %composer.endReplaceGroup()
+      } else {
+        %composer.startReplaceGroup(<>)
+        %composer.endReplaceGroup()
       }
-      %composer.endReplaceGroup()
     }
     if (isTraceInProgress()) {
       traceEventEnd()
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testGroupsInLoops[useFir = true].txt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testGroupsInLoops[useFir = true].txt
index 57e056852d9ad..51b535e889763 100644
--- a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testGroupsInLoops[useFir = true].txt	
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testGroupsInLoops[useFir = true].txt	
@@ -47,9 +47,9 @@ private fun KeyContent1(items: List<Int>, %composer: Composer?, %changed: Int) {
       traceEventStart(<>, %dirty, -1, <>)
     }
     items.forEach { item: Int ->
-      %composer.startReplaceGroup(<>)
-      sourceInformation(%composer, "")
       if (item > -1) {
+        %composer.startReplaceGroup(<>)
+        sourceInformation(%composer, "")
         %composer.startMovableGroup(<>, item)
         sourceInformation(%composer, "<rememb...>")
         sourceInformationMarkerStart(%composer, <>, "CC(remember):Test.kt#9igjgp")
@@ -58,8 +58,11 @@ private fun KeyContent1(items: List<Int>, %composer: Composer?, %changed: Int) {
         }
         sourceInformationMarkerEnd(%composer)
         %composer.endMovableGroup()
+        %composer.endReplaceGroup()
+      } else {
+        %composer.startReplaceGroup(<>)
+        %composer.endReplaceGroup()
       }
-      %composer.endReplaceGroup()
     }
     if (isTraceInProgress()) {
       traceEventEnd()
@@ -86,9 +89,9 @@ private fun KeyContent2(items: List<Int>, %composer: Composer?, %changed: Int) {
     val <iterator> = items.iterator()
     while (<iterator>.hasNext()) {
       val item = <iterator>.next()
-      %composer.startReplaceGroup(<>)
-      sourceInformation(%composer, "")
       if (item > -1) {
+        %composer.startReplaceGroup(<>)
+        sourceInformation(%composer, "")
         %composer.startMovableGroup(<>, item)
         sourceInformation(%composer, "<rememb...>")
         val tmp0 = <block>{
@@ -101,8 +104,11 @@ private fun KeyContent2(items: List<Int>, %composer: Composer?, %changed: Int) {
         }
         %composer.endMovableGroup()
         tmp0
+        %composer.endReplaceGroup()
+      } else {
+        %composer.startReplaceGroup(<>)
+        %composer.endReplaceGroup()
       }
-      %composer.endReplaceGroup()
     }
     if (isTraceInProgress()) {
       traceEventEnd()
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testIfElseWithCallsInConditions[useFir = false].txt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testIfElseWithCallsInConditions[useFir = false].txt
index 93b4d577e4151..a4d532c650786 100644
--- a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testIfElseWithCallsInConditions[useFir = false].txt	
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testIfElseWithCallsInConditions[useFir = false].txt	
@@ -35,8 +35,6 @@ fun Example(x: Int, %composer: Composer?, %changed: Int) {
   if (isTraceInProgress()) {
     traceEventStart(<>, %changed, -1, <>)
   }
-  %composer.startReplaceGroup(<>)
-  sourceInformation(%composer, "")
   if (%composer.startReplaceGroup(<>)
   sourceInformation(%composer, "<B(a)>")
   val tmp0_group = B(a, %composer, 0)
@@ -49,10 +47,11 @@ fun Example(x: Int, %composer: Composer?, %changed: Int) {
   %composer.endReplaceGroup()
   tmp1_group) {
     NA()
-  } else {
+  } else if (%composer.startReplaceGroup(<>)
+  %composer.endReplaceGroup()
+  true) {
     NA()
   }
-  %composer.endReplaceGroup()
   if (isTraceInProgress()) {
     traceEventEnd()
   }
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testIfElseWithCallsInConditions[useFir = true].txt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testIfElseWithCallsInConditions[useFir = true].txt
index 93b4d577e4151..a4d532c650786 100644
--- a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testIfElseWithCallsInConditions[useFir = true].txt	
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testIfElseWithCallsInConditions[useFir = true].txt	
@@ -35,8 +35,6 @@ fun Example(x: Int, %composer: Composer?, %changed: Int) {
   if (isTraceInProgress()) {
     traceEventStart(<>, %changed, -1, <>)
   }
-  %composer.startReplaceGroup(<>)
-  sourceInformation(%composer, "")
   if (%composer.startReplaceGroup(<>)
   sourceInformation(%composer, "<B(a)>")
   val tmp0_group = B(a, %composer, 0)
@@ -49,10 +47,11 @@ fun Example(x: Int, %composer: Composer?, %changed: Int) {
   %composer.endReplaceGroup()
   tmp1_group) {
     NA()
-  } else {
+  } else if (%composer.startReplaceGroup(<>)
+  %composer.endReplaceGroup()
+  true) {
     NA()
   }
-  %composer.endReplaceGroup()
   if (isTraceInProgress()) {
     traceEventEnd()
   }
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testIfWithCallInCondition[useFir = false].txt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testIfWithCallInCondition[useFir = false].txt
index afa4db1c6900d..1696718afbb88 100644
--- a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testIfWithCallInCondition[useFir = false].txt	
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testIfWithCallInCondition[useFir = false].txt	
@@ -27,13 +27,19 @@ fun Example(x: Int) {
 @NonRestartableComposable
 @Composable
 fun Example(x: Int, %composer: Composer?, %changed: Int) {
-  sourceInformationMarkerStart(%composer, <>, "C(Example)<B()>:Test.kt")
+  sourceInformationMarkerStart(%composer, <>, "C(Example):Test.kt")
   if (isTraceInProgress()) {
     traceEventStart(<>, %changed, -1, <>)
   }
-  if (B(%composer, 0)) {
+  if (%composer.startReplaceGroup(<>)
+  sourceInformation(%composer, "<B()>")
+  val tmp0_group = B(%composer, 0)
+  %composer.endReplaceGroup()
+  tmp0_group) {
     NA()
-  } else {
+  } else if (%composer.startReplaceGroup(<>)
+  %composer.endReplaceGroup()
+  true) {
     NA()
   }
   if (isTraceInProgress()) {
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testIfWithCallInCondition[useFir = true].txt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testIfWithCallInCondition[useFir = true].txt
index afa4db1c6900d..1696718afbb88 100644
--- a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testIfWithCallInCondition[useFir = true].txt	
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testIfWithCallInCondition[useFir = true].txt	
@@ -27,13 +27,19 @@ fun Example(x: Int) {
 @NonRestartableComposable
 @Composable
 fun Example(x: Int, %composer: Composer?, %changed: Int) {
-  sourceInformationMarkerStart(%composer, <>, "C(Example)<B()>:Test.kt")
+  sourceInformationMarkerStart(%composer, <>, "C(Example):Test.kt")
   if (isTraceInProgress()) {
     traceEventStart(<>, %changed, -1, <>)
   }
-  if (B(%composer, 0)) {
+  if (%composer.startReplaceGroup(<>)
+  sourceInformation(%composer, "<B()>")
+  val tmp0_group = B(%composer, 0)
+  %composer.endReplaceGroup()
+  tmp0_group) {
     NA()
-  } else {
+  } else if (%composer.startReplaceGroup(<>)
+  %composer.endReplaceGroup()
+  true) {
     NA()
   }
   if (isTraceInProgress()) {
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testIfWithCallsInBranch[useFir = false].txt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testIfWithCallsInBranch[useFir = false].txt
index 59f08b038401e..82a1cf4573c5d 100644
--- a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testIfWithCallsInBranch[useFir = false].txt	
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testIfWithCallsInBranch[useFir = false].txt	
@@ -28,12 +28,15 @@ fun Example(x: Int, %composer: Composer?, %changed: Int) {
   if (isTraceInProgress()) {
     traceEventStart(<>, %changed, -1, <>)
   }
-  %composer.startReplaceGroup(<>)
-  sourceInformation(%composer, "<A()>")
   if (x > 0) {
+    %composer.startReplaceGroup(<>)
+    sourceInformation(%composer, "<A()>")
     A(%composer, 0)
+    %composer.endReplaceGroup()
+  } else {
+    %composer.startReplaceGroup(<>)
+    %composer.endReplaceGroup()
   }
-  %composer.endReplaceGroup()
   if (isTraceInProgress()) {
     traceEventEnd()
   }
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testIfWithCallsInBranch[useFir = true].txt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testIfWithCallsInBranch[useFir = true].txt
index 59f08b038401e..82a1cf4573c5d 100644
--- a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testIfWithCallsInBranch[useFir = true].txt	
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testIfWithCallsInBranch[useFir = true].txt	
@@ -28,12 +28,15 @@ fun Example(x: Int, %composer: Composer?, %changed: Int) {
   if (isTraceInProgress()) {
     traceEventStart(<>, %changed, -1, <>)
   }
-  %composer.startReplaceGroup(<>)
-  sourceInformation(%composer, "<A()>")
   if (x > 0) {
+    %composer.startReplaceGroup(<>)
+    sourceInformation(%composer, "<A()>")
     A(%composer, 0)
+    %composer.endReplaceGroup()
+  } else {
+    %composer.startReplaceGroup(<>)
+    %composer.endReplaceGroup()
   }
-  %composer.endReplaceGroup()
   if (isTraceInProgress()) {
     traceEventEnd()
   }
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testIfWithEarlyReturnInsideInlineLambda[useFir = false].txt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testIfWithEarlyReturnInsideInlineLambda[useFir = false].txt
index fe1b30e362370..72f0fd74e7297 100644
--- a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testIfWithEarlyReturnInsideInlineLambda[useFir = false].txt	
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testIfWithEarlyReturnInsideInlineLambda[useFir = false].txt	
@@ -28,17 +28,17 @@ fun Test(%composer: Composer?, %changed: Int) {
       traceEventStart(<>, %changed, -1, <>)
     }
     run {
-      %composer.startReplaceGroup(<>)
-      sourceInformation(%composer, "<Test()>")
       if (true) {
+        %composer.startReplaceGroup(<>)
         %composer.endReplaceGroup()
         return@run
       } else {
+        %composer.startReplaceGroup(<>)
+        sourceInformation(%composer, "<Test()>")
         Test(%composer, 0)
         %composer.endReplaceGroup()
         return@run
       }
-      %composer.endReplaceGroup()
     }
     if (isTraceInProgress()) {
       traceEventEnd()
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testIfWithEarlyReturnInsideInlineLambda[useFir = true].txt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testIfWithEarlyReturnInsideInlineLambda[useFir = true].txt
index fe1b30e362370..72f0fd74e7297 100644
--- a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testIfWithEarlyReturnInsideInlineLambda[useFir = true].txt	
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testIfWithEarlyReturnInsideInlineLambda[useFir = true].txt	
@@ -28,17 +28,17 @@ fun Test(%composer: Composer?, %changed: Int) {
       traceEventStart(<>, %changed, -1, <>)
     }
     run {
-      %composer.startReplaceGroup(<>)
-      sourceInformation(%composer, "<Test()>")
       if (true) {
+        %composer.startReplaceGroup(<>)
         %composer.endReplaceGroup()
         return@run
       } else {
+        %composer.startReplaceGroup(<>)
+        sourceInformation(%composer, "<Test()>")
         Test(%composer, 0)
         %composer.endReplaceGroup()
         return@run
       }
-      %composer.endReplaceGroup()
     }
     if (isTraceInProgress()) {
       traceEventEnd()
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testKeyInIfAndCallsAfter[useFir = false].txt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testKeyInIfAndCallsAfter[useFir = false].txt
index 0141b323dd875..3de72378e964a 100644
--- a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testKeyInIfAndCallsAfter[useFir = false].txt	
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testKeyInIfAndCallsAfter[useFir = false].txt	
@@ -29,16 +29,19 @@ fun Example(x: Int, %composer: Composer?, %changed: Int) {
   if (isTraceInProgress()) {
     traceEventStart(<>, %changed, -1, <>)
   }
-  %composer.startReplaceGroup(<>)
-  sourceInformation(%composer, "<A(b)>")
   if (x > 0) {
+    %composer.startReplaceGroup(<>)
+    sourceInformation(%composer, "<A(b)>")
     %composer.startMovableGroup(<>, x)
     sourceInformation(%composer, "<A(a)>")
     A(a, %composer, 0)
     %composer.endMovableGroup()
     A(b, %composer, 0)
+    %composer.endReplaceGroup()
+  } else {
+    %composer.startReplaceGroup(<>)
+    %composer.endReplaceGroup()
   }
-  %composer.endReplaceGroup()
   if (isTraceInProgress()) {
     traceEventEnd()
   }
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testKeyInIfAndCallsAfter[useFir = true].txt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testKeyInIfAndCallsAfter[useFir = true].txt
index 0141b323dd875..3de72378e964a 100644
--- a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testKeyInIfAndCallsAfter[useFir = true].txt	
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testKeyInIfAndCallsAfter[useFir = true].txt	
@@ -29,16 +29,19 @@ fun Example(x: Int, %composer: Composer?, %changed: Int) {
   if (isTraceInProgress()) {
     traceEventStart(<>, %changed, -1, <>)
   }
-  %composer.startReplaceGroup(<>)
-  sourceInformation(%composer, "<A(b)>")
   if (x > 0) {
+    %composer.startReplaceGroup(<>)
+    sourceInformation(%composer, "<A(b)>")
     %composer.startMovableGroup(<>, x)
     sourceInformation(%composer, "<A(a)>")
     A(a, %composer, 0)
     %composer.endMovableGroup()
     A(b, %composer, 0)
+    %composer.endReplaceGroup()
+  } else {
+    %composer.startReplaceGroup(<>)
+    %composer.endReplaceGroup()
   }
-  %composer.endReplaceGroup()
   if (isTraceInProgress()) {
     traceEventEnd()
   }
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testKeyInIfAndCallsBefore[useFir = false].txt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testKeyInIfAndCallsBefore[useFir = false].txt
index 0f63c6b120f0a..5af7dcec807aa 100644
--- a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testKeyInIfAndCallsBefore[useFir = false].txt	
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testKeyInIfAndCallsBefore[useFir = false].txt	
@@ -29,16 +29,19 @@ fun Example(x: Int, %composer: Composer?, %changed: Int) {
   if (isTraceInProgress()) {
     traceEventStart(<>, %changed, -1, <>)
   }
-  %composer.startReplaceGroup(<>)
-  sourceInformation(%composer, "<A(a)>")
   if (x > 0) {
+    %composer.startReplaceGroup(<>)
+    sourceInformation(%composer, "<A(a)>")
     A(a, %composer, 0)
     %composer.startMovableGroup(<>, x)
     sourceInformation(%composer, "<A(b)>")
     A(b, %composer, 0)
     %composer.endMovableGroup()
+    %composer.endReplaceGroup()
+  } else {
+    %composer.startReplaceGroup(<>)
+    %composer.endReplaceGroup()
   }
-  %composer.endReplaceGroup()
   if (isTraceInProgress()) {
     traceEventEnd()
   }
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testKeyInIfAndCallsBefore[useFir = true].txt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testKeyInIfAndCallsBefore[useFir = true].txt
index 0f63c6b120f0a..5af7dcec807aa 100644
--- a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testKeyInIfAndCallsBefore[useFir = true].txt	
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testKeyInIfAndCallsBefore[useFir = true].txt	
@@ -29,16 +29,19 @@ fun Example(x: Int, %composer: Composer?, %changed: Int) {
   if (isTraceInProgress()) {
     traceEventStart(<>, %changed, -1, <>)
   }
-  %composer.startReplaceGroup(<>)
-  sourceInformation(%composer, "<A(a)>")
   if (x > 0) {
+    %composer.startReplaceGroup(<>)
+    sourceInformation(%composer, "<A(a)>")
     A(a, %composer, 0)
     %composer.startMovableGroup(<>, x)
     sourceInformation(%composer, "<A(b)>")
     A(b, %composer, 0)
     %composer.endMovableGroup()
+    %composer.endReplaceGroup()
+  } else {
+    %composer.startReplaceGroup(<>)
+    %composer.endReplaceGroup()
   }
-  %composer.endReplaceGroup()
   if (isTraceInProgress()) {
     traceEventEnd()
   }
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testKeyInIf[useFir = false].txt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testKeyInIf[useFir = false].txt
index ccfb16c8c7576..5918c0ac636f6 100644
--- a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testKeyInIf[useFir = false].txt	
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testKeyInIf[useFir = false].txt	
@@ -28,15 +28,18 @@ fun Example(x: Int, %composer: Composer?, %changed: Int) {
   if (isTraceInProgress()) {
     traceEventStart(<>, %changed, -1, <>)
   }
-  %composer.startReplaceGroup(<>)
-  sourceInformation(%composer, "")
   if (x > 0) {
+    %composer.startReplaceGroup(<>)
+    sourceInformation(%composer, "")
     %composer.startMovableGroup(<>, x)
     sourceInformation(%composer, "<A()>")
     A(%composer, 0)
     %composer.endMovableGroup()
+    %composer.endReplaceGroup()
+  } else {
+    %composer.startReplaceGroup(<>)
+    %composer.endReplaceGroup()
   }
-  %composer.endReplaceGroup()
   if (isTraceInProgress()) {
     traceEventEnd()
   }
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testKeyInIf[useFir = true].txt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testKeyInIf[useFir = true].txt
index ccfb16c8c7576..5918c0ac636f6 100644
--- a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testKeyInIf[useFir = true].txt	
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testKeyInIf[useFir = true].txt	
@@ -28,15 +28,18 @@ fun Example(x: Int, %composer: Composer?, %changed: Int) {
   if (isTraceInProgress()) {
     traceEventStart(<>, %changed, -1, <>)
   }
-  %composer.startReplaceGroup(<>)
-  sourceInformation(%composer, "")
   if (x > 0) {
+    %composer.startReplaceGroup(<>)
+    sourceInformation(%composer, "")
     %composer.startMovableGroup(<>, x)
     sourceInformation(%composer, "<A()>")
     A(%composer, 0)
     %composer.endMovableGroup()
+    %composer.endReplaceGroup()
+  } else {
+    %composer.startReplaceGroup(<>)
+    %composer.endReplaceGroup()
   }
-  %composer.endReplaceGroup()
   if (isTraceInProgress()) {
     traceEventEnd()
   }
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testLetWithComposableCalls[useFir = false].txt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testLetWithComposableCalls[useFir = false].txt
index 5c540b6fca124..87a6eb1bbd15e 100644
--- a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testLetWithComposableCalls[useFir = false].txt	
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testLetWithComposableCalls[useFir = false].txt	
@@ -36,26 +36,30 @@ fun Example(x: Int?, %composer: Composer?, %changed: Int) {
       traceEventStart(<>, %dirty, -1, <>)
     }
     val tmp0_safe_receiver = x
-    %composer.startReplaceGroup(<>)
-    sourceInformation(%composer, "*<A(b)>")
-    val tmp0_group = when {
+    when {
       tmp0_safe_receiver == null -> {
+        %composer.startReplaceGroup(<>)
+        %composer.endReplaceGroup()
         null
       }
       else -> {
+        %composer.startReplaceGroup(<>)
+        sourceInformation(%composer, "*<A(b)>")
         tmp0_safe_receiver.let { it: Int ->
-          %composer.startReplaceGroup(<>)
-          sourceInformation(%composer, "<A(a)>")
           if (it > 0) {
+            %composer.startReplaceGroup(<>)
+            sourceInformation(%composer, "<A(a)>")
             A(a, %composer, 0)
+            %composer.endReplaceGroup()
+          } else {
+            %composer.startReplaceGroup(<>)
+            %composer.endReplaceGroup()
           }
-          %composer.endReplaceGroup()
           A(b, %composer, 0)
         }
+        %composer.endReplaceGroup()
       }
     }
-    %composer.endReplaceGroup()
-    tmp0_group
     A(c, %composer, 0)
     if (isTraceInProgress()) {
       traceEventEnd()
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testLetWithComposableCalls[useFir = true].txt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testLetWithComposableCalls[useFir = true].txt
index 5c540b6fca124..87a6eb1bbd15e 100644
--- a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testLetWithComposableCalls[useFir = true].txt	
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testLetWithComposableCalls[useFir = true].txt	
@@ -36,26 +36,30 @@ fun Example(x: Int?, %composer: Composer?, %changed: Int) {
       traceEventStart(<>, %dirty, -1, <>)
     }
     val tmp0_safe_receiver = x
-    %composer.startReplaceGroup(<>)
-    sourceInformation(%composer, "*<A(b)>")
-    val tmp0_group = when {
+    when {
       tmp0_safe_receiver == null -> {
+        %composer.startReplaceGroup(<>)
+        %composer.endReplaceGroup()
         null
       }
       else -> {
+        %composer.startReplaceGroup(<>)
+        sourceInformation(%composer, "*<A(b)>")
         tmp0_safe_receiver.let { it: Int ->
-          %composer.startReplaceGroup(<>)
-          sourceInformation(%composer, "<A(a)>")
           if (it > 0) {
+            %composer.startReplaceGroup(<>)
+            sourceInformation(%composer, "<A(a)>")
             A(a, %composer, 0)
+            %composer.endReplaceGroup()
+          } else {
+            %composer.startReplaceGroup(<>)
+            %composer.endReplaceGroup()
           }
-          %composer.endReplaceGroup()
           A(b, %composer, 0)
         }
+        %composer.endReplaceGroup()
       }
     }
-    %composer.endReplaceGroup()
-    tmp0_group
     A(c, %composer, 0)
     if (isTraceInProgress()) {
       traceEventEnd()
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testOR[useFir = false].txt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testOR[useFir = false].txt
index 0355a3b25c49d..a2933cfb7e21e 100644
--- a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testOR[useFir = false].txt	
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testOR[useFir = false].txt	
@@ -26,10 +26,14 @@ fun Example(%composer: Composer?, %changed: Int) {
     traceEventStart(<>, %changed, -1, <>)
   }
   %composer.startReplaceGroup(<>)
-  sourceInformation(%composer, "<B()>,<B()>")
-  val tmp0_group = B(%composer, 0) || B(%composer, 0)
+  sourceInformation(%composer, "<B()>")
+  val tmp0_group = B(%composer, 0)
   %composer.endReplaceGroup()
-  tmp0_group
+  tmp0_group || %composer.startReplaceGroup(<>)
+  sourceInformation(%composer, "<B()>")
+  val tmp1_group = B(%composer, 0)
+  %composer.endReplaceGroup()
+  tmp1_group
   if (isTraceInProgress()) {
     traceEventEnd()
   }
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testOR[useFir = true].txt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testOR[useFir = true].txt
index 0355a3b25c49d..a2933cfb7e21e 100644
--- a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testOR[useFir = true].txt	
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testOR[useFir = true].txt	
@@ -26,10 +26,14 @@ fun Example(%composer: Composer?, %changed: Int) {
     traceEventStart(<>, %changed, -1, <>)
   }
   %composer.startReplaceGroup(<>)
-  sourceInformation(%composer, "<B()>,<B()>")
-  val tmp0_group = B(%composer, 0) || B(%composer, 0)
+  sourceInformation(%composer, "<B()>")
+  val tmp0_group = B(%composer, 0)
   %composer.endReplaceGroup()
-  tmp0_group
+  tmp0_group || %composer.startReplaceGroup(<>)
+  sourceInformation(%composer, "<B()>")
+  val tmp1_group = B(%composer, 0)
+  %composer.endReplaceGroup()
+  tmp1_group
   if (isTraceInProgress()) {
     traceEventEnd()
   }
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testRememberInConditionalCallArgument[useFir = false].txt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testRememberInConditionalCallArgument[useFir = false].txt
index 936c9e6fcf02d..bf5b7c0354eaa 100644
--- a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testRememberInConditionalCallArgument[useFir = false].txt	
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testRememberInConditionalCallArgument[useFir = false].txt	
@@ -31,21 +31,23 @@ private fun Test(param: String?, %composer: Composer?, %changed: Int) {
     if (isTraceInProgress()) {
       traceEventStart(<>, %dirty, -1, <>)
     }
-    Test(<block>{
+    Test(if (param == null) {
       %composer.startReplaceGroup(<>)
       sourceInformation(%composer, "<rememb...>")
-      val tmp1_group = if (param == null) {
+      val tmp1_group = <block>{
         sourceInformationMarkerStart(%composer, <>, "CC(remember):Test.kt#9igjgp")
         val tmp0_group = %composer.cache(false) {
           ""
         }
         sourceInformationMarkerEnd(%composer)
         tmp0_group
-      } else {
-        null
       }
       %composer.endReplaceGroup()
       tmp1_group
+    } else {
+      %composer.startReplaceGroup(<>)
+      %composer.endReplaceGroup()
+      null
     }, %composer, 0)
     if (isTraceInProgress()) {
       traceEventEnd()
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testRememberInConditionalCallArgument[useFir = true].txt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testRememberInConditionalCallArgument[useFir = true].txt
index 936c9e6fcf02d..bf5b7c0354eaa 100644
--- a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testRememberInConditionalCallArgument[useFir = true].txt	
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testRememberInConditionalCallArgument[useFir = true].txt	
@@ -31,21 +31,23 @@ private fun Test(param: String?, %composer: Composer?, %changed: Int) {
     if (isTraceInProgress()) {
       traceEventStart(<>, %dirty, -1, <>)
     }
-    Test(<block>{
+    Test(if (param == null) {
       %composer.startReplaceGroup(<>)
       sourceInformation(%composer, "<rememb...>")
-      val tmp1_group = if (param == null) {
+      val tmp1_group = <block>{
         sourceInformationMarkerStart(%composer, <>, "CC(remember):Test.kt#9igjgp")
         val tmp0_group = %composer.cache(false) {
           ""
         }
         sourceInformationMarkerEnd(%composer)
         tmp0_group
-      } else {
-        null
       }
       %composer.endReplaceGroup()
       tmp1_group
+    } else {
+      %composer.startReplaceGroup(<>)
+      %composer.endReplaceGroup()
+      null
     }, %composer, 0)
     if (isTraceInProgress()) {
       traceEventEnd()
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testRememberInNestedConditionalCallArgument[useFir = false].txt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testRememberInNestedConditionalCallArgument[useFir = false].txt
index f07fc50bed6ae..08063378c76e5 100644
--- a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testRememberInNestedConditionalCallArgument[useFir = false].txt	
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testRememberInNestedConditionalCallArgument[useFir = false].txt	
@@ -31,31 +31,35 @@ private fun Test(param: String?, %composer: Composer?, %changed: Int): String? {
   if (isTraceInProgress()) {
     traceEventStart(<>, %changed, -1, <>)
   }
-  val tmp0 = Test(<block>{
+  val tmp0 = Test(if (param == null) {
     %composer.startReplaceGroup(<>)
     sourceInformation(%composer, "<Test(>")
-    val tmp3_group = if (param == null) {
-      Test(<block>{
+    val tmp3_group = <block>{
+      Test(if (param == null) {
         %composer.startReplaceGroup(<>)
         sourceInformation(%composer, "<rememb...>")
-        val tmp2_group = if (param == null) {
+        val tmp2_group = <block>{
           sourceInformationMarkerStart(%composer, <>, "CC(remember):Test.kt#9igjgp")
           val tmp1_group = %composer.cache(false) {
             ""
           }
           sourceInformationMarkerEnd(%composer)
           tmp1_group
-        } else {
-          null
         }
         %composer.endReplaceGroup()
         tmp2_group
+      } else {
+        %composer.startReplaceGroup(<>)
+        %composer.endReplaceGroup()
+        null
       }, %composer, 0)
-    } else {
-      null
     }
     %composer.endReplaceGroup()
     tmp3_group
+  } else {
+    %composer.startReplaceGroup(<>)
+    %composer.endReplaceGroup()
+    null
   }, %composer, 0)
   if (isTraceInProgress()) {
     traceEventEnd()
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testRememberInNestedConditionalCallArgument[useFir = true].txt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testRememberInNestedConditionalCallArgument[useFir = true].txt
index f07fc50bed6ae..08063378c76e5 100644
--- a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testRememberInNestedConditionalCallArgument[useFir = true].txt	
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testRememberInNestedConditionalCallArgument[useFir = true].txt	
@@ -31,31 +31,35 @@ private fun Test(param: String?, %composer: Composer?, %changed: Int): String? {
   if (isTraceInProgress()) {
     traceEventStart(<>, %changed, -1, <>)
   }
-  val tmp0 = Test(<block>{
+  val tmp0 = Test(if (param == null) {
     %composer.startReplaceGroup(<>)
     sourceInformation(%composer, "<Test(>")
-    val tmp3_group = if (param == null) {
-      Test(<block>{
+    val tmp3_group = <block>{
+      Test(if (param == null) {
         %composer.startReplaceGroup(<>)
         sourceInformation(%composer, "<rememb...>")
-        val tmp2_group = if (param == null) {
+        val tmp2_group = <block>{
           sourceInformationMarkerStart(%composer, <>, "CC(remember):Test.kt#9igjgp")
           val tmp1_group = %composer.cache(false) {
             ""
           }
           sourceInformationMarkerEnd(%composer)
           tmp1_group
-        } else {
-          null
         }
         %composer.endReplaceGroup()
         tmp2_group
+      } else {
+        %composer.startReplaceGroup(<>)
+        %composer.endReplaceGroup()
+        null
       }, %composer, 0)
-    } else {
-      null
     }
     %composer.endReplaceGroup()
     tmp3_group
+  } else {
+    %composer.startReplaceGroup(<>)
+    %composer.endReplaceGroup()
+    null
   }, %composer, 0)
   if (isTraceInProgress()) {
     traceEventEnd()
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testSafeCall[useFir = false].txt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testSafeCall[useFir = false].txt
index eed4dc6fb0b93..8f8c302c27d78 100644
--- a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testSafeCall[useFir = false].txt	
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testSafeCall[useFir = false].txt	
@@ -28,18 +28,19 @@ fun Example(x: Int?, %composer: Composer?, %changed: Int) {
     traceEventStart(<>, %changed, -1, <>)
   }
   val tmp0_safe_receiver = x
-  %composer.startReplaceGroup(<>)
-  sourceInformation(%composer, "<A()>")
-  val tmp0_group = when {
+  when {
     tmp0_safe_receiver == null -> {
+      %composer.startReplaceGroup(<>)
+      %composer.endReplaceGroup()
       null
     }
     else -> {
+      %composer.startReplaceGroup(<>)
+      sourceInformation(%composer, "<A()>")
       tmp0_safe_receiver.A(%composer, 0b1110 and %changed)
+      %composer.endReplaceGroup()
     }
   }
-  %composer.endReplaceGroup()
-  tmp0_group
   if (isTraceInProgress()) {
     traceEventEnd()
   }
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testSafeCall[useFir = true].txt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testSafeCall[useFir = true].txt
index eed4dc6fb0b93..8f8c302c27d78 100644
--- a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testSafeCall[useFir = true].txt	
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testSafeCall[useFir = true].txt	
@@ -28,18 +28,19 @@ fun Example(x: Int?, %composer: Composer?, %changed: Int) {
     traceEventStart(<>, %changed, -1, <>)
   }
   val tmp0_safe_receiver = x
-  %composer.startReplaceGroup(<>)
-  sourceInformation(%composer, "<A()>")
-  val tmp0_group = when {
+  when {
     tmp0_safe_receiver == null -> {
+      %composer.startReplaceGroup(<>)
+      %composer.endReplaceGroup()
       null
     }
     else -> {
+      %composer.startReplaceGroup(<>)
+      sourceInformation(%composer, "<A()>")
       tmp0_safe_receiver.A(%composer, 0b1110 and %changed)
+      %composer.endReplaceGroup()
     }
   }
-  %composer.endReplaceGroup()
-  tmp0_group
   if (isTraceInProgress()) {
     traceEventEnd()
   }
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testWhenWithCallsInConditionsAndCallAfter[useFir = false].txt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testWhenWithCallsInConditionsAndCallAfter[useFir = false].txt
index abdb1fbb33297..5eb44294ceb2d 100644
--- a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testWhenWithCallsInConditionsAndCallAfter[useFir = false].txt	
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testWhenWithCallsInConditionsAndCallAfter[useFir = false].txt	
@@ -32,8 +32,6 @@ fun Example(x: Int, %composer: Composer?, %changed: Int) {
   if (isTraceInProgress()) {
     traceEventStart(<>, %changed, -1, <>)
   }
-  %composer.startReplaceGroup(<>)
-  sourceInformation(%composer, "")
   when {
     %composer.startReplaceGroup(<>)
     sourceInformation(%composer, "<R(a)>")
@@ -49,11 +47,12 @@ fun Example(x: Int, %composer: Composer?, %changed: Int) {
     tmp1_group -> {
       NA()
     }
-    else -> {
+    %composer.startReplaceGroup(<>)
+    %composer.endReplaceGroup()
+    true -> {
       NA()
     }
   }
-  %composer.endReplaceGroup()
   A(%composer, 0)
   if (isTraceInProgress()) {
     traceEventEnd()
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testWhenWithCallsInConditionsAndCallAfter[useFir = true].txt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testWhenWithCallsInConditionsAndCallAfter[useFir = true].txt
index abdb1fbb33297..5eb44294ceb2d 100644
--- a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testWhenWithCallsInConditionsAndCallAfter[useFir = true].txt	
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testWhenWithCallsInConditionsAndCallAfter[useFir = true].txt	
@@ -32,8 +32,6 @@ fun Example(x: Int, %composer: Composer?, %changed: Int) {
   if (isTraceInProgress()) {
     traceEventStart(<>, %changed, -1, <>)
   }
-  %composer.startReplaceGroup(<>)
-  sourceInformation(%composer, "")
   when {
     %composer.startReplaceGroup(<>)
     sourceInformation(%composer, "<R(a)>")
@@ -49,11 +47,12 @@ fun Example(x: Int, %composer: Composer?, %changed: Int) {
     tmp1_group -> {
       NA()
     }
-    else -> {
+    %composer.startReplaceGroup(<>)
+    %composer.endReplaceGroup()
+    true -> {
       NA()
     }
   }
-  %composer.endReplaceGroup()
   A(%composer, 0)
   if (isTraceInProgress()) {
     traceEventEnd()
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testWhenWithCallsInConditions[useFir = false].txt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testWhenWithCallsInConditions[useFir = false].txt
index 13e1ae3997389..c76de7f642d5f 100644
--- a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testWhenWithCallsInConditions[useFir = false].txt	
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testWhenWithCallsInConditions[useFir = false].txt	
@@ -32,8 +32,6 @@ fun Example(x: Int, %composer: Composer?, %changed: Int) {
   if (isTraceInProgress()) {
     traceEventStart(<>, %changed, -1, <>)
   }
-  %composer.startReplaceGroup(<>)
-  sourceInformation(%composer, "")
   when {
     %composer.startReplaceGroup(<>)
     sourceInformation(%composer, "<R(a)>")
@@ -49,11 +47,12 @@ fun Example(x: Int, %composer: Composer?, %changed: Int) {
     tmp1_group -> {
       NA()
     }
-    else -> {
+    %composer.startReplaceGroup(<>)
+    %composer.endReplaceGroup()
+    true -> {
       NA()
     }
   }
-  %composer.endReplaceGroup()
   if (isTraceInProgress()) {
     traceEventEnd()
   }
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testWhenWithCallsInConditions[useFir = true].txt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testWhenWithCallsInConditions[useFir = true].txt
index 13e1ae3997389..c76de7f642d5f 100644
--- a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testWhenWithCallsInConditions[useFir = true].txt	
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testWhenWithCallsInConditions[useFir = true].txt	
@@ -32,8 +32,6 @@ fun Example(x: Int, %composer: Composer?, %changed: Int) {
   if (isTraceInProgress()) {
     traceEventStart(<>, %changed, -1, <>)
   }
-  %composer.startReplaceGroup(<>)
-  sourceInformation(%composer, "")
   when {
     %composer.startReplaceGroup(<>)
     sourceInformation(%composer, "<R(a)>")
@@ -49,11 +47,12 @@ fun Example(x: Int, %composer: Composer?, %changed: Int) {
     tmp1_group -> {
       NA()
     }
-    else -> {
+    %composer.startReplaceGroup(<>)
+    %composer.endReplaceGroup()
+    true -> {
       NA()
     }
   }
-  %composer.endReplaceGroup()
   if (isTraceInProgress()) {
     traceEventEnd()
   }
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testWhileInsideIfAndCallAfter[useFir = false].txt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testWhileInsideIfAndCallAfter[useFir = false].txt
index ff05c971ba04c..0e12fefe3ed49 100644
--- a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testWhileInsideIfAndCallAfter[useFir = false].txt	
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testWhileInsideIfAndCallAfter[useFir = false].txt	
@@ -29,9 +29,9 @@ fun Example(x: Int, %composer: Composer?, %changed: Int) {
   if (isTraceInProgress()) {
     traceEventStart(<>, %changed, -1, <>)
   }
-  %composer.startReplaceGroup(<>)
-  sourceInformation(%composer, "<A()>")
   if (x > 0) {
+    %composer.startReplaceGroup(<>)
+    sourceInformation(%composer, "<A()>")
     %composer.startReplaceGroup(<>)
     sourceInformation(%composer, "*<A()>")
     while (x > 0) {
@@ -39,8 +39,11 @@ fun Example(x: Int, %composer: Composer?, %changed: Int) {
     }
     %composer.endReplaceGroup()
     A(%composer, 0)
+    %composer.endReplaceGroup()
+  } else {
+    %composer.startReplaceGroup(<>)
+    %composer.endReplaceGroup()
   }
-  %composer.endReplaceGroup()
   if (isTraceInProgress()) {
     traceEventEnd()
   }
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testWhileInsideIfAndCallAfter[useFir = true].txt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testWhileInsideIfAndCallAfter[useFir = true].txt
index ff05c971ba04c..0e12fefe3ed49 100644
--- a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testWhileInsideIfAndCallAfter[useFir = true].txt	
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testWhileInsideIfAndCallAfter[useFir = true].txt	
@@ -29,9 +29,9 @@ fun Example(x: Int, %composer: Composer?, %changed: Int) {
   if (isTraceInProgress()) {
     traceEventStart(<>, %changed, -1, <>)
   }
-  %composer.startReplaceGroup(<>)
-  sourceInformation(%composer, "<A()>")
   if (x > 0) {
+    %composer.startReplaceGroup(<>)
+    sourceInformation(%composer, "<A()>")
     %composer.startReplaceGroup(<>)
     sourceInformation(%composer, "*<A()>")
     while (x > 0) {
@@ -39,8 +39,11 @@ fun Example(x: Int, %composer: Composer?, %changed: Int) {
     }
     %composer.endReplaceGroup()
     A(%composer, 0)
+    %composer.endReplaceGroup()
+  } else {
+    %composer.startReplaceGroup(<>)
+    %composer.endReplaceGroup()
   }
-  %composer.endReplaceGroup()
   if (isTraceInProgress()) {
     traceEventEnd()
   }
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testWhileInsideIfAndCallBefore[useFir = false].txt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testWhileInsideIfAndCallBefore[useFir = false].txt
index 7f98ae59985dc..8d50a740aff12 100644
--- a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testWhileInsideIfAndCallBefore[useFir = false].txt	
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testWhileInsideIfAndCallBefore[useFir = false].txt	
@@ -29,15 +29,18 @@ fun Example(x: Int, %composer: Composer?, %changed: Int) {
   if (isTraceInProgress()) {
     traceEventStart(<>, %changed, -1, <>)
   }
-  %composer.startReplaceGroup(<>)
-  sourceInformation(%composer, "<A()>,*<A()>")
   if (x > 0) {
+    %composer.startReplaceGroup(<>)
+    sourceInformation(%composer, "<A()>,*<A()>")
     A(%composer, 0)
     while (x > 0) {
       A(%composer, 0)
     }
+    %composer.endReplaceGroup()
+  } else {
+    %composer.startReplaceGroup(<>)
+    %composer.endReplaceGroup()
   }
-  %composer.endReplaceGroup()
   if (isTraceInProgress()) {
     traceEventEnd()
   }
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testWhileInsideIfAndCallBefore[useFir = true].txt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testWhileInsideIfAndCallBefore[useFir = true].txt
index 7f98ae59985dc..8d50a740aff12 100644
--- a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testWhileInsideIfAndCallBefore[useFir = true].txt	
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testWhileInsideIfAndCallBefore[useFir = true].txt	
@@ -29,15 +29,18 @@ fun Example(x: Int, %composer: Composer?, %changed: Int) {
   if (isTraceInProgress()) {
     traceEventStart(<>, %changed, -1, <>)
   }
-  %composer.startReplaceGroup(<>)
-  sourceInformation(%composer, "<A()>,*<A()>")
   if (x > 0) {
+    %composer.startReplaceGroup(<>)
+    sourceInformation(%composer, "<A()>,*<A()>")
     A(%composer, 0)
     while (x > 0) {
       A(%composer, 0)
     }
+    %composer.endReplaceGroup()
+  } else {
+    %composer.startReplaceGroup(<>)
+    %composer.endReplaceGroup()
   }
-  %composer.endReplaceGroup()
   if (isTraceInProgress()) {
     traceEventEnd()
   }
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testWhileInsideIf[useFir = false].txt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testWhileInsideIf[useFir = false].txt
index 4cf7057d48adb..27dbf17a8ba9e 100644
--- a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testWhileInsideIf[useFir = false].txt	
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testWhileInsideIf[useFir = false].txt	
@@ -28,14 +28,17 @@ fun Example(x: Int, %composer: Composer?, %changed: Int) {
   if (isTraceInProgress()) {
     traceEventStart(<>, %changed, -1, <>)
   }
-  %composer.startReplaceGroup(<>)
-  sourceInformation(%composer, "*<A()>")
   if (x > 0) {
+    %composer.startReplaceGroup(<>)
+    sourceInformation(%composer, "*<A()>")
     while (x > 0) {
       A(%composer, 0)
     }
+    %composer.endReplaceGroup()
+  } else {
+    %composer.startReplaceGroup(<>)
+    %composer.endReplaceGroup()
   }
-  %composer.endReplaceGroup()
   if (isTraceInProgress()) {
     traceEventEnd()
   }
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testWhileInsideIf[useFir = true].txt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testWhileInsideIf[useFir = true].txt
index 4cf7057d48adb..27dbf17a8ba9e 100644
--- a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testWhileInsideIf[useFir = true].txt	
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testWhileInsideIf[useFir = true].txt	
@@ -28,14 +28,17 @@ fun Example(x: Int, %composer: Composer?, %changed: Int) {
   if (isTraceInProgress()) {
     traceEventStart(<>, %changed, -1, <>)
   }
-  %composer.startReplaceGroup(<>)
-  sourceInformation(%composer, "*<A()>")
   if (x > 0) {
+    %composer.startReplaceGroup(<>)
+    sourceInformation(%composer, "*<A()>")
     while (x > 0) {
       A(%composer, 0)
     }
+    %composer.endReplaceGroup()
+  } else {
+    %composer.startReplaceGroup(<>)
+    %composer.endReplaceGroup()
   }
-  %composer.endReplaceGroup()
   if (isTraceInProgress()) {
     traceEventEnd()
   }
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/test_CM1_CCM1_RetFun[useFir = false].txt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/test_CM1_CCM1_RetFun[useFir = false].txt
index f895e8b82b821..f30afffba0788 100644
--- a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/test_CM1_CCM1_RetFun[useFir = false].txt	
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/test_CM1_CCM1_RetFun[useFir = false].txt	
@@ -43,9 +43,9 @@ fun test_CM1_CCM1_RetFun(condition: Boolean, %composer: Composer?, %changed: Int
       %composer.startReplaceGroup(<>)
       sourceInformation(%composer, "C<Text("...>,<Text("...>:Test.kt")
       Text("M1 - begin", %composer, 0b0110)
-      %composer.startReplaceGroup(<>)
-      sourceInformation(%composer, "<Text("...>,<M1>")
       if (condition) {
+        %composer.startReplaceGroup(<>)
+        sourceInformation(%composer, "<Text("...>,<M1>")
         Text("if - begin", %composer, 0b0110)
         M1({ %composer: Composer?, %changed: Int ->
           %composer.startReplaceGroup(<>)
@@ -61,8 +61,11 @@ fun test_CM1_CCM1_RetFun(condition: Boolean, %composer: Composer?, %changed: Int
           return
           %composer.endReplaceGroup()
         }, %composer, 0)
+        %composer.endReplaceGroup()
+      } else {
+        %composer.startReplaceGroup(<>)
+        %composer.endReplaceGroup()
       }
-      %composer.endReplaceGroup()
       Text("M1 - end", %composer, 0b0110)
       %composer.endReplaceGroup()
     }, %composer, 0)
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/test_CM1_CCM1_RetFun[useFir = true].txt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/test_CM1_CCM1_RetFun[useFir = true].txt
index f895e8b82b821..f30afffba0788 100644
--- a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/test_CM1_CCM1_RetFun[useFir = true].txt	
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/test_CM1_CCM1_RetFun[useFir = true].txt	
@@ -43,9 +43,9 @@ fun test_CM1_CCM1_RetFun(condition: Boolean, %composer: Composer?, %changed: Int
       %composer.startReplaceGroup(<>)
       sourceInformation(%composer, "C<Text("...>,<Text("...>:Test.kt")
       Text("M1 - begin", %composer, 0b0110)
-      %composer.startReplaceGroup(<>)
-      sourceInformation(%composer, "<Text("...>,<M1>")
       if (condition) {
+        %composer.startReplaceGroup(<>)
+        sourceInformation(%composer, "<Text("...>,<M1>")
         Text("if - begin", %composer, 0b0110)
         M1({ %composer: Composer?, %changed: Int ->
           %composer.startReplaceGroup(<>)
@@ -61,8 +61,11 @@ fun test_CM1_CCM1_RetFun(condition: Boolean, %composer: Composer?, %changed: Int
           return
           %composer.endReplaceGroup()
         }, %composer, 0)
+        %composer.endReplaceGroup()
+      } else {
+        %composer.startReplaceGroup(<>)
+        %composer.endReplaceGroup()
       }
-      %composer.endReplaceGroup()
       Text("M1 - end", %composer, 0b0110)
       %composer.endReplaceGroup()
     }, %composer, 0)
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.FunctionBodySkippingTransformTests/testFunInterfaces2[useFir = false].txt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.FunctionBodySkippingTransformTests/testFunInterfaces2[useFir = false].txt
index 6de1c6dafe230..67206bc64b7f7 100644
--- a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.FunctionBodySkippingTransformTests/testFunInterfaces2[useFir = false].txt	
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.FunctionBodySkippingTransformTests/testFunInterfaces2[useFir = false].txt	
@@ -64,13 +64,19 @@ fun Test(%composer: Composer?, %changed: Int) {
         @Composable
         override fun getColor(%composer: Composer?, %changed: Int): Color {
           %composer.startReplaceGroup(<>)
-          sourceInformation(%composer, "C(getColor)<condit...>:Test.kt")
+          sourceInformation(%composer, "C(getColor):Test.kt")
           if (isTraceInProgress()) {
             traceEventStart(<>, %changed, -1, <>)
           }
-          val tmp0 = if (condition(%composer, 0)) {
+          val tmp0 = if (%composer.startReplaceGroup(<>)
+          sourceInformation(%composer, "<condit...>")
+          val tmp1_group = condition(%composer, 0)
+          %composer.endReplaceGroup()
+          tmp1_group) {
             Companion.Red
-          } else {
+          } else if (%composer.startReplaceGroup(<>)
+          %composer.endReplaceGroup()
+          true) {
             Companion.Blue
           }
           if (isTraceInProgress()) {
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.FunctionBodySkippingTransformTests/testFunInterfaces2[useFir = true].txt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.FunctionBodySkippingTransformTests/testFunInterfaces2[useFir = true].txt
index 6de1c6dafe230..67206bc64b7f7 100644
--- a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.FunctionBodySkippingTransformTests/testFunInterfaces2[useFir = true].txt	
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.FunctionBodySkippingTransformTests/testFunInterfaces2[useFir = true].txt	
@@ -64,13 +64,19 @@ fun Test(%composer: Composer?, %changed: Int) {
         @Composable
         override fun getColor(%composer: Composer?, %changed: Int): Color {
           %composer.startReplaceGroup(<>)
-          sourceInformation(%composer, "C(getColor)<condit...>:Test.kt")
+          sourceInformation(%composer, "C(getColor):Test.kt")
           if (isTraceInProgress()) {
             traceEventStart(<>, %changed, -1, <>)
           }
-          val tmp0 = if (condition(%composer, 0)) {
+          val tmp0 = if (%composer.startReplaceGroup(<>)
+          sourceInformation(%composer, "<condit...>")
+          val tmp1_group = condition(%composer, 0)
+          %composer.endReplaceGroup()
+          tmp1_group) {
             Companion.Red
-          } else {
+          } else if (%composer.startReplaceGroup(<>)
+          %composer.endReplaceGroup()
+          true) {
             Companion.Blue
           }
           if (isTraceInProgress()) {
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.FunctionBodySkippingTransformTests/testSiblingIfsWithoutElseHaveUniqueKeys[useFir = false].txt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.FunctionBodySkippingTransformTests/testSiblingIfsWithoutElseHaveUniqueKeys[useFir = false].txt
index 0c0f57147149b..195eb4cd6c785 100644
--- a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.FunctionBodySkippingTransformTests/testSiblingIfsWithoutElseHaveUniqueKeys[useFir = false].txt	
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.FunctionBodySkippingTransformTests/testSiblingIfsWithoutElseHaveUniqueKeys[useFir = false].txt	
@@ -24,7 +24,7 @@ fun Test(cond: Boolean) {
 @Composable
 fun Test(cond: Boolean, %composer: Composer?, %changed: Int) {
   %composer = %composer.startRestartGroup(<>)
-  sourceInformation(%composer, "C(Test)<B()>:Test.kt")
+  sourceInformation(%composer, "C(Test):Test.kt")
   val %dirty = %changed
   if (%changed and 0b0110 == 0) {
     %dirty = %dirty or if (%composer.changed(cond)) 0b0100 else 0b0010
@@ -33,14 +33,23 @@ fun Test(cond: Boolean, %composer: Composer?, %changed: Int) {
     if (isTraceInProgress()) {
       traceEventStart(<>, %dirty, -1, <>)
     }
-    %composer.startReplaceGroup(<>)
-    sourceInformation(%composer, "<A()>")
     if (cond) {
+      %composer.startReplaceGroup(<>)
+      sourceInformation(%composer, "<A()>")
       A(%composer, 0)
+      %composer.endReplaceGroup()
+    } else {
+      %composer.startReplaceGroup(<>)
+      %composer.endReplaceGroup()
     }
-    %composer.endReplaceGroup()
     if (cond) {
+      %composer.startReplaceGroup(<>)
+      sourceInformation(%composer, "<B()>")
       B(%composer, 0)
+      %composer.endReplaceGroup()
+    } else {
+      %composer.startReplaceGroup(<>)
+      %composer.endReplaceGroup()
     }
     if (isTraceInProgress()) {
       traceEventEnd()
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.FunctionBodySkippingTransformTests/testSiblingIfsWithoutElseHaveUniqueKeys[useFir = true].txt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.FunctionBodySkippingTransformTests/testSiblingIfsWithoutElseHaveUniqueKeys[useFir = true].txt
index 0c0f57147149b..195eb4cd6c785 100644
--- a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.FunctionBodySkippingTransformTests/testSiblingIfsWithoutElseHaveUniqueKeys[useFir = true].txt	
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.FunctionBodySkippingTransformTests/testSiblingIfsWithoutElseHaveUniqueKeys[useFir = true].txt	
@@ -24,7 +24,7 @@ fun Test(cond: Boolean) {
 @Composable
 fun Test(cond: Boolean, %composer: Composer?, %changed: Int) {
   %composer = %composer.startRestartGroup(<>)
-  sourceInformation(%composer, "C(Test)<B()>:Test.kt")
+  sourceInformation(%composer, "C(Test):Test.kt")
   val %dirty = %changed
   if (%changed and 0b0110 == 0) {
     %dirty = %dirty or if (%composer.changed(cond)) 0b0100 else 0b0010
@@ -33,14 +33,23 @@ fun Test(cond: Boolean, %composer: Composer?, %changed: Int) {
     if (isTraceInProgress()) {
       traceEventStart(<>, %dirty, -1, <>)
     }
-    %composer.startReplaceGroup(<>)
-    sourceInformation(%composer, "<A()>")
     if (cond) {
+      %composer.startReplaceGroup(<>)
+      sourceInformation(%composer, "<A()>")
       A(%composer, 0)
+      %composer.endReplaceGroup()
+    } else {
+      %composer.startReplaceGroup(<>)
+      %composer.endReplaceGroup()
     }
-    %composer.endReplaceGroup()
     if (cond) {
+      %composer.startReplaceGroup(<>)
+      sourceInformation(%composer, "<B()>")
       B(%composer, 0)
+      %composer.endReplaceGroup()
+    } else {
+      %composer.startReplaceGroup(<>)
+      %composer.endReplaceGroup()
     }
     if (isTraceInProgress()) {
       traceEventEnd()
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.FunctionBodySkippingTransformTests/test_groupAroundIfComposeCallInIfConditionWithShortCircuit[useFir = false].txt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.FunctionBodySkippingTransformTests/test_groupAroundIfComposeCallInIfConditionWithShortCircuit[useFir = false].txt
index d57b4db7e29c2..e846eb825aa44 100644
--- a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.FunctionBodySkippingTransformTests/test_groupAroundIfComposeCallInIfConditionWithShortCircuit[useFir = false].txt	
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.FunctionBodySkippingTransformTests/test_groupAroundIfComposeCallInIfConditionWithShortCircuit[useFir = false].txt	
@@ -21,9 +21,9 @@ fun Test(%composer: Composer?, %changed: Int) {
     if (isTraceInProgress()) {
       traceEventStart(<>, %changed, -1, <>)
     }
-    ReceiveValue(if (%composer.startReplaceGroup(<>)
+    ReceiveValue(if (state && %composer.startReplaceGroup(<>)
     sourceInformation(%composer, "<getCon...>")
-    val tmp0_group = state && getCondition(%composer, 0)
+    val tmp0_group = getCondition(%composer, 0)
     %composer.endReplaceGroup()
     tmp0_group) 0 else 1, %composer, 0)
     if (isTraceInProgress()) {
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.FunctionBodySkippingTransformTests/test_groupAroundIfComposeCallInIfConditionWithShortCircuit[useFir = true].txt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.FunctionBodySkippingTransformTests/test_groupAroundIfComposeCallInIfConditionWithShortCircuit[useFir = true].txt
index d57b4db7e29c2..e846eb825aa44 100644
--- a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.FunctionBodySkippingTransformTests/test_groupAroundIfComposeCallInIfConditionWithShortCircuit[useFir = true].txt	
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.FunctionBodySkippingTransformTests/test_groupAroundIfComposeCallInIfConditionWithShortCircuit[useFir = true].txt	
@@ -21,9 +21,9 @@ fun Test(%composer: Composer?, %changed: Int) {
     if (isTraceInProgress()) {
       traceEventStart(<>, %changed, -1, <>)
     }
-    ReceiveValue(if (%composer.startReplaceGroup(<>)
+    ReceiveValue(if (state && %composer.startReplaceGroup(<>)
     sourceInformation(%composer, "<getCon...>")
-    val tmp0_group = state && getCondition(%composer, 0)
+    val tmp0_group = getCondition(%composer, 0)
     %composer.endReplaceGroup()
     tmp0_group) 0 else 1, %composer, 0)
     if (isTraceInProgress()) {
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.FunctionBodySkippingTransformTestsNoSource/testIfStatementGroups[useFir = false].txt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.FunctionBodySkippingTransformTestsNoSource/testIfStatementGroups[useFir = false].txt
new file mode 100644
index 0000000000000..29f9b813eb1a1
--- /dev/null
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.FunctionBodySkippingTransformTestsNoSource/testIfStatementGroups[useFir = false].txt	
@@ -0,0 +1,75 @@
+//
+// Source
+// ------------------------------------------
+
+import androidx.compose.runtime.*
+
+@Composable
+fun Test(level: Int) {
+    Wrap {
+        if (level > 0) {
+            used(remember { "Before" })
+            Wrap {
+                used(remember { "Middle" })
+            }
+            used(remember { "End" })
+        }
+    }
+}
+
+//
+// Transformed IR
+// ------------------------------------------
+
+@Composable
+fun Test(level: Int, %composer: Composer?, %changed: Int) {
+  %composer = %composer.startRestartGroup(<>)
+  val %dirty = %changed
+  if (%changed and 0b0110 == 0) {
+    %dirty = %dirty or if (%composer.changed(level)) 0b0100 else 0b0010
+  }
+  if (%dirty and 0b0011 != 0b0010 || !%composer.skipping) {
+    Wrap(rememberComposableLambda(<>, true, { %composer: Composer?, %changed: Int ->
+      if (%changed and 0b0011 != 0b0010 || !%composer.skipping) {
+        if (level > 0) {
+          %composer.startReplaceGroup(<>)
+          used(<block>{
+            %composer.cache(false) {
+              "Before"
+            }
+          })
+          Wrap(ComposableSingletons%TestKt.lambda-1, %composer, 0b0110)
+          used(<block>{
+            %composer.cache(false) {
+              "End"
+            }
+          })
+          %composer.endReplaceGroup()
+        } else {
+          %composer.startReplaceGroup(<>)
+          %composer.endReplaceGroup()
+        }
+      } else {
+        %composer.skipToGroupEnd()
+      }
+    }, %composer, 0b00110110), %composer, 0b0110)
+  } else {
+    %composer.skipToGroupEnd()
+  }
+  %composer.endRestartGroup()?.updateScope { %composer: Composer?, %force: Int ->
+    Test(level, %composer, updateChangedFlags(%changed or 0b0001))
+  }
+}
+internal object ComposableSingletons%TestKt {
+  val lambda-1: Function2<Composer, Int, Unit> = composableLambdaInstance(<>, false) { %composer: Composer?, %changed: Int ->
+    if (%changed and 0b0011 != 0b0010 || !%composer.skipping) {
+      used(<block>{
+        %composer.cache(false) {
+          "Middle"
+        }
+      })
+    } else {
+      %composer.skipToGroupEnd()
+    }
+  }
+}
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.FunctionBodySkippingTransformTestsNoSource/testIfStatementGroups[useFir = true].txt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.FunctionBodySkippingTransformTestsNoSource/testIfStatementGroups[useFir = true].txt
new file mode 100644
index 0000000000000..29f9b813eb1a1
--- /dev/null
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.FunctionBodySkippingTransformTestsNoSource/testIfStatementGroups[useFir = true].txt	
@@ -0,0 +1,75 @@
+//
+// Source
+// ------------------------------------------
+
+import androidx.compose.runtime.*
+
+@Composable
+fun Test(level: Int) {
+    Wrap {
+        if (level > 0) {
+            used(remember { "Before" })
+            Wrap {
+                used(remember { "Middle" })
+            }
+            used(remember { "End" })
+        }
+    }
+}
+
+//
+// Transformed IR
+// ------------------------------------------
+
+@Composable
+fun Test(level: Int, %composer: Composer?, %changed: Int) {
+  %composer = %composer.startRestartGroup(<>)
+  val %dirty = %changed
+  if (%changed and 0b0110 == 0) {
+    %dirty = %dirty or if (%composer.changed(level)) 0b0100 else 0b0010
+  }
+  if (%dirty and 0b0011 != 0b0010 || !%composer.skipping) {
+    Wrap(rememberComposableLambda(<>, true, { %composer: Composer?, %changed: Int ->
+      if (%changed and 0b0011 != 0b0010 || !%composer.skipping) {
+        if (level > 0) {
+          %composer.startReplaceGroup(<>)
+          used(<block>{
+            %composer.cache(false) {
+              "Before"
+            }
+          })
+          Wrap(ComposableSingletons%TestKt.lambda-1, %composer, 0b0110)
+          used(<block>{
+            %composer.cache(false) {
+              "End"
+            }
+          })
+          %composer.endReplaceGroup()
+        } else {
+          %composer.startReplaceGroup(<>)
+          %composer.endReplaceGroup()
+        }
+      } else {
+        %composer.skipToGroupEnd()
+      }
+    }, %composer, 0b00110110), %composer, 0b0110)
+  } else {
+    %composer.skipToGroupEnd()
+  }
+  %composer.endRestartGroup()?.updateScope { %composer: Composer?, %force: Int ->
+    Test(level, %composer, updateChangedFlags(%changed or 0b0001))
+  }
+}
+internal object ComposableSingletons%TestKt {
+  val lambda-1: Function2<Composer, Int, Unit> = composableLambdaInstance(<>, false) { %composer: Composer?, %changed: Int ->
+    if (%changed and 0b0011 != 0b0010 || !%composer.skipping) {
+      used(<block>{
+        %composer.cache(false) {
+          "Middle"
+        }
+      })
+    } else {
+      %composer.skipToGroupEnd()
+    }
+  }
+}
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.LambdaMemoizationTransformTests/testCrossinlineLambda[useFir = false].txt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.LambdaMemoizationTransformTests/testCrossinlineLambda[useFir = false].txt
index 66c51caa6494e..e2322f0e4e4b0 100644
--- a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.LambdaMemoizationTransformTests/testCrossinlineLambda[useFir = false].txt	
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.LambdaMemoizationTransformTests/testCrossinlineLambda[useFir = false].txt	
@@ -18,7 +18,7 @@ fun Test() {
 @Composable
 fun Test(%composer: Composer?, %changed: Int) {
   %composer = %composer.startRestartGroup(<>)
-  sourceInformation(%composer, "C(Test)*<it()>:Test.kt")
+  sourceInformation(%composer, "C(Test):Test.kt")
   if (%changed != 0 || !%composer.skipping) {
     if (isTraceInProgress()) {
       traceEventStart(<>, %changed, -1, <>)
@@ -41,17 +41,21 @@ fun Test(%composer: Composer?, %changed: Int) {
       }
     }
     val tmp0_safe_receiver = lambda
-    val tmp0_group = when {
+    when {
       tmp0_safe_receiver == null -> {
+        %composer.startReplaceGroup(<>)
+        %composer.endReplaceGroup()
         null
       }
       else -> {
+        %composer.startReplaceGroup(<>)
+        sourceInformation(%composer, "*<it()>")
         tmp0_safe_receiver.let { it: Function2<Composer, Int, Unit> ->
           it(%composer, 0)
         }
+        %composer.endReplaceGroup()
       }
     }
-    tmp0_group
     if (isTraceInProgress()) {
       traceEventEnd()
     }
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.LambdaMemoizationTransformTests/testCrossinlineLambda[useFir = true].txt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.LambdaMemoizationTransformTests/testCrossinlineLambda[useFir = true].txt
index 66c51caa6494e..e2322f0e4e4b0 100644
--- a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.LambdaMemoizationTransformTests/testCrossinlineLambda[useFir = true].txt	
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.LambdaMemoizationTransformTests/testCrossinlineLambda[useFir = true].txt	
@@ -18,7 +18,7 @@ fun Test() {
 @Composable
 fun Test(%composer: Composer?, %changed: Int) {
   %composer = %composer.startRestartGroup(<>)
-  sourceInformation(%composer, "C(Test)*<it()>:Test.kt")
+  sourceInformation(%composer, "C(Test):Test.kt")
   if (%changed != 0 || !%composer.skipping) {
     if (isTraceInProgress()) {
       traceEventStart(<>, %changed, -1, <>)
@@ -41,17 +41,21 @@ fun Test(%composer: Composer?, %changed: Int) {
       }
     }
     val tmp0_safe_receiver = lambda
-    val tmp0_group = when {
+    when {
       tmp0_safe_receiver == null -> {
+        %composer.startReplaceGroup(<>)
+        %composer.endReplaceGroup()
         null
       }
       else -> {
+        %composer.startReplaceGroup(<>)
+        sourceInformation(%composer, "*<it()>")
         tmp0_safe_receiver.let { it: Function2<Composer, Int, Unit> ->
           it(%composer, 0)
         }
+        %composer.endReplaceGroup()
       }
     }
-    tmp0_group
     if (isTraceInProgress()) {
       traceEventEnd()
     }
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.LambdaMemoizationTransformTests/testMemoizingFunctionInIf[useFir = false].txt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.LambdaMemoizationTransformTests/testMemoizingFunctionInIf[useFir = false].txt
index 853879194a9a0..650550a4192d0 100644
--- a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.LambdaMemoizationTransformTests/testMemoizingFunctionInIf[useFir = false].txt	
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.LambdaMemoizationTransformTests/testMemoizingFunctionInIf[useFir = false].txt	
@@ -31,10 +31,10 @@ fun Something(param: Function0<String>?, %composer: Composer?, %changed: Int) {
     if (isTraceInProgress()) {
       traceEventStart(<>, %dirty, -1, <>)
     }
-    Something(<block>{
+    Something(if (param != null) {
       %composer.startReplaceGroup(<>)
       sourceInformation(%composer, "<{>")
-      val tmp1_group = if (param != null) {
+      val tmp1_group = <block>{
         sourceInformationMarkerStart(%composer, <>, "CC(remember):Test.kt#9igjgp")
         val tmp0_group = %composer.cache(%dirty and 0b1110 == 0b0100) {
           {
@@ -43,11 +43,13 @@ fun Something(param: Function0<String>?, %composer: Composer?, %changed: Int) {
         }
         sourceInformationMarkerEnd(%composer)
         tmp0_group
-      } else {
-        null
       }
       %composer.endReplaceGroup()
       tmp1_group
+    } else {
+      %composer.startReplaceGroup(<>)
+      %composer.endReplaceGroup()
+      null
     }, %composer, 0)
     if (isTraceInProgress()) {
       traceEventEnd()
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.LambdaMemoizationTransformTests/testMemoizingFunctionInIf[useFir = true].txt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.LambdaMemoizationTransformTests/testMemoizingFunctionInIf[useFir = true].txt
index 853879194a9a0..650550a4192d0 100644
--- a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.LambdaMemoizationTransformTests/testMemoizingFunctionInIf[useFir = true].txt	
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.LambdaMemoizationTransformTests/testMemoizingFunctionInIf[useFir = true].txt	
@@ -31,10 +31,10 @@ fun Something(param: Function0<String>?, %composer: Composer?, %changed: Int) {
     if (isTraceInProgress()) {
       traceEventStart(<>, %dirty, -1, <>)
     }
-    Something(<block>{
+    Something(if (param != null) {
       %composer.startReplaceGroup(<>)
       sourceInformation(%composer, "<{>")
-      val tmp1_group = if (param != null) {
+      val tmp1_group = <block>{
         sourceInformationMarkerStart(%composer, <>, "CC(remember):Test.kt#9igjgp")
         val tmp0_group = %composer.cache(%dirty and 0b1110 == 0b0100) {
           {
@@ -43,11 +43,13 @@ fun Something(param: Function0<String>?, %composer: Composer?, %changed: Int) {
         }
         sourceInformationMarkerEnd(%composer)
         tmp0_group
-      } else {
-        null
       }
       %composer.endReplaceGroup()
       tmp1_group
+    } else {
+      %composer.startReplaceGroup(<>)
+      %composer.endReplaceGroup()
+      null
     }, %composer, 0)
     if (isTraceInProgress()) {
       traceEventEnd()
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.RememberIntrinsicTransformTests/testElidedRememberInsideIfDeoptsRememberAfterIf[useFir = false].txt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.RememberIntrinsicTransformTests/testElidedRememberInsideIfDeoptsRememberAfterIf[useFir = false].txt
index 4df31fbec7236..23062f624f303 100644
--- a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.RememberIntrinsicTransformTests/testElidedRememberInsideIfDeoptsRememberAfterIf[useFir = false].txt	
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.RememberIntrinsicTransformTests/testElidedRememberInsideIfDeoptsRememberAfterIf[useFir = false].txt	
@@ -26,21 +26,23 @@ fun app(x: Boolean, %composer: Composer?, %changed: Int) {
   if (isTraceInProgress()) {
     traceEventStart(<>, %changed, -1, <>)
   }
-  val a = <block>{
+  val a = if (x) {
     %composer.startReplaceGroup(<>)
     sourceInformation(%composer, "<rememb...>")
-    val tmp1_group = if (x) {
+    val tmp1_group = <block>{
       sourceInformationMarkerStart(%composer, <>, "CC(remember):Test.kt#9igjgp")
       val tmp0_group = %composer.cache(false) {
         1
       }
       sourceInformationMarkerEnd(%composer)
       tmp0_group
-    } else {
-      2
     }
     %composer.endReplaceGroup()
     tmp1_group
+  } else {
+    %composer.startReplaceGroup(<>)
+    %composer.endReplaceGroup()
+    2
   }
   val b = <block>{
     sourceInformationMarkerStart(%composer, <>, "CC(remember):Test.kt#9igjgp")
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.RememberIntrinsicTransformTests/testElidedRememberInsideIfDeoptsRememberAfterIf[useFir = true].txt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.RememberIntrinsicTransformTests/testElidedRememberInsideIfDeoptsRememberAfterIf[useFir = true].txt
index 4df31fbec7236..23062f624f303 100644
--- a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.RememberIntrinsicTransformTests/testElidedRememberInsideIfDeoptsRememberAfterIf[useFir = true].txt	
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.RememberIntrinsicTransformTests/testElidedRememberInsideIfDeoptsRememberAfterIf[useFir = true].txt	
@@ -26,21 +26,23 @@ fun app(x: Boolean, %composer: Composer?, %changed: Int) {
   if (isTraceInProgress()) {
     traceEventStart(<>, %changed, -1, <>)
   }
-  val a = <block>{
+  val a = if (x) {
     %composer.startReplaceGroup(<>)
     sourceInformation(%composer, "<rememb...>")
-    val tmp1_group = if (x) {
+    val tmp1_group = <block>{
       sourceInformationMarkerStart(%composer, <>, "CC(remember):Test.kt#9igjgp")
       val tmp0_group = %composer.cache(false) {
         1
       }
       sourceInformationMarkerEnd(%composer)
       tmp0_group
-    } else {
-      2
     }
     %composer.endReplaceGroup()
     tmp1_group
+  } else {
+    %composer.startReplaceGroup(<>)
+    %composer.endReplaceGroup()
+    2
   }
   val b = <block>{
     sourceInformationMarkerStart(%composer, <>, "CC(remember):Test.kt#9igjgp")
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.RememberIntrinsicTransformTests/testIntrinsicRememberOfLambdaInIfBlock[useFir = false].txt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.RememberIntrinsicTransformTests/testIntrinsicRememberOfLambdaInIfBlock[useFir = false].txt
index 6ab60d155be07..eb18a7e8be952 100644
--- a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.RememberIntrinsicTransformTests/testIntrinsicRememberOfLambdaInIfBlock[useFir = false].txt	
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.RememberIntrinsicTransformTests/testIntrinsicRememberOfLambdaInIfBlock[useFir = false].txt	
@@ -28,7 +28,7 @@ fun Test(a: Boolean, visible: Boolean, onDismiss: () -> Unit) {
 @Composable
 fun Test(a: Boolean, visible: Boolean, onDismiss: Function0<Unit>, %composer: Composer?, %changed: Int) {
   %composer = %composer.startRestartGroup(<>)
-  sourceInformation(%composer, "C(Test)P(!1,2)<someCo...>,<{>:Test.kt")
+  sourceInformation(%composer, "C(Test)P(!1,2):Test.kt")
   val %dirty = %changed
   if (%changed and 0b0110 == 0) {
     %dirty = %dirty or if (%composer.changed(a)) 0b0100 else 0b0010
@@ -44,11 +44,15 @@ fun Test(a: Boolean, visible: Boolean, onDismiss: Function0<Unit>, %composer: Co
       traceEventStart(<>, %dirty, -1, <>)
     }
     if (a) {
+      %composer.startReplaceGroup(<>)
+      sourceInformation(%composer, "<someCo...>")
       val a = someComposableValue(%composer, 0)
       used(a)
       val m = Modifier()
-      val dismissModifier = <block>{
-        val tmp1_group = if (visible) {
+      val dismissModifier = if (visible) {
+        %composer.startReplaceGroup(<>)
+        sourceInformation(%composer, "<{>")
+        val tmp1_group = <block>{
           m.pointerInput(Unit, <block>{
             sourceInformationMarkerStart(%composer, <>, "CC(remember):Test.kt#9igjgp")
             val tmp0_group = %composer.cache(%dirty and 0b001110000000 == 0b000100000000) {
@@ -61,12 +65,19 @@ fun Test(a: Boolean, visible: Boolean, onDismiss: Function0<Unit>, %composer: Co
             sourceInformationMarkerEnd(%composer)
             tmp0_group
           })
-        } else {
-          m
         }
+        %composer.endReplaceGroup()
         tmp1_group
+      } else {
+        %composer.startReplaceGroup(<>)
+        %composer.endReplaceGroup()
+        m
       }
       used(dismissModifier)
+      %composer.endReplaceGroup()
+    } else {
+      %composer.startReplaceGroup(<>)
+      %composer.endReplaceGroup()
     }
     if (isTraceInProgress()) {
       traceEventEnd()
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.RememberIntrinsicTransformTests/testIntrinsicRememberOfLambdaInIfBlock[useFir = true].txt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.RememberIntrinsicTransformTests/testIntrinsicRememberOfLambdaInIfBlock[useFir = true].txt
index 6ab60d155be07..eb18a7e8be952 100644
--- a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.RememberIntrinsicTransformTests/testIntrinsicRememberOfLambdaInIfBlock[useFir = true].txt	
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.RememberIntrinsicTransformTests/testIntrinsicRememberOfLambdaInIfBlock[useFir = true].txt	
@@ -28,7 +28,7 @@ fun Test(a: Boolean, visible: Boolean, onDismiss: () -> Unit) {
 @Composable
 fun Test(a: Boolean, visible: Boolean, onDismiss: Function0<Unit>, %composer: Composer?, %changed: Int) {
   %composer = %composer.startRestartGroup(<>)
-  sourceInformation(%composer, "C(Test)P(!1,2)<someCo...>,<{>:Test.kt")
+  sourceInformation(%composer, "C(Test)P(!1,2):Test.kt")
   val %dirty = %changed
   if (%changed and 0b0110 == 0) {
     %dirty = %dirty or if (%composer.changed(a)) 0b0100 else 0b0010
@@ -44,11 +44,15 @@ fun Test(a: Boolean, visible: Boolean, onDismiss: Function0<Unit>, %composer: Co
       traceEventStart(<>, %dirty, -1, <>)
     }
     if (a) {
+      %composer.startReplaceGroup(<>)
+      sourceInformation(%composer, "<someCo...>")
       val a = someComposableValue(%composer, 0)
       used(a)
       val m = Modifier()
-      val dismissModifier = <block>{
-        val tmp1_group = if (visible) {
+      val dismissModifier = if (visible) {
+        %composer.startReplaceGroup(<>)
+        sourceInformation(%composer, "<{>")
+        val tmp1_group = <block>{
           m.pointerInput(Unit, <block>{
             sourceInformationMarkerStart(%composer, <>, "CC(remember):Test.kt#9igjgp")
             val tmp0_group = %composer.cache(%dirty and 0b001110000000 == 0b000100000000) {
@@ -61,12 +65,19 @@ fun Test(a: Boolean, visible: Boolean, onDismiss: Function0<Unit>, %composer: Co
             sourceInformationMarkerEnd(%composer)
             tmp0_group
           })
-        } else {
-          m
         }
+        %composer.endReplaceGroup()
         tmp1_group
+      } else {
+        %composer.startReplaceGroup(<>)
+        %composer.endReplaceGroup()
+        m
       }
       used(dismissModifier)
+      %composer.endReplaceGroup()
+    } else {
+      %composer.startReplaceGroup(<>)
+      %composer.endReplaceGroup()
     }
     if (isTraceInProgress()) {
       traceEventEnd()
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.RememberIntrinsicTransformTests/testRememberInsideOfIfWithComposableCallBefore[useFir = false].txt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.RememberIntrinsicTransformTests/testRememberInsideOfIfWithComposableCallBefore[useFir = false].txt
index 14b6bb82d35dd..40ef7c195e5e7 100644
--- a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.RememberIntrinsicTransformTests/testRememberInsideOfIfWithComposableCallBefore[useFir = false].txt	
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.RememberIntrinsicTransformTests/testRememberInsideOfIfWithComposableCallBefore[useFir = false].txt	
@@ -21,7 +21,7 @@ fun Test(condition: Boolean) {
 @Composable
 fun Test(condition: Boolean, %composer: Composer?, %changed: Int) {
   %composer = %composer.startRestartGroup(<>)
-  sourceInformation(%composer, "C(Test)<A()>,<rememb...>:Test.kt")
+  sourceInformation(%composer, "C(Test):Test.kt")
   val %dirty = %changed
   if (%changed and 0b0110 == 0) {
     %dirty = %dirty or if (%composer.changed(condition)) 0b0100 else 0b0010
@@ -31,6 +31,8 @@ fun Test(condition: Boolean, %composer: Composer?, %changed: Int) {
       traceEventStart(<>, %dirty, -1, <>)
     }
     if (condition) {
+      %composer.startReplaceGroup(<>)
+      sourceInformation(%composer, "<A()>,<rememb...>")
       A(%composer, 0)
       val foo = <block>{
         sourceInformationMarkerStart(%composer, <>, "CC(remember):Test.kt#9igjgp")
@@ -40,6 +42,10 @@ fun Test(condition: Boolean, %composer: Composer?, %changed: Int) {
         sourceInformationMarkerEnd(%composer)
         tmp0_group
       }
+      %composer.endReplaceGroup()
+    } else {
+      %composer.startReplaceGroup(<>)
+      %composer.endReplaceGroup()
     }
     if (isTraceInProgress()) {
       traceEventEnd()
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.RememberIntrinsicTransformTests/testRememberInsideOfIfWithComposableCallBefore[useFir = true].txt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.RememberIntrinsicTransformTests/testRememberInsideOfIfWithComposableCallBefore[useFir = true].txt
index 14b6bb82d35dd..40ef7c195e5e7 100644
--- a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.RememberIntrinsicTransformTests/testRememberInsideOfIfWithComposableCallBefore[useFir = true].txt	
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.RememberIntrinsicTransformTests/testRememberInsideOfIfWithComposableCallBefore[useFir = true].txt	
@@ -21,7 +21,7 @@ fun Test(condition: Boolean) {
 @Composable
 fun Test(condition: Boolean, %composer: Composer?, %changed: Int) {
   %composer = %composer.startRestartGroup(<>)
-  sourceInformation(%composer, "C(Test)<A()>,<rememb...>:Test.kt")
+  sourceInformation(%composer, "C(Test):Test.kt")
   val %dirty = %changed
   if (%changed and 0b0110 == 0) {
     %dirty = %dirty or if (%composer.changed(condition)) 0b0100 else 0b0010
@@ -31,6 +31,8 @@ fun Test(condition: Boolean, %composer: Composer?, %changed: Int) {
       traceEventStart(<>, %dirty, -1, <>)
     }
     if (condition) {
+      %composer.startReplaceGroup(<>)
+      sourceInformation(%composer, "<A()>,<rememb...>")
       A(%composer, 0)
       val foo = <block>{
         sourceInformationMarkerStart(%composer, <>, "CC(remember):Test.kt#9igjgp")
@@ -40,6 +42,10 @@ fun Test(condition: Boolean, %composer: Composer?, %changed: Int) {
         sourceInformationMarkerEnd(%composer)
         tmp0_group
       }
+      %composer.endReplaceGroup()
+    } else {
+      %composer.startReplaceGroup(<>)
+      %composer.endReplaceGroup()
     }
     if (isTraceInProgress()) {
       traceEventEnd()
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.RememberIntrinsicTransformTests/testRememberInsideOfIf[useFir = false].txt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.RememberIntrinsicTransformTests/testRememberInsideOfIf[useFir = false].txt
index 7020b47fb16ab..1a4d50d21b58f 100644
--- a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.RememberIntrinsicTransformTests/testRememberInsideOfIf[useFir = false].txt	
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.RememberIntrinsicTransformTests/testRememberInsideOfIf[useFir = false].txt	
@@ -21,7 +21,7 @@ fun Test(condition: Boolean) {
 @Composable
 fun Test(condition: Boolean, %composer: Composer?, %changed: Int) {
   %composer = %composer.startRestartGroup(<>)
-  sourceInformation(%composer, "C(Test)<A()>,<rememb...>:Test.kt")
+  sourceInformation(%composer, "C(Test)<A()>:Test.kt")
   val %dirty = %changed
   if (%changed and 0b0110 == 0) {
     %dirty = %dirty or if (%composer.changed(condition)) 0b0100 else 0b0010
@@ -32,6 +32,8 @@ fun Test(condition: Boolean, %composer: Composer?, %changed: Int) {
     }
     A(%composer, 0)
     if (condition) {
+      %composer.startReplaceGroup(<>)
+      sourceInformation(%composer, "<rememb...>")
       val foo = <block>{
         sourceInformationMarkerStart(%composer, <>, "CC(remember):Test.kt#9igjgp")
         val tmp0_group = %composer.cache(false) {
@@ -40,6 +42,10 @@ fun Test(condition: Boolean, %composer: Composer?, %changed: Int) {
         sourceInformationMarkerEnd(%composer)
         tmp0_group
       }
+      %composer.endReplaceGroup()
+    } else {
+      %composer.startReplaceGroup(<>)
+      %composer.endReplaceGroup()
     }
     if (isTraceInProgress()) {
       traceEventEnd()
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.RememberIntrinsicTransformTests/testRememberInsideOfIf[useFir = true].txt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.RememberIntrinsicTransformTests/testRememberInsideOfIf[useFir = true].txt
index 7020b47fb16ab..1a4d50d21b58f 100644
--- a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.RememberIntrinsicTransformTests/testRememberInsideOfIf[useFir = true].txt	
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.RememberIntrinsicTransformTests/testRememberInsideOfIf[useFir = true].txt	
@@ -21,7 +21,7 @@ fun Test(condition: Boolean) {
 @Composable
 fun Test(condition: Boolean, %composer: Composer?, %changed: Int) {
   %composer = %composer.startRestartGroup(<>)
-  sourceInformation(%composer, "C(Test)<A()>,<rememb...>:Test.kt")
+  sourceInformation(%composer, "C(Test)<A()>:Test.kt")
   val %dirty = %changed
   if (%changed and 0b0110 == 0) {
     %dirty = %dirty or if (%composer.changed(condition)) 0b0100 else 0b0010
@@ -32,6 +32,8 @@ fun Test(condition: Boolean, %composer: Composer?, %changed: Int) {
     }
     A(%composer, 0)
     if (condition) {
+      %composer.startReplaceGroup(<>)
+      sourceInformation(%composer, "<rememb...>")
       val foo = <block>{
         sourceInformationMarkerStart(%composer, <>, "CC(remember):Test.kt#9igjgp")
         val tmp0_group = %composer.cache(false) {
@@ -40,6 +42,10 @@ fun Test(condition: Boolean, %composer: Composer?, %changed: Int) {
         sourceInformationMarkerEnd(%composer)
         tmp0_group
       }
+      %composer.endReplaceGroup()
+    } else {
+      %composer.startReplaceGroup(<>)
+      %composer.endReplaceGroup()
     }
     if (isTraceInProgress()) {
       traceEventEnd()
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.RememberIntrinsicTransformTests/testVarargsIntrinsicRemember[useFir = false].txt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.RememberIntrinsicTransformTests/testVarargsIntrinsicRemember[useFir = false].txt
index ac97240a474d3..d0ae40167940a 100644
--- a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.RememberIntrinsicTransformTests/testVarargsIntrinsicRemember[useFir = false].txt	
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.RememberIntrinsicTransformTests/testVarargsIntrinsicRemember[useFir = false].txt	
@@ -19,7 +19,7 @@ fun Test(vararg strings: String) {
 @Composable
 fun Test(strings: Array<out String>, %composer: Composer?, %changed: Int) {
   %composer = %composer.startRestartGroup(<>)
-  sourceInformation(%composer, "C(Test)<rememb...>,<Text("...>:Test.kt")
+  sourceInformation(%composer, "C(Test)<rememb...>:Test.kt")
   val %dirty = %changed
   %composer.startMovableGroup(<>, strings.size)
   val <iterator> = strings.iterator()
@@ -46,7 +46,13 @@ fun Test(strings: Array<out String>, %composer: Composer?, %changed: Int) {
       tmp0_group
     }
     if (show.value) {
+      %composer.startReplaceGroup(<>)
+      sourceInformation(%composer, "<Text("...>")
       Text("Showing", %composer, 0b0110)
+      %composer.endReplaceGroup()
+    } else {
+      %composer.startReplaceGroup(<>)
+      %composer.endReplaceGroup()
     }
     if (isTraceInProgress()) {
       traceEventEnd()
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.RememberIntrinsicTransformTests/testVarargsIntrinsicRemember[useFir = true].txt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.RememberIntrinsicTransformTests/testVarargsIntrinsicRemember[useFir = true].txt
index ac97240a474d3..d0ae40167940a 100644
--- a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.RememberIntrinsicTransformTests/testVarargsIntrinsicRemember[useFir = true].txt	
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.RememberIntrinsicTransformTests/testVarargsIntrinsicRemember[useFir = true].txt	
@@ -19,7 +19,7 @@ fun Test(vararg strings: String) {
 @Composable
 fun Test(strings: Array<out String>, %composer: Composer?, %changed: Int) {
   %composer = %composer.startRestartGroup(<>)
-  sourceInformation(%composer, "C(Test)<rememb...>,<Text("...>:Test.kt")
+  sourceInformation(%composer, "C(Test)<rememb...>:Test.kt")
   val %dirty = %changed
   %composer.startMovableGroup(<>, strings.size)
   val <iterator> = strings.iterator()
@@ -46,7 +46,13 @@ fun Test(strings: Array<out String>, %composer: Composer?, %changed: Int) {
       tmp0_group
     }
     if (show.value) {
+      %composer.startReplaceGroup(<>)
+      sourceInformation(%composer, "<Text("...>")
       Text("Showing", %composer, 0b0110)
+      %composer.endReplaceGroup()
+    } else {
+      %composer.startReplaceGroup(<>)
+      %composer.endReplaceGroup()
     }
     if (isTraceInProgress()) {
       traceEventEnd()
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.TargetAnnotationsTransformTests/testLetIt[useFir = false].txt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.TargetAnnotationsTransformTests/testLetIt[useFir = false].txt
index 639daabf1ca07..0fa8a2e75cc14 100644
--- a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.TargetAnnotationsTransformTests/testLetIt[useFir = false].txt	
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.TargetAnnotationsTransformTests/testLetIt[useFir = false].txt	
@@ -17,7 +17,7 @@ fun Test(content: (@Composable () -> Unit?)) {
 @ComposableInferredTarget(scheme = "[0[0]]")
 fun Test(content: Function2<Composer, Int, Unit?>, %composer: Composer?, %changed: Int) {
   %composer = %composer.startRestartGroup(<>)
-  sourceInformation(%composer, "C(Test)*<it()>:Test.kt")
+  sourceInformation(%composer, "C(Test):Test.kt")
   val %dirty = %changed
   if (%changed and 0b0110 == 0) {
     %dirty = %dirty or if (%composer.changedInstance(content)) 0b0100 else 0b0010
@@ -27,19 +27,26 @@ fun Test(content: Function2<Composer, Int, Unit?>, %composer: Composer?, %change
       traceEventStart(<>, %dirty, -1, <>)
     }
     val tmp0_safe_receiver = content
-    val tmp1_group = when {
+    when {
       tmp0_safe_receiver == null -> {
+        %composer.startReplaceGroup(<>)
+        %composer.endReplaceGroup()
         null
       }
       else -> {
-        val tmp0_group = tmp0_safe_receiver.let { it: Function2<Composer, Int, Unit?> ->
-          val tmp0_return = it(%composer, 0)
-          tmp0_return
+        %composer.startReplaceGroup(<>)
+        sourceInformation(%composer, "*<it()>")
+        val tmp1_group = <block>{
+          val tmp0_group = tmp0_safe_receiver.let { it: Function2<Composer, Int, Unit?> ->
+            val tmp0_return = it(%composer, 0)
+            tmp0_return
+          }
+          tmp0_group
         }
-        tmp0_group
+        %composer.endReplaceGroup()
+        tmp1_group
       }
     }
-    tmp1_group
     if (isTraceInProgress()) {
       traceEventEnd()
     }
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.TargetAnnotationsTransformTests/testLetIt[useFir = true].txt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.TargetAnnotationsTransformTests/testLetIt[useFir = true].txt
index 639daabf1ca07..0fa8a2e75cc14 100644
--- a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.TargetAnnotationsTransformTests/testLetIt[useFir = true].txt	
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.TargetAnnotationsTransformTests/testLetIt[useFir = true].txt	
@@ -17,7 +17,7 @@ fun Test(content: (@Composable () -> Unit?)) {
 @ComposableInferredTarget(scheme = "[0[0]]")
 fun Test(content: Function2<Composer, Int, Unit?>, %composer: Composer?, %changed: Int) {
   %composer = %composer.startRestartGroup(<>)
-  sourceInformation(%composer, "C(Test)*<it()>:Test.kt")
+  sourceInformation(%composer, "C(Test):Test.kt")
   val %dirty = %changed
   if (%changed and 0b0110 == 0) {
     %dirty = %dirty or if (%composer.changedInstance(content)) 0b0100 else 0b0010
@@ -27,19 +27,26 @@ fun Test(content: Function2<Composer, Int, Unit?>, %composer: Composer?, %change
       traceEventStart(<>, %dirty, -1, <>)
     }
     val tmp0_safe_receiver = content
-    val tmp1_group = when {
+    when {
       tmp0_safe_receiver == null -> {
+        %composer.startReplaceGroup(<>)
+        %composer.endReplaceGroup()
         null
       }
       else -> {
-        val tmp0_group = tmp0_safe_receiver.let { it: Function2<Composer, Int, Unit?> ->
-          val tmp0_return = it(%composer, 0)
-          tmp0_return
+        %composer.startReplaceGroup(<>)
+        sourceInformation(%composer, "*<it()>")
+        val tmp1_group = <block>{
+          val tmp0_group = tmp0_safe_receiver.let { it: Function2<Composer, Int, Unit?> ->
+            val tmp0_return = it(%composer, 0)
+            tmp0_return
+          }
+          tmp0_group
         }
-        tmp0_group
+        %composer.endReplaceGroup()
+        tmp1_group
       }
     }
-    tmp1_group
     if (isTraceInProgress()) {
       traceEventEnd()
     }
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.TargetAnnotationsTransformTests/testOptionalParameters[useFir = false].txt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.TargetAnnotationsTransformTests/testOptionalParameters[useFir = false].txt
index 0f427c95e91fa..2ddd3a3b4990e 100644
--- a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.TargetAnnotationsTransformTests/testOptionalParameters[useFir = false].txt	
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.TargetAnnotationsTransformTests/testOptionalParameters[useFir = false].txt	
@@ -161,69 +161,76 @@ fun Optional(one: Function2<Composer, Int, Unit>?, two: Function2<Composer, Int,
     }
     one(%composer, 0b1110 and %dirty)
     val tmp0_safe_receiver = two
-    %composer.startReplaceGroup(<>)
-    sourceInformation(%composer, "<invoke...>")
-    val tmp0_group = when {
+    when {
       tmp0_safe_receiver == null -> {
+        %composer.startReplaceGroup(<>)
+        %composer.endReplaceGroup()
         null
       }
       else -> {
+        %composer.startReplaceGroup(<>)
+        sourceInformation(%composer, "<invoke...>")
         tmp0_safe_receiver(%composer, 0b1110 and %dirty shr 0b0011)
+        %composer.endReplaceGroup()
       }
     }
-    %composer.endReplaceGroup()
-    tmp0_group
     val tmp1_safe_receiver = three
-    %composer.startReplaceGroup(<>)
-    sourceInformation(%composer, "*<it()>")
-    val tmp1_group = when {
+    when {
       tmp1_safe_receiver == null -> {
+        %composer.startReplaceGroup(<>)
+        %composer.endReplaceGroup()
         null
       }
       else -> {
+        %composer.startReplaceGroup(<>)
+        sourceInformation(%composer, "*<it()>")
         tmp1_safe_receiver.let { it: Function2<Composer, Int, Unit> ->
           it(%composer, 0)
         }
+        %composer.endReplaceGroup()
       }
     }
-    %composer.endReplaceGroup()
-    tmp1_group
     val tmp2_safe_receiver = four
-    %composer.startReplaceGroup(<>)
-    sourceInformation(%composer, "*<four()>")
-    val tmp2_group = when {
+    when {
       tmp2_safe_receiver == null -> {
+        %composer.startReplaceGroup(<>)
+        %composer.endReplaceGroup()
         null
       }
       else -> {
+        %composer.startReplaceGroup(<>)
+        sourceInformation(%composer, "*<four()>")
         tmp2_safe_receiver.let { it: Function2<Composer, Int, Unit> ->
           four(%composer, 0b1110 and %dirty shr 0b1001)
         }
+        %composer.endReplaceGroup()
       }
     }
-    %composer.endReplaceGroup()
-    tmp2_group
-    %composer.startReplaceGroup(<>)
-    sourceInformation(%composer, "<five()>")
     if (five != null) {
+      %composer.startReplaceGroup(<>)
+      sourceInformation(%composer, "<five()>")
       five(%composer, 0b1110 and %dirty shr 0b1100)
+      %composer.endReplaceGroup()
+    } else {
+      %composer.startReplaceGroup(<>)
+      %composer.endReplaceGroup()
     }
-    %composer.endReplaceGroup()
     val tmp3_safe_receiver = six
-    %composer.startReplaceGroup(<>)
-    sourceInformation(%composer, "*<Wrappe...>")
-    val tmp3_group = when {
+    when {
       tmp3_safe_receiver == null -> {
+        %composer.startReplaceGroup(<>)
+        %composer.endReplaceGroup()
         null
       }
       else -> {
+        %composer.startReplaceGroup(<>)
+        sourceInformation(%composer, "*<Wrappe...>")
         tmp3_safe_receiver.let { it: Function2<Composer, Int, Unit> ->
           Wrapper(it, %composer, 0)
         }
+        %composer.endReplaceGroup()
       }
     }
-    %composer.endReplaceGroup()
-    tmp3_group
     content(%composer, 0b1110 and %dirty shr 0b00010010)
     if (isTraceInProgress()) {
       traceEventEnd()
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.TargetAnnotationsTransformTests/testOptionalParameters[useFir = true].txt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.TargetAnnotationsTransformTests/testOptionalParameters[useFir = true].txt
index 0f427c95e91fa..2ddd3a3b4990e 100644
--- a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.TargetAnnotationsTransformTests/testOptionalParameters[useFir = true].txt	
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.TargetAnnotationsTransformTests/testOptionalParameters[useFir = true].txt	
@@ -161,69 +161,76 @@ fun Optional(one: Function2<Composer, Int, Unit>?, two: Function2<Composer, Int,
     }
     one(%composer, 0b1110 and %dirty)
     val tmp0_safe_receiver = two
-    %composer.startReplaceGroup(<>)
-    sourceInformation(%composer, "<invoke...>")
-    val tmp0_group = when {
+    when {
       tmp0_safe_receiver == null -> {
+        %composer.startReplaceGroup(<>)
+        %composer.endReplaceGroup()
         null
       }
       else -> {
+        %composer.startReplaceGroup(<>)
+        sourceInformation(%composer, "<invoke...>")
         tmp0_safe_receiver(%composer, 0b1110 and %dirty shr 0b0011)
+        %composer.endReplaceGroup()
       }
     }
-    %composer.endReplaceGroup()
-    tmp0_group
     val tmp1_safe_receiver = three
-    %composer.startReplaceGroup(<>)
-    sourceInformation(%composer, "*<it()>")
-    val tmp1_group = when {
+    when {
       tmp1_safe_receiver == null -> {
+        %composer.startReplaceGroup(<>)
+        %composer.endReplaceGroup()
         null
       }
       else -> {
+        %composer.startReplaceGroup(<>)
+        sourceInformation(%composer, "*<it()>")
         tmp1_safe_receiver.let { it: Function2<Composer, Int, Unit> ->
           it(%composer, 0)
         }
+        %composer.endReplaceGroup()
       }
     }
-    %composer.endReplaceGroup()
-    tmp1_group
     val tmp2_safe_receiver = four
-    %composer.startReplaceGroup(<>)
-    sourceInformation(%composer, "*<four()>")
-    val tmp2_group = when {
+    when {
       tmp2_safe_receiver == null -> {
+        %composer.startReplaceGroup(<>)
+        %composer.endReplaceGroup()
         null
       }
       else -> {
+        %composer.startReplaceGroup(<>)
+        sourceInformation(%composer, "*<four()>")
         tmp2_safe_receiver.let { it: Function2<Composer, Int, Unit> ->
           four(%composer, 0b1110 and %dirty shr 0b1001)
         }
+        %composer.endReplaceGroup()
       }
     }
-    %composer.endReplaceGroup()
-    tmp2_group
-    %composer.startReplaceGroup(<>)
-    sourceInformation(%composer, "<five()>")
     if (five != null) {
+      %composer.startReplaceGroup(<>)
+      sourceInformation(%composer, "<five()>")
       five(%composer, 0b1110 and %dirty shr 0b1100)
+      %composer.endReplaceGroup()
+    } else {
+      %composer.startReplaceGroup(<>)
+      %composer.endReplaceGroup()
     }
-    %composer.endReplaceGroup()
     val tmp3_safe_receiver = six
-    %composer.startReplaceGroup(<>)
-    sourceInformation(%composer, "*<Wrappe...>")
-    val tmp3_group = when {
+    when {
       tmp3_safe_receiver == null -> {
+        %composer.startReplaceGroup(<>)
+        %composer.endReplaceGroup()
         null
       }
       else -> {
+        %composer.startReplaceGroup(<>)
+        sourceInformation(%composer, "*<Wrappe...>")
         tmp3_safe_receiver.let { it: Function2<Composer, Int, Unit> ->
           Wrapper(it, %composer, 0)
         }
+        %composer.endReplaceGroup()
       }
     }
-    %composer.endReplaceGroup()
-    tmp3_group
     content(%composer, 0b1110 and %dirty shr 0b00010010)
     if (isTraceInProgress()) {
       traceEventEnd()
diff --git a/plugins/compose/compiler-hosted/src/main/java/androidx/compose/compiler/plugins/kotlin/lower/ComposableFunctionBodyTransformer.kt b/plugins/compose/compiler-hosted/src/main/java/androidx/compose/compiler/plugins/kotlin/lower/ComposableFunctionBodyTransformer.kt
index dfe4c005058de..e46aca7b31ec5 100644
--- a/plugins/compose/compiler-hosted/src/main/java/androidx/compose/compiler/plugins/kotlin/lower/ComposableFunctionBodyTransformer.kt
+++ b/plugins/compose/compiler-hosted/src/main/java/androidx/compose/compiler/plugins/kotlin/lower/ComposableFunctionBodyTransformer.kt
@@ -3710,6 +3710,7 @@ class ComposableFunctionBodyTransformer(
         // conditionally executed.
         var needsWrappingGroup = false
         var resultsWithCalls = 0
+        var conditionsWithCalls = 0
         var hasElseBranch = false
 
         val transformed = IrWhenImpl(
@@ -3751,12 +3752,19 @@ class ComposableFunctionBodyTransformer(
                     condScopes.add(condScope)
                     resultScopes.add(resultScope)
 
-                    // the first condition is always executed so if it has a composable call in it,
-                    // it doesn't necessitate a group
-                    needsWrappingGroup =
-                        needsWrappingGroup || (index != 0 && condScope.hasComposableCalls)
-                    if (resultScope.hasComposableCalls)
+                    // If we are optimizing non-skipping groups we never use a wrapping group, each condition
+                    // is wrapped in a group instead.
+                    if (!FeatureFlag.OptimizeNonSkippingGroups.enabled) {
+                        // the first condition is always executed so if it has a composable call in it,
+                        // it doesn't necessitate a group. However, non-skipping group optimization is
+                        // enabled, we need a wrapping group if any conditions have a composable call.
+                        needsWrappingGroup =
+                            needsWrappingGroup || ((index != 0) && condScope.hasComposableCalls)
+                    }
+                    if (resultScope.hasComposableCalls && !it.result.isGroupBalanced())
                         resultsWithCalls++
+                    if (condScope.hasComposableCalls && !it.condition.isGroupBalanced())
+                        conditionsWithCalls++
 
                     transformed.branches.add(
                         IrBranchImpl(
@@ -3770,13 +3778,30 @@ class ComposableFunctionBodyTransformer(
             }
         }
 
+        // If we are optimizing the non-skipping functions we always need the
+        // same number of groups if any of the results have composable functions
+        // and it needs to be the same number even if only one branch requires a
+        // group.
+        val needsResultGroups = if (FeatureFlag.OptimizeNonSkippingGroups.enabled) {
+            resultsWithCalls > 0
+        } else {
+            resultsWithCalls > 1 && !needsWrappingGroup
+        }
+
+        val needsConditionGroups = if (FeatureFlag.OptimizeNonSkippingGroups.enabled) {
+            conditionsWithCalls > 0
+        } else {
+            // A wrapping group is used instead.
+            false
+        }
+
         // If we are putting groups around the result branches, we need to guarantee that exactly
         // one result branch is executed. We do this by adding an else branch if it there is not
         // one already. Note that we only need to do this if we aren't going to wrap the if
         // statement in a group entirely, which we will do if the conditions have calls in them.
         // NOTE: we might also be able to assume that the when is exhaustive if it has a non-unit
         // resulting type, since the type system should enforce that.
-        if (!hasElseBranch && resultsWithCalls > 1 && !needsWrappingGroup) {
+        if (!hasElseBranch && needsResultGroups) {
             condScopes.add(Scope.BranchScope())
             resultScopes.add(Scope.BranchScope())
             transformed.branches.add(
@@ -3815,17 +3840,23 @@ class ComposableFunctionBodyTransformer(
                 }
             }
 
+            // if no wrapping group but more than we need branch groups, we have to have every
+            // result be a group so that we have a consistent number of groups during execution
             if (
-                // if no wrapping group but more than one result have calls, we have to have every
-                // result be a group so that we have a consistent number of groups during execution
-                (resultsWithCalls > 1 && !needsWrappingGroup) ||
+                needsResultGroups ||
                 // if we are wrapping the if with a group, then we only need to add a group when
-                // the block has composable calls
-                (needsWrappingGroup && resultScope.hasComposableCalls)
+                // the block has composable calls. The check of the feature flag check here is redundant
+                // as needsBranchGroups will be true if any result scope has composable calls but it
+                // is here redundantly so when this flag is removed this code will be updated.
+                !FeatureFlag.OptimizeNonSkippingGroups.enabled && (needsWrappingGroup && resultScope.hasComposableCalls)
             ) {
                 it.result = it.result.asReplaceGroup(resultScope)
             }
 
+            if (needsConditionGroups) {
+                it.condition = it.condition.asReplaceGroup(condScope)
+            }
+
             if (resultsWithCalls == 1 && resultScope.hasComposableCalls) {
                 // Realize all groups in the branch result with a conditional call - making sure
                 // that nested control structures are wrapped correctly.
@@ -3834,11 +3865,20 @@ class ComposableFunctionBodyTransformer(
         }
 
         return when {
-            resultsWithCalls == 1 || needsWrappingGroup -> transformed.asCoalescableGroup(whenScope)
+            !FeatureFlag.OptimizeNonSkippingGroups.enabled &&
+                    (resultsWithCalls == 1 || needsWrappingGroup) -> transformed.asCoalescableGroup(whenScope)
             else -> transformed
         }
     }
 
+    // Returns true if the number of groups added are required to be fix and a group is inserted  to balance the groups if they are not.
+    // Currently this is only guaranteed for IrWhen nodes when the group non-skipping group optimization is enabled. This avoids
+    // inserting a redundant group to balance an already balanced set of groups.
+    private fun IrExpression.isGroupBalanced(): Boolean = when(this) {
+        is IrWhen -> FeatureFlag.OptimizeNonSkippingGroups.enabled
+        else -> false
+    }
+
     sealed class Scope(val name: String) {
         var parent: Scope? = null
         var level: Int = 0