Update DeferredJsonMerger to take pending and completed into account.

Author: BoD

libraries/apollo-runtime/src/commonMain/kotlin/com/apollographql/apollo/internal/DeferredJsonMerger.kt

@@ -15,19 +15,25 @@ private typealias MutableJsonMap = MutableMap<String, Any?>
  * Each call to [merge] will merge the given chunk into the [merged] Map, and will also update the [mergedFragmentIds] Set with the
  * value of its `path` and `label` field.
  *
- * The fields in `data` are merged into the node found in [merged] at `path` (for the first call to [merge], the payload is
- * copied to [merged] as-is).
+ * The fields in `data` are merged into the node found in [merged] at the path known by looking at the `id` field (for the first call to
+ * [merge], the payload is copied to [merged] as-is).
  *
- * `errors` in incremental items (if present) are merged together in an array and then set to the `errors` field of the [merged] Map,
- * at each call to [merge].
- * `extensions` in incremental items (if present) are merged together in an array and then set to the `extensions/incremental` field of the
+ * `errors` in incremental and completed items (if present) are merged together in an array and then set to the `errors` field of the
  * [merged] Map, at each call to [merge].
+ * `extensions` in incremental items (if present) are merged together in an array and then set to the `extensions` field of the [merged]
+ * Map, at each call to [merge].
  */
 @ApolloInternal
+@Suppress("UNCHECKED_CAST")
 class DeferredJsonMerger {
   private val _merged: MutableJsonMap = mutableMapOf()
   val merged: JsonMap = _merged
 
+  /**
+   * Map of identifiers to their corresponding DeferredFragmentIdentifier, found in `pending`.
+   */
+  private val idsToDeferredFragmentIdentifiers = mutableMapOf<String, DeferredFragmentIdentifier>()
+
   private val _mergedFragmentIds = mutableSetOf<DeferredFragmentIdentifier>()
   val mergedFragmentIds: Set<DeferredFragmentIdentifier> = _mergedFragmentIds
 
@@ -47,11 +53,12 @@ class DeferredJsonMerger {
     return merge(payloadMap)
   }
 
-  @Suppress("UNCHECKED_CAST")
   fun merge(payload: JsonMap): JsonMap {
     if (merged.isEmpty()) {
-      // Initial payload, no merging needed
-      _merged += payload
+      // Initial payload, no merging needed (strip some fields that should not appear in the final result)
+      _merged += payload - "hasNext" - "pending"
+      handlePending(payload)
+      handleCompleted(payload)
       return merged
     }
 
@@ -60,48 +67,68 @@ class DeferredJsonMerger {
       isEmptyPayload = true
     } else {
       isEmptyPayload = false
-      val mergedErrors = mutableListOf<JsonMap>()
-      val mergedExtensions = mutableListOf<JsonMap>()
       for (incrementalItem in incrementalList) {
-        mergeData(incrementalItem)
-        // Merge errors and extensions (if any) of the incremental list
-        (incrementalItem["errors"] as? List<JsonMap>)?.let { mergedErrors += it }
-        (incrementalItem["extensions"] as? JsonMap)?.let { mergedExtensions += it }
-      }
-      // Keep only this payload's errors and extensions, if any
-      if (mergedErrors.isNotEmpty()) {
-        _merged["errors"] = mergedErrors
-      } else {
-        _merged.remove("errors")
-      }
-      if (mergedExtensions.isNotEmpty()) {
-        _merged["extensions"] = mapOf("incremental" to mergedExtensions)
-      } else {
-        _merged.remove("extensions")
+        mergeIncrementalData(incrementalItem)
+        // Merge errors (if any) of the incremental item
+        (incrementalItem["errors"] as? List<JsonMap>)?.let { getOrPutMergedErrors() += it }
       }
     }
 
     hasNext = payload["hasNext"] as Boolean? ?: false
 
+    handlePending(payload)
+    handleCompleted(payload)
+
+    (payload["extensions"] as? JsonMap)?.let { getOrPutExtensions() += it }
+
     return merged
   }
 
-  @Suppress("UNCHECKED_CAST")
-  private fun mergeData(incrementalItem: JsonMap) {
-    val data = incrementalItem["data"] as JsonMap?
-    val path = incrementalItem["path"] as List<Any>
-    val mergedData = merged["data"] as JsonMap
+  private fun getOrPutMergedErrors() = _merged.getOrPut("errors") { mutableListOf<JsonMap>() } as MutableList<JsonMap>
 
-    // payloadData can be null if there are errors
-    if (data != null) {
-      val nodeToMergeInto = nodeAtPath(mergedData, path) as MutableJsonMap
-      deepMerge(nodeToMergeInto, data)
+  private fun getOrPutExtensions() = _merged.getOrPut("extensions") { mutableMapOf<String, Any?>() } as MutableJsonMap
 
-      _mergedFragmentIds += DeferredFragmentIdentifier(path = path, label = incrementalItem["label"] as String?)
+  private fun handlePending(payload: JsonMap) {
+    val pending = payload["pending"] as? List<JsonMap>
+    if (pending != null) {
+      for (pendingItem in pending) {
+        val id = pendingItem["id"] as String
+        val path = pendingItem["path"] as List<Any>
+        val label = pendingItem["label"] as String?
+        idsToDeferredFragmentIdentifiers[id] = DeferredFragmentIdentifier(path = path, label = label)
+      }
     }
   }
 
-  @Suppress("UNCHECKED_CAST")
+  private fun handleCompleted(payload: JsonMap) {
+    val completed = payload["completed"] as? List<JsonMap>
+    if (completed != null) {
+      for (completedItem in completed) {
+        val errors = completedItem["errors"] as? List<JsonMap>
+        if (errors != null) {
+          // Merge errors (if any) of the completed item
+          getOrPutMergedErrors() += errors
+        } else {
+          // No errors: we have merged all the fields of the fragment so it can be parsed
+          val id = completedItem["id"] as String
+          val deferredFragmentIdentifier = idsToDeferredFragmentIdentifiers.remove(id)
+              ?: error("Id '$id' not found in pending results")
+          _mergedFragmentIds += deferredFragmentIdentifier
+        }
+      }
+    }
+  }
+
+  private fun mergeIncrementalData(incrementalItem: JsonMap) {
+    val id = incrementalItem["id"] as String? ?: error("No id found in incremental item")
+    val data = incrementalItem["data"] as JsonMap? ?: error("No data found in incremental item")
+    val subPath = incrementalItem["subPath"] as List<Any>? ?: emptyList()
+    val path = (idsToDeferredFragmentIdentifiers[id]?.path ?: error("Id '$id' not found in pending results")) + subPath
+    val mergedData = merged["data"] as JsonMap
+    val nodeToMergeInto = nodeAtPath(mergedData, path) as MutableJsonMap
+    deepMerge(nodeToMergeInto, data)
+  }
+
   private fun deepMerge(destination: MutableJsonMap, map: JsonMap) {
     for ((key, value) in map) {
       if (destination.containsKey(key) && destination[key] is MutableMap<*, *>) {
@@ -116,7 +143,6 @@ class DeferredJsonMerger {
     }
   }
 
-  @Suppress("UNCHECKED_CAST")
   private fun jsonToMap(json: BufferedSource): JsonMap = BufferedSourceJsonReader(json).readAny() as JsonMap
 
 
@@ -130,7 +156,6 @@ class DeferredJsonMerger {
       node = if (node is List<*>) {
         node[key as Int]
       } else {
-        @Suppress("UNCHECKED_CAST")
         node as JsonMap
         node[key]
       }