Fixed stack animation not updating to new child component instance when active configuration is unchanged

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

@@ -49,22 +49,31 @@ internal class DefaultStackAnimation<C : Any, T : Any>(
     ) {
         var currentStack by remember { mutableStateOf(stack) }
         var items by remember { mutableStateOf(getAnimationItems(newStack = currentStack)) }
+        var nextItems: Map<Any, AnimationItem<C, T>>? by remember { mutableStateOf(null) }
         val stackKeys = remember(stack) { stack.items.map { it.key } }
         val currentStackKeys = remember(currentStack) { currentStack.items.map { it.key } }
 
-        if (stackKeys != currentStackKeys) {
+        if (stack != currentStack) {
             val oldStack = currentStack
             currentStack = stack
 
             val updateItems =
                 when {
-                    stack.active.key == oldStack.active.key -> items.keys.singleOrNull() != stack.active.key
+                    stack.active.key == oldStack.active.key ->
+                        (items.keys.singleOrNull() != stack.active.key) ||
+                            (items.values.singleOrNull()?.child?.instance != stack.active.instance)
+
                     items.size == 1 -> items.keys.single() != stack.active.key
                     else -> items.keys.toList() != stackKeys
                 }
 
             if (updateItems) {
-                items = getAnimationItems(newStack = currentStack, oldStack = oldStack)
+                val newItems = getAnimationItems(newStack = currentStack, oldStack = oldStack)
+                if (items.size == 1) {
+                    items = newItems
+                } else {
+                    nextItems = newItems
+                }
             }
         }
 
@@ -82,12 +91,21 @@ internal class DefaultStackAnimation<C : Any, T : Any>(
                         },
                         content = content,
                     )
+
+                    if (item.direction.isExit) {
+                        DisposableEffect(Unit) {
+                            onDispose {
+                                nextItems?.also { items = it }
+                                nextItems = null
+                            }
+                        }
+                    }
                 }
             }
 
             // A workaround until https://issuetracker.google.com/issues/214231672.
             // Normally only the exiting child should be disabled.
-            if (disableInputDuringAnimation && (items.size > 1)) {
+            if (disableInputDuringAnimation && ((items.size > 1) || (nextItems != null))) {
                 Overlay(modifier = Modifier.matchParentSize())
             }
         }
@@ -129,7 +147,7 @@ internal class DefaultStackAnimation<C : Any, T : Any>(
 
     private fun getAnimationItems(newStack: ChildStack<C, T>, oldStack: ChildStack<C, T>? = null): Map<Any, AnimationItem<C, T>> =
         when {
-            oldStack == null ->
+            (oldStack == null) || (newStack.active.key == oldStack.active.key) ->
                 keyedItemsOf(
                     AnimationItem(
                         child = newStack.active,

extensions-compose-experimental/src/jvmTest/kotlin/com/arkivanov/decompose/extensions/compose/experimental/TestUtils.kt

@@ -1,6 +1,7 @@
 package com.arkivanov.decompose.extensions.compose.experimental
 
 import androidx.compose.animation.EnterExitState
+import androidx.compose.animation.core.AnimationConstants
 import androidx.compose.animation.core.LinearEasing
 import androidx.compose.animation.core.Transition
 import androidx.compose.animation.core.animateFloat
@@ -15,15 +16,25 @@ import androidx.compose.ui.test.SemanticsNodeInteraction
 import kotlin.test.fail
 
 @Composable
-internal fun Transition<EnterExitState>.animateFloat(): State<Float> =
-    animateFloat(transitionSpec = { tween(easing = LinearEasing) }) { state ->
+internal fun Transition<EnterExitState>.animateFloat(durationMillis: Int = AnimationConstants.DefaultDurationMillis): State<Float> =
+    animateFloat(transitionSpec = { tween(easing = LinearEasing, durationMillis = durationMillis) }) { state ->
         when (state) {
             EnterExitState.PreEnter -> 0F
             EnterExitState.Visible -> 1F
             EnterExitState.PostExit -> 0F
         }
     }
 
+internal fun <T> List<T>.takeSorted(comparator: Comparator<T>): List<T> =
+    takeWhileIndexed { index, item ->
+        (index == 0) || (comparator.compare(item, get(index - 1)) >= 0)
+    }
+
+internal fun <T> Iterable<T>.takeWhileIndexed(predicate: (Int, T) -> Boolean): List<T> =
+    withIndex()
+        .takeWhile { (index, item) -> predicate(index, item) }
+        .map { it.value }
+
 internal fun SemanticsNodeInteraction.assertTestTagToRootExists(testTag: String) {
     val count = collectTestTagsToRoot().filter { it == testTag }.size
     if (count != 1) {

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

@@ -1,5 +1,6 @@
 package com.arkivanov.decompose.extensions.compose.experimental.stack
 
+import androidx.compose.animation.AnimatedVisibilityScope
 import androidx.compose.foundation.clickable
 import androidx.compose.foundation.text.BasicText
 import androidx.compose.runtime.Composable
@@ -29,12 +30,13 @@ import org.junit.Rule
 import org.junit.Test
 import org.junit.runner.RunWith
 import org.junit.runners.Parameterized
+import kotlin.test.assertSame
 
 @OptIn(ExperimentalDecomposeApi::class)
 @Suppress("TestFunctionName")
 @RunWith(Parameterized::class)
 class ChildStackTest(
-    private val animation: StackAnimation<Config, Config>?,
+    private val animation: StackAnimation<Config, Any>?,
 ) {
 
     @get:Rule
@@ -92,6 +94,18 @@ class ChildStackTest(
         composeRule.onNodeWithText(text = "ChildA1", substring = true).assertExists()
     }
 
+    @Test
+    fun GIVEN_child_displayed_WHEN_new_child_instance_with_the_same_key_THEN_new_child_instance_displayed() {
+        val state = mutableStateOf(routerState(Child.Created(configuration = Config.A, instance = Any(), key = "A")))
+        var lastInstance: Any? = null
+        setContent(state) { lastInstance = it.instance }
+
+        val instance2 = Any()
+        state.setValueOnIdle(routerState(Child.Created(configuration = Config.A, instance = instance2, key = "A")))
+
+        assertSame(instance2, lastInstance)
+    }
+
     @Test
     fun GIVEN_child_B_displayed_and_child_A_in_back_stack_WHEN_pop_child_B_THEN_state_restored_for_child_A() {
         val state = mutableStateOf(routerState(Config.A))
@@ -186,11 +200,14 @@ class ChildStackTest(
         composeRule.onNodeWithText(text = "ChildA2", substring = true).assertDoesNotExist()
     }
 
-    private fun setContent(state: State<ChildStack<Config, Config>>) {
+    private fun setContent(
+        state: State<ChildStack<Config, Any>>,
+        content: @Composable AnimatedVisibilityScope.(Child.Created<Config, Any>) -> Unit = {
+            Child(name = it.key.toString())
+        },
+    ) {
         composeRule.setContent {
-            ChildStack(stack = state.value, animation = animation) { child ->
-                Child(name = child.key.toString())
-            }
+            ChildStack(stack = state.value, animation = animation, content = content)
         }
 
         composeRule.runOnIdle {}
@@ -211,6 +228,12 @@ class ChildStackTest(
             backStack = stack.dropLast(1).map { it.toChild() },
         )
 
+    private fun routerState(vararg stack: Child.Created<Config, Any>): ChildStack<Config, Any> =
+        ChildStack(
+            active = stack.last(),
+            backStack = stack.dropLast(1),
+        )
+
     private fun Pair<Any, Config>.toChild(): Child.Created<Config, Config> =
         Child.Created(configuration = second, instance = second, key = first)
 
@@ -235,7 +258,7 @@ class ChildStackTest(
         fun parameters(): List<Array<out Any?>> =
             getParameters().map { arrayOf(it) }
 
-        private fun getParameters(): List<StackAnimation<Config, Config>?> {
+        private fun getParameters(): List<StackAnimation<Config, Any>?> {
             val predictiveBackParams1 =
                 PredictiveBackParams(
                     backHandler = BackDispatcher(),

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

@@ -0,0 +1,121 @@
+package com.arkivanov.decompose.extensions.compose.experimental.stack.animation
+
+import androidx.compose.runtime.LaunchedEffect
+import androidx.compose.runtime.getValue
+import androidx.compose.runtime.mutableStateOf
+import androidx.compose.runtime.setValue
+import androidx.compose.ui.Modifier
+import androidx.compose.ui.test.MainTestClock
+import androidx.compose.ui.test.junit4.createComposeRule
+import com.arkivanov.decompose.Child
+import com.arkivanov.decompose.extensions.compose.experimental.animateFloat
+import com.arkivanov.decompose.extensions.compose.experimental.takeSorted
+import com.arkivanov.decompose.router.stack.ChildStack
+import org.junit.Rule
+import kotlin.test.Test
+import kotlin.test.assertEquals
+import kotlin.test.assertTrue
+
+@Suppress("TestFunctionName")
+class DefaultStackAnimationTest {
+
+    @get:Rule
+    val composeRule = createComposeRule()
+
+    @Test
+    fun WHEN_animating_push_and_stack_popped_during_animation_THEN_animated_push_and_pop_fully() {
+        val anim =
+            DefaultStackAnimation<Int, Any>(
+                disableInputDuringAnimation = true,
+                predictiveBackParams = { null },
+                selector = { _, _, _ -> null },
+            )
+
+        var stack by mutableStateOf(stack(1))
+        val values1 = ArrayList<Pair<Long, Float>>()
+        val values2 = ArrayList<Pair<Long, Float>>()
+
+        composeRule.setContent {
+            anim(stack = stack, modifier = Modifier) {
+                val value by transition.animateFloat(durationMillis = 1000)
+                val pair = composeRule.mainClock.currentTime to value
+
+                when (it.key) {
+                    1 -> values1 += pair
+                    2 -> values2 += pair
+                }
+            }
+        }
+
+        stack = stack(1, 2)
+        composeRule.mainClock.advanceFramesBy(millis = 500L)
+        stack = stack(1)
+        composeRule.waitForIdle()
+
+        val v11 = values1.takeSorted(compareByDescending { it.second })
+        val v21 = values2.takeSorted(compareBy { it.second })
+        assertTrue(v11.size > 10)
+        assertTrue(v21.size > 10)
+        assertEquals(1F, v11.first().second)
+        assertEquals(0F, v11.last().second)
+        assertEquals(0F, v21.first().second)
+        assertEquals(1F, v21.last().second)
+
+        val v12 = values1.drop(v11.size - 1).takeSorted(compareBy { it.second })
+        val v22 = values2.drop(v21.size - 1).takeSorted(compareByDescending { it.second })
+        assertTrue(v12.size > 10)
+        assertTrue(v22.size > 10)
+        assertEquals(0F, v12.first().second)
+        assertEquals(1F, v12.last().second)
+        assertEquals(1F, v22.first().second)
+        assertEquals(0F, v22.last().second)
+    }
+
+    @Test
+    fun WHEN_animating_push_and_stack_popped_during_animation_THEN_first_child_restarted() {
+        val anim =
+            DefaultStackAnimation<Int, Any>(
+                disableInputDuringAnimation = true,
+                predictiveBackParams = { null },
+                selector = { _, _, _ -> null },
+            )
+
+        var stack by mutableStateOf(stack(1))
+        var counter = 0
+
+        composeRule.setContent {
+            anim(stack = stack, modifier = Modifier) {
+                transition.animateFloat(durationMillis = 1000)
+
+                if (it.key == 1) {
+                    LaunchedEffect(Unit) {
+                        counter++
+                    }
+                }
+            }
+        }
+
+        stack = stack(1, 2)
+        composeRule.mainClock.advanceFramesBy(millis = 500L)
+        stack = stack(1)
+        composeRule.waitForIdle()
+
+        assertEquals(2, counter)
+    }
+
+    private fun MainTestClock.advanceFramesBy(millis: Long) {
+        val endTime = currentTime + millis
+        while (currentTime < endTime) {
+            advanceTimeByFrame()
+        }
+    }
+
+    private fun child(config: Int): Child.Created<Int, Any> =
+        Child.Created(configuration = config, instance = Any())
+
+    private fun stack(vararg stack: Int): ChildStack<Int, Any> =
+        ChildStack(
+            active = child(stack.last()),
+            backStack = stack.dropLast(1).map(::child),
+        )
+}

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

@@ -2,6 +2,7 @@ package com.arkivanov.decompose.extensions.compose.stack.animation
 
 import androidx.compose.foundation.layout.Box
 import androidx.compose.runtime.Composable
+import androidx.compose.runtime.DisposableEffect
 import androidx.compose.runtime.getValue
 import androidx.compose.runtime.key
 import androidx.compose.runtime.mutableStateOf
@@ -29,11 +30,18 @@ internal abstract class AbstractStackAnimation<C : Any, T : Any>(
     override operator fun invoke(stack: ChildStack<C, T>, modifier: Modifier, content: @Composable (child: Child.Created<C, T>) -> Unit) {
         var currentStack by remember { mutableStateOf(stack) }
         var items by remember { mutableStateOf(getAnimationItems(newStack = currentStack, oldStack = null)) }
+        var nextItems: Map<Any, AnimationItem<C, T>>? by remember { mutableStateOf(null) }
 
-        if (stack.active.key != currentStack.active.key) {
+        if (stack.active.key != currentStack.active.key || stack.active.instance != currentStack.active.instance) {
             val oldStack = currentStack
             currentStack = stack
-            items = getAnimationItems(newStack = currentStack, oldStack = oldStack)
+
+            val newItems = getAnimationItems(newStack = currentStack, oldStack = oldStack)
+            if (items.size == 1) {
+                items = newItems
+            } else {
+                nextItems = newItems
+            }
         }
 
         Box(modifier = modifier) {
@@ -50,20 +58,29 @@ internal abstract class AbstractStackAnimation<C : Any, T : Any>(
                         },
                         content = content,
                     )
+
+                    if (item.direction.isExit) {
+                        DisposableEffect(Unit) {
+                            onDispose {
+                                nextItems?.also { items = it }
+                                nextItems = null
+                            }
+                        }
+                    }
                 }
             }
 
             // A workaround until https://issuetracker.google.com/issues/214231672.
             // Normally only the exiting child should be disabled.
-            if (disableInputDuringAnimation && (items.size > 1)) {
+            if (disableInputDuringAnimation && ((items.size > 1) || (nextItems != null))) {
                 InputConsumingOverlay(modifier = Modifier.matchParentSize())
             }
         }
     }
 
     private fun getAnimationItems(newStack: ChildStack<C, T>, oldStack: ChildStack<C, T>?): Map<Any, AnimationItem<C, T>> =
         when {
-            oldStack == null ->
+            (oldStack == null) || (newStack.active.key == oldStack.active.key) ->
                 listOf(AnimationItem(child = newStack.active, direction = Direction.ENTER_FRONT, isInitial = true))
 
             (newStack.size < oldStack.size) && (newStack.active.key in oldStack.backStack) ->

extensions-compose/src/jvmTest/kotlin/com/arkivanov/decompose/extensions/compose/Utils.kt

@@ -0,0 +1,11 @@
+package com.arkivanov.decompose.extensions.compose
+
+internal fun <T> List<T>.takeSorted(comparator: Comparator<T>): List<T> =
+    takeWhileIndexed { index, item ->
+        (index == 0) || (comparator.compare(item, get(index - 1)) >= 0)
+    }
+
+internal fun <T> Iterable<T>.takeWhileIndexed(predicate: (Int, T) -> Boolean): List<T> =
+    withIndex()
+        .takeWhile { (index, item) -> predicate(index, item) }
+        .map { it.value }