Add entry animations and BackHandler into SwipeDismissableNavHost

Bug: 313564642
Test: Manual testing and SwipeDismissableNavHostTest.kt

Relnote: "Added animations visible while entering a new screen in
SwipeDismissableNavHost. These animations will be turned off when
REDUCE_MOTION is turned on. Also integrated BackHandler into the host
instead of setting OnBackPressedDispatcher to the host ourselves."

Change-Id: Ia473b2996a779866867d1d02ca477737f6074473
diff --git a/wear/compose/compose-navigation/src/main/java/androidx/wear/compose/navigation/SwipeDismissableNavHost.kt b/wear/compose/compose-navigation/src/main/java/androidx/wear/compose/navigation/SwipeDismissableNavHost.kt
index cac68b5..794c74f 100644
--- a/wear/compose/compose-navigation/src/main/java/androidx/wear/compose/navigation/SwipeDismissableNavHost.kt
+++ b/wear/compose/compose-navigation/src/main/java/androidx/wear/compose/navigation/SwipeDismissableNavHost.kt
@@ -17,8 +17,16 @@
 package androidx.wear.compose.navigation
 
 import android.util.Log
-import androidx.activity.compose.LocalOnBackPressedDispatcherOwner
+import androidx.activity.compose.BackHandler
+import androidx.compose.animation.core.Animatable
+import androidx.compose.animation.core.AnimationVector1D
+import androidx.compose.animation.core.CubicBezierEasing
+import androidx.compose.animation.core.LinearEasing
+import androidx.compose.animation.core.tween
+import androidx.compose.foundation.Canvas
 import androidx.compose.foundation.layout.Box
+import androidx.compose.foundation.layout.fillMaxSize
+import androidx.compose.foundation.shape.CircleShape
 import androidx.compose.runtime.Composable
 import androidx.compose.runtime.DisposableEffect
 import androidx.compose.runtime.LaunchedEffect
@@ -30,7 +38,12 @@
 import androidx.compose.runtime.saveable.rememberSaveableStateHolder
 import androidx.compose.runtime.setValue
 import androidx.compose.ui.Modifier
+import androidx.compose.ui.graphics.Color
+import androidx.compose.ui.graphics.RectangleShape
+import androidx.compose.ui.graphics.graphicsLayer
+import androidx.compose.ui.platform.LocalConfiguration
 import androidx.compose.ui.platform.LocalLifecycleOwner
+import androidx.compose.ui.util.lerp
 import androidx.lifecycle.Lifecycle
 import androidx.lifecycle.LifecycleEventObserver
 import androidx.lifecycle.viewmodel.compose.LocalViewModelStoreOwner
@@ -44,11 +57,12 @@
 import androidx.navigation.createGraph
 import androidx.navigation.get
 import androidx.wear.compose.foundation.BasicSwipeToDismissBox
+import androidx.wear.compose.foundation.ExperimentalWearFoundationApi
+import androidx.wear.compose.foundation.LocalReduceMotion
 import androidx.wear.compose.foundation.LocalSwipeToDismissBackgroundScrimColor
 import androidx.wear.compose.foundation.LocalSwipeToDismissContentScrimColor
 import androidx.wear.compose.foundation.SwipeToDismissBoxState
 import androidx.wear.compose.foundation.SwipeToDismissKeys
-import androidx.wear.compose.foundation.SwipeToDismissValue
 import androidx.wear.compose.foundation.edgeSwipeToDismiss
 import androidx.wear.compose.foundation.rememberSwipeToDismissBoxState
 
@@ -91,14 +105,14 @@
     route: String? = null,
     builder: NavGraphBuilder.() -> Unit
 ) = SwipeDismissableNavHost(
-        navController,
-        remember(route, startDestination, builder) {
-            navController.createGraph(startDestination, route, builder)
-        },
-        modifier,
-        userSwipeEnabled,
-        state = state,
-    )
+    navController,
+    remember(route, startDestination, builder) {
+        navController.createGraph(startDestination, route, builder)
+    },
+    modifier,
+    userSwipeEnabled,
+    state = state,
+)
 
 /**
  * Provides a place in the Compose hierarchy for self-contained navigation to occur,
@@ -129,6 +143,7 @@
  *
  * @throws IllegalArgumentException if no WearNavigation.Destination is on the navigation backstack.
  */
+@OptIn(ExperimentalWearFoundationApi::class)
 @Composable
 public fun SwipeDismissableNavHost(
     navController: NavHostController,
@@ -142,29 +157,12 @@
         "SwipeDismissableNavHost requires a ViewModelStoreOwner to be provided " +
             "via LocalViewModelStoreOwner"
     }
-    val onBackPressedDispatcherOwner = LocalOnBackPressedDispatcherOwner.current
-    val onBackPressedDispatcher = onBackPressedDispatcherOwner?.onBackPressedDispatcher
 
-    // Setup the navController with proper owners
-    navController.setLifecycleOwner(lifecycleOwner)
     navController.setViewModelStore(viewModelStoreOwner.viewModelStore)
-    if (onBackPressedDispatcher != null) {
-        navController.setOnBackPressedDispatcher(onBackPressedDispatcher)
-    }
-    // Ensure that the NavController only receives back events while
-    // the NavHost is in composition
-    DisposableEffect(navController) {
-        navController.enableOnBackPressed(true)
-        onDispose {
-            navController.enableOnBackPressed(false)
-        }
-    }
 
     // Then set the graph
     navController.graph = graph
 
-    val stateHolder = rememberSaveableStateHolder()
-
     // Find the WearNavigator, returning early if it isn't found
     // (such as is the case when using TestNavHostController).
     val wearNavigator = navController.navigatorProvider.get>(
@@ -172,9 +170,21 @@
     ) as? WearNavigator ?: return
 
     val backStack by wearNavigator.backStack.collectAsState()
+
+    val navigateBack: () -> Unit = { navController.popBackStack() }
+    BackHandler(enabled = backStack.size > 1, onBack = navigateBack)
+
     val transitionsInProgress by wearNavigator.transitionsInProgress.collectAsState()
     var initialContent by remember { mutableStateOf(true) }
 
+    DisposableEffect(lifecycleOwner) {
+        // Setup the navController with proper owners
+        navController.setLifecycleOwner(lifecycleOwner)
+        onDispose { }
+    }
+
+    val stateHolder = rememberSaveableStateHolder()
+
     val previous = if (backStack.size <= 1) null else backStack[backStack.lastIndex - 1]
     // Get the current navigation backstack entry. If the backstack is empty, it could be because
     // no WearNavigator.Destinations were added to the navigation backstack (be sure to build
@@ -199,40 +209,93 @@
     }
 
     val swipeState = state.swipeToDismissBoxState
-    LaunchedEffect(swipeState.currentValue) {
-        // This effect operates when the swipe gesture is complete:
-        // 1) Resets the screen offset (otherwise, the next destination is draw off-screen)
-        // 2) Pops the navigation back stack to return to the previous level
-        if (swipeState.currentValue == SwipeToDismissValue.Dismissed) {
-            swipeState.snapTo(SwipeToDismissValue.Default)
-            navController.popBackStack()
-        }
-    }
     LaunchedEffect(swipeState.isAnimationRunning) {
         // This effect marks the transitions completed when swipe animations finish,
         // so that the navigation backstack entries can go to Lifecycle.State.RESUMED.
-        if (swipeState.isAnimationRunning == false) {
+        if (!swipeState.isAnimationRunning) {
             transitionsInProgress.forEach { entry ->
                 wearNavigator.onTransitionComplete(entry)
             }
         }
     }
 
+    val isRoundDevice = isRoundDevice()
+    val reduceMotionEnabled = LocalReduceMotion.current.enabled()
+
+    val animationProgress = remember(current) {
+        if (!wearNavigator.isPop.value) {
+            Animatable(0f)
+        } else {
+            Animatable(1f)
+        }
+    }
+
+    LaunchedEffect(animationProgress.isRunning) {
+        if (!animationProgress.isRunning) {
+            transitionsInProgress.forEach { entry ->
+                wearNavigator.onTransitionComplete(entry)
+            }
+        }
+    }
+
+    LaunchedEffect(current) {
+        if (!wearNavigator.isPop.value) {
+            if (reduceMotionEnabled) {
+                animationProgress.snapTo(1f)
+            } else {
+                animationProgress.animateTo(
+                    targetValue = 1f,
+                    animationSpec = tween(
+                        durationMillis = NAV_HOST_ENTER_TRANSITION_DURATION_MEDIUM +
+                            NAV_HOST_ENTER_TRANSITION_DURATION_SHORT,
+                        easing = LinearEasing
+                    )
+                )
+            }
+        }
+    }
+
     BasicSwipeToDismissBox(
+        onDismissed = navigateBack,
         state = swipeState,
         modifier = Modifier,
         userSwipeEnabled = userSwipeEnabled && previous != null,
         backgroundKey = previous?.id ?: SwipeToDismissKeys.Background,
-        contentKey = current?.id ?: SwipeToDismissKeys.Content,
-        content = { isBackground ->
-            BoxedStackEntryContent(if (isBackground) previous else current, stateHolder, modifier)
-        }
-    )
+        contentKey = current?.id ?: SwipeToDismissKeys.Content
+    ) { isBackground ->
+        BoxedStackEntryContent(
+            entry = if (isBackground) previous else current,
+            saveableStateHolder = stateHolder,
+            modifier = modifier.then(
+                if (isBackground) {
+                    Modifier // Not applying graphicsLayer on the background.
+                } else {
+                    Modifier.graphicsLayer {
+                        val scaleProgression = NAV_HOST_ENTER_TRANSITION_EASING_STANDARD
+                            .transform((animationProgress.value / 0.75f).coerceAtMost(1f))
+                        val opacityProgression = NAV_HOST_ENTER_TRANSITION_EASING_STANDARD
+                            .transform((animationProgress.value / 0.25f).coerceAtMost(1f))
+                        val scale = lerp(0.75f, 1f, scaleProgression)
+                        val opacity = lerp(0.1f, 1f, opacityProgression)
+                        scaleX = scale
+                        scaleY = scale
+                        alpha = opacity
+                        clip = true
+                        shape = if (isRoundDevice) CircleShape else RectangleShape
+                    }
+                }
+            ),
+            layerColor = if (isBackground || wearNavigator.isPop.value) {
+                Color.Unspecified
+            } else FLASH_COLOR,
+            animatable = animationProgress
+        )
+    }
 
     DisposableEffect(previous, current) {
         if (initialContent) {
-            // There are no animations for showing the initial content, so mark transitions complete,
-            // allowing the navigation backstack entry to go to Lifecycle.State.RESUMED.
+            // There are no animations for showing the initial content, so mark transitions
+            // complete, allowing the navigation backstack entry to go to Lifecycle.State.RESUMED.
             transitionsInProgress.forEach { entry ->
                 wearNavigator.onTransitionComplete(entry)
             }
@@ -401,8 +464,11 @@
     entry: NavBackStackEntry?,
     saveableStateHolder: SaveableStateHolder,
     modifier: Modifier = Modifier,
+    layerColor: Color,
+    animatable: Animatable
 ) {
     if (entry != null) {
+        val isRoundDevice = isRoundDevice()
         var lifecycleState by remember {
             mutableStateOf(entry.lifecycle.currentState)
         }
@@ -421,9 +487,36 @@
                 entry.LocalOwnersProvider(saveableStateHolder) {
                     destination.content(entry)
                 }
+                // Adding a flash effect when a new screen opens
+                if (layerColor != Color.Unspecified) {
+                    Canvas(Modifier.fillMaxSize()) {
+                        val absoluteProgression =
+                            ((animatable.value - 0.25f).coerceAtLeast(0f) / 0.75f).coerceAtMost(1f)
+                        val easedProgression = NAV_HOST_ENTER_TRANSITION_EASING_STANDARD
+                            .transform(absoluteProgression)
+                        val alpha = lerp(0.07f, 0f, easedProgression)
+                        if (isRoundDevice) {
+                            drawCircle(color = layerColor.copy(alpha))
+                        } else {
+                            drawRect(color = layerColor.copy(alpha))
+                        }
+                    }
+                }
             }
         }
     }
 }
 
+@Composable
+private fun isRoundDevice(): Boolean {
+    val configuration = LocalConfiguration.current
+    return remember(configuration) {
+        configuration.isScreenRound
+    }
+}
+
 private const val TAG = "SwipeDismissableNavHost"
+private const val NAV_HOST_ENTER_TRANSITION_DURATION_SHORT = 100
+private const val NAV_HOST_ENTER_TRANSITION_DURATION_MEDIUM = 300
+private val NAV_HOST_ENTER_TRANSITION_EASING_STANDARD = CubicBezierEasing(0.4f, 0f, 0.2f, 1f)
+private val FLASH_COLOR = Color.White
diff --git a/wear/compose/compose-navigation/src/main/java/androidx/wear/compose/navigation/WearNavigator.kt b/wear/compose/compose-navigation/src/main/java/androidx/wear/compose/navigation/WearNavigator.kt
index d6adcc6..11b3934 100644
--- a/wear/compose/compose-navigation/src/main/java/androidx/wear/compose/navigation/WearNavigator.kt
+++ b/wear/compose/compose-navigation/src/main/java/androidx/wear/compose/navigation/WearNavigator.kt
@@ -17,6 +17,7 @@
 package androidx.wear.compose.navigation
 
 import androidx.compose.runtime.Composable
+import androidx.compose.runtime.mutableStateOf
 import androidx.navigation.NavBackStackEntry
 import androidx.navigation.NavDestination
 import androidx.navigation.NavOptions
@@ -39,20 +40,27 @@
      */
     internal val backStack get() = state.backStack
 
+    /**
+     * Indicates if an entry is being popped from [backStack].
+     */
+    internal val isPop = mutableStateOf(false)
+
     override fun navigate(
         entries: List,
         navOptions: NavOptions?,
         navigatorExtras: Extras?
     ) {
         entries.forEach { entry ->
-            state.push(entry)
+            state.pushWithTransition(entry)
         }
+        isPop.value = false
     }
 
     override fun createDestination() = Destination(this) {}
 
     override fun popBackStack(popUpTo: NavBackStackEntry, savedState: Boolean) {
         state.popWithTransition(popUpTo, savedState)
+        isPop.value = true
     }
 
     /**