From 4ed19c5e36a1295cd9dea6eac64605f01d9508f5 Mon Sep 17 00:00:00 2001
From: ogaclejapan <ogaclejapan@gmail.com>
Date: Sun, 28 Jul 2024 16:08:47 +0900
Subject: [PATCH] Improve the results of synchronously executed Mutations.

When `isOneShot` in MutationOptions is set to false, multiple calls might not receive the intended results
(there was a possibility of referencing the state of the next execution).

By using a callback, the execution results of the Mutation can be awaited on the calling side,
ensuring that the results of the call are reliably referenced.
---
 .../kotlin/soil/query/MutationClient.kt         |  1 +
 .../kotlin/soil/query/MutationCommand.kt        |  6 +++++-
 .../kotlin/soil/query/MutationCommands.kt       |  7 +++++--
 .../commonMain/kotlin/soil/query/MutationRef.kt | 17 +++++------------
 .../query/internal/CompletableDeferredExt.kt    | 14 ++++++++++++++
 5 files changed, 30 insertions(+), 15 deletions(-)
 create mode 100644 soil-query-core/src/commonMain/kotlin/soil/query/internal/CompletableDeferredExt.kt

diff --git a/soil-query-core/src/commonMain/kotlin/soil/query/MutationClient.kt b/soil-query-core/src/commonMain/kotlin/soil/query/MutationClient.kt
index 93cc0e7..555f160 100644
--- a/soil-query-core/src/commonMain/kotlin/soil/query/MutationClient.kt
+++ b/soil-query-core/src/commonMain/kotlin/soil/query/MutationClient.kt
@@ -22,3 +22,4 @@ interface MutationClient {
 }
 
 typealias MutationOptionsOverride = (MutationOptions) -> MutationOptions
+typealias MutationCallback<T> = (Result<T>) -> Unit
diff --git a/soil-query-core/src/commonMain/kotlin/soil/query/MutationCommand.kt b/soil-query-core/src/commonMain/kotlin/soil/query/MutationCommand.kt
index 7b8a99c..18e69ce 100644
--- a/soil-query-core/src/commonMain/kotlin/soil/query/MutationCommand.kt
+++ b/soil-query-core/src/commonMain/kotlin/soil/query/MutationCommand.kt
@@ -84,7 +84,8 @@ suspend fun <T, S> MutationCommand.Context<T>.mutate(
  */
 suspend inline fun <T, S> MutationCommand.Context<T>.dispatchMutateResult(
     key: MutationKey<T, S>,
-    variable: S
+    variable: S,
+    noinline callback: MutationCallback<T>?
 ) {
     mutate(key, variable)
         .onSuccess { data ->
@@ -98,6 +99,9 @@ suspend inline fun <T, S> MutationCommand.Context<T>.dispatchMutateResult(
         }
         .onFailure { dispatch(MutationAction.MutateFailure(it)) }
         .onFailure { options.onError?.invoke(it, state, key.id) }
+        .also {
+            callback?.invoke(it)
+        }
 }
 
 internal fun <T> MutationCommand.Context<T>.onRetryCallback(
diff --git a/soil-query-core/src/commonMain/kotlin/soil/query/MutationCommands.kt b/soil-query-core/src/commonMain/kotlin/soil/query/MutationCommands.kt
index 684e377..004544b 100644
--- a/soil-query-core/src/commonMain/kotlin/soil/query/MutationCommands.kt
+++ b/soil-query-core/src/commonMain/kotlin/soil/query/MutationCommands.kt
@@ -4,6 +4,7 @@
 package soil.query
 
 import soil.query.internal.vvv
+import kotlin.coroutines.cancellation.CancellationException
 
 /**
  * Mutation commands are used to update the [mutation state][MutationState].
@@ -24,16 +25,18 @@ sealed class MutationCommands<T> : MutationCommand<T> {
     data class Mutate<T, S>(
         val key: MutationKey<T, S>,
         val variable: S,
-        val revision: String
+        val revision: String,
+        val callback: MutationCallback<T>? = null
     ) : MutationCommands<T>() {
 
         override suspend fun handle(ctx: MutationCommand.Context<T>) {
             if (!ctx.shouldMutate(revision)) {
                 ctx.options.vvv(key.id) { "skip mutation(shouldMutate=false)" }
+                callback?.invoke(Result.failure(CancellationException("skip mutation")))
                 return
             }
             ctx.dispatch(MutationAction.Mutating)
-            ctx.dispatchMutateResult(key, variable)
+            ctx.dispatchMutateResult(key, variable, callback)
         }
     }
 
diff --git a/soil-query-core/src/commonMain/kotlin/soil/query/MutationRef.kt b/soil-query-core/src/commonMain/kotlin/soil/query/MutationRef.kt
index 8ca0389..d5156cb 100644
--- a/soil-query-core/src/commonMain/kotlin/soil/query/MutationRef.kt
+++ b/soil-query-core/src/commonMain/kotlin/soil/query/MutationRef.kt
@@ -3,8 +3,8 @@
 
 package soil.query
 
-import kotlinx.coroutines.flow.dropWhile
-import kotlinx.coroutines.flow.first
+import kotlinx.coroutines.CompletableDeferred
+import soil.query.internal.toResultCallback
 
 /**
  * A reference to a [Mutation] for [MutationKey].
@@ -28,16 +28,9 @@ class MutationRef<T, S>(
      * @return The result of the mutation.
      */
     suspend fun mutate(variable: S): T {
-        mutateAsync(variable)
-        val submittedAt = state.value.submittedAt
-        val result = state.dropWhile { it.submittedAt <= submittedAt }.first()
-        if (result.isSuccess) {
-            return result.data!!
-        } else if (result.isFailure) {
-            throw result.error!!
-        } else {
-            error("Unexpected ${result.status}")
-        }
+        val deferred = CompletableDeferred<T>()
+        command.send(MutationCommands.Mutate(key, variable, state.value.revision, deferred.toResultCallback()))
+        return deferred.await()
     }
 
     /**
diff --git a/soil-query-core/src/commonMain/kotlin/soil/query/internal/CompletableDeferredExt.kt b/soil-query-core/src/commonMain/kotlin/soil/query/internal/CompletableDeferredExt.kt
new file mode 100644
index 0000000..b8a23e5
--- /dev/null
+++ b/soil-query-core/src/commonMain/kotlin/soil/query/internal/CompletableDeferredExt.kt
@@ -0,0 +1,14 @@
+// Copyright 2024 Soil Contributors
+// SPDX-License-Identifier: Apache-2.0
+
+package soil.query.internal
+
+import kotlinx.coroutines.CompletableDeferred
+
+internal fun <T> CompletableDeferred<T>.toResultCallback(): (Result<T>) -> Unit {
+    return { result ->
+        result
+            .onSuccess { complete(it) }
+            .onFailure { completeExceptionally(it) }
+    }
+}