Use task approach

Author: gusty

src/FSharpPlus/Extensions/ValueTask.fs

@@ -243,22 +243,20 @@ module ValueTask =
 
     /// Flattens two nested ValueTask into one.
     let join (source: ValueTask<ValueTask<'T>>) : ValueTask<'T> =
-        let tcs = TaskCompletionSource<'T> ()
-        continueTask tcs source (fun x ->
-            continueTask tcs x (fun x ->
-                tcs.SetResult x))
-        tcs.Task |> ValueTask<'T>
+        if source.IsCompleted then
+            match source with
+            | Succeeded inner -> inner
+            | Faulted aex -> raise aex
+            | Canceled -> canceled
+        else
+            let tcs = TaskCompletionSource<'T> ()
+            continueTask tcs source (fun inner ->
+                continueTask tcs inner tcs.SetResult)
+            ValueTask<'T> tcs.Task
 
 
     /// <summary>Creates a ValueTask workflow from 'source' workflow, mapping and flattening its result with 'f'.</summary>
-    let bind (f: 'T -> ValueTask<'U>) (source: ValueTask<'T>) : ValueTask<'U> =
-        let tcs = TaskCompletionSource<'U> ()
-        continueTask tcs source (fun x ->
-            try 
-                continueTask tcs (f x) (fun fx ->
-                    tcs.SetResult fx)
-            with e -> tcs.SetException e)
-        tcs.Task |> ValueTask<'U>
+    let bind (f: 'T -> ValueTask<'U>) (source: ValueTask<'T>) : ValueTask<'U> = source |> Unchecked.nonNull |> map f |> join
 
     /// <summary>Creates a ValueTask that ignores the result of the source ValueTask.</summary>
     /// <remarks>It can be used to convert non-generic ValueTask to unit ValueTask.</remarks>
@@ -319,23 +317,10 @@ module ValueTask =
                 compensation ()
         try
             let task = body ()
-            if task.IsCompleted then 
-                try
-                    compensation ()
-                    task
-                with e -> raise e
+            if task.IsCompleted then compensation (); task
             else
-                let tcs = TaskCompletionSource<'T> ()
-                let f = function
-                    | Succeeded r -> tcs.SetResult r
-                    | Faulted aex -> tcs.SetException aex.InnerExceptions
-                    | Canceled    -> tcs.SetCanceled ()
-                task.ConfigureAwait(false).GetAwaiter().UnsafeOnCompleted (fun () ->
-                    try
-                        compensation ()
-                        f task
-                    with e -> tcs.SetException e)
-                ValueTask<'T> tcs.Task
+                let t = task.AsTask()
+                t.ContinueWith(fun (_: Task<'T>) -> compensation (); t).Unwrap () |> ValueTask<'T>
         with _ ->
             compensation ()
             reraise ()