Fix predictive back completion on Android 14 with 3-button nav

extensions-compose-experimental/src/commonMain/kotlin/com/arkivanov/decompose/extensions/compose/experimental/stack/animation/DefaultStackAnimation.kt

@@ -15,6 +15,7 @@ import androidx.compose.runtime.mutableStateOf
 import androidx.compose.runtime.remember
 import androidx.compose.runtime.rememberCoroutineScope
 import androidx.compose.runtime.setValue
+import androidx.compose.runtime.withFrameNanos
 import androidx.compose.ui.Modifier
 import androidx.compose.ui.input.pointer.pointerInput
 import com.arkivanov.decompose.Child
@@ -29,7 +30,9 @@ import com.arkivanov.decompose.extensions.compose.stack.animation.predictiveback
 import com.arkivanov.decompose.router.stack.ChildStack
 import com.arkivanov.essenty.backhandler.BackCallback
 import com.arkivanov.essenty.backhandler.BackEvent
+import kotlinx.coroutines.CompletableDeferred
 import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.Deferred
 import kotlinx.coroutines.Job
 import kotlinx.coroutines.cancel
 import kotlinx.coroutines.launch
@@ -139,6 +142,10 @@ internal class DefaultStackAnimation<C : Any, T : Any>(
         val transition = rememberTransition(item.transitionState)
         val isTransitionIdle = item.transitionState.isIdle()
 
+        LaunchedEffect(item.transitionComposed) {
+            item.transitionComposed?.complete(Unit)
+        }
+
         if (isTransitionIdle) {
             LaunchedEffect(Unit) {
                 onFinished()
@@ -230,13 +237,15 @@ internal class DefaultStackAnimation<C : Any, T : Any>(
         direction: Direction,
         transitionState: TransitionState<EnterExitState>,
         otherChild: Child.Created<C, T>,
+        transitionComposed: CompletableDeferred<Unit>? = null,
         isPredictiveBack: Boolean = false,
         predictiveBackAnimator: StackAnimator? = null,
     ): AnimationItem<C, T> =
         AnimationItem(
             child = child,
             direction = direction,
             transitionState = transitionState,
+            transitionComposed = transitionComposed,
             animator = predictiveBackAnimator ?: selector(child, otherChild, direction, isPredictiveBack),
         )
 
@@ -257,7 +266,7 @@ internal class DefaultStackAnimation<C : Any, T : Any>(
                 currentState.scope.cancel()
                 val newScope = parentScope + Job(parentScope.coroutineContext[Job])
                 state = State.Started(currentState.lastBackEvent, newScope)
-                onBackProgressed(currentState.lastBackEvent)
+                onBackProgressed(backEvent)
             }
         }
 
@@ -286,6 +295,7 @@ internal class DefaultStackAnimation<C : Any, T : Any>(
                         direction = Direction.ENTER_BACK,
                         transitionState = animationHandler.enterTransitionState,
                         otherChild = exitChild,
+                        transitionComposed = animationHandler.enterTransitionComposed,
                         isPredictiveBack = true,
                         predictiveBackAnimator = animationHandler.animatable?.let { anim -> SimpleStackAnimator { anim.enterModifier } },
                     ),
@@ -294,6 +304,7 @@ internal class DefaultStackAnimation<C : Any, T : Any>(
                         direction = Direction.EXIT_FRONT,
                         transitionState = animationHandler.exitTransitionState,
                         otherChild = enterChild,
+                        transitionComposed = animationHandler.exitTransitionComposed,
                         isPredictiveBack = true,
                         predictiveBackAnimator = animationHandler.animatable?.let { anim -> SimpleStackAnimator { anim.exitModifier } },
                     ),
@@ -377,7 +388,11 @@ internal class DefaultStackAnimation<C : Any, T : Any>(
     private class AnimationHandler(
         val animatable: PredictiveBackAnimatable?,
     ) {
+
+        val exitTransitionComposed: CompletableDeferred<Unit> = CompletableDeferred()
         val exitTransitionState: SeekableTransitionState<EnterExitState> = SeekableTransitionState(EnterExitState.Visible)
+
+        val enterTransitionComposed: CompletableDeferred<Unit> = CompletableDeferred()
         val enterTransitionState: SeekableTransitionState<EnterExitState> = SeekableTransitionState(EnterExitState.PreEnter)
 
         suspend fun progress(backEvent: BackEvent) {
@@ -402,11 +417,26 @@ internal class DefaultStackAnimation<C : Any, T : Any>(
 
         suspend fun finish() {
             awaitAll(
-                { exitTransitionState.animateTo(EnterExitState.PostExit) },
-                { enterTransitionState.animateTo(EnterExitState.Visible) },
+                {
+                    exitTransitionComposed.awaitCompletionForAFewFrames()
+                    exitTransitionState.animateTo(EnterExitState.PostExit)
+                },
+                {
+                    enterTransitionComposed.awaitCompletionForAFewFrames()
+                    enterTransitionState.animateTo(EnterExitState.Visible)
+                },
                 { animatable?.finish() },
             )
         }
+
+        // Wait a few frames for the SeekableTransitionState to receive its Transition
+        private suspend fun Deferred<*>.awaitCompletionForAFewFrames() {
+            repeat(3) {
+                if (!isCompleted) {
+                    withFrameNanos {}
+                }
+            }
+        }
     }
 }
 
@@ -429,6 +459,7 @@ private data class AnimationItem<out C : Any, out T : Any>(
     val child: Child.Created<C, T>,
     val direction: Direction,
     val transitionState: TransitionState<EnterExitState>,
+    val transitionComposed: CompletableDeferred<Unit>? = null,
     val animator: StackAnimator? = null,
 )
 

extensions-compose-experimental/src/jvmTest/kotlin/com/arkivanov/decompose/extensions/compose/experimental/stack/animation/PredictiveBackGestureTest.kt

@@ -3,6 +3,7 @@ package com.arkivanov.decompose.extensions.compose.experimental.stack.animation
 import androidx.compose.animation.EnterExitState
 import androidx.compose.animation.core.animateFloat
 import androidx.compose.material.Text
+import androidx.compose.runtime.LaunchedEffect
 import androidx.compose.runtime.getValue
 import androidx.compose.runtime.mutableStateOf
 import androidx.compose.runtime.setValue
@@ -16,6 +17,7 @@ import com.arkivanov.decompose.extensions.compose.experimental.animateFloat
 import com.arkivanov.decompose.extensions.compose.experimental.assertTestTagToRootDoesNotExist
 import com.arkivanov.decompose.extensions.compose.experimental.assertTestTagToRootExists
 import com.arkivanov.decompose.extensions.compose.experimental.stack.dropLast
+import com.arkivanov.decompose.extensions.compose.stack.animation.Direction
 import com.arkivanov.decompose.extensions.compose.stack.animation.predictiveback.PredictiveBackAnimatable
 import com.arkivanov.decompose.router.stack.ChildStack
 import com.arkivanov.essenty.backhandler.BackDispatcher
@@ -622,6 +624,103 @@ class PredictiveBackGestureTest {
         composeRule.onNodeWithText("2").assertDoesNotExist()
     }
 
+    @Test
+    fun GIVEN_gesture_progressed_WHEN_cancelled_and_restarted_THEN_animatable_created_with_last_known_progress() {
+        var stack by mutableStateOf(stack("1", "2"))
+        var animatable: TestAnimatable? = null
+
+        val animation =
+            DefaultStackAnimation(
+                predictiveBackAnimatable = {
+                    animatable = TestAnimatable(initialBackEvent = it, cancel = suspendForever())
+                    animatable
+                },
+                onBack = { stack = stack.dropLast() },
+            )
+
+        composeRule.setContent {
+            animation(stack, Modifier) {
+                Text(text = it.configuration)
+            }
+        }
+
+        backDispatcher.startPredictiveBack(BackEvent(progress = 0F))
+        composeRule.waitForIdle()
+        backDispatcher.progressPredictiveBack(BackEvent(progress = 0.5F))
+        composeRule.waitForIdle()
+
+        backDispatcher.cancelPredictiveBack()
+        backDispatcher.startPredictiveBack(BackEvent(progress = 0F))
+        composeRule.waitForIdle()
+
+        assertEquals(0.5F, animatable?.initialBackEvent?.progress)
+    }
+
+    @Test
+    fun GIVEN_gesture_progressed_WHEN_cancelled_and_restarted_THEN_gesture_restarted_and_stack_unchanged() {
+        var stack by mutableStateOf(stack("1", "2"))
+
+        val animation =
+            DefaultStackAnimation(
+                predictiveBackAnimatable = { TestAnimatable(initialBackEvent = it, cancel = suspendForever()) },
+                onBack = { stack = stack.dropLast() },
+            )
+
+        composeRule.setContent {
+            animation(stack, Modifier) {
+                Text(text = it.configuration)
+            }
+        }
+
+        backDispatcher.startPredictiveBack(BackEvent(progress = 0F))
+        composeRule.waitForIdle()
+        backDispatcher.progressPredictiveBack(BackEvent(progress = 0.5F))
+        composeRule.waitForIdle()
+
+        backDispatcher.cancelPredictiveBack()
+        backDispatcher.startPredictiveBack(BackEvent(progress = 0F))
+        composeRule.waitForIdle()
+
+        assertEquals(stack("1", "2"), stack)
+        composeRule.onNodeWithText("1").assertTestTagToRootExists(enterTestTag(progress = 0F))
+        composeRule.onNodeWithText("2").assertTestTagToRootExists(exitTestTag(progress = 0F))
+    }
+
+    @Test
+    fun WHEN_gesture_started_progress_and_finished_in_the_same_frame_THEN_transitions_animated() {
+        var stack by mutableStateOf(stack("1", "2"))
+        val childDirections = HashMap<String, Direction?>()
+        val childAnimations = HashMap<String, Float>()
+
+        val anim =
+            DefaultStackAnimation(
+                predictiveBackAnimatable = { null },
+                onBack = { stack = stack.dropLast() },
+            )
+
+        composeRule.setContent {
+            anim(stack = stack, modifier = Modifier) {
+                val animationValue by transition.animateFloat(durationMillis = 1000)
+
+                LaunchedEffect(stackAnimationDirection, animationValue) {
+                    childDirections[it.configuration] = stackAnimationDirection
+                    childAnimations[it.configuration] = animationValue
+                }
+            }
+        }
+        composeRule.waitForIdle()
+
+        backDispatcher.startPredictiveBack(BackEvent(progress = 0F))
+        backDispatcher.progressPredictiveBack(BackEvent(progress = 0F))
+        backDispatcher.back()
+        composeRule.mainClock.advanceTimeBy(500L)
+
+        assertEquals(Direction.ENTER_BACK, childDirections["1"])
+        assertEquals(expected = 0.5F, actual = childAnimations["1"] ?: -1F, absoluteTolerance = 0.05F)
+        assertEquals(Direction.EXIT_FRONT, childDirections["2"])
+        assertEquals(expected = 0.5F, actual = childAnimations["2"] ?: -1F, absoluteTolerance = 0.05F)
+    }
+
     private fun DefaultStackAnimation(
         predictiveBackAnimatable: (initialBackEvent: BackEvent) -> PredictiveBackAnimatable? = ::TestAnimatable,
         animator: StackAnimator? = null,
@@ -667,7 +766,7 @@ class PredictiveBackGestureTest {
     }
 
     private class TestAnimatable(
-        initialBackEvent: BackEvent,
+        val initialBackEvent: BackEvent,
         private val finish: suspend () -> Unit = {},
         private val cancel: suspend () -> Unit = {},
     ) : PredictiveBackAnimatable {