Save Expandable state across navigation

To save and restore Expandable state across navigation, we are using
rememberSaveable instead of remember. This ensures that the data is
stored while navigating and restore when coming back to screen.

Bug: 289991224
Test: Existing tests
Relnote: """
Use rememberSaveable to save and restore Expandable state across
navigation.
"""

Change-Id: I77285e6d2e8bfa0ed6c3992f75d575a5e349e426
diff --git a/wear/compose/compose-foundation/api/current.txt b/wear/compose/compose-foundation/api/current.txt
index 1db9c05..729cabc 100644
--- a/wear/compose/compose-foundation/api/current.txt
+++ b/wear/compose/compose-foundation/api/current.txt
@@ -210,6 +210,11 @@
     method public void setExpanded(boolean);
     property public final float expandProgress;
     property public final boolean expanded;
+    field public static final androidx.wear.compose.foundation.ExpandableState.Companion Companion;
+  }
+
+  public static final class ExpandableState.Companion {
+    method @androidx.compose.runtime.Composable public androidx.compose.runtime.saveable.Saver<androidx.wear.compose.foundation.ExpandableState,java.lang.Boolean> saver(androidx.compose.animation.core.AnimationSpec<java.lang.Float> expandAnimationSpec, androidx.compose.animation.core.AnimationSpec<java.lang.Float> collapseAnimationSpec);
   }
 
   @SuppressCompatibility @androidx.wear.compose.foundation.ExperimentalWearFoundationApi public final class ExpandableStateMapping<T> {
diff --git a/wear/compose/compose-foundation/api/restricted_current.txt b/wear/compose/compose-foundation/api/restricted_current.txt
index 1db9c05..729cabc 100644
--- a/wear/compose/compose-foundation/api/restricted_current.txt
+++ b/wear/compose/compose-foundation/api/restricted_current.txt
@@ -210,6 +210,11 @@
     method public void setExpanded(boolean);
     property public final float expandProgress;
     property public final boolean expanded;
+    field public static final androidx.wear.compose.foundation.ExpandableState.Companion Companion;
+  }
+
+  public static final class ExpandableState.Companion {
+    method @androidx.compose.runtime.Composable public androidx.compose.runtime.saveable.Saver<androidx.wear.compose.foundation.ExpandableState,java.lang.Boolean> saver(androidx.compose.animation.core.AnimationSpec<java.lang.Float> expandAnimationSpec, androidx.compose.animation.core.AnimationSpec<java.lang.Float> collapseAnimationSpec);
   }
 
   @SuppressCompatibility @androidx.wear.compose.foundation.ExperimentalWearFoundationApi public final class ExpandableStateMapping<T> {
diff --git a/wear/compose/compose-foundation/src/androidTest/kotlin/androidx/wear/compose/foundation/ExpandableTest.kt b/wear/compose/compose-foundation/src/androidTest/kotlin/androidx/wear/compose/foundation/ExpandableTest.kt
index b79b17d..773663d 100644
--- a/wear/compose/compose-foundation/src/androidTest/kotlin/androidx/wear/compose/foundation/ExpandableTest.kt
+++ b/wear/compose/compose-foundation/src/androidTest/kotlin/androidx/wear/compose/foundation/ExpandableTest.kt
@@ -33,6 +33,7 @@
 import androidx.compose.ui.platform.testTag
 import androidx.compose.ui.test.captureToImage
 import androidx.compose.ui.test.junit4.ComposeTestRule
+import androidx.compose.ui.test.junit4.StateRestorationTester
 import androidx.compose.ui.test.junit4.createComposeRule
 import androidx.compose.ui.test.onNodeWithTag
 import androidx.compose.ui.test.performClick
@@ -49,6 +50,8 @@
     @get:Rule
     val rule = createComposeRule()
 
+    private val restorationTester = StateRestorationTester(rule)
+
     @RequiresApi(Build.VERSION_CODES.O)
     @Test
     fun initially_collapsed() =
@@ -103,6 +106,24 @@
     @Test
     fun expanded_click() = verifyClick(true)
 
+    @Test
+    fun restoreState_after_recomposition() {
+        var expandableState: ExpandableState? = null
+        restorationTester.setContent {
+            expandableState = rememberExpandableState() // initially set expanded to false
+            expandableState?.expanded = true
+        }
+
+        rule.runOnUiThread {
+            // set to null which signifies recomposition
+            expandableState = null
+        }
+
+        restorationTester.emulateSavedInstanceStateRestore()
+
+        assertEquals(expandableState?.expanded, true)
+    }
+
     @RequiresApi(Build.VERSION_CODES.O)
     private fun verifyClick(initiallyExpanded: Boolean) {
         val clicked = mutableListOf<Boolean>()
diff --git a/wear/compose/compose-foundation/src/main/java/androidx/wear/compose/foundation/Expandable.kt b/wear/compose/compose-foundation/src/main/java/androidx/wear/compose/foundation/Expandable.kt
index 84fd745..4a6c60c 100644
--- a/wear/compose/compose-foundation/src/main/java/androidx/wear/compose/foundation/Expandable.kt
+++ b/wear/compose/compose-foundation/src/main/java/androidx/wear/compose/foundation/Expandable.kt
@@ -25,6 +25,8 @@
 import androidx.compose.runtime.mutableStateMapOf
 import androidx.compose.runtime.remember
 import androidx.compose.runtime.rememberCoroutineScope
+import androidx.compose.runtime.saveable.Saver
+import androidx.compose.runtime.saveable.rememberSaveable
 import androidx.compose.ui.Modifier
 import androidx.compose.ui.draw.clipToBounds
 import androidx.compose.ui.layout.Layout
@@ -55,7 +57,12 @@
     collapseAnimationSpec: AnimationSpec<Float> = ExpandableItemsDefaults.collapseAnimationSpec,
 ): ExpandableState {
     val scope = rememberCoroutineScope()
-    return remember {
+    return rememberSaveable(
+        saver = ExpandableState.saver(
+            expandAnimationSpec = expandAnimationSpec,
+            collapseAnimationSpec = collapseAnimationSpec,
+        )
+    ) {
         ExpandableState(initiallyExpanded, scope, expandAnimationSpec, collapseAnimationSpec)
     }
 }
@@ -258,6 +265,30 @@
                 }
             }
         }
+
+    companion object {
+        /**
+         * The default [Saver] implementation for [ExpandableState].
+         */
+        @Composable
+        fun saver(
+            expandAnimationSpec: AnimationSpec<Float>,
+            collapseAnimationSpec: AnimationSpec<Float>,
+        ): Saver<ExpandableState, Boolean> {
+            val coroutineScope = rememberCoroutineScope()
+            return Saver(
+                save = { it.expanded },
+                restore = {
+                    ExpandableState(
+                        initiallyExpanded = it,
+                        expandAnimationSpec = expandAnimationSpec,
+                        collapseAnimationSpec = collapseAnimationSpec,
+                        coroutineScope = coroutineScope
+                    )
+                }
+            )
+        }
+    }
 }
 
 /**