Realize groups in when expr if a wrapping group is omitted

Test: compiler tests
Fixes: 336571300 ( https://issuetracker.google.com/issues/336571300 )
Change-Id: I1ce2de38d9b13528a2039528d69e7c10854329fa ( https://android-review.googlesource.com/q/I1ce2de38d9b13528a2039528d69e7c10854329fa )

Moved from: androidx/androidx@9ab1db3

Author: ShikaSD

plugins/compose/compiler-hosted/integration-tests/src/androidUnitTest/kotlin/androidx/compose/compiler/plugins/kotlin/ComposeBytecodeCodegenTest.kt

@@ -811,4 +811,24 @@ class ComposeBytecodeCodegenTest(useFir: Boolean) : AbstractCodegenTest(useFir)
                 }
             },
         )
+
+    @Test // regression test for 336571300
+    fun test_groupAroundIfComposeCallInIfConditionWithShortCircuit() = testCompile(
+        source = """
+            import androidx.compose.runtime.*
+
+            @Composable
+            fun Test() {
+                ReceiveValue(if (state && getCondition()) 0 else 1)
+            }
+
+            val state by mutableStateOf(true)
+
+            @Composable
+            fun getCondition() = remember { mutableStateOf(false) }.value
+
+            @Composable
+            fun ReceiveValue(value: Int) { }
+        """
+    )
 }

plugins/compose/compiler-hosted/integration-tests/src/jvmTest/kotlin/androidx/compose/compiler/plugins/kotlin/FunctionBodySkippingTransformTests.kt

@@ -1281,6 +1281,30 @@ class FunctionBodySkippingTransformTests(
             fun Text(value: String) {}
         """
     )
+
+    @Test // regression test for 336571300
+    fun test_groupAroundIfComposeCallInIfConditionWithShortCircuit() =
+        verifyGoldenComposeIrTransform(
+            source = """
+                import androidx.compose.runtime.*
+
+                @Composable
+                fun Test() {
+                    ReceiveValue(if (state && getCondition()) 0 else 1)
+                }
+            """,
+            """
+                import androidx.compose.runtime.*
+
+                val state by mutableStateOf(true)
+
+                @Composable
+                fun getCondition() = remember { mutableStateOf(false) }.value
+
+                @Composable
+                fun ReceiveValue(value: Int) { }
+            """
+        )
 }
 
 class FunctionBodySkippingTransformTestsNoSource(

plugins/compose/compiler-hosted/integration-tests/src/test/resources/golden/androidx.compose.compiler.plugins.kotlin.FunctionBodySkippingTransformTests/test_groupAroundIfComposeCallInIfConditionWithShortCircuit[useFir = false].txt

@@ -0,0 +1,38 @@
+//
+// Source
+// ------------------------------------------
+
+import androidx.compose.runtime.*
+
+@Composable
+fun Test() {
+    ReceiveValue(if (state && getCondition()) 0 else 1)
+}
+
+//
+// Transformed IR
+// ------------------------------------------
+
+@Composable
+fun Test(%composer: Composer?, %changed: Int) {
+  %composer = %composer.startRestartGroup(<>)
+  sourceInformation(%composer, "C(Test)<Receiv...>:Test.kt")
+  if (%changed != 0 || !%composer.skipping) {
+    if (isTraceInProgress()) {
+      traceEventStart(<>, %changed, -1, <>)
+    }
+    ReceiveValue(if (%composer.startReplaceGroup(<>)
+    sourceInformation(%composer, "<getCon...>")
+    val tmp0_group = state && getCondition(%composer, 0)
+    %composer.endReplaceGroup()
+    tmp0_group) 0 else 1, %composer, 0)
+    if (isTraceInProgress()) {
+      traceEventEnd()
+    }
+  } else {
+    %composer.skipToGroupEnd()
+  }
+  %composer.endRestartGroup()?.updateScope { %composer: Composer?, %force: Int ->
+    Test(%composer, updateChangedFlags(%changed or 0b0001))
+  }
+}

plugins/compose/compiler-hosted/integration-tests/src/test/resources/golden/androidx.compose.compiler.plugins.kotlin.FunctionBodySkippingTransformTests/test_groupAroundIfComposeCallInIfConditionWithShortCircuit[useFir = true].txt

@@ -0,0 +1,38 @@
+//
+// Source
+// ------------------------------------------
+
+import androidx.compose.runtime.*
+
+@Composable
+fun Test() {
+    ReceiveValue(if (state && getCondition()) 0 else 1)
+}
+
+//
+// Transformed IR
+// ------------------------------------------
+
+@Composable
+fun Test(%composer: Composer?, %changed: Int) {
+  %composer = %composer.startRestartGroup(<>)
+  sourceInformation(%composer, "C(Test)<Receiv...>:Test.kt")
+  if (%changed != 0 || !%composer.skipping) {
+    if (isTraceInProgress()) {
+      traceEventStart(<>, %changed, -1, <>)
+    }
+    ReceiveValue(if (%composer.startReplaceGroup(<>)
+    sourceInformation(%composer, "<getCon...>")
+    val tmp0_group = state && getCondition(%composer, 0)
+    %composer.endReplaceGroup()
+    tmp0_group) 0 else 1, %composer, 0)
+    if (isTraceInProgress()) {
+      traceEventEnd()
+    }
+  } else {
+    %composer.skipToGroupEnd()
+  }
+  %composer.endRestartGroup()?.updateScope { %composer: Composer?, %force: Int ->
+    Test(%composer, updateChangedFlags(%changed or 0b0001))
+  }
+}

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

@@ -17,6 +17,7 @@
 package androidx.compose.compiler.test
 
 import androidx.compose.runtime.Composable
+import androidx.compose.runtime.NonRestartableComposable
 import androidx.compose.runtime.getValue
 import androidx.compose.runtime.mock.Text
 import androidx.compose.runtime.mock.compositionTest
@@ -26,6 +27,7 @@ import androidx.compose.runtime.remember
 import androidx.compose.runtime.rememberUpdatedState
 import androidx.compose.runtime.setValue
 import kotlin.test.Test
+import kotlin.test.assertEquals
 
 class CompositionTests {
     @Test
@@ -80,6 +82,37 @@ class CompositionTests {
             DefaultValueClass()
         }
     }
+
+    @Test
+    fun groupAroundIfComposeCallInIfConditionWithShortCircuit() = compositionTest {
+        var state by mutableStateOf(true)
+        compose {
+            ReceiveValue(if (state && getCondition()) 0 else 1)
+
+            ReceiveValue(
+                when {
+                    state -> when {
+                        state -> getCondition()
+                        else -> false
+                    }.let { if (it) 0 else 1 }
+                    else -> 1
+                }
+            )
+        }
+
+        state = false
+        advance()
+    }
+}
+
+@Composable
+fun getCondition() = remember { false }
+
+@NonRestartableComposable
+@Composable
+fun ReceiveValue(value: Int) {
+    val string = remember { "$value" }
+    assertEquals(1, string.length)
 }
 
 class CrossInlineState(content: @Composable () -> Unit = { }) {

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

@@ -3892,12 +3892,19 @@ class ComposableFunctionBodyTransformer(
         }
 
         forEachWith(transformed.branches, condScopes, resultScopes) { it, condScope, resultScope ->
-            // If the conditional block doesn't have a composable call in it, we don't need
-            // to generate a group around it because we will be generating one around the entire
-            // if statement
-            if (needsWrappingGroup && condScope.hasComposableCalls) {
-                it.condition = it.condition.asReplaceGroup(condScope)
+            if (condScope.hasComposableCalls) {
+                if (needsWrappingGroup) {
+                    // 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)
+                } else {
+                    // Ensure that the inner structure of condition is correct if the wrapping group
+                    // is not required by realizing groups in condition scope.
+                    condScope.realizeAllDirectChildren()
+                    condScope.realizeCoalescableGroup()
+                }
             }
+
             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