[Compose] Fix code generation for group optimization Fixes a code generation issue added with fc48dbdba that caused the generation of && and || that contained composable functions to fail IR validation. Fixes: [346821372](https://issuetracker.google.com/346821372)

diff --git "a/plugins/compose/compiler-hosted/integration-tests/src/test/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testAND\133useFir = false\135.txt" "b/plugins/compose/compiler-hosted/integration-tests/src/test/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testAND\133useFir = false\135.txt"
index 9040a49..dfa0afc 100644
--- "a/plugins/compose/compiler-hosted/integration-tests/src/test/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testAND\133useFir = false\135.txt"
+++ "b/plugins/compose/compiler-hosted/integration-tests/src/test/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testAND\133useFir = false\135.txt"

@@ -21,19 +21,24 @@
 @NonRestartableComposable
 @Composable
 fun Example(%composer: Composer?, %changed: Int) {
-  sourceInformationMarkerStart(%composer, <>, "C(Example):Test.kt")
+  sourceInformationMarkerStart(%composer, <>, "C(Example)<B()>:Test.kt")
   if (isTraceInProgress()) {
     traceEventStart(<>, %changed, -1, <>)
   }
-  %composer.startReplaceGroup(<>)
-  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
+  when {
+    B(%composer, 0) -> {
+      %composer.startReplaceGroup(<>)
+      sourceInformation(%composer, "<B()>")
+      val tmp0_group = B(%composer, 0)
+      %composer.endReplaceGroup()
+      tmp0_group
+    }
+    else -> {
+      %composer.startReplaceGroup(<>)
+      %composer.endReplaceGroup()
+      false
+    }
+  }
   if (isTraceInProgress()) {
     traceEventEnd()
   }

diff --git "a/plugins/compose/compiler-hosted/integration-tests/src/test/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testAND\133useFir = true\135.txt" "b/plugins/compose/compiler-hosted/integration-tests/src/test/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testAND\133useFir = true\135.txt"
index 9040a49..dfa0afc 100644
--- "a/plugins/compose/compiler-hosted/integration-tests/src/test/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testAND\133useFir = true\135.txt"
+++ "b/plugins/compose/compiler-hosted/integration-tests/src/test/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testAND\133useFir = true\135.txt"

@@ -21,19 +21,24 @@
 @NonRestartableComposable
 @Composable
 fun Example(%composer: Composer?, %changed: Int) {
-  sourceInformationMarkerStart(%composer, <>, "C(Example):Test.kt")
+  sourceInformationMarkerStart(%composer, <>, "C(Example)<B()>:Test.kt")
   if (isTraceInProgress()) {
     traceEventStart(<>, %changed, -1, <>)
   }
-  %composer.startReplaceGroup(<>)
-  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
+  when {
+    B(%composer, 0) -> {
+      %composer.startReplaceGroup(<>)
+      sourceInformation(%composer, "<B()>")
+      val tmp0_group = B(%composer, 0)
+      %composer.endReplaceGroup()
+      tmp0_group
+    }
+    else -> {
+      %composer.startReplaceGroup(<>)
+      %composer.endReplaceGroup()
+      false
+    }
+  }
   if (isTraceInProgress()) {
     traceEventEnd()
   }

diff --git "a/plugins/compose/compiler-hosted/integration-tests/src/test/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testDynamicWrappingGroupWithReturnValue\133useFir = false\135.txt" "b/plugins/compose/compiler-hosted/integration-tests/src/test/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testDynamicWrappingGroupWithReturnValue\133useFir = false\135.txt"
index 7d41355..7ee900d 100644
--- "a/plugins/compose/compiler-hosted/integration-tests/src/test/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testDynamicWrappingGroupWithReturnValue\133useFir = false\135.txt"
+++ "b/plugins/compose/compiler-hosted/integration-tests/src/test/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testDynamicWrappingGroupWithReturnValue\133useFir = false\135.txt"

@@ -30,22 +30,13 @@
   }
   val tmp0 = if (x > 0) {
     %composer.startReplaceGroup(<>)
-    sourceInformation(%composer, "")
-    val tmp3_group = <block>{
-      if (%composer.startReplaceGroup(<>)
-      sourceInformation(%composer, "<B()>")
-      val tmp1_group = B(%composer, 0)
-      %composer.endReplaceGroup()
-      tmp1_group) 1 else if (%composer.startReplaceGroup(<>)
-      sourceInformation(%composer, "<B()>")
-      val tmp2_group = B(%composer, 0)
-      %composer.endReplaceGroup()
-      tmp2_group) 2 else if (%composer.startReplaceGroup(<>)
-      %composer.endReplaceGroup()
-      true) 3
+    sourceInformation(%composer, "<B()>,<B()>")
+    val tmp2_group = <block>{
+      val tmp1_group = if (B(%composer, 0)) 1 else if (B(%composer, 0)) 2 else 3
+      tmp1_group
     }
     %composer.endReplaceGroup()
-    tmp3_group
+    tmp2_group
   } else {
     %composer.startReplaceGroup(<>)
     %composer.endReplaceGroup()

diff --git "a/plugins/compose/compiler-hosted/integration-tests/src/test/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testDynamicWrappingGroupWithReturnValue\133useFir = true\135.txt" "b/plugins/compose/compiler-hosted/integration-tests/src/test/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testDynamicWrappingGroupWithReturnValue\133useFir = true\135.txt"
index 7d41355..7ee900d 100644
--- "a/plugins/compose/compiler-hosted/integration-tests/src/test/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testDynamicWrappingGroupWithReturnValue\133useFir = true\135.txt"
+++ "b/plugins/compose/compiler-hosted/integration-tests/src/test/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testDynamicWrappingGroupWithReturnValue\133useFir = true\135.txt"

@@ -30,22 +30,13 @@
   }
   val tmp0 = if (x > 0) {
     %composer.startReplaceGroup(<>)
-    sourceInformation(%composer, "")
-    val tmp3_group = <block>{
-      if (%composer.startReplaceGroup(<>)
-      sourceInformation(%composer, "<B()>")
-      val tmp1_group = B(%composer, 0)
-      %composer.endReplaceGroup()
-      tmp1_group) 1 else if (%composer.startReplaceGroup(<>)
-      sourceInformation(%composer, "<B()>")
-      val tmp2_group = B(%composer, 0)
-      %composer.endReplaceGroup()
-      tmp2_group) 2 else if (%composer.startReplaceGroup(<>)
-      %composer.endReplaceGroup()
-      true) 3
+    sourceInformation(%composer, "<B()>,<B()>")
+    val tmp2_group = <block>{
+      val tmp1_group = if (B(%composer, 0)) 1 else if (B(%composer, 0)) 2 else 3
+      tmp1_group
     }
     %composer.endReplaceGroup()
-    tmp3_group
+    tmp2_group
   } else {
     %composer.startReplaceGroup(<>)
     %composer.endReplaceGroup()

diff --git "a/plugins/compose/compiler-hosted/integration-tests/src/test/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testIfElseWithCallsInConditions\133useFir = false\135.txt" "b/plugins/compose/compiler-hosted/integration-tests/src/test/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testIfElseWithCallsInConditions\133useFir = false\135.txt"
index a4d532c..6dfd5e6 100644
--- "a/plugins/compose/compiler-hosted/integration-tests/src/test/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testIfElseWithCallsInConditions\133useFir = false\135.txt"
+++ "b/plugins/compose/compiler-hosted/integration-tests/src/test/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testIfElseWithCallsInConditions\133useFir = false\135.txt"

@@ -35,23 +35,16 @@
   if (isTraceInProgress()) {
     traceEventStart(<>, %changed, -1, <>)
   }
-  if (%composer.startReplaceGroup(<>)
-  sourceInformation(%composer, "<B(a)>")
-  val tmp0_group = B(a, %composer, 0)
-  %composer.endReplaceGroup()
-  tmp0_group) {
+  %composer.startReplaceGroup(<>)
+  sourceInformation(%composer, "<B(a)>,<B(b)>")
+  if (B(a, %composer, 0)) {
     NA()
-  } else if (%composer.startReplaceGroup(<>)
-  sourceInformation(%composer, "<B(b)>")
-  val tmp1_group = B(b, %composer, 0)
-  %composer.endReplaceGroup()
-  tmp1_group) {
+  } else if (B(b, %composer, 0)) {
     NA()
-  } else if (%composer.startReplaceGroup(<>)
-  %composer.endReplaceGroup()
-  true) {
+  } else {
     NA()
   }
+  %composer.endReplaceGroup()
   if (isTraceInProgress()) {
     traceEventEnd()
   }

diff --git "a/plugins/compose/compiler-hosted/integration-tests/src/test/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testIfElseWithCallsInConditions\133useFir = true\135.txt" "b/plugins/compose/compiler-hosted/integration-tests/src/test/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testIfElseWithCallsInConditions\133useFir = true\135.txt"
index a4d532c..6dfd5e6 100644
--- "a/plugins/compose/compiler-hosted/integration-tests/src/test/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testIfElseWithCallsInConditions\133useFir = true\135.txt"
+++ "b/plugins/compose/compiler-hosted/integration-tests/src/test/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testIfElseWithCallsInConditions\133useFir = true\135.txt"

@@ -35,23 +35,16 @@
   if (isTraceInProgress()) {
     traceEventStart(<>, %changed, -1, <>)
   }
-  if (%composer.startReplaceGroup(<>)
-  sourceInformation(%composer, "<B(a)>")
-  val tmp0_group = B(a, %composer, 0)
-  %composer.endReplaceGroup()
-  tmp0_group) {
+  %composer.startReplaceGroup(<>)
+  sourceInformation(%composer, "<B(a)>,<B(b)>")
+  if (B(a, %composer, 0)) {
     NA()
-  } else if (%composer.startReplaceGroup(<>)
-  sourceInformation(%composer, "<B(b)>")
-  val tmp1_group = B(b, %composer, 0)
-  %composer.endReplaceGroup()
-  tmp1_group) {
+  } else if (B(b, %composer, 0)) {
     NA()
-  } else if (%composer.startReplaceGroup(<>)
-  %composer.endReplaceGroup()
-  true) {
+  } else {
     NA()
   }
+  %composer.endReplaceGroup()
   if (isTraceInProgress()) {
     traceEventEnd()
   }

diff --git "a/plugins/compose/compiler-hosted/integration-tests/src/test/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testIfWithCallInCondition\133useFir = false\135.txt" "b/plugins/compose/compiler-hosted/integration-tests/src/test/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testIfWithCallInCondition\133useFir = false\135.txt"
index 1696718..afa4db1 100644
--- "a/plugins/compose/compiler-hosted/integration-tests/src/test/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testIfWithCallInCondition\133useFir = false\135.txt"
+++ "b/plugins/compose/compiler-hosted/integration-tests/src/test/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testIfWithCallInCondition\133useFir = false\135.txt"

@@ -27,19 +27,13 @@
 @NonRestartableComposable
 @Composable
 fun Example(x: Int, %composer: Composer?, %changed: Int) {
-  sourceInformationMarkerStart(%composer, <>, "C(Example):Test.kt")
+  sourceInformationMarkerStart(%composer, <>, "C(Example)<B()>:Test.kt")
   if (isTraceInProgress()) {
     traceEventStart(<>, %changed, -1, <>)
   }
-  if (%composer.startReplaceGroup(<>)
-  sourceInformation(%composer, "<B()>")
-  val tmp0_group = B(%composer, 0)
-  %composer.endReplaceGroup()
-  tmp0_group) {
+  if (B(%composer, 0)) {
     NA()
-  } else if (%composer.startReplaceGroup(<>)
-  %composer.endReplaceGroup()
-  true) {
+  } else {
     NA()
   }
   if (isTraceInProgress()) {

diff --git "a/plugins/compose/compiler-hosted/integration-tests/src/test/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testIfWithCallInCondition\133useFir = true\135.txt" "b/plugins/compose/compiler-hosted/integration-tests/src/test/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testIfWithCallInCondition\133useFir = true\135.txt"
index 1696718..afa4db1 100644
--- "a/plugins/compose/compiler-hosted/integration-tests/src/test/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testIfWithCallInCondition\133useFir = true\135.txt"
+++ "b/plugins/compose/compiler-hosted/integration-tests/src/test/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testIfWithCallInCondition\133useFir = true\135.txt"

@@ -27,19 +27,13 @@
 @NonRestartableComposable
 @Composable
 fun Example(x: Int, %composer: Composer?, %changed: Int) {
-  sourceInformationMarkerStart(%composer, <>, "C(Example):Test.kt")
+  sourceInformationMarkerStart(%composer, <>, "C(Example)<B()>:Test.kt")
   if (isTraceInProgress()) {
     traceEventStart(<>, %changed, -1, <>)
   }
-  if (%composer.startReplaceGroup(<>)
-  sourceInformation(%composer, "<B()>")
-  val tmp0_group = B(%composer, 0)
-  %composer.endReplaceGroup()
-  tmp0_group) {
+  if (B(%composer, 0)) {
     NA()
-  } else if (%composer.startReplaceGroup(<>)
-  %composer.endReplaceGroup()
-  true) {
+  } else {
     NA()
   }
   if (isTraceInProgress()) {

diff --git "a/plugins/compose/compiler-hosted/integration-tests/src/test/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testOR\133useFir = false\135.txt" "b/plugins/compose/compiler-hosted/integration-tests/src/test/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testOR\133useFir = false\135.txt"
index a2933cf..54c4439 100644
--- "a/plugins/compose/compiler-hosted/integration-tests/src/test/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testOR\133useFir = false\135.txt"
+++ "b/plugins/compose/compiler-hosted/integration-tests/src/test/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testOR\133useFir = false\135.txt"

@@ -21,19 +21,24 @@
 @NonRestartableComposable
 @Composable
 fun Example(%composer: Composer?, %changed: Int) {
-  sourceInformationMarkerStart(%composer, <>, "C(Example):Test.kt")
+  sourceInformationMarkerStart(%composer, <>, "C(Example)<B()>:Test.kt")
   if (isTraceInProgress()) {
     traceEventStart(<>, %changed, -1, <>)
   }
-  %composer.startReplaceGroup(<>)
-  sourceInformation(%composer, "<B()>")
-  val tmp0_group = B(%composer, 0)
-  %composer.endReplaceGroup()
-  tmp0_group || %composer.startReplaceGroup(<>)
-  sourceInformation(%composer, "<B()>")
-  val tmp1_group = B(%composer, 0)
-  %composer.endReplaceGroup()
-  tmp1_group
+  when {
+    B(%composer, 0) -> {
+      %composer.startReplaceGroup(<>)
+      %composer.endReplaceGroup()
+      true
+    }
+    else -> {
+      %composer.startReplaceGroup(<>)
+      sourceInformation(%composer, "<B()>")
+      val tmp0_group = B(%composer, 0)
+      %composer.endReplaceGroup()
+      tmp0_group
+    }
+  }
   if (isTraceInProgress()) {
     traceEventEnd()
   }

diff --git "a/plugins/compose/compiler-hosted/integration-tests/src/test/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testOR\133useFir = true\135.txt" "b/plugins/compose/compiler-hosted/integration-tests/src/test/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testOR\133useFir = true\135.txt"
index a2933cf..54c4439 100644
--- "a/plugins/compose/compiler-hosted/integration-tests/src/test/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testOR\133useFir = true\135.txt"
+++ "b/plugins/compose/compiler-hosted/integration-tests/src/test/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testOR\133useFir = true\135.txt"

@@ -21,19 +21,24 @@
 @NonRestartableComposable
 @Composable
 fun Example(%composer: Composer?, %changed: Int) {
-  sourceInformationMarkerStart(%composer, <>, "C(Example):Test.kt")
+  sourceInformationMarkerStart(%composer, <>, "C(Example)<B()>:Test.kt")
   if (isTraceInProgress()) {
     traceEventStart(<>, %changed, -1, <>)
   }
-  %composer.startReplaceGroup(<>)
-  sourceInformation(%composer, "<B()>")
-  val tmp0_group = B(%composer, 0)
-  %composer.endReplaceGroup()
-  tmp0_group || %composer.startReplaceGroup(<>)
-  sourceInformation(%composer, "<B()>")
-  val tmp1_group = B(%composer, 0)
-  %composer.endReplaceGroup()
-  tmp1_group
+  when {
+    B(%composer, 0) -> {
+      %composer.startReplaceGroup(<>)
+      %composer.endReplaceGroup()
+      true
+    }
+    else -> {
+      %composer.startReplaceGroup(<>)
+      sourceInformation(%composer, "<B()>")
+      val tmp0_group = B(%composer, 0)
+      %composer.endReplaceGroup()
+      tmp0_group
+    }
+  }
   if (isTraceInProgress()) {
     traceEventEnd()
   }

diff --git "a/plugins/compose/compiler-hosted/integration-tests/src/test/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testWhenWithCallsInConditionsAndCallAfter\133useFir = false\135.txt" "b/plugins/compose/compiler-hosted/integration-tests/src/test/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testWhenWithCallsInConditionsAndCallAfter\133useFir = false\135.txt"
index 5eb4429..f5ad2c5 100644
--- "a/plugins/compose/compiler-hosted/integration-tests/src/test/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testWhenWithCallsInConditionsAndCallAfter\133useFir = false\135.txt"
+++ "b/plugins/compose/compiler-hosted/integration-tests/src/test/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testWhenWithCallsInConditionsAndCallAfter\133useFir = false\135.txt"

@@ -32,27 +32,20 @@
   if (isTraceInProgress()) {
     traceEventStart(<>, %changed, -1, <>)
   }
+  %composer.startReplaceGroup(<>)
+  sourceInformation(%composer, "<R(a)>,<R(b)>")
   when {
-    %composer.startReplaceGroup(<>)
-    sourceInformation(%composer, "<R(a)>")
-    val tmp0_group = x == R(a, %composer, 0)
-    %composer.endReplaceGroup()
-    tmp0_group -> {
+    x == R(a, %composer, 0) -> {
       NA()
     }
-    %composer.startReplaceGroup(<>)
-    sourceInformation(%composer, "<R(b)>")
-    val tmp1_group = x > R(b, %composer, 0)
-    %composer.endReplaceGroup()
-    tmp1_group -> {
+    x > R(b, %composer, 0) -> {
       NA()
     }
-    %composer.startReplaceGroup(<>)
-    %composer.endReplaceGroup()
-    true -> {
+    else -> {
       NA()
     }
   }
+  %composer.endReplaceGroup()
   A(%composer, 0)
   if (isTraceInProgress()) {
     traceEventEnd()

diff --git "a/plugins/compose/compiler-hosted/integration-tests/src/test/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testWhenWithCallsInConditionsAndCallAfter\133useFir = true\135.txt" "b/plugins/compose/compiler-hosted/integration-tests/src/test/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testWhenWithCallsInConditionsAndCallAfter\133useFir = true\135.txt"
index 5eb4429..f5ad2c5 100644
--- "a/plugins/compose/compiler-hosted/integration-tests/src/test/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testWhenWithCallsInConditionsAndCallAfter\133useFir = true\135.txt"
+++ "b/plugins/compose/compiler-hosted/integration-tests/src/test/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testWhenWithCallsInConditionsAndCallAfter\133useFir = true\135.txt"

@@ -32,27 +32,20 @@
   if (isTraceInProgress()) {
     traceEventStart(<>, %changed, -1, <>)
   }
+  %composer.startReplaceGroup(<>)
+  sourceInformation(%composer, "<R(a)>,<R(b)>")
   when {
-    %composer.startReplaceGroup(<>)
-    sourceInformation(%composer, "<R(a)>")
-    val tmp0_group = x == R(a, %composer, 0)
-    %composer.endReplaceGroup()
-    tmp0_group -> {
+    x == R(a, %composer, 0) -> {
       NA()
     }
-    %composer.startReplaceGroup(<>)
-    sourceInformation(%composer, "<R(b)>")
-    val tmp1_group = x > R(b, %composer, 0)
-    %composer.endReplaceGroup()
-    tmp1_group -> {
+    x > R(b, %composer, 0) -> {
       NA()
     }
-    %composer.startReplaceGroup(<>)
-    %composer.endReplaceGroup()
-    true -> {
+    else -> {
       NA()
     }
   }
+  %composer.endReplaceGroup()
   A(%composer, 0)
   if (isTraceInProgress()) {
     traceEventEnd()

diff --git "a/plugins/compose/compiler-hosted/integration-tests/src/test/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testWhenWithCallsInConditions\133useFir = false\135.txt" "b/plugins/compose/compiler-hosted/integration-tests/src/test/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testWhenWithCallsInConditions\133useFir = false\135.txt"
index c76de7f..1e9d59a 100644
--- "a/plugins/compose/compiler-hosted/integration-tests/src/test/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testWhenWithCallsInConditions\133useFir = false\135.txt"
+++ "b/plugins/compose/compiler-hosted/integration-tests/src/test/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testWhenWithCallsInConditions\133useFir = false\135.txt"

@@ -32,27 +32,20 @@
   if (isTraceInProgress()) {
     traceEventStart(<>, %changed, -1, <>)
   }
+  %composer.startReplaceGroup(<>)
+  sourceInformation(%composer, "<R(a)>,<R(b)>")
   when {
-    %composer.startReplaceGroup(<>)
-    sourceInformation(%composer, "<R(a)>")
-    val tmp0_group = x == R(a, %composer, 0)
-    %composer.endReplaceGroup()
-    tmp0_group -> {
+    x == R(a, %composer, 0) -> {
       NA()
     }
-    %composer.startReplaceGroup(<>)
-    sourceInformation(%composer, "<R(b)>")
-    val tmp1_group = x > R(b, %composer, 0)
-    %composer.endReplaceGroup()
-    tmp1_group -> {
+    x > R(b, %composer, 0) -> {
       NA()
     }
-    %composer.startReplaceGroup(<>)
-    %composer.endReplaceGroup()
-    true -> {
+    else -> {
       NA()
     }
   }
+  %composer.endReplaceGroup()
   if (isTraceInProgress()) {
     traceEventEnd()
   }

diff --git "a/plugins/compose/compiler-hosted/integration-tests/src/test/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testWhenWithCallsInConditions\133useFir = true\135.txt" "b/plugins/compose/compiler-hosted/integration-tests/src/test/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testWhenWithCallsInConditions\133useFir = true\135.txt"
index c76de7f..1e9d59a 100644
--- "a/plugins/compose/compiler-hosted/integration-tests/src/test/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testWhenWithCallsInConditions\133useFir = true\135.txt"
+++ "b/plugins/compose/compiler-hosted/integration-tests/src/test/resources/golden/androidx.compose.compiler.plugins.kotlin.ControlFlowTransformTests/testWhenWithCallsInConditions\133useFir = true\135.txt"

@@ -32,27 +32,20 @@
   if (isTraceInProgress()) {
     traceEventStart(<>, %changed, -1, <>)
   }
+  %composer.startReplaceGroup(<>)
+  sourceInformation(%composer, "<R(a)>,<R(b)>")
   when {
-    %composer.startReplaceGroup(<>)
-    sourceInformation(%composer, "<R(a)>")
-    val tmp0_group = x == R(a, %composer, 0)
-    %composer.endReplaceGroup()
-    tmp0_group -> {
+    x == R(a, %composer, 0) -> {
       NA()
     }
-    %composer.startReplaceGroup(<>)
-    sourceInformation(%composer, "<R(b)>")
-    val tmp1_group = x > R(b, %composer, 0)
-    %composer.endReplaceGroup()
-    tmp1_group -> {
+    x > R(b, %composer, 0) -> {
       NA()
     }
-    %composer.startReplaceGroup(<>)
-    %composer.endReplaceGroup()
-    true -> {
+    else -> {
       NA()
     }
   }
+  %composer.endReplaceGroup()
   if (isTraceInProgress()) {
     traceEventEnd()
   }

diff --git "a/plugins/compose/compiler-hosted/integration-tests/src/test/resources/golden/androidx.compose.compiler.plugins.kotlin.FunctionBodySkippingTransformTests/testFunInterfaces2\133useFir = false\135.txt" "b/plugins/compose/compiler-hosted/integration-tests/src/test/resources/golden/androidx.compose.compiler.plugins.kotlin.FunctionBodySkippingTransformTests/testFunInterfaces2\133useFir = false\135.txt"
index 67206bc..6de1c6d 100644
--- "a/plugins/compose/compiler-hosted/integration-tests/src/test/resources/golden/androidx.compose.compiler.plugins.kotlin.FunctionBodySkippingTransformTests/testFunInterfaces2\133useFir = false\135.txt"
+++ "b/plugins/compose/compiler-hosted/integration-tests/src/test/resources/golden/androidx.compose.compiler.plugins.kotlin.FunctionBodySkippingTransformTests/testFunInterfaces2\133useFir = false\135.txt"

@@ -64,19 +64,13 @@
         @Composable
         override fun getColor(%composer: Composer?, %changed: Int): Color {
           %composer.startReplaceGroup(<>)
-          sourceInformation(%composer, "C(getColor):Test.kt")
+          sourceInformation(%composer, "C(getColor)<condit...>:Test.kt")
           if (isTraceInProgress()) {
             traceEventStart(<>, %changed, -1, <>)
           }
-          val tmp0 = if (%composer.startReplaceGroup(<>)
-          sourceInformation(%composer, "<condit...>")
-          val tmp1_group = condition(%composer, 0)
-          %composer.endReplaceGroup()
-          tmp1_group) {
+          val tmp0 = if (condition(%composer, 0)) {
             Companion.Red
-          } else if (%composer.startReplaceGroup(<>)
-          %composer.endReplaceGroup()
-          true) {
+          } else {
             Companion.Blue
           }
           if (isTraceInProgress()) {

diff --git "a/plugins/compose/compiler-hosted/integration-tests/src/test/resources/golden/androidx.compose.compiler.plugins.kotlin.FunctionBodySkippingTransformTests/testFunInterfaces2\133useFir = true\135.txt" "b/plugins/compose/compiler-hosted/integration-tests/src/test/resources/golden/androidx.compose.compiler.plugins.kotlin.FunctionBodySkippingTransformTests/testFunInterfaces2\133useFir = true\135.txt"
index 67206bc..6de1c6d 100644
--- "a/plugins/compose/compiler-hosted/integration-tests/src/test/resources/golden/androidx.compose.compiler.plugins.kotlin.FunctionBodySkippingTransformTests/testFunInterfaces2\133useFir = true\135.txt"
+++ "b/plugins/compose/compiler-hosted/integration-tests/src/test/resources/golden/androidx.compose.compiler.plugins.kotlin.FunctionBodySkippingTransformTests/testFunInterfaces2\133useFir = true\135.txt"

@@ -64,19 +64,13 @@
         @Composable
         override fun getColor(%composer: Composer?, %changed: Int): Color {
           %composer.startReplaceGroup(<>)
-          sourceInformation(%composer, "C(getColor):Test.kt")
+          sourceInformation(%composer, "C(getColor)<condit...>:Test.kt")
           if (isTraceInProgress()) {
             traceEventStart(<>, %changed, -1, <>)
           }
-          val tmp0 = if (%composer.startReplaceGroup(<>)
-          sourceInformation(%composer, "<condit...>")
-          val tmp1_group = condition(%composer, 0)
-          %composer.endReplaceGroup()
-          tmp1_group) {
+          val tmp0 = if (condition(%composer, 0)) {
             Companion.Red
-          } else if (%composer.startReplaceGroup(<>)
-          %composer.endReplaceGroup()
-          true) {
+          } else {
             Companion.Blue
           }
           if (isTraceInProgress()) {

diff --git "a/plugins/compose/compiler-hosted/integration-tests/src/test/resources/golden/androidx.compose.compiler.plugins.kotlin.FunctionBodySkippingTransformTests/test_groupAroundIfComposeCallInIfConditionWithShortCircuit\133useFir = false\135.txt" "b/plugins/compose/compiler-hosted/integration-tests/src/test/resources/golden/androidx.compose.compiler.plugins.kotlin.FunctionBodySkippingTransformTests/test_groupAroundIfComposeCallInIfConditionWithShortCircuit\133useFir = false\135.txt"
index e846eb8..97c3c4d 100644
--- "a/plugins/compose/compiler-hosted/integration-tests/src/test/resources/golden/androidx.compose.compiler.plugins.kotlin.FunctionBodySkippingTransformTests/test_groupAroundIfComposeCallInIfConditionWithShortCircuit\133useFir = false\135.txt"
+++ "b/plugins/compose/compiler-hosted/integration-tests/src/test/resources/golden/androidx.compose.compiler.plugins.kotlin.FunctionBodySkippingTransformTests/test_groupAroundIfComposeCallInIfConditionWithShortCircuit\133useFir = false\135.txt"

@@ -21,11 +21,21 @@
     if (isTraceInProgress()) {
       traceEventStart(<>, %changed, -1, <>)
     }
-    ReceiveValue(if (state && %composer.startReplaceGroup(<>)
-    sourceInformation(%composer, "<getCon...>")
-    val tmp0_group = getCondition(%composer, 0)
-    %composer.endReplaceGroup()
-    tmp0_group) 0 else 1, %composer, 0)
+    ReceiveValue(if (when {
+      state -> {
+        %composer.startReplaceGroup(<>)
+        sourceInformation(%composer, "<getCon...>")
+        val tmp0_group = getCondition(%composer, 0)
+        %composer.endReplaceGroup()
+        tmp0_group
+      }
+      else -> {
+        %composer.startReplaceGroup(<>)
+        %composer.endReplaceGroup()
+        false
+      }
+    }
+    ) 0 else 1, %composer, 0)
     if (isTraceInProgress()) {
       traceEventEnd()
     }

diff --git "a/plugins/compose/compiler-hosted/integration-tests/src/test/resources/golden/androidx.compose.compiler.plugins.kotlin.FunctionBodySkippingTransformTests/test_groupAroundIfComposeCallInIfConditionWithShortCircuit\133useFir = true\135.txt" "b/plugins/compose/compiler-hosted/integration-tests/src/test/resources/golden/androidx.compose.compiler.plugins.kotlin.FunctionBodySkippingTransformTests/test_groupAroundIfComposeCallInIfConditionWithShortCircuit\133useFir = true\135.txt"
index e846eb8..97c3c4d 100644
--- "a/plugins/compose/compiler-hosted/integration-tests/src/test/resources/golden/androidx.compose.compiler.plugins.kotlin.FunctionBodySkippingTransformTests/test_groupAroundIfComposeCallInIfConditionWithShortCircuit\133useFir = true\135.txt"
+++ "b/plugins/compose/compiler-hosted/integration-tests/src/test/resources/golden/androidx.compose.compiler.plugins.kotlin.FunctionBodySkippingTransformTests/test_groupAroundIfComposeCallInIfConditionWithShortCircuit\133useFir = true\135.txt"

@@ -21,11 +21,21 @@
     if (isTraceInProgress()) {
       traceEventStart(<>, %changed, -1, <>)
     }
-    ReceiveValue(if (state && %composer.startReplaceGroup(<>)
-    sourceInformation(%composer, "<getCon...>")
-    val tmp0_group = getCondition(%composer, 0)
-    %composer.endReplaceGroup()
-    tmp0_group) 0 else 1, %composer, 0)
+    ReceiveValue(if (when {
+      state -> {
+        %composer.startReplaceGroup(<>)
+        sourceInformation(%composer, "<getCon...>")
+        val tmp0_group = getCondition(%composer, 0)
+        %composer.endReplaceGroup()
+        tmp0_group
+      }
+      else -> {
+        %composer.startReplaceGroup(<>)
+        %composer.endReplaceGroup()
+        false
+      }
+    }
+    ) 0 else 1, %composer, 0)
     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 575c950..1e33713 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

@@ -3830,7 +3830,6 @@
         // conditionally executed.
         var needsWrappingGroup = false
         var resultsWithCalls = 0
-        var conditionsWithCalls = 0
         var hasElseBranch = false
 
         val transformed = IrWhenImpl(
@@ -3872,19 +3871,13 @@
                     condScopes.add(condScope)
                     resultScopes.add(resultScope)
 
-                    // 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)
-                    }
+                    // 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(
@@ -3908,13 +3901,6 @@
             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
@@ -3948,7 +3934,7 @@
 
         forEachWith(transformed.branches, condScopes, resultScopes) { it, condScope, resultScope ->
             if (condScope.hasComposableCalls) {
-                if (needsWrappingGroup) {
+                if (needsWrappingGroup && !FeatureFlag.OptimizeNonSkippingGroups.enabled) {
                     // Generate a group around the conditional block when it has a composable call
                     // in it and we are generating a group around when block.
                     it.condition = it.condition.asReplaceGroup(condScope)
@@ -3973,10 +3959,6 @@
                 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.
@@ -3984,9 +3966,20 @@
             }
         }
 
+        if (
+            FeatureFlag.OptimizeNonSkippingGroups.enabled && needsResultGroups && (
+                    transformed.origin == IrStatementOrigin.ANDAND || transformed.origin == IrStatementOrigin.OROR
+                    )
+        ) {
+            // When a IrWhen has a ANDAND or OROR origin it is required they also have a specific shape such as for ANDAND requires a
+            // `true -> false` clause at the end.  As we violate this by adding a wrapping group around all results, this origin is removed
+            // down-stream lowerings will no longer special case this IrWhen.
+            transformed.origin = IrStatementOrigin.WHEN
+        }
+
         return when {
-            !FeatureFlag.OptimizeNonSkippingGroups.enabled &&
-                    (resultsWithCalls == 1 || needsWrappingGroup) -> transformed.asCoalescableGroup(whenScope)
+            ((!FeatureFlag.OptimizeNonSkippingGroups.enabled && resultsWithCalls == 1) || needsWrappingGroup) ->
+                transformed.asCoalescableGroup(whenScope)
             else -> transformed
         }
     }