Improve performance of DataCollectionFragment

app/src/main/java/com/google/android/ground/ui/datacollection/DataCollectionFragment.kt

@@ -102,7 +102,6 @@ class DataCollectionFragment : AbstractFragment(), BackPressListener {
       }
     )
 
-    viewModel.init()
     lifecycleScope.launch { viewModel.uiState.filterNotNull().collect { updateUI(it) } }
   }
 

app/src/main/java/com/google/android/ground/ui/datacollection/DataCollectionViewModel.kt

@@ -27,7 +27,6 @@
 import com.google.android.ground.model.submission.TaskData
 import com.google.android.ground.model.submission.ValueDelta
 import com.google.android.ground.model.submission.isNotNullOrEmpty
-import com.google.android.ground.model.task.Condition
 import com.google.android.ground.model.task.Task
 import com.google.android.ground.persistence.local.room.converter.SubmissionDeltasConverter
 import com.google.android.ground.persistence.uuid.OfflineUuidGenerator
@@ -58,6 +57,7 @@
 import kotlinx.coroutines.flow.SharingStarted
 import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.asStateFlow
+import kotlinx.coroutines.flow.collectLatest
 import kotlinx.coroutines.flow.filterNotNull
 import kotlinx.coroutines.flow.first
 import kotlinx.coroutines.flow.flow
@@ -140,15 +140,30 @@
     MutableStateFlow(mutableMapOf())
 
   private val taskDataHandler = TaskDataHandler()
+  private val taskSequenceHandler = TaskSequenceHandler(tasks, taskDataHandler)
 
   // Tracks the current task ID to compute the position in the list of tasks for the current job.
-  private val currentTaskId: StateFlow<String> =
-    savedStateHandle.getStateFlow(TASK_POSITION_ID, tasks.firstOrNull()?.id ?: "")
+  private val currentTaskId: StateFlow<String> = savedStateHandle.getStateFlow(TASK_POSITION_ID, "")
 
-  lateinit var submissionId: String
+  init {
+    // Set current task's ID for new task submissions.
+    if (currentTaskId.value == "") {
+      savedStateHandle[TASK_POSITION_ID] = taskSequenceHandler.getTaskSequence().first().id
+    }
+
+    _uiState.update { UiState.TaskListAvailable(tasks, getTaskPosition(getCurrentTaskId())) }
 
-  fun init() {
-    _uiState.update { UiState.TaskListAvailable(tasks, getTaskPosition(currentTaskId.value)) }
+    // Update the computed task sequence cache if the task's data is updated.
+    viewModelScope.launch {
+      taskDataHandler.dataState.collectLatest { taskSequenceHandler.refreshTaskSequence() }
+    }
+  }
+
+  /** Returns the ID of the user visible task. */
+  private fun getCurrentTaskId(): String {
+    val taskId = currentTaskId.value
+    check(taskId.isNotBlank()) { "Task ID can't be blank" }
+    return taskId
   }
 
   fun setLoiName(name: String) {
@@ -245,7 +260,7 @@
 
   /** Persists the current UI state locally which can be resumed whenever the app re-opens. */
   fun saveCurrentState() {
-    val taskId = currentTaskId.value
+    val taskId = getCurrentTaskId()
     val viewModel = getTaskViewModel(taskId) ?: error("ViewModel not found for task $taskId")
 
     viewModel.validate()?.let {
@@ -260,8 +275,8 @@
 
   /** Retrieves a list of [ValueDelta] for tasks that are part of the current sequence. */
   private fun getDeltas(): List<ValueDelta> =
-    // Filter deltas to valid tasks.
-    getTaskSequence()
+    taskSequenceHandler
+      .getTaskSequence()
       .mapNotNull { task ->
         taskDataHandler.getData(task)?.let { ValueDelta(task.id, task.type, it) }
       }
@@ -275,9 +290,6 @@
     }
   }
 
-  private fun getAbsolutePosition(taskId: String): Int =
-    if (taskId == "") 0 else tasks.indexOf(tasks.first { it.id == taskId })
-
   /** Persists the collected data as draft to local storage. */
   private fun saveDraft(taskId: String) {
     val deltas = getDeltas()
@@ -304,112 +316,47 @@
     externalScope.launch(ioDispatcher) { submissionRepository.deleteDraftSubmission() }
   }
 
-  /**
-   * Get the current index within the computed task sequence, and the number of tasks in the
-   * sequence, e.g (0, 2) means the first task of 2.
-   */
-  private fun getPositionInTaskSequence(taskId: String): Pair<Int, Int> {
-    var currentIndex = 0
-    var size = 0
-    getTaskSequence().forEachIndexed { index, task ->
-      if (task.id == taskId) {
-        currentIndex = index
-      }
-      size++
-    }
-    return currentIndex to size
-  }
-
-  /** Returns the index of the task ID, or -1 if null or not found. */
-  private fun getIndexOfTask(taskId: String?) =
-    if (taskId == null) {
-      -1
-    } else {
-      tasks.indexOfFirst { it.id == taskId }
-    }
-
-  /**
-   * Retrieves the current task sequence given the inputs and conditions set on the tasks. Setting a
-   * start ID will always generate a sequence with the start ID as the first element, and if
-   * reversed is set, will generate the previous tasks from there.
-   */
-  private fun getTaskSequence(
-    startId: String? = null,
-    reversed: Boolean = false,
-    taskValueOverride: Pair<String, TaskData?>? = null,
-  ): Sequence<Task> {
-    if (tasks.isEmpty()) {
-      error("Can't generate sequence for empty task list")
-    }
-    val startIndex =
-      getIndexOfTask(startId).let {
-        if (it < 0) {
-          // Default to 0 if startId is not found or is null.
-          if (startId != null) Timber.w("startId, $startId, was not found. Defaulting to 0")
-          0
-        } else {
-          it
-        }
-      }
-    return if (reversed) {
-        tasks.subList(0, startIndex + 1).reversed()
-      } else {
-        tasks.subList(startIndex, tasks.size)
-      }
-      .let { tasks ->
-        tasks.asSequence().filter {
-          it.condition == null || evaluateCondition(it.condition, taskValueOverride)
-        }
-      }
-  }
-
   private fun moveToNextTask() {
-    step(false)
+    val taskId = taskSequenceHandler.getNextTask(getCurrentTaskId())
+    moveToTask(taskId)
   }
 
   fun moveToPreviousTask() {
-    step(true)
+    val taskId = taskSequenceHandler.getPreviousTask(getCurrentTaskId())
+    moveToTask(taskId)
   }
 
-  /** Displays the task at the relative position to the current one. */
-  private fun step(reversed: Boolean) {
-    val task = getTaskSequence(startId = currentTaskId.value, reversed).take(2).last()
-    savedStateHandle[TASK_POSITION_ID] = task.id
+  private fun moveToTask(taskId: String) {
+    savedStateHandle[TASK_POSITION_ID] = taskId
 
     // Save collected data as draft
     clearDraft()
-    saveDraft(task.id)
+    saveDraft(taskId)
 
-    _uiState.update { UiState.TaskUpdated(getTaskPosition(task.id)) }
+    _uiState.update { UiState.TaskUpdated(getTaskPosition(taskId)) }
   }
 
-  private fun getTaskPosition(taskId: String): TaskPosition {
-    val (index, size) = getPositionInTaskSequence(taskId)
-    return TaskPosition(
-      absoluteIndex = getAbsolutePosition(taskId),
-      relativeIndex = index,
-      sequenceSize = size,
-    )
-  }
+  private fun getTaskPosition(taskId: String) = taskSequenceHandler.getTaskPosition(taskId)
 
-  /** Returns true if the given [taskId] is first in the sequence of displayed tasks. */
-  fun isFirstPosition(taskId: String): Boolean = taskId == getTaskSequence().first().id
+  /** Returns true if the given [taskId] is first task in the sequence of displayed tasks. */
+  fun isFirstPosition(taskId: String): Boolean = taskSequenceHandler.isFirstPosition(taskId)
+
+  /** Returns true if the given [taskId] is last task in the sequence of displayed tasks. */
+  fun isLastPosition(taskId: String): Boolean = taskSequenceHandler.isLastPosition(taskId)
 
   /**
-   * Returns true if the given [taskId] with task data would be last in sequence. Useful for
-   * handling conditional tasks, see #2394.
+   * Returns true if the given [taskId] with [newValue] would be last in the sequence of displayed
+   * tasks. Required for handling conditional tasks, see #2394.
    */
-  fun checkLastPositionWithTaskData(taskId: String, value: TaskData?): Boolean =
-    taskId == getTaskSequence(taskValueOverride = taskId to value).last().id
-
-  /** Returns true if the given [taskId] is last task in sequence. */
-  fun isLastPosition(taskId: String): Boolean = taskId == getTaskSequence().last().id
+  fun isLastPositionWithValue(taskId: String, newValue: TaskData?): Boolean {
+    if (taskDataHandler.getData(taskId) == newValue) {
+      // Reuse the existing task sequence if the value has already been saved (i.e. after pressing
+      // "Next" and going back).
+      return isLastPosition(taskId)
+    }
 
-  /** Evaluates the task condition against the current inputs. */
-  private fun evaluateCondition(
-    condition: Condition,
-    taskValueOverride: Pair<String, TaskData?>? = null,
-  ): Boolean = condition.fulfilledBy(taskDataHandler.getTaskSelections(taskValueOverride))
+    return taskSequenceHandler.checkIfTaskIsLastWithValue(taskValueOverride = taskId to newValue)
+  }
 
   companion object {
     private const val TASK_JOB_ID_KEY = "jobId"

app/src/main/java/com/google/android/ground/ui/datacollection/TaskDataHandler.kt

@@ -54,8 +54,14 @@ class TaskDataHandler {
    * @param task The [Task] to retrieve data for.
    * @return The [TaskData] associated with the task, or `null` if no data is found.
    */
+  @Deprecated("Use getData(taskId: String) instead")
   fun getData(task: Task): TaskData? = _dataState.value[task]
 
+  // TODO: Refactor task data handler to use TaskId as key instead of Task.
+  // Issue URL: https://github.com/google/ground-android/issues/3009
+  fun getData(taskId: String): TaskData? =
+    _dataState.value.filterKeys { it.id == taskId }.values.firstOrNull()
+
   /**
    * Returns a [TaskSelections] map representing the current state of task data.
    *

app/src/main/java/com/google/android/ground/ui/datacollection/TaskSequenceHandler.kt

@@ -15,11 +15,12 @@
  */
 package com.google.android.ground.ui.datacollection
 
+import com.google.android.ground.model.submission.TaskData
 import com.google.android.ground.model.task.Task
 import com.google.android.ground.model.task.TaskSelections
 
 /**
- * Manages operations related to a sequence of tasks.
+ * Manages state and operations related to a sequence of tasks.
  *
  * This class provides methods to retrieve, navigate, and query the position of tasks within a
  * sequence. The sequence is derived from a list of [Task] objects, filtered based on a provided
@@ -43,17 +44,27 @@ class TaskSequenceHandler(
     require(taskId.isNotBlank()) { "Task ID can't be blank." }
   }
 
-  /** Generates the task sequence based on conditions and overrides. */
-  fun generateTaskSequence(taskSelections: TaskSelections? = null): Sequence<Task> {
-    val selections = taskSelections ?: taskDataHandler.getTaskSelections()
+  /** Generates the task sequence based on task's conditions. */
+  fun generateTaskSequence(): Sequence<Task> {
+    val selections = taskDataHandler.getTaskSelections()
     return tasks.filter { shouldIncludeTaskInSequence(it, selections) }.asSequence()
   }
 
+  /** Returns true if the specified task would be last in the sequence with the given value. */
+  fun checkIfTaskIsLastWithValue(taskValueOverride: Pair<String, TaskData?>): Boolean {
+    val overriddenTaskId = taskValueOverride.first
+    validateTaskId(overriddenTaskId)
+
+    val selections = taskDataHandler.getTaskSelections(taskValueOverride)
+    val lastTask = tasks.last { shouldIncludeTaskInSequence(it, selections) }
+    return lastTask.id == overriddenTaskId
+  }
+
   /** Determines if a task should be included with the given overrides. */
   private fun shouldIncludeTaskInSequence(task: Task, taskSelections: TaskSelections): Boolean =
     task.condition == null || task.condition.fulfilledBy(taskSelections)
 
-  /** Lazily retrieves the task sequence. */
+  /** Returns the pre-computed task sequence or generates and caches it first, if missing. */
   fun getTaskSequence(): Sequence<Task> {
     if (!isSequenceInitialized) {
       taskSequence = generateTaskSequence()
@@ -62,8 +73,15 @@ class TaskSequenceHandler(
     return taskSequence
   }
 
-  // TODO: Add a method to update the cached sequence whenever necessary.
-  // Issue URL: https://github.com/google/ground-android/issues/2993
+  /**
+   * Re-computes the task sequence and updates the local cache.
+   *
+   * Note: This is a heavy computing operation. So, it should only be triggered if the task's value
+   * has been updated.
+   */
+  fun refreshTaskSequence() {
+    taskSequence = generateTaskSequence()
+  }
 
   /**
    * Checks if the specified task is the first task in the displayed sequence.

app/src/main/java/com/google/android/ground/ui/datacollection/tasks/AbstractTaskFragment.kt

@@ -155,7 +155,7 @@ abstract class AbstractTaskFragment<T : AbstractTaskViewModel> : AbstractFragmen
           button.showIfTrue(value.isNotNullOrEmpty())
         }
         button.enableIfTrue(value.isNotNullOrEmpty())
-        val isLastPosition = checkLastPositionWithTaskData(value)
+        val isLastPosition = dataCollectionViewModel.isLastPositionWithValue(taskId, value)
         button.toggleDone(done = isLastPosition)
       }
       .disable()
@@ -239,13 +239,6 @@ abstract class AbstractTaskFragment<T : AbstractTaskViewModel> : AbstractFragmen
   /** Returns true if the current task is in the last position in the sequence. */
   private fun isLastPosition() = dataCollectionViewModel.isLastPosition(taskId)
 
-  /**
-   * Returns true if the current task with the given task data would be last in sequence. Useful for
-   * handling conditional tasks, see #2394.
-   */
-  private fun checkLastPositionWithTaskData(value: TaskData?) =
-    dataCollectionViewModel.checkLastPositionWithTaskData(taskId, value)
-
   private fun getTask(): Task = viewModel.task
 
   fun getCurrentValue(): TaskData? = viewModel.taskTaskData.value

app/src/test/java/com/google/android/ground/ui/datacollection/TaskDataHandlerTest.kt

@@ -108,6 +108,24 @@ class TaskDataHandlerTest {
     assertThat(handler.getData(task)).isNull()
   }
 
+  @Test
+  fun `getData with taskId returns correct data`() = runTest {
+    val handler = TaskDataHandler()
+    val task = createTask("task1")
+    val taskData = createTaskData("data1")
+
+    handler.setData(task, taskData)
+
+    assertThat(handler.getData(task.id)).isEqualTo(taskData)
+  }
+
+  @Test
+  fun `getData with taskId returns null for unknown task`() = runTest {
+    val handler = TaskDataHandler()
+
+    assertThat(handler.getData("task1")).isNull()
+  }
+
   @Test
   fun `getTaskSelections returns correct values`() = runTest {
     val handler = TaskDataHandler()