Android Wear paging improvements

app/src/googleplay/java/org/tasks/wear/WearService.kt

@@ -35,6 +35,7 @@ class WearService(
             completedAtBottom = preferences.completedTasksAtBottom,
         )
         return Tasks.newBuilder()
+            .setTotalItems(payload.size)
             .addAllItems(
                 payload
                     .subList(position, position + limit)
@@ -42,6 +43,7 @@ class WearService(
                         when (item) {
                             is UiItem.Header ->
                                 GrpcProto.UiItem.newBuilder()
+                                    .setId(item.value)
                                     .setType(GrpcProto.UiItemType.Header)
                                     .setTitle(headerFormatter.headerString(item.value))
                                     .build()

wear-datalayer/src/main/proto/grpc.proto

@@ -20,7 +20,8 @@ message UiItem {
 }
 
 message Tasks {
-  repeated UiItem items = 1;
+  uint32 totalItems = 1;
+  repeated UiItem items = 2;
 }
 
 message LastUpdate {

wear/src/main/java/org/tasks/presentation/MainActivity.kt

@@ -51,15 +51,16 @@ class MainActivity : ComponentActivity() {
                     modifier = Modifier.background(MaterialTheme.colors.background),
                 ) {
                     val navController = rememberSwipeDismissableNavController()
+                    val taskListViewModel: TaskListViewModel = viewModel()
+                    val taskListItems = taskListViewModel.uiItems.collectAsLazyPagingItems()
                     SwipeDismissableNavHost(
                         startDestination = "task_list",
                         navController = navController,
                     ) {
-                        composable("task_list") { navBackStackEntry ->
-                            val viewModel: TaskListViewModel = viewModel(navBackStackEntry)
+                        composable("task_list") {
                             TaskListScreen(
-                                uiItems = viewModel.uiItems.collectAsLazyPagingItems(),
-                                onComplete = { viewModel.completeTask(it) },
+                                uiItems = taskListItems,
+                                onComplete = { taskListViewModel.completeTask(it) },
                                 onClick = { navController.navigate("task_edit/$it") },
                             )
                         }

wear/src/main/java/org/tasks/presentation/MyPagingSource.kt

@@ -1,36 +1,85 @@
 package org.tasks.presentation
 
-import android.util.Log
 import androidx.paging.PagingSource
+import androidx.paging.PagingSource.LoadParams
+import androidx.paging.PagingSource.LoadParams.Append
+import androidx.paging.PagingSource.LoadParams.Prepend
+import androidx.paging.PagingSource.LoadParams.Refresh
 import androidx.paging.PagingState
-import kotlinx.coroutines.Dispatchers
-import kotlinx.coroutines.withContext
+
+private const val INITIAL_ITEM_COUNT = -1
 
 class MyPagingSource<T : Any>(
-    private val fetch: suspend (position: Int, limit: Int) -> List<T>?,
+    private val fetch: suspend (position: Int, limit: Int) -> Pair<Int, List<T>>,
 ) : PagingSource<Int, T>() {
 
-    override suspend fun load(params: LoadParams<Int>): LoadResult<Int, T> {
-        val position = params.key ?: 0
-        val limit = params.loadSize
+    private var itemCount = INITIAL_ITEM_COUNT
 
+    override suspend fun load(params: LoadParams<Int>): LoadResult<Int, T> {
         return try {
-            val items = withContext (Dispatchers.IO) {
-                fetch(position, limit) ?: emptyList()
+            val key = params.key ?: 0
+            val limit = getLimit(params, key)
+            val offset = getOffset(params, key, itemCount)
+            val (newCount, data) = fetch(offset, limit)
+            if (itemCount == INITIAL_ITEM_COUNT) {
+                itemCount = newCount
             }
-
-            LoadResult.Page(
-                data = items,
-                prevKey = if (position <= 0) null else position - limit,
-                nextKey = if (items.isEmpty()) null else position + limit
+            val nextPosToLoad = offset + data.size
+            val nextKey =
+                if (data.isEmpty() || data.size < limit || nextPosToLoad >= itemCount) {
+                    null
+                } else {
+                    nextPosToLoad
+                }
+            val prevKey = if (offset <= 0 || data.isEmpty()) null else offset
+            return LoadResult.Page(
+                data = data,
+                prevKey = prevKey,
+                nextKey = nextKey,
+                itemsBefore = offset,
+                itemsAfter = maxOf(0, itemCount - nextPosToLoad)
             )
         } catch (e: Exception) {
-            Log.e("MyPagingSource", "${e.message}\n${e.stackTrace}")
             LoadResult.Error(e)
         }
     }
 
-    override fun getRefreshKey(state: PagingState<Int, T>): Int {
-        return ((state.anchorPosition ?: 0) - state.config.initialLoadSize / 2).coerceAtLeast(0)
+    override fun getRefreshKey(state: PagingState<Int, T>): Int? = state.getClippedRefreshKey()
+}
+
+private fun <Value : Any> PagingState<Int, Value>.getClippedRefreshKey(): Int? {
+    return when (val anchorPosition = anchorPosition) {
+        null -> null
+        else -> maxOf(0, anchorPosition - (config.initialLoadSize / 2))
+    }
+}
+
+fun getLimit(params: LoadParams<Int>, key: Int): Int {
+    return when (params) {
+        is Prepend ->
+            if (key < params.loadSize) {
+                key
+            } else {
+                params.loadSize
+            }
+        else -> params.loadSize
+    }
+}
+
+fun getOffset(params: LoadParams<Int>, key: Int, itemCount: Int): Int {
+    return when (params) {
+        is Prepend ->
+            if (key < params.loadSize) {
+                0
+            } else {
+                key - params.loadSize
+            }
+        is Append -> key
+        is Refresh ->
+            if (itemCount != INITIAL_ITEM_COUNT && key >= itemCount) {
+                maxOf(0, itemCount - params.loadSize)
+            } else {
+                key
+            }
     }
 }

wear/src/main/java/org/tasks/presentation/screens/TaskListScreen.kt

@@ -10,7 +10,6 @@ import androidx.compose.material.icons.outlined.CheckBox
 import androidx.compose.material.icons.outlined.CheckBoxOutlineBlank
 import androidx.compose.material.icons.outlined.Repeat
 import androidx.compose.runtime.Composable
-import androidx.compose.runtime.key
 import androidx.compose.ui.Alignment
 import androidx.compose.ui.Modifier
 import androidx.compose.ui.graphics.Color
@@ -28,6 +27,7 @@ import com.google.android.horologist.annotations.ExperimentalHorologistApi
 import com.google.android.horologist.compose.layout.ScalingLazyColumn
 import com.google.android.horologist.compose.layout.ScreenScaffold
 import com.google.android.horologist.compose.layout.rememberResponsiveColumnState
+import com.google.android.horologist.compose.paging.items
 import org.tasks.GrpcProto
 import org.tasks.kmp.org.tasks.themes.ColorProvider
 
@@ -41,19 +41,28 @@ fun TaskListScreen(
     val columnState = rememberResponsiveColumnState()
     ScreenScaffold(
         scrollState = columnState,
-        positionIndicator = {},
     ) {
         ScalingLazyColumn(
             modifier = Modifier.fillMaxSize(),
             columnState = columnState,
         ) {
-            items(uiItems.itemCount) { index ->
-                val item = uiItems[index] ?: return@items
-                key(item.id) {
+            items(
+                items = uiItems,
+                key = { item -> "${item.type}_${item.id}" },
+            ) { item ->
+                if (item == null) {
+                    TaskCard(
+                        task = GrpcProto.UiItem.getDefaultInstance(),
+                        showCheckbox = false,
+                        onComplete = {},
+                        onClick = {},
+                    )
+                } else {
                     when (item.type) {
                         GrpcProto.UiItemType.Task ->
                             TaskCard(
                                 task = item,
+                                showCheckbox = true,
                                 onComplete = { onComplete(item.id) },
                                 onClick = { onClick(item.id) },
                             )
@@ -87,6 +96,7 @@ fun GroupSeparator(
 @Composable
 fun TaskCard(
     task: GrpcProto.UiItem,
+    showCheckbox: Boolean,
     onComplete: () -> Unit,
     onClick: () -> Unit,
 ) {
@@ -98,25 +108,27 @@ fun TaskCard(
         Row(
             verticalAlignment = Alignment.CenterVertically,
         ) {
-            Button(
-                onClick = onComplete,
-                colors = ButtonDefaults.iconButtonColors(),
-            ) {
-                Icon(
-                    imageVector = when {
-                        task.completed -> Icons.Outlined.CheckBox
-                        task.repeating -> Icons.Outlined.Repeat
-                        else -> Icons.Outlined.CheckBoxOutlineBlank
-                    },
-                    tint = Color(
-                        ColorProvider.priorityColor(
-                            task.priority,
-                            isDarkMode = true,
-                            desaturate = true
-                        )
-                    ),
-                    contentDescription = null,
-                )
+            if (showCheckbox) {
+                Button(
+                    onClick = onComplete,
+                    colors = ButtonDefaults.iconButtonColors(),
+                ) {
+                    Icon(
+                        imageVector = when {
+                            task.completed -> Icons.Outlined.CheckBox
+                            task.repeating -> Icons.Outlined.Repeat
+                            else -> Icons.Outlined.CheckBoxOutlineBlank
+                        },
+                        tint = Color(
+                            ColorProvider.priorityColor(
+                                task.priority,
+                                isDarkMode = true,
+                                desaturate = true
+                            )
+                        ),
+                        contentDescription = null,
+                    )
+                }
             }
             Text(
                 text = task.title,

wear/src/main/java/org/tasks/presentation/screens/TaskListViewModel.kt

@@ -6,6 +6,7 @@ import androidx.lifecycle.viewModelScope
 import androidx.paging.Pager
 import androidx.paging.PagingConfig
 import androidx.paging.PagingData
+import androidx.paging.cachedIn
 import com.google.android.horologist.annotations.ExperimentalHorologistApi
 import com.google.android.horologist.data.ProtoDataStoreHelper.protoFlow
 import com.google.android.horologist.data.TargetNodeId
@@ -40,12 +41,14 @@ class TaskListViewModel(
                             .setLimit(limit)
                             .build()
                     )
-                    .itemsList
+                    .let { Pair(it.totalItems, it.itemsList) }
             }
                 .also { pagingSource = it }
         }
     )
         .flow
+        .cachedIn(viewModelScope)
+
     private val wearDataLayerRegistry = WearDataLayerRegistry.fromContext(
         application = application,
         coroutineScope = viewModelScope,