@@ -55,59 +55,75 @@ module CancellableTasks =
5555 let initialResumptionFunc =
5656 CancellableTaskBaseResumptionFunc< 'T, _>( fun sm -> code.Invoke(& sm))
5757
58- let resumptionInfo () =
59- let mutable state = InitialYield
58+ let bounceAllowed = Trampoline.Current.IsAwaited()
59+
60+ let maybeBounce state =
61+ if
62+ bounceAllowed
63+ && Trampoline.Current.ShouldBounce
64+ then
65+ Bounce state
66+ else
67+ Immediate state
68+
69+ let resumptionInfo =
70+ let initialState = maybeBounce Running
6071
61- { new CancellableTaskBaseResumptionDynamicInfo< 'T, _>( initialResumptionFunc) with
72+ { new CancellableTaskBaseResumptionDynamicInfo< 'T, _>( initialResumptionFunc,
73+ ResumptionData = initialState) with
6274 member info.MoveNext ( sm ) =
63- let current = state
64- let mutable continuation = Stop
6575
66- match current with
67- | InitialYield ->
68- state <- Running
76+ let getCurrent () =
77+ match info.ResumptionData with
78+ | null -> failwith " Invalid state: ResumptionData is null"
79+ | state -> unbox state
6980
70- continuation <-
71- if BindContext.CheckWhenIsBind() then Bounce else Immediate
72- | Running ->
73- try
74- let step = info.ResumptionFunc.Invoke(& sm)
81+ let setState state = info.ResumptionData <- state
7582
76- if step then
77- state <- SetResult
83+ match getCurrent () with
84+ | Immediate state ->
85+ setState state
86+ info.MoveNext & sm
87+ | Running ->
88+ let mutable keepGoing = true
7889
79- continuation <-
80- if BindContext.Check() then Bounce else Immediate
90+ try
91+ if info.ResumptionFunc.Invoke(& sm) then
92+ setState ( maybeBounce SetResult)
8193 else
82- continuation <-
83- Await( downcast sm.ResumptionDynamicInfo.ResumptionData)
94+ keepGoing <-
95+ match getCurrent () with
96+ | Awaiting _ -> true
97+ | _ -> false
8498 with exn ->
85- state <- SetException( ExceptionCache.CaptureOrRetrieve exn)
86- continuation <- if BindContext.Check() then Bounce else Immediate
87- | SetResult ->
88- MethodBuilder.SetResult(& sm.Data.MethodBuilder, sm.Data.Result)
89- | SetException edi ->
90- MethodBuilder.SetException(& sm.Data.MethodBuilder, edi.SourceException)
91-
92- let continuation = continuation
99+ setState (
100+ maybeBounce
101+ <| SetException( ExceptionCache.CaptureOrRetrieve exn)
102+ )
103+
104+ if keepGoing then
105+ info.MoveNext & sm
106+ | Awaiting awaiter ->
107+ setState Running
108+ let mutable awaiter = awaiter
93109
94- match continuation with
95- | Stop -> ()
96- | Immediate -> info.MoveNext(& sm)
97- | Bounce ->
98- MethodBuilder.AwaitOnCompleted(
110+ MethodBuilder.AwaitUnsafeOnCompleted(
99111 & sm.Data.MethodBuilder,
100- Trampoline.Current.Ref ,
112+ & awaiter ,
101113 & sm
102114 )
103- | Await awaiter ->
104- let mutable awaiter = awaiter
115+ | Bounce next ->
116+ setState next
105117
106- MethodBuilder.AwaitUnsafeOnCompleted (
118+ MethodBuilder.AwaitOnCompleted (
107119 & sm.Data.MethodBuilder,
108- & awaiter ,
120+ Trampoline.Current.Ref ,
109121 & sm
110122 )
123+ | SetResult ->
124+ MethodBuilder.SetResult(& sm.Data.MethodBuilder, sm.Data.Result)
125+ | SetException edi ->
126+ MethodBuilder.SetException(& sm.Data.MethodBuilder, edi.SourceException)
111127
112128 member _.SetStateMachine ( sm , state ) =
113129 MethodBuilder.SetStateMachine(& sm.Data.MethodBuilder, state)
@@ -120,7 +136,7 @@ module CancellableTasks =
120136 Task.FromCanceled<_>( ct)
121137 else
122138 sm.Data.CancellationToken <- ct
123- sm.ResumptionDynamicInfo <- resumptionInfo ()
139+ sm.ResumptionDynamicInfo <- resumptionInfo
124140 sm.Data.MethodBuilder <- AsyncTaskMethodBuilder< 'T>. Create()
125141 sm.Data.MethodBuilder.Start(& sm)
126142 sm.Data.MethodBuilder.Task
@@ -132,29 +148,36 @@ module CancellableTasks =
132148 __ stateMachine< CancellableTaskBaseStateMachineData< 'T, _>, CancellableTask< 'T>>
133149 ( MoveNextMethodImpl<_>( fun sm ->
134150 __ resumeAt sm.ResumptionPoint
151+
135152 let mutable error = ValueNone
136153
137- let __stack_go1 = yieldOnBindLimitWhenIsBind() .Invoke(& sm)
154+ let noBounce = not ( Trampoline.Current.IsAwaited())
155+
156+ let __stack_go1 =
157+ noBounce
158+ || yieldOnBindLimit() .Invoke(& sm)
138159
139160 if __ stack_ go1 then
140161 try
141162 let __stack_code_fin = code.Invoke(& sm)
142163
143164 if __ stack_ code_ fin then
144- let __stack_go2 = yieldOnBindLimit() .Invoke(& sm)
165+ let __stack_go2 =
166+ noBounce
167+ || yieldOnBindLimit() .Invoke(& sm)
145168
146169 if __ stack_ go2 then
147170 MethodBuilder.SetResult(
148171 & sm.Data.MethodBuilder,
149172 sm.Data.Result
150173 )
151174 with exn ->
152- error <-
153- ValueSome
154- <| ExceptionCache.CaptureOrRetrieve exn
175+ error <- ValueSome( ExceptionCache.CaptureOrRetrieve exn)
155176
156177 if error.IsSome then
157- let __stack_go2 = yieldOnBindLimit() .Invoke(& sm)
178+ let __stack_go2 =
179+ noBounce
180+ || yieldOnBindLimit() .Invoke(& sm)
158181
159182 if __ stack_ go2 then
160183 MethodBuilder.SetException(
@@ -185,7 +208,7 @@ module CancellableTasks =
185208 member inline _.Source
186209 ( [<InlineIfLambda>] x : CancellationToken -> Task < _ >)
187210 : CancellationToken -> Awaiter < TaskAwaiter < _ >, _ > =
188- fun ct -> Awaitable.GetTaskAwaiter( BindContext.SetIsBind x ct)
211+ fun ct -> Awaitable.GetTaskAwaiter( Trampoline.Allow x ct)
189212
190213 [<NoEagerConstraintApplication>]
191214 member inline this.MergeSources
@@ -293,29 +316,36 @@ module CancellableTasks =
293316 __ stateMachine< CancellableTaskBaseStateMachineData< 'T, _>, CancellableTask< 'T>>
294317 ( MoveNextMethodImpl<_>( fun sm ->
295318 __ resumeAt sm.ResumptionPoint
319+
296320 let mutable error = ValueNone
297321
298- let __stack_go1 = yieldOnBindLimitWhenIsBind() .Invoke(& sm)
322+ let noBounce = not ( Trampoline.Current.IsAwaited())
323+
324+ let __stack_go1 =
325+ noBounce
326+ || yieldOnBindLimit() .Invoke(& sm)
299327
300328 if __ stack_ go1 then
301329 try
302330 let __stack_code_fin = code.Invoke(& sm)
303331
304332 if __ stack_ code_ fin then
305- let __stack_go2 = yieldOnBindLimit() .Invoke(& sm)
333+ let __stack_go2 =
334+ noBounce
335+ || yieldOnBindLimit() .Invoke(& sm)
306336
307337 if __ stack_ go2 then
308338 MethodBuilder.SetResult(
309339 & sm.Data.MethodBuilder,
310340 sm.Data.Result
311341 )
312342 with exn ->
313- error <-
314- ValueSome
315- <| ExceptionCache.CaptureOrRetrieve exn
343+ error <- ValueSome( ExceptionCache.CaptureOrRetrieve exn)
316344
317345 if error.IsSome then
318- let __stack_go2 = yieldOnBindLimit() .Invoke(& sm)
346+ let __stack_go2 =
347+ noBounce
348+ || yieldOnBindLimit() .Invoke(& sm)
319349
320350 if __ stack_ go2 then
321351 MethodBuilder.SetException(
0 commit comments