Skip to content

Commit 39cd1ca

Browse files
committed
update
1 parent e3d2023 commit 39cd1ca

File tree

6 files changed

+403
-351
lines changed

6 files changed

+403
-351
lines changed

src/IcedTasks/CancellablePoolingValueTask.fs

Lines changed: 68 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -60,59 +60,75 @@ module CancellablePoolingValueTasks =
6060
let initialResumptionFunc =
6161
CancellableTaskBaseResumptionFunc<'T, _>(fun sm -> code.Invoke(&sm))
6262

63-
let resumptionInfo () =
64-
let mutable state = InitialYield
63+
let bounceAllowed = Trampoline.Current.IsAwaited()
64+
65+
let maybeBounce state =
66+
if
67+
bounceAllowed
68+
&& Trampoline.Current.ShouldBounce
69+
then
70+
Bounce state
71+
else
72+
Immediate state
73+
74+
let resumptionInfo =
75+
let initialState = maybeBounce Running
6576

66-
{ new CancellableTaskBaseResumptionDynamicInfo<'T, _>(initialResumptionFunc) with
77+
{ new CancellableTaskBaseResumptionDynamicInfo<'T, _>(initialResumptionFunc,
78+
ResumptionData = initialState) with
6779
member info.MoveNext(sm) =
68-
let current = state
69-
let mutable continuation = Stop
7080

71-
match current with
72-
| InitialYield ->
73-
state <- Running
81+
let getCurrent () =
82+
match info.ResumptionData with
83+
| :? DynamicState as state -> state
84+
| _ -> failwith "Invalid resumption data"
7485

75-
continuation <-
76-
if BindContext.CheckWhenIsBind() then Bounce else Immediate
77-
| Running ->
78-
try
79-
let step = info.ResumptionFunc.Invoke(&sm)
86+
let setState state = info.ResumptionData <- state
8087

81-
if step then
82-
state <- SetResult
88+
match getCurrent () with
89+
| Immediate state ->
90+
setState state
91+
info.MoveNext &sm
92+
| Running ->
93+
let mutable keepGoing = true
8394

84-
continuation <-
85-
if BindContext.Check() then Bounce else Immediate
95+
try
96+
if info.ResumptionFunc.Invoke(&sm) then
97+
setState (maybeBounce SetResult)
8698
else
87-
continuation <-
88-
Await(downcast sm.ResumptionDynamicInfo.ResumptionData)
99+
keepGoing <-
100+
match getCurrent () with
101+
| Awaiting _ -> true
102+
| _ -> false
89103
with exn ->
90-
state <- SetException(ExceptionCache.CaptureOrRetrieve exn)
91-
continuation <- if BindContext.Check() then Bounce else Immediate
92-
| SetResult ->
93-
MethodBuilder.SetResult(&sm.Data.MethodBuilder, sm.Data.Result)
94-
| SetException edi ->
95-
MethodBuilder.SetException(&sm.Data.MethodBuilder, edi.SourceException)
96-
97-
let continuation = continuation
104+
setState (
105+
maybeBounce
106+
<| SetException(ExceptionCache.CaptureOrRetrieve exn)
107+
)
108+
109+
if keepGoing then
110+
info.MoveNext &sm
111+
| Awaiting awaiter ->
112+
setState Running
113+
let mutable awaiter = awaiter
98114

99-
match continuation with
100-
| Stop -> ()
101-
| Immediate -> info.MoveNext(&sm)
102-
| Bounce ->
103-
MethodBuilder.AwaitOnCompleted(
115+
MethodBuilder.AwaitUnsafeOnCompleted(
104116
&sm.Data.MethodBuilder,
105-
Trampoline.Current.Ref,
117+
&awaiter,
106118
&sm
107119
)
108-
| Await awaiter ->
109-
let mutable awaiter = awaiter
120+
| Bounce next ->
121+
setState next
110122

111-
MethodBuilder.AwaitUnsafeOnCompleted(
123+
MethodBuilder.AwaitOnCompleted(
112124
&sm.Data.MethodBuilder,
113-
&awaiter,
125+
Trampoline.Current.Ref,
114126
&sm
115127
)
128+
| SetResult ->
129+
MethodBuilder.SetResult(&sm.Data.MethodBuilder, sm.Data.Result)
130+
| SetException edi ->
131+
MethodBuilder.SetException(&sm.Data.MethodBuilder, edi.SourceException)
116132

117133
member _.SetStateMachine(sm, state) =
118134
MethodBuilder.SetStateMachine(&sm.Data.MethodBuilder, state)
@@ -123,7 +139,7 @@ module CancellablePoolingValueTasks =
123139
ValueTask.FromCanceled<_>(ct)
124140
else
125141
sm.Data.CancellationToken <- ct
126-
sm.ResumptionDynamicInfo <- resumptionInfo ()
142+
sm.ResumptionDynamicInfo <- resumptionInfo
127143
sm.Data.MethodBuilder <- PoolingAsyncValueTaskMethodBuilder<'T>.Create()
128144
sm.Data.MethodBuilder.Start(&sm)
129145
sm.Data.MethodBuilder.Task
@@ -134,29 +150,36 @@ module CancellablePoolingValueTasks =
134150
__stateMachine<CancellableTaskBaseStateMachineData<'T, _>, CancellableValueTask<'T>>
135151
(MoveNextMethodImpl<_>(fun sm ->
136152
__resumeAt sm.ResumptionPoint
153+
137154
let mutable error = ValueNone
138155

139-
let __stack_go1 = yieldOnBindLimitWhenIsBind().Invoke(&sm)
156+
let noBounce = not (Trampoline.Current.IsAwaited())
157+
158+
let __stack_go1 =
159+
noBounce
160+
|| yieldOnBindLimit().Invoke(&sm)
140161

141162
if __stack_go1 then
142163
try
143164
let __stack_code_fin = code.Invoke(&sm)
144165

145166
if __stack_code_fin then
146-
let __stack_go2 = yieldOnBindLimit().Invoke(&sm)
167+
let __stack_go2 =
168+
noBounce
169+
|| yieldOnBindLimit().Invoke(&sm)
147170

148171
if __stack_go2 then
149172
MethodBuilder.SetResult(
150173
&sm.Data.MethodBuilder,
151174
sm.Data.Result
152175
)
153176
with exn ->
154-
error <-
155-
ValueSome
156-
<| ExceptionCache.CaptureOrRetrieve exn
177+
error <- ValueSome(ExceptionCache.CaptureOrRetrieve exn)
157178

158179
if error.IsSome then
159-
let __stack_go2 = yieldOnBindLimit().Invoke(&sm)
180+
let __stack_go2 =
181+
noBounce
182+
|| yieldOnBindLimit().Invoke(&sm)
160183

161184
if __stack_go2 then
162185
MethodBuilder.SetException(
@@ -192,7 +215,7 @@ module CancellablePoolingValueTasks =
192215
([<InlineIfLambda>] x: CancellationToken -> ValueTask<_>)
193216
: CancellationToken -> Awaiter<ValueTaskAwaiter<_>, _> =
194217
fun ct ->
195-
BindContext.SetIsBind x ct
218+
Trampoline.Allow x ct
196219
|> Awaitable.GetAwaiter
197220

198221
[<NoEagerConstraintApplication>]

src/IcedTasks/CancellableTask.fs

Lines changed: 81 additions & 51 deletions
Original file line numberDiff line numberDiff line change
@@ -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

Comments
 (0)