public void AwaitUnsafeOnCompleted <TAwaiter, TStateMachine>(ref TAwaiter awaiter, ref TStateMachine stateMachine)
where TAwaiter : ICriticalNotifyCompletion
where TStateMachine : IAsyncStateMachine
{
if (awaiter is CancellationTokenInjector cti)
{
this.cancellationTokenInjector = cti;
stateMachine.MoveNext();
return;
}
if (this.cancellationTokenInjector.CancellationToken.CanBeCanceled)
{
var moveNext = new Action(stateMachine.MoveNext);
awaiter.UnsafeOnCompleted(() =>
{
var ct = this.cancellationTokenInjector.CancellationToken;
if (ct.IsCancellationRequested)
{
this.completionSource.SetException(new OperationCanceledException(ct));
return;
}
if (this.cancellationTokenInjector.RestoreInjection)
{
_ = YTask.InjectToStatic(ct);
}
moveNext();
});
}
else
{
awaiter.UnsafeOnCompleted(stateMachine.MoveNext);
}
}
public static SyntaxNode Transform(string source)
{
var tree = CSharpSyntaxTree.ParseText(source);
var rewriter = new CancellationTokenInjector(tree.GetRoot());
return(rewriter.Visit(tree.GetRoot()));
}