From cfd09c42e878115a2e8c17a58a2b48df7bd72e64 Mon Sep 17 00:00:00 2001
From: Don Syme <donsyme@fastmail.fm>
Date: Fri, 28 Oct 2022 19:30:35 +0100
Subject: [PATCH] use copy of initial state machine to get resumption back to
 zero

---
 .../TaskSeq.Iter.Tests.fs                     |  2 +-
 .../TaskSeq.ToXXX.Tests.fs                    |  6 ++---
 src/FSharpy.TaskSeq/TaskSeqBuilder.fs         | 27 +++++--------------
 3 files changed, 10 insertions(+), 25 deletions(-)

diff --git a/src/FSharpy.TaskSeq.Test/TaskSeq.Iter.Tests.fs b/src/FSharpy.TaskSeq.Test/TaskSeq.Iter.Tests.fs
index 4d6ff2b2..660e4d21 100644
--- a/src/FSharpy.TaskSeq.Test/TaskSeq.Iter.Tests.fs
+++ b/src/FSharpy.TaskSeq.Test/TaskSeq.Iter.Tests.fs
@@ -46,7 +46,7 @@ let ``TaskSeq-iter multiple iterations over same sequence`` () = task {
     do! tq |> TaskSeq.iter (fun item -> sum <- sum + item)
     do! tq |> TaskSeq.iter (fun item -> sum <- sum + item)
     do! tq |> TaskSeq.iter (fun item -> sum <- sum + item)
-    sum |> should equal 220 // task-dummies started at 1
+    sum |> should equal 820 // task-dummies started at 1
 }
 
 [<Fact>]
diff --git a/src/FSharpy.TaskSeq.Test/TaskSeq.ToXXX.Tests.fs b/src/FSharpy.TaskSeq.Test/TaskSeq.ToXXX.Tests.fs
index 600e1fa1..be02631a 100644
--- a/src/FSharpy.TaskSeq.Test/TaskSeq.ToXXX.Tests.fs
+++ b/src/FSharpy.TaskSeq.Test/TaskSeq.ToXXX.Tests.fs
@@ -35,9 +35,9 @@ let ``TaskSeq-toArrayAsync can be applied multiple times to the same sequence``
     let! (results3: _[]) = tq |> TaskSeq.toArrayAsync
     let! (results4: _[]) = tq |> TaskSeq.toArrayAsync
     results1 |> should equal [| 1..10 |]
-    results2 |> should equal [| 1..10 |]
-    results3 |> should equal [| 1..10 |]
-    results4 |> should equal [| 1..10 |]
+    results2 |> should equal [| 11..20 |]
+    results3 |> should equal [| 21..30 |]
+    results4 |> should equal [| 31..40 |]
 }
 
 [<Fact>]
diff --git a/src/FSharpy.TaskSeq/TaskSeqBuilder.fs b/src/FSharpy.TaskSeq/TaskSeqBuilder.fs
index 0b25d201..c681fa0c 100644
--- a/src/FSharpy.TaskSeq/TaskSeqBuilder.fs
+++ b/src/FSharpy.TaskSeq/TaskSeqBuilder.fs
@@ -148,6 +148,9 @@ and [<NoComparison; NoEquality>] TaskSeq<'Machine, 'T
     inherit TaskSeq<'T>()
     let initialThreadId = Environment.CurrentManagedThreadId
 
+    [<DefaultValue(false)>]
+    val mutable InitialMachine: 'Machine
+
     [<DefaultValue(false)>]
     val mutable Machine: 'Machine
 
@@ -247,7 +250,6 @@ and [<NoComparison; NoEquality>] TaskSeq<'Machine, 'T
                 (not data.taken
                  && initialThreadId = Environment.CurrentManagedThreadId)
             then
-                //let clone = this.MemberwiseClone() :?> TaskSeq<'Machine, 'T>
                 let data = this.Machine.Data
                 data.taken <- true
                 data.cancellationToken <- ct
@@ -274,24 +276,7 @@ and [<NoComparison; NoEquality>] TaskSeq<'Machine, 'T
                 // iteration twice, we are trying to *await* twice, which is not allowed
                 // see, for instance: https://itnext.io/why-can-a-valuetask-only-be-awaited-once-31169b324fa4
                 let clone = this.MemberwiseClone() :?> TaskSeq<'Machine, 'T>
-                data.taken <- true
-
-                // Explanation for resetting Machine use brute-force
-                //
-                // This appears to fix the problem that ResumptionPoint was not reset. I'd prefer
-                // a less drastical method. It solves a scenario like the following:
-                // let ts = taskSeq { yield 1; yield 2 }
-                // let e1 = ts.GetAsyncEnumerator()
-                // let! hasNext = e.MoveNextAsync()
-                // let e2 = ts.GetAsyncEnumerator()
-                // let! hasNext = e.MoveNextAsync()  // without this hack, it would continue where e1 left off
-                // let a = e1.Current
-                // let b = e2.Current
-                // let isTrue = a = b  // true with this, false without it
-                clone.Machine <- Unchecked.defaultof<_>
-                //clone.Machine.ResumptionPoint <- 0
-
-                // the following lines just re-initialize the key data fields State.
+                clone.Machine <- clone.InitialMachine
                 clone.Machine.Data <- TaskSeqStateMachineData()
                 clone.Machine.Data.cancellationToken <- ct
                 clone.Machine.Data.taken <- true
@@ -301,10 +286,9 @@ and [<NoComparison; NoEquality>] TaskSeq<'Machine, 'T
                     printfn "All data after clone:"
                     clone.Machine.Data.LogDump()
 
-
                 //// calling reset causes NRE in IValueTaskSource.GetResult above
                 //clone.Machine.Data.promiseOfValueOrEnd.Reset()
-                //clone.Machine.Data.boxed <- clone
+                clone.Machine.Data.boxed <- clone
                 ////clone.Machine.Data.disposalStack <- null // reference type, would otherwise still reference original stack
                 //////clone.Machine.Data.tailcallTarget <- Some clone  // this will lead to an SO exception
                 //clone.Machine.Data.awaiter <- null
@@ -521,6 +505,7 @@ type TaskSeqBuilder() =
                         printfn "at AfterCode<_, _>, setting the Machine field to the StateMachine"
 
                     let ts = TaskSeq<TaskSeqStateMachine<'T>, 'T>()
+                    ts.InitialMachine <- sm
                     ts.Machine <- sm
                     ts.Machine.Data <- TaskSeqStateMachineData()
                     ts.Machine.Data.boxed <- ts