Skip to content
Draft
126 changes: 82 additions & 44 deletions src/IcedTasks/CancellablePoolingValueTask.fs
Original file line number Diff line number Diff line change
Expand Up @@ -54,35 +54,59 @@ module CancellablePoolingValueTasks =
let initialResumptionFunc =
CancellableTaskBaseResumptionFunc<'T, _>(fun sm -> code.Invoke(&sm))

let resumptionInfo =
let resumptionInfo () =
let mutable state = InitialYield

{ new CancellableTaskBaseResumptionDynamicInfo<'T, _>(initialResumptionFunc) with
member info.MoveNext(sm) =
let mutable savedExn = null

try
sm.ResumptionDynamicInfo.ResumptionData <- null
let step = info.ResumptionFunc.Invoke(&sm)

if step then
MethodBuilder.SetResult(&sm.Data.MethodBuilder, sm.Data.Result)
else
match sm.ResumptionDynamicInfo.ResumptionData with
| :? ICriticalNotifyCompletion as awaiter ->
let mutable awaiter = awaiter
// assert not (isNull awaiter)
MethodBuilder.AwaitOnCompleted(
&sm.Data.MethodBuilder,
&awaiter,
&sm
)
| awaiter -> assert not (isNull awaiter)

with exn ->
savedExn <- exn
// Run SetException outside the stack unwind, see https://github.com/dotnet/roslyn/issues/26567
match savedExn with
| null -> ()
| exn -> MethodBuilder.SetException(&sm.Data.MethodBuilder, exn)
let current = state
let mutable continuation = Stop

match current with
| InitialYield ->
state <- Running

continuation <-
if BindContext.CheckWhenIsBind() then Bounce else Immediate
| Running ->
try
let step = info.ResumptionFunc.Invoke(&sm)

if step then
state <- SetResult

continuation <-
if BindContext.Check() then Bounce else Immediate
else
continuation <-
Await(downcast sm.ResumptionDynamicInfo.ResumptionData)
with exn ->
state <- SetException(ExceptionCache.CaptureOrRetrieve exn)
continuation <- if BindContext.Check() then Bounce else Immediate
| SetResult ->
MethodBuilder.SetResult(&sm.Data.MethodBuilder, sm.Data.Result)
| SetException edi ->
MethodBuilder.SetException(&sm.Data.MethodBuilder, edi.SourceException)

let continuation = continuation

match continuation with
| Stop -> ()
| Immediate -> info.MoveNext(&sm)
| Bounce ->
MethodBuilder.AwaitOnCompleted(
&sm.Data.MethodBuilder,
Trampoline.AwaiterRef,
&sm
)
| Await awaiter ->
let mutable awaiter = awaiter

MethodBuilder.AwaitUnsafeOnCompleted(
&sm.Data.MethodBuilder,
&awaiter,
&sm
)

member _.SetStateMachine(sm, state) =
MethodBuilder.SetStateMachine(&sm.Data.MethodBuilder, state)
Expand All @@ -93,7 +117,7 @@ module CancellablePoolingValueTasks =
ValueTask.FromCanceled<_>(ct)
else
sm.Data.CancellationToken <- ct
sm.ResumptionDynamicInfo <- resumptionInfo
sm.ResumptionDynamicInfo <- resumptionInfo ()
sm.Data.MethodBuilder <- PoolingAsyncValueTaskMethodBuilder<'T>.Create()
sm.Data.MethodBuilder.Start(&sm)
sm.Data.MethodBuilder.Task
Expand All @@ -103,22 +127,36 @@ module CancellablePoolingValueTasks =
if __useResumableCode then
__stateMachine<CancellableTaskBaseStateMachineData<'T, _>, CancellableValueTask<'T>>
(MoveNextMethodImpl<_>(fun sm ->
//-- RESUMABLE CODE START
__resumeAt sm.ResumptionPoint
let mutable __stack_exn = null

try
let __stack_code_fin = code.Invoke(&sm)

if __stack_code_fin then
MethodBuilder.SetResult(&sm.Data.MethodBuilder, sm.Data.Result)
with exn ->
__stack_exn <- exn
// Run SetException outside the stack unwind, see https://github.com/dotnet/roslyn/issues/26567
match __stack_exn with
| null -> ()
| exn -> MethodBuilder.SetException(&sm.Data.MethodBuilder, exn)
//-- RESUMABLE CODE END
let mutable error = ValueNone

let __stack_go1 = yieldOnBindLimitWhenIsBind().Invoke(&sm)

if __stack_go1 then
try
let __stack_code_fin = code.Invoke(&sm)

if __stack_code_fin then
let __stack_go2 = yieldOnBindLimit().Invoke(&sm)

if __stack_go2 then
MethodBuilder.SetResult(
&sm.Data.MethodBuilder,
sm.Data.Result
)
with exn ->
error <-
ValueSome
<| ExceptionCache.CaptureOrRetrieve exn

if error.IsSome then
let __stack_go2 = yieldOnBindLimit().Invoke(&sm)

if __stack_go2 then
MethodBuilder.SetException(
&sm.Data.MethodBuilder,
error.Value.SourceException
)
))
(SetStateMachineMethodImpl<_>(fun sm state ->
MethodBuilder.SetStateMachine(&sm.Data.MethodBuilder, state)
Expand Down Expand Up @@ -148,7 +186,7 @@ module CancellablePoolingValueTasks =
([<InlineIfLambda>] x: CancellationToken -> ValueTask<_>)
: CancellationToken -> Awaiter<ValueTaskAwaiter<_>, _> =
fun ct ->
(x ct)
BindContext.SetIsBind x ct
|> Awaitable.GetAwaiter

[<NoEagerConstraintApplication>]
Expand Down
Loading
Loading