Bug 1843536 - Fix tab reordering animation issues

app/src/main/java/org/mozilla/fenix/tabstray/TabsTrayTabLayouts.kt

@@ -14,9 +14,9 @@ import androidx.compose.foundation.lazy.LazyColumn
 import androidx.compose.foundation.lazy.grid.GridCells
 import androidx.compose.foundation.lazy.grid.GridItemSpan
 import androidx.compose.foundation.lazy.grid.LazyVerticalGrid
-import androidx.compose.foundation.lazy.grid.items
+import androidx.compose.foundation.lazy.grid.itemsIndexed
 import androidx.compose.foundation.lazy.grid.rememberLazyGridState
-import androidx.compose.foundation.lazy.items
+import androidx.compose.foundation.lazy.itemsIndexed
 import androidx.compose.foundation.lazy.rememberLazyListState
 import androidx.compose.runtime.Composable
 import androidx.compose.runtime.LaunchedEffect
@@ -131,7 +131,7 @@ fun TabLayout(
 }
 
 @OptIn(ExperimentalFoundationApi::class)
-@Suppress("LongParameterList")
+@Suppress("LongParameterList", "LongMethod")
 @Composable
 private fun TabGrid(
     tabs: List<TabSessionState>,
@@ -197,11 +197,15 @@ private fun TabGrid(
             }
         }
 
-        items(
+        itemsIndexed(
             items = tabs,
-            key = { tab -> tab.id },
-        ) { tab ->
-            DragItemContainer(state = reorderState, key = tab.id) {
+            key = { _, tab -> tab.id },
+        ) { index, tab ->
+            DragItemContainer(
+                state = reorderState,
+                position = index + if (header != null) 1 else 0,
+                key = tab.id,
+            ) {
                 TabGridItem(
                     tab = tab,
                     thumbnailSize = tabThumbnailSize,
@@ -288,11 +292,15 @@ private fun TabList(
             }
         }
 
-        items(
+        itemsIndexed(
             items = tabs,
-            key = { tab -> tab.id },
-        ) { tab ->
-            DragItemContainer(state = reorderState, key = tab.id) {
+            key = { _, tab -> tab.id },
+        ) { index, tab ->
+            DragItemContainer(
+                state = reorderState,
+                position = index + if (header != null) 1 else 0,
+                key = tab.id,
+            ) {
                 TabListItem(
                     tab = tab,
                     thumbnailSize = tabThumbnailSize,

app/src/main/java/org/mozilla/fenix/tabstray/browser/compose/ReorderableGrid.kt

@@ -102,11 +102,17 @@ class GridReorderState internal constructor(
     private var draggingItemCumulatedOffset by mutableStateOf(Offset.Zero)
     private var draggingItemInitialOffset by mutableStateOf(Offset.Zero)
     internal var moved by mutableStateOf(false)
-    internal val draggingItemOffset: Offset
+    private val draggingItemOffset: Offset
         get() = draggingItemLayoutInfo?.let { item ->
             draggingItemInitialOffset + draggingItemCumulatedOffset - item.offset.toOffset()
         } ?: Offset.Zero
 
+    internal fun computeItemOffset(index: Int): Offset {
+        val itemAtIndex = gridState.layoutInfo.visibleItemsInfo.firstOrNull { info -> info.index == index }
+            ?: return Offset.Zero
+        return draggingItemInitialOffset + draggingItemCumulatedOffset - itemAtIndex.offset.toOffset()
+    }
+
     private val draggingItemLayoutInfo: LazyGridItemInfo?
         get() = gridState.layoutInfo.visibleItemsInfo.firstOrNull { it.key == draggingItemKey }
 
@@ -202,6 +208,7 @@ class GridReorderState internal constructor(
  *
  * @param state State of the lazy grid.
  * @param key Key of the item to be displayed.
+ * @param position Position in the grid of the item to be displayed.
  * @param content Content of the item to be displayed.
  */
 @ExperimentalFoundationApi
@@ -209,6 +216,7 @@ class GridReorderState internal constructor(
 fun LazyGridItemScope.DragItemContainer(
     state: GridReorderState,
     key: Any,
+    position: Int,
     content: @Composable () -> Unit,
 ) {
     val modifier = when (key) {
@@ -216,8 +224,8 @@ fun LazyGridItemScope.DragItemContainer(
             Modifier
                 .zIndex(1f)
                 .graphicsLayer {
-                    translationX = state.draggingItemOffset.x
-                    translationY = state.draggingItemOffset.y
+                    translationX = state.computeItemOffset(position).x
+                    translationY = state.computeItemOffset(position).y
                 }
         }
 

app/src/main/java/org/mozilla/fenix/tabstray/browser/compose/ReorderableList.kt

@@ -17,6 +17,7 @@ import androidx.compose.foundation.lazy.LazyListItemInfo
 import androidx.compose.foundation.lazy.LazyListState
 import androidx.compose.runtime.Composable
 import androidx.compose.runtime.getValue
+import androidx.compose.runtime.mutableFloatStateOf
 import androidx.compose.runtime.mutableStateOf
 import androidx.compose.runtime.remember
 import androidx.compose.runtime.rememberCoroutineScope
@@ -93,14 +94,20 @@ class ListReorderState internal constructor(
     var draggingItemKey by mutableStateOf<Any?>(null)
         private set
 
-    private var draggingItemCumulatedOffset by mutableStateOf(0f)
-    private var draggingItemInitialOffset by mutableStateOf(0f)
+    private var draggingItemCumulatedOffset by mutableFloatStateOf(0f)
+    private var draggingItemInitialOffset by mutableFloatStateOf(0f)
     internal var moved by mutableStateOf(false)
-    internal val draggingItemOffset: Float
+    private val draggingItemOffset: Float
         get() = draggingItemLayoutInfo?.let { item ->
             draggingItemInitialOffset + draggingItemCumulatedOffset - item.offset
         } ?: 0f
 
+    internal fun computeItemOffset(index: Int): Float {
+        val itemAtIndex = listState.layoutInfo.visibleItemsInfo.firstOrNull { info -> info.index == index }
+            ?: return draggingItemOffset
+        return draggingItemInitialOffset + draggingItemCumulatedOffset - itemAtIndex.offset
+    }
+
     private val draggingItemLayoutInfo: LazyListItemInfo?
         get() = listState.layoutInfo.visibleItemsInfo.firstOrNull { it.key == draggingItemKey }
 
@@ -192,6 +199,7 @@ class ListReorderState internal constructor(
  *
  * @param state List reordering state.
  * @param key Key of the item to be displayed.
+ * @param position Position in the list of the item to be displayed.
  * @param content Content of the item to be displayed.
  */
 @ExperimentalFoundationApi
@@ -199,6 +207,7 @@ class ListReorderState internal constructor(
 fun LazyItemScope.DragItemContainer(
     state: ListReorderState,
     key: Any,
+    position: Int,
     content: @Composable () -> Unit,
 ) {
     val modifier = when (key) {
@@ -206,7 +215,7 @@ fun LazyItemScope.DragItemContainer(
             Modifier
                 .zIndex(1f)
                 .graphicsLayer {
-                    translationY = state.draggingItemOffset
+                    translationY = state.computeItemOffset(position)
                 }
         }