Ensure the ComposeNavigator back stack is restored

We need to make sure that the ComposeNavigator back stack is restored
through config changes and process death.

RelNote: "Fixed an issue where `popBackStack()` and `navigateUp()` would not work after a configuration change or process death and recreation."
Test: ComposeNavigatorTest
Bug: 173281473

Change-Id: Icea477e53970d8fcd22a698ed2feca1734503896

diff --git a/navigation/navigation-compose/src/androidTest/java/androidx/navigation/compose/ComposeNavigatorTest.kt b/navigation/navigation-compose/src/androidTest/java/androidx/navigation/compose/ComposeNavigatorTest.kt
index c11367b..1079ad8 100644
--- a/navigation/navigation-compose/src/androidTest/java/androidx/navigation/compose/ComposeNavigatorTest.kt
+++ b/navigation/navigation-compose/src/androidTest/java/androidx/navigation/compose/ComposeNavigatorTest.kt

@@ -16,21 +16,40 @@
 
 package androidx.navigation.compose
 
-import androidx.compose.ui.test.junit4.createComposeRule
 import androidx.navigation.navOptions
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.LargeTest
 import com.google.common.truth.Truth.assertThat
 import com.google.common.truth.Truth.assertWithMessage
-import org.junit.Rule
 import org.junit.Test
 import org.junit.runner.RunWith
 
 @LargeTest
 @RunWith(AndroidJUnit4::class)
 class ComposeNavigatorTest {
-    @get:Rule
-    val composeTestRule = createComposeRule()
+
+    @Test
+    fun testNavigateConfigChangeThenPop() {
+        val navigator = ComposeNavigator()
+        val destination = navigator.createDestination()
+        destination.id = FIRST_DESTINATION_ID
+
+        assertThat(navigator.navigate(destination, null, null, null))
+            .isEqualTo(destination)
+
+        destination.id = SECOND_DESTINATION_ID
+        assertThat(navigator.navigate(destination, null, null, null))
+            .isEqualTo(destination)
+
+        val savedState = navigator.onSaveState()!!
+        val restoredNavigator = ComposeNavigator()
+
+        restoredNavigator.onRestoreState(savedState)
+
+        assertWithMessage("ComposeNavigator should return true when popping the second destination")
+            .that(navigator.popBackStack())
+            .isTrue()
+    }
 
     @Test
     fun testNavigateWithPopUpToThenPop() {

diff --git a/navigation/navigation-compose/src/main/java/androidx/navigation/compose/ComposeNavigator.kt b/navigation/navigation-compose/src/main/java/androidx/navigation/compose/ComposeNavigator.kt
index 4b8f293..4903062 100644
--- a/navigation/navigation-compose/src/main/java/androidx/navigation/compose/ComposeNavigator.kt
+++ b/navigation/navigation-compose/src/main/java/androidx/navigation/compose/ComposeNavigator.kt

@@ -18,6 +18,7 @@
 
 import android.os.Bundle
 import androidx.compose.runtime.Composable
+import androidx.core.os.bundleOf
 import androidx.navigation.NavBackStackEntry
 import androidx.navigation.NavDestination
 import androidx.navigation.NavOptions
@@ -55,6 +56,18 @@
         return backstack.removeLastOrNull() != null
     }
 
+    override fun onSaveState(): Bundle? {
+        return bundleOf(KEY_BACK_STACK_IDS to backstack.toIntArray())
+    }
+
+    override fun onRestoreState(savedState: Bundle) {
+        val restoredBackStack = savedState.getIntArray(KEY_BACK_STACK_IDS)
+        if (restoredBackStack != null) {
+            backstack.clear()
+            backstack.addAll(restoredBackStack.asList())
+        }
+    }
+
     /**
      * NavDestination specific to [ComposeNavigator]
      */
@@ -63,4 +76,8 @@
         navigator: ComposeNavigator,
         internal val content: @Composable (NavBackStackEntry) -> Unit
     ) : NavDestination(navigator)
+
+    private companion object {
+        private const val KEY_BACK_STACK_IDS = "androidx-nav-compose:navigator:backStackIds"
+    }
 }