[Compose] Add compiler support for pausable composition

Pausable composition is an experimental runtime feature that allows
a composition to split across multiple frames.

Relnote: Added the PausableComposition feature flags
Bug: 328817808

KT-71227 Fixed

(cherry picked from commit ee9217f8f0f37967684fbfe4a568c2b3c8707507)
diff --git a/libraries/tools/kotlin-compose-compiler/src/common/kotlin/org/jetbrains/kotlin/compose/compiler/gradle/ComposeFeatureFlags.kt b/libraries/tools/kotlin-compose-compiler/src/common/kotlin/org/jetbrains/kotlin/compose/compiler/gradle/ComposeFeatureFlags.kt
index 0a9d0ca..956ab40 100644
--- a/libraries/tools/kotlin-compose-compiler/src/common/kotlin/org/jetbrains/kotlin/compose/compiler/gradle/ComposeFeatureFlags.kt
+++ b/libraries/tools/kotlin-compose-compiler/src/common/kotlin/org/jetbrains/kotlin/compose/compiler/gradle/ComposeFeatureFlags.kt
@@ -66,6 +66,7 @@
         StrongSkipping("StrongSkipping"),
         IntrinsicRemember("IntrinsicRemember"),
         OptimizeNonSkippingGroups("OptimizeNonSkippingGroups"),
+        PausableComposition("PausableComposition"),
     }
 
     companion object {
@@ -124,5 +125,22 @@
          */
         @JvmField
         val OptimizeNonSkippingGroups: ComposeFeatureFlag = Enabled(Feature.OptimizeNonSkippingGroups)
+
+        /**
+         * Change the code generation of composable function to enable pausing when part of pausable composition.
+         *
+         * Pausable composition is an experimental runtime feature. Experiments with this feature can be run by enabling this feature flag
+         * and using a runtime version that supports pausable composition. If the runtime used does not support pausable composition, no
+         * change is made to the code generation.
+         *
+         * This feature is still considered experimental and is thus disabled by default. It can be enabled by adding,
+         *```
+         * composeCompiler {
+         *   featureFlag = setOf(ComposeFeatureFlag.PausableComposition)
+         * }
+         * ```
+         */
+        @JvmField
+        val PausableComposition: ComposeFeatureFlag = Enabled(Feature.PausableComposition)
     }
 }
diff --git a/libraries/tools/kotlin-compose-compiler/src/functionalTest/kotlin/org/jetbrains/kotlin/compose/compiler/gradle/ExtensionConfigurationTest.kt b/libraries/tools/kotlin-compose-compiler/src/functionalTest/kotlin/org/jetbrains/kotlin/compose/compiler/gradle/ExtensionConfigurationTest.kt
index 660cf71..36b8ee3 100644
--- a/libraries/tools/kotlin-compose-compiler/src/functionalTest/kotlin/org/jetbrains/kotlin/compose/compiler/gradle/ExtensionConfigurationTest.kt
+++ b/libraries/tools/kotlin-compose-compiler/src/functionalTest/kotlin/org/jetbrains/kotlin/compose/compiler/gradle/ExtensionConfigurationTest.kt
@@ -117,6 +117,13 @@
     }
 
     @Test
+    fun enablePausableComposition() {
+        testComposeFeatureFlags(listOf("PausableComposition")) { extension ->
+            extension.featureFlags.value(setOf(ComposeFeatureFlag.PausableComposition))
+        }
+    }
+
+    @Test
     fun disableIntrinsicRememberCompatibility() {
         testComposeFeatureFlags(listOf("-IntrinsicRemember")) { extension ->
             @Suppress("DEPRECATION")
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/kotlin/androidx/compose/compiler/plugins/kotlin/ComposeModuleMetricsTests.kt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/kotlin/androidx/compose/compiler/plugins/kotlin/ComposeModuleMetricsTests.kt
index 9e7ad71..2da112c 100644
--- a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/kotlin/androidx/compose/compiler/plugins/kotlin/ComposeModuleMetricsTests.kt
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/kotlin/androidx/compose/compiler/plugins/kotlin/ComposeModuleMetricsTests.kt
@@ -151,7 +151,8 @@
               "featureFlags": {
                 "StrongSkipping": false,
                 "IntrinsicRemember": true,
-                "OptimizeNonSkippingGroups": false
+                "OptimizeNonSkippingGroups": false,
+                "PausableComposition": false
               }
             }
         """
diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/kotlin/androidx/compose/compiler/plugins/kotlin/ComposePausableCompositionTests.kt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/kotlin/androidx/compose/compiler/plugins/kotlin/ComposePausableCompositionTests.kt
new file mode 100644
index 0000000..6182e7d
--- /dev/null
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/kotlin/androidx/compose/compiler/plugins/kotlin/ComposePausableCompositionTests.kt
@@ -0,0 +1,86 @@
+/*
+ * Copyright 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package androidx.compose.compiler.plugins.kotlin
+
+import org.jetbrains.kotlin.config.CompilerConfiguration
+import org.junit.runners.Parameterized
+import kotlin.test.Test
+
+class ComposePausableCompositionTests(
+    useFir: Boolean,
+    private val pausableEnabled: Boolean
+) : AbstractControlFlowTransformTests(useFir) {
+    companion object {
+        @JvmStatic
+        @Parameterized.Parameters(name = "useFir = {0}, pausableEnabled = {1}")
+        fun data() = arrayOf<Any>(
+            arrayOf(false, false),
+            arrayOf(false, true),
+            arrayOf(true, false),
+            arrayOf(true, true)
+        )
+    }
+
+    override fun CompilerConfiguration.updateConfiguration() {
+        put(ComposeConfiguration.SOURCE_INFORMATION_ENABLED_KEY, false)
+        put(ComposeConfiguration.TRACE_MARKERS_ENABLED_KEY, false)
+        put(
+            ComposeConfiguration.FEATURE_FLAGS,
+            listOf(FeatureFlag.PausableComposition.name(pausableEnabled))
+        )
+    }
+
+    @Test
+    fun testRestartableComposableFunction() = verifyGoldenComposeIrTransform(
+        source = """
+            import androidx.compose.runtime.*
+
+            @Composable
+            fun Test(a: Int, b: String, c: Float) {
+                use(a)
+                use(b)
+                use(c)
+            }
+        """,
+        extra = """
+            fun use(value: Any?) { println(value) } 
+        """
+    )
+
+    @Test
+    fun testRestartableComposableLambda() = verifyGoldenComposeIrTransform(
+        source = """
+            import androidx.compose.runtime.*
+
+            @Composable
+            fun Test(a: Int, b: String, c: Float) {
+                Wrap {
+                    use(a)
+                    use(b)
+                    use(c)
+                }
+            }
+        """,
+        extra = """
+            import androidx.compose.runtime.*
+
+            fun use(value: Any?) { println(value) }
+            @Composable fun Wrap(content: @Composable () -> Unit) = content()
+        """
+    )
+
+}
\ No newline at end of file
diff --git "a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ComposePausableCompositionTests/testRestartableComposableFunction\133useFir = false, pausableEnabled = false\135.txt" "b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ComposePausableCompositionTests/testRestartableComposableFunction\133useFir = false, pausableEnabled = false\135.txt"
new file mode 100644
index 0000000..66860c0
--- /dev/null
+++ "b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ComposePausableCompositionTests/testRestartableComposableFunction\133useFir = false, pausableEnabled = false\135.txt"
@@ -0,0 +1,41 @@
+//
+// Source
+// ------------------------------------------
+
+import androidx.compose.runtime.*
+
+@Composable
+fun Test(a: Int, b: String, c: Float) {
+    use(a)
+    use(b)
+    use(c)
+}
+
+//
+// Transformed IR
+// ------------------------------------------
+
+@Composable
+fun Test(a: Int, b: String, c: Float, %composer: Composer?, %changed: Int) {
+  %composer = %composer.startRestartGroup(<>)
+  val %dirty = %changed
+  if (%changed and 0b0110 == 0) {
+    %dirty = %dirty or if (%composer.changed(a)) 0b0100 else 0b0010
+  }
+  if (%changed and 0b00110000 == 0) {
+    %dirty = %dirty or if (%composer.changed(b)) 0b00100000 else 0b00010000
+  }
+  if (%changed and 0b000110000000 == 0) {
+    %dirty = %dirty or if (%composer.changed(c)) 0b000100000000 else 0b10000000
+  }
+  if (%dirty and 0b10010011 != 0b10010010 || !%composer.skipping) {
+    use(a)
+    use(b)
+    use(c)
+  } else {
+    %composer.skipToGroupEnd()
+  }
+  %composer.endRestartGroup()?.updateScope { %composer: Composer?, %force: Int ->
+    Test(a, b, c, %composer, updateChangedFlags(%changed or 0b0001))
+  }
+}
diff --git "a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ComposePausableCompositionTests/testRestartableComposableFunction\133useFir = false, pausableEnabled = true\135.txt" "b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ComposePausableCompositionTests/testRestartableComposableFunction\133useFir = false, pausableEnabled = true\135.txt"
new file mode 100644
index 0000000..96b143e
--- /dev/null
+++ "b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ComposePausableCompositionTests/testRestartableComposableFunction\133useFir = false, pausableEnabled = true\135.txt"
@@ -0,0 +1,41 @@
+//
+// Source
+// ------------------------------------------
+
+import androidx.compose.runtime.*
+
+@Composable
+fun Test(a: Int, b: String, c: Float) {
+    use(a)
+    use(b)
+    use(c)
+}
+
+//
+// Transformed IR
+// ------------------------------------------
+
+@Composable
+fun Test(a: Int, b: String, c: Float, %composer: Composer?, %changed: Int) {
+  %composer = %composer.startRestartGroup(<>)
+  val %dirty = %changed
+  if (%changed and 0b0110 == 0) {
+    %dirty = %dirty or if (%composer.changed(a)) 0b0100 else 0b0010
+  }
+  if (%changed and 0b00110000 == 0) {
+    %dirty = %dirty or if (%composer.changed(b)) 0b00100000 else 0b00010000
+  }
+  if (%changed and 0b000110000000 == 0) {
+    %dirty = %dirty or if (%composer.changed(c)) 0b000100000000 else 0b10000000
+  }
+  if (%composer.shouldExecute(%dirty and 0b10010011 != 0b10010010, %dirty and 0b0001)) {
+    use(a)
+    use(b)
+    use(c)
+  } else {
+    %composer.skipToGroupEnd()
+  }
+  %composer.endRestartGroup()?.updateScope { %composer: Composer?, %force: Int ->
+    Test(a, b, c, %composer, updateChangedFlags(%changed or 0b0001))
+  }
+}
diff --git "a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ComposePausableCompositionTests/testRestartableComposableFunction\133useFir = true, pausableEnabled = false\135.txt" "b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ComposePausableCompositionTests/testRestartableComposableFunction\133useFir = true, pausableEnabled = false\135.txt"
new file mode 100644
index 0000000..66860c0
--- /dev/null
+++ "b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ComposePausableCompositionTests/testRestartableComposableFunction\133useFir = true, pausableEnabled = false\135.txt"
@@ -0,0 +1,41 @@
+//
+// Source
+// ------------------------------------------
+
+import androidx.compose.runtime.*
+
+@Composable
+fun Test(a: Int, b: String, c: Float) {
+    use(a)
+    use(b)
+    use(c)
+}
+
+//
+// Transformed IR
+// ------------------------------------------
+
+@Composable
+fun Test(a: Int, b: String, c: Float, %composer: Composer?, %changed: Int) {
+  %composer = %composer.startRestartGroup(<>)
+  val %dirty = %changed
+  if (%changed and 0b0110 == 0) {
+    %dirty = %dirty or if (%composer.changed(a)) 0b0100 else 0b0010
+  }
+  if (%changed and 0b00110000 == 0) {
+    %dirty = %dirty or if (%composer.changed(b)) 0b00100000 else 0b00010000
+  }
+  if (%changed and 0b000110000000 == 0) {
+    %dirty = %dirty or if (%composer.changed(c)) 0b000100000000 else 0b10000000
+  }
+  if (%dirty and 0b10010011 != 0b10010010 || !%composer.skipping) {
+    use(a)
+    use(b)
+    use(c)
+  } else {
+    %composer.skipToGroupEnd()
+  }
+  %composer.endRestartGroup()?.updateScope { %composer: Composer?, %force: Int ->
+    Test(a, b, c, %composer, updateChangedFlags(%changed or 0b0001))
+  }
+}
diff --git "a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ComposePausableCompositionTests/testRestartableComposableFunction\133useFir = true, pausableEnabled = true\135.txt" "b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ComposePausableCompositionTests/testRestartableComposableFunction\133useFir = true, pausableEnabled = true\135.txt"
new file mode 100644
index 0000000..96b143e
--- /dev/null
+++ "b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ComposePausableCompositionTests/testRestartableComposableFunction\133useFir = true, pausableEnabled = true\135.txt"
@@ -0,0 +1,41 @@
+//
+// Source
+// ------------------------------------------
+
+import androidx.compose.runtime.*
+
+@Composable
+fun Test(a: Int, b: String, c: Float) {
+    use(a)
+    use(b)
+    use(c)
+}
+
+//
+// Transformed IR
+// ------------------------------------------
+
+@Composable
+fun Test(a: Int, b: String, c: Float, %composer: Composer?, %changed: Int) {
+  %composer = %composer.startRestartGroup(<>)
+  val %dirty = %changed
+  if (%changed and 0b0110 == 0) {
+    %dirty = %dirty or if (%composer.changed(a)) 0b0100 else 0b0010
+  }
+  if (%changed and 0b00110000 == 0) {
+    %dirty = %dirty or if (%composer.changed(b)) 0b00100000 else 0b00010000
+  }
+  if (%changed and 0b000110000000 == 0) {
+    %dirty = %dirty or if (%composer.changed(c)) 0b000100000000 else 0b10000000
+  }
+  if (%composer.shouldExecute(%dirty and 0b10010011 != 0b10010010, %dirty and 0b0001)) {
+    use(a)
+    use(b)
+    use(c)
+  } else {
+    %composer.skipToGroupEnd()
+  }
+  %composer.endRestartGroup()?.updateScope { %composer: Composer?, %force: Int ->
+    Test(a, b, c, %composer, updateChangedFlags(%changed or 0b0001))
+  }
+}
diff --git "a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ComposePausableCompositionTests/testRestartableComposableLambda\133useFir = false, pausableEnabled = false\135.txt" "b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ComposePausableCompositionTests/testRestartableComposableLambda\133useFir = false, pausableEnabled = false\135.txt"
new file mode 100644
index 0000000..11fc4ab
--- /dev/null
+++ "b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ComposePausableCompositionTests/testRestartableComposableLambda\133useFir = false, pausableEnabled = false\135.txt"
@@ -0,0 +1,49 @@
+//
+// Source
+// ------------------------------------------
+
+import androidx.compose.runtime.*
+
+@Composable
+fun Test(a: Int, b: String, c: Float) {
+    Wrap {
+        use(a)
+        use(b)
+        use(c)
+    }
+}
+
+//
+// Transformed IR
+// ------------------------------------------
+
+@Composable
+fun Test(a: Int, b: String, c: Float, %composer: Composer?, %changed: Int) {
+  %composer = %composer.startRestartGroup(<>)
+  val %dirty = %changed
+  if (%changed and 0b0110 == 0) {
+    %dirty = %dirty or if (%composer.changed(a)) 0b0100 else 0b0010
+  }
+  if (%changed and 0b00110000 == 0) {
+    %dirty = %dirty or if (%composer.changed(b)) 0b00100000 else 0b00010000
+  }
+  if (%changed and 0b000110000000 == 0) {
+    %dirty = %dirty or if (%composer.changed(c)) 0b000100000000 else 0b10000000
+  }
+  if (%dirty and 0b10010011 != 0b10010010 || !%composer.skipping) {
+    Wrap(rememberComposableLambda(<>, true, { %composer: Composer?, %changed: Int ->
+      if (%changed and 0b0011 != 0b0010 || !%composer.skipping) {
+        use(a)
+        use(b)
+        use(c)
+      } else {
+        %composer.skipToGroupEnd()
+      }
+    }, %composer, 0b00110110), %composer, 0b0110)
+  } else {
+    %composer.skipToGroupEnd()
+  }
+  %composer.endRestartGroup()?.updateScope { %composer: Composer?, %force: Int ->
+    Test(a, b, c, %composer, updateChangedFlags(%changed or 0b0001))
+  }
+}
diff --git "a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ComposePausableCompositionTests/testRestartableComposableLambda\133useFir = false, pausableEnabled = true\135.txt" "b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ComposePausableCompositionTests/testRestartableComposableLambda\133useFir = false, pausableEnabled = true\135.txt"
new file mode 100644
index 0000000..3ccce03
--- /dev/null
+++ "b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ComposePausableCompositionTests/testRestartableComposableLambda\133useFir = false, pausableEnabled = true\135.txt"
@@ -0,0 +1,49 @@
+//
+// Source
+// ------------------------------------------
+
+import androidx.compose.runtime.*
+
+@Composable
+fun Test(a: Int, b: String, c: Float) {
+    Wrap {
+        use(a)
+        use(b)
+        use(c)
+    }
+}
+
+//
+// Transformed IR
+// ------------------------------------------
+
+@Composable
+fun Test(a: Int, b: String, c: Float, %composer: Composer?, %changed: Int) {
+  %composer = %composer.startRestartGroup(<>)
+  val %dirty = %changed
+  if (%changed and 0b0110 == 0) {
+    %dirty = %dirty or if (%composer.changed(a)) 0b0100 else 0b0010
+  }
+  if (%changed and 0b00110000 == 0) {
+    %dirty = %dirty or if (%composer.changed(b)) 0b00100000 else 0b00010000
+  }
+  if (%changed and 0b000110000000 == 0) {
+    %dirty = %dirty or if (%composer.changed(c)) 0b000100000000 else 0b10000000
+  }
+  if (%composer.shouldExecute(%dirty and 0b10010011 != 0b10010010, %dirty and 0b0001)) {
+    Wrap(rememberComposableLambda(<>, true, { %composer: Composer?, %changed: Int ->
+      if (%composer.shouldExecute(%changed and 0b0011 != 0b0010, %changed and 0b0001)) {
+        use(a)
+        use(b)
+        use(c)
+      } else {
+        %composer.skipToGroupEnd()
+      }
+    }, %composer, 0b00110110), %composer, 0b0110)
+  } else {
+    %composer.skipToGroupEnd()
+  }
+  %composer.endRestartGroup()?.updateScope { %composer: Composer?, %force: Int ->
+    Test(a, b, c, %composer, updateChangedFlags(%changed or 0b0001))
+  }
+}
diff --git "a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ComposePausableCompositionTests/testRestartableComposableLambda\133useFir = true, pausableEnabled = false\135.txt" "b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ComposePausableCompositionTests/testRestartableComposableLambda\133useFir = true, pausableEnabled = false\135.txt"
new file mode 100644
index 0000000..11fc4ab
--- /dev/null
+++ "b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ComposePausableCompositionTests/testRestartableComposableLambda\133useFir = true, pausableEnabled = false\135.txt"
@@ -0,0 +1,49 @@
+//
+// Source
+// ------------------------------------------
+
+import androidx.compose.runtime.*
+
+@Composable
+fun Test(a: Int, b: String, c: Float) {
+    Wrap {
+        use(a)
+        use(b)
+        use(c)
+    }
+}
+
+//
+// Transformed IR
+// ------------------------------------------
+
+@Composable
+fun Test(a: Int, b: String, c: Float, %composer: Composer?, %changed: Int) {
+  %composer = %composer.startRestartGroup(<>)
+  val %dirty = %changed
+  if (%changed and 0b0110 == 0) {
+    %dirty = %dirty or if (%composer.changed(a)) 0b0100 else 0b0010
+  }
+  if (%changed and 0b00110000 == 0) {
+    %dirty = %dirty or if (%composer.changed(b)) 0b00100000 else 0b00010000
+  }
+  if (%changed and 0b000110000000 == 0) {
+    %dirty = %dirty or if (%composer.changed(c)) 0b000100000000 else 0b10000000
+  }
+  if (%dirty and 0b10010011 != 0b10010010 || !%composer.skipping) {
+    Wrap(rememberComposableLambda(<>, true, { %composer: Composer?, %changed: Int ->
+      if (%changed and 0b0011 != 0b0010 || !%composer.skipping) {
+        use(a)
+        use(b)
+        use(c)
+      } else {
+        %composer.skipToGroupEnd()
+      }
+    }, %composer, 0b00110110), %composer, 0b0110)
+  } else {
+    %composer.skipToGroupEnd()
+  }
+  %composer.endRestartGroup()?.updateScope { %composer: Composer?, %force: Int ->
+    Test(a, b, c, %composer, updateChangedFlags(%changed or 0b0001))
+  }
+}
diff --git "a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ComposePausableCompositionTests/testRestartableComposableLambda\133useFir = true, pausableEnabled = true\135.txt" "b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ComposePausableCompositionTests/testRestartableComposableLambda\133useFir = true, pausableEnabled = true\135.txt"
new file mode 100644
index 0000000..3ccce03
--- /dev/null
+++ "b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/resources/golden/androidx.compose.compiler.plugins.kotlin.ComposePausableCompositionTests/testRestartableComposableLambda\133useFir = true, pausableEnabled = true\135.txt"
@@ -0,0 +1,49 @@
+//
+// Source
+// ------------------------------------------
+
+import androidx.compose.runtime.*
+
+@Composable
+fun Test(a: Int, b: String, c: Float) {
+    Wrap {
+        use(a)
+        use(b)
+        use(c)
+    }
+}
+
+//
+// Transformed IR
+// ------------------------------------------
+
+@Composable
+fun Test(a: Int, b: String, c: Float, %composer: Composer?, %changed: Int) {
+  %composer = %composer.startRestartGroup(<>)
+  val %dirty = %changed
+  if (%changed and 0b0110 == 0) {
+    %dirty = %dirty or if (%composer.changed(a)) 0b0100 else 0b0010
+  }
+  if (%changed and 0b00110000 == 0) {
+    %dirty = %dirty or if (%composer.changed(b)) 0b00100000 else 0b00010000
+  }
+  if (%changed and 0b000110000000 == 0) {
+    %dirty = %dirty or if (%composer.changed(c)) 0b000100000000 else 0b10000000
+  }
+  if (%composer.shouldExecute(%dirty and 0b10010011 != 0b10010010, %dirty and 0b0001)) {
+    Wrap(rememberComposableLambda(<>, true, { %composer: Composer?, %changed: Int ->
+      if (%composer.shouldExecute(%changed and 0b0011 != 0b0010, %changed and 0b0001)) {
+        use(a)
+        use(b)
+        use(c)
+      } else {
+        %composer.skipToGroupEnd()
+      }
+    }, %composer, 0b00110110), %composer, 0b0110)
+  } else {
+    %composer.skipToGroupEnd()
+  }
+  %composer.endRestartGroup()?.updateScope { %composer: Composer?, %force: Int ->
+    Test(a, b, c, %composer, updateChangedFlags(%changed or 0b0001))
+  }
+}
diff --git a/plugins/compose/compiler-hosted/src/main/java/androidx/compose/compiler/plugins/kotlin/ComposeNames.kt b/plugins/compose/compiler-hosted/src/main/java/androidx/compose/compiler/plugins/kotlin/ComposeNames.kt
index 922c1b9..1f0c442 100644
--- a/plugins/compose/compiler-hosted/src/main/java/androidx/compose/compiler/plugins/kotlin/ComposeNames.kt
+++ b/plugins/compose/compiler-hosted/src/main/java/androidx/compose/compiler/plugins/kotlin/ComposeNames.kt
@@ -27,4 +27,5 @@
     val REMEMBER_COMPOSABLE_LAMBDA = "rememberComposableLambda"
     val REMEMBER_COMPOSABLE_LAMBDAN = "rememberComposableLambdaN"
     val DEFAULT_IMPLS = Name.identifier("ComposeDefaultImpls")
+    val SHOULD_EXECUTE = Name.identifier("shouldExecute")
 }
diff --git a/plugins/compose/compiler-hosted/src/main/java/androidx/compose/compiler/plugins/kotlin/ComposePlugin.kt b/plugins/compose/compiler-hosted/src/main/java/androidx/compose/compiler/plugins/kotlin/ComposePlugin.kt
index 4c768b9..36f7ada 100644
--- a/plugins/compose/compiler-hosted/src/main/java/androidx/compose/compiler/plugins/kotlin/ComposePlugin.kt
+++ b/plugins/compose/compiler-hosted/src/main/java/androidx/compose/compiler/plugins/kotlin/ComposePlugin.kt
@@ -355,7 +355,7 @@
  * AbstractComposeLowering by using the FeatureFlag.enabled extension property. For example
  * testing if StrongSkipping is enabled can be checked by checking
  *
- *   FeatureFlag.StrongSkipping.enabled
+ *    FeatureFlag.StrongSkipping.enabled
  *
  * The `default` field is the source of truth for the default of the property. Turning it
  * to `true` here will make it default on even if the value was previous enabled through
@@ -371,7 +371,9 @@
 enum class FeatureFlag(val featureName: String, val default: Boolean) {
     StrongSkipping("StrongSkipping", default = true),
     IntrinsicRemember("IntrinsicRemember", default = true),
-    OptimizeNonSkippingGroups("OptimizeNonSkippingGroups", default = false);
+    OptimizeNonSkippingGroups("OptimizeNonSkippingGroups", default = false),
+    PausableComposition("PausableComposition", default = false),
+    ;
 
     val disabledName get() = "-$featureName"
     fun name(enabled: Boolean) = if (enabled) featureName else disabledName
@@ -383,8 +385,8 @@
                 featureName.startsWith("-") -> featureName.substring(1) to false
                 else -> featureName to true
             }
-            return FeatureFlag.values().firstOrNull {
-                featureToSearch.trim().compareTo(it.featureName, ignoreCase = true) == 0
+            return FeatureFlag.entries.firstOrNull {
+                featureToSearch.trim().equals(it.featureName, ignoreCase = true)
             } to enabled
         }
     }
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 2a1b6a3..c90b01a 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
@@ -77,6 +77,7 @@
      * look at it.
      */
     Uncertain(0b000),
+
     /**
      * This indicates that the value is known to be the same since the last time the function was
      * executed. There is no need to store the value in the slot table in this case because the
@@ -84,6 +85,7 @@
      * in the previous execution.
      */
     Same(0b001),
+
     /**
      * This indicates that the value is known to be different since the last time the function
      * was executed. There is no need to store the value in the slot table in this case because
@@ -91,6 +93,7 @@
      * was in the previous execution.
      */
     Different(0b010),
+
     /**
      * This indicates that the value is known to *never change* for the duration of the running
      * program.
@@ -176,7 +179,7 @@
     thisParams: Int = 0,
 ): Int {
     return 1 + // composer param
-        changedParamCount(realValueParams, thisParams)
+            changedParamCount(realValueParams, thisParams)
 }
 
 @JvmDefaultWithCompatibility
@@ -187,16 +190,19 @@
     fun irIsolateBitsAtSlot(slot: Int, includeStableBit: Boolean): IrExpression
     fun irSlotAnd(slot: Int, bits: Int): IrExpression
     fun irHasDifferences(usedParams: BooleanArray): IrExpression
+    fun irRestartFlags(): IrExpression
     fun irCopyToTemporary(
         nameHint: String? = null,
         isVar: Boolean = false,
         exactName: Boolean = false
     ): IrChangedBitMaskVariable
+
     fun putAsValueArgumentInWithLowBit(
         fn: IrFunctionAccessExpression,
         startIndex: Int,
         lowBit: Boolean
     )
+
     fun irShiftBits(fromSlot: Int, toSlot: Int): IrExpression
     fun irStableBitAtSlot(slot: Int): IrExpression
 }
@@ -475,6 +481,18 @@
             }
     }
 
+    private val shouldExecuteFunction by guardedLazy {
+        if (FeatureFlag.PausableComposition.enabled)
+            composerIrClass
+                .functions
+                .firstOrNull {
+                    it.name == ComposeNames.SHOULD_EXECUTE && it.valueParameters.size == 2 &&
+                            it.valueParameters.first().type.isBoolean() &&
+                            it.valueParameters.drop(1).first().type.isInt()
+                }
+        else null
+    }
+
     private val sourceInformationFunction by guardedLazy {
         getTopLevelFunction(ComposeCallableIds.sourceInformation).owner
     }
@@ -987,11 +1005,10 @@
             // 1. if any of the stable parameters have *differences* from last execution.
             // 2. if the composer.skipping call returns false
             // 3. function is inline
-            val shouldExecute = irOrOr(
+            val shouldExecute = irShouldExecute(
                 dirtyForSkipping.irHasDifferences(scope.usedParams),
-                irNot(irIsSkipping())
+                dirtyForSkipping.irRestartFlags(),
             )
-
             val transformedBody = irIfThenElse(
                 condition = shouldExecute,
                 thenPart = irBlock(
@@ -1157,9 +1174,9 @@
 
             // (3) is only necessary to check if we actually have unstable params, so we only
             // generate that check if we need to.
-            var shouldExecute = irOrOr(
+            var shouldExecute = irShouldExecute(
                 dirtyForSkipping.irHasDifferences(scope.usedParams),
-                irNot(irIsSkipping())
+                dirtyForSkipping.irRestartFlags(),
             )
 
             // boolean array mapped to parameters. true indicates that the type is unstable
@@ -1761,6 +1778,22 @@
 
     private fun irIsSkipping() =
         irMethodCall(irCurrentComposer(), isSkippingFunction.getter!!)
+
+    private fun irShouldExecute(parametersChanged: IrExpression, flags: IrExpression): IrExpression {
+        val shouldExecuteFunction = shouldExecuteFunction
+        return if (shouldExecuteFunction != null) {
+            irMethodCall(irCurrentComposer(), shouldExecuteFunction).apply {
+                putValueArgument(0, parametersChanged)
+                putValueArgument(1, flags)
+            }
+        } else {
+            irOrOr(
+                parametersChanged,
+                irNot(irIsSkipping())
+            )
+        }
+    }
+
     private fun irDefaultsInvalid() =
         irMethodCall(irCurrentComposer(), defaultsInvalidFunction.getter!!)
 
@@ -4513,6 +4546,9 @@
             )
         }
 
+        // The restart flag is always in the first parameter flags (or the implied changed parameter for 0 parameters)
+        override fun irRestartFlags(): IrExpression = irAnd(irGet(params[0]), irConst(1))
+
         override fun irHasDifferences(
             usedParams: BooleanArray
         ): IrExpression {