pengtaihao
/
UniTask
mirror of https://github.com/Cysharp/UniTask
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Add overload in UniTask.WaitUntil, UniTask.WaitWhile and UniTask.Defer to avoid closure allocation

pull/619/head
Kochoyume 2024-09-24 23:37:44 +09:00
parent 83d8a2b424
commit dfe5ee43c2
3 changed files with 398 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

103
src/UniTask/Assets/Plugins/UniTask/Runtime/UniTask.Factory.cs
View File

@ -202,6 +202,22 @@ namespace Cysharp.Threading.Tasks
return new UniTask<T>(new DeferPromise<T>(factory), 0);
}
/// <summary>
/// Defer the task creation just before call await.
/// </summary>
public static UniTask Defer<TState>(TState state, Func<TState, UniTask> factory)
{
return new UniTask(new DeferPromiseWithState<TState>(state, factory), 0);
}
/// <summary>
/// Defer the task creation just before call await.
/// </summary>
public static UniTask<TResult> Defer<TState, TResult>(TState state, Func<TState, UniTask<TResult>> factory)
{
return new UniTask<TResult>(new DeferPromiseWithState<TState, TResult>(state, factory), 0);
}
/// <summary>
/// Never complete.
/// </summary>
@ -465,6 +481,93 @@ namespace Cysharp.Threading.Tasks
}
}
sealed class DeferPromiseWithState<TState> : IUniTaskSource
{
Func<TState, UniTask> factory;
TState argument;
UniTask task;
UniTask.Awaiter awaiter;
public DeferPromiseWithState(TState argument, Func<TState, UniTask> factory)
{
this.argument = argument;
this.factory = factory;
}
public void GetResult(short token)
{
awaiter.GetResult();
}
public UniTaskStatus GetStatus(short token)
{
var f = Interlocked.Exchange(ref factory, null);
if (f != null)
{
task = f(argument);
awaiter = task.GetAwaiter();
}
return task.Status;
}
public void OnCompleted(Action<object> continuation, object state, short token)
{
awaiter.SourceOnCompleted(continuation, state);
}
public UniTaskStatus UnsafeGetStatus()
{
return task.Status;
}
}
sealed class DeferPromiseWithState<TState, TResult> : IUniTaskSource<TResult>
{
Func<TState, UniTask<TResult>> factory;
TState argument;
UniTask<TResult> task;
UniTask<TResult>.Awaiter awaiter;
public DeferPromiseWithState(TState argument, Func<TState, UniTask<TResult>> factory)
{
this.argument = argument;
this.factory = factory;
}
public TResult GetResult(short token)
{
return awaiter.GetResult();
}
void IUniTaskSource.GetResult(short token)
{
awaiter.GetResult();
}
public UniTaskStatus GetStatus(short token)
{
var f = Interlocked.Exchange(ref factory, null);
if (f != null)
{
task = f(argument);
awaiter = task.GetAwaiter();
}
return task.Status;
}
public void OnCompleted(Action<object> continuation, object state, short token)
{
awaiter.SourceOnCompleted(continuation, state);
}
public UniTaskStatus UnsafeGetStatus()
{
return task.Status;
}
}
sealed class NeverPromise<T> : IUniTaskSource<T>
{
static readonly Action<object> cancellationCallback = CancellationCallback;

265
src/UniTask/Assets/Plugins/UniTask/Runtime/UniTask.WaitUntil.cs
View File

@ -15,11 +15,21 @@ namespace Cysharp.Threading.Tasks
return new UniTask(WaitUntilPromise.Create(predicate, timing, cancellationToken, cancelImmediately, out var token), token);
}
public static UniTask WaitUntil<T>(T state, Func<T, bool> predicate, PlayerLoopTiming timing = PlayerLoopTiming.Update, CancellationToken cancellationToken = default(CancellationToken), bool cancelImmediately = false)
{
return new UniTask(WaitUntilPromise<T>.Create(state, predicate, timing, cancellationToken, cancelImmediately, out var token), token);
}
public static UniTask WaitWhile(Func<bool> predicate, PlayerLoopTiming timing = PlayerLoopTiming.Update, CancellationToken cancellationToken = default(CancellationToken), bool cancelImmediately = false)
{
return new UniTask(WaitWhilePromise.Create(predicate, timing, cancellationToken, cancelImmediately, out var token), token);
}
public static UniTask WaitWhile<T>(T state, Func<T, bool> predicate, PlayerLoopTiming timing = PlayerLoopTiming.Update, CancellationToken cancellationToken = default(CancellationToken), bool cancelImmediately = false)
{
return new UniTask(WaitWhilePromise<T>.Create(state, predicate, timing, cancellationToken, cancelImmediately, out var token), token);
}
public static UniTask WaitUntilCanceled(CancellationToken cancellationToken, PlayerLoopTiming timing = PlayerLoopTiming.Update, bool completeImmediately = false)
{
return new UniTask(WaitUntilCanceledPromise.Create(cancellationToken, timing, completeImmediately, out var token), token);
@ -162,6 +172,131 @@ namespace Cysharp.Threading.Tasks
}
}
sealed class WaitUntilPromise<T> : IUniTaskSource, IPlayerLoopItem, ITaskPoolNode<WaitUntilPromise<T>>
{
static TaskPool<WaitUntilPromise<T>> pool;
WaitUntilPromise<T> nextNode;
public ref WaitUntilPromise<T> NextNode => ref nextNode;
static WaitUntilPromise()
{
TaskPool.RegisterSizeGetter(typeof(WaitUntilPromise<T>), () => pool.Size);
}
Func<T, bool> predicate;
T argument;
CancellationToken cancellationToken;
CancellationTokenRegistration cancellationTokenRegistration;
bool cancelImmediately;
UniTaskCompletionSourceCore<object> core;
WaitUntilPromise()
{
}
public static IUniTaskSource Create(T argument, Func<T, bool> predicate, PlayerLoopTiming timing, CancellationToken cancellationToken, bool cancelImmediately, out short token)
{
if (cancellationToken.IsCancellationRequested)
{
return AutoResetUniTaskCompletionSource.CreateFromCanceled(cancellationToken, out token);
}
if (!pool.TryPop(out var result))
{
result = new WaitUntilPromise<T>();
}
result.predicate = predicate;
result.argument = argument;
result.cancellationToken = cancellationToken;
result.cancelImmediately = cancelImmediately;
if (cancelImmediately && cancellationToken.CanBeCanceled)
{
result.cancellationTokenRegistration = cancellationToken.RegisterWithoutCaptureExecutionContext(state =>
{
var promise = (WaitUntilPromise<T>)state;
promise.core.TrySetCanceled(promise.cancellationToken);
}, result);
}
TaskTracker.TrackActiveTask(result, 3);
PlayerLoopHelper.AddAction(timing, result);
token = result.core.Version;
return result;
}
public void GetResult(short token)
{
try
{
core.GetResult(token);
}
finally
{
if (!(cancelImmediately && cancellationToken.IsCancellationRequested))
{
TryReturn();
}
}
}
public UniTaskStatus GetStatus(short token)
{
return core.GetStatus(token);
}
public UniTaskStatus UnsafeGetStatus()
{
return core.UnsafeGetStatus();
}
public void OnCompleted(Action<object> continuation, object state, short token)
{
core.OnCompleted(continuation, state, token);
}
public bool MoveNext()
{
if (cancellationToken.IsCancellationRequested)
{
core.TrySetCanceled(cancellationToken);
return false;
}
try
{
if (!predicate(argument))
{
return true;
}
}
catch (Exception ex)
{
core.TrySetException(ex);
return false;
}
core.TrySetResult(null);
return false;
}
bool TryReturn()
{
TaskTracker.RemoveTracking(this);
core.Reset();
predicate = default;
argument = default;
cancellationToken = default;
cancellationTokenRegistration.Dispose();
cancelImmediately = default;
return pool.TryPush(this);
}
}
sealed class WaitWhilePromise : IUniTaskSource, IPlayerLoopItem, ITaskPoolNode<WaitWhilePromise>
{
static TaskPool<WaitWhilePromise> pool;
@ -199,7 +334,7 @@ namespace Cysharp.Threading.Tasks
result.predicate = predicate;
result.cancellationToken = cancellationToken;
result.cancelImmediately = cancelImmediately;
if (cancelImmediately && cancellationToken.CanBeCanceled)
{
result.cancellationTokenRegistration = cancellationToken.RegisterWithoutCaptureExecutionContext(state =>
@ -288,6 +423,130 @@ namespace Cysharp.Threading.Tasks
}
}
sealed class WaitWhilePromise<T> : IUniTaskSource, IPlayerLoopItem, ITaskPoolNode<WaitWhilePromise<T>>
{
static TaskPool<WaitWhilePromise<T>> pool;
WaitWhilePromise<T> nextNode;
public ref WaitWhilePromise<T> NextNode => ref nextNode;
static WaitWhilePromise()
{
TaskPool.RegisterSizeGetter(typeof(WaitWhilePromise<T>), () => pool.Size);
}
Func<T, bool> predicate;
T argument;
CancellationToken cancellationToken;
CancellationTokenRegistration cancellationTokenRegistration;
bool cancelImmediately;
UniTaskCompletionSourceCore<object> core;
WaitWhilePromise()
{
}
public static IUniTaskSource Create(T argument, Func<T, bool> predicate, PlayerLoopTiming timing, CancellationToken cancellationToken, bool cancelImmediately, out short token)
{
if (cancellationToken.IsCancellationRequested)
{
return AutoResetUniTaskCompletionSource.CreateFromCanceled(cancellationToken, out token);
}
if (!pool.TryPop(out var result))
{
result = new WaitWhilePromise<T>();
}
result.predicate = predicate;
result.argument = argument;
result.cancellationToken = cancellationToken;
if (cancelImmediately && cancellationToken.CanBeCanceled)
{
result.cancellationTokenRegistration = cancellationToken.RegisterWithoutCaptureExecutionContext(state =>
{
var promise = (WaitWhilePromise<T>)state;
promise.core.TrySetCanceled(promise.cancellationToken);
}, result);
}
TaskTracker.TrackActiveTask(result, 3);
PlayerLoopHelper.AddAction(timing, result);
token = result.core.Version;
return result;
}
public void GetResult(short token)
{
try
{
core.GetResult(token);
}
finally
{
if (!(cancelImmediately && cancellationToken.IsCancellationRequested))
{
TryReturn();
}
}
}
public UniTaskStatus GetStatus(short token)
{
return core.GetStatus(token);
}
public UniTaskStatus UnsafeGetStatus()
{
return core.UnsafeGetStatus();
}
public void OnCompleted(Action<object> continuation, object state, short token)
{
core.OnCompleted(continuation, state, token);
}
public bool MoveNext()
{
if (cancellationToken.IsCancellationRequested)
{
core.TrySetCanceled(cancellationToken);
return false;
}
try
{
if (predicate(argument))
{
return true;
}
}
catch (Exception ex)
{
core.TrySetException(ex);
return false;
}
core.TrySetResult(null);
return false;
}
bool TryReturn()
{
TaskTracker.RemoveTracking(this);
core.Reset();
predicate = default;
argument = default;
cancellationToken = default;
cancellationTokenRegistration.Dispose();
cancelImmediately = default;
return pool.TryPush(this);
}
}
sealed class WaitUntilCanceledPromise : IUniTaskSource, IPlayerLoopItem, ITaskPoolNode<WaitUntilCanceledPromise>
{
static TaskPool<WaitUntilCanceledPromise> pool;
@ -443,7 +702,7 @@ namespace Cysharp.Threading.Tasks
result.equalityComparer = equalityComparer ?? UnityEqualityComparer.GetDefault<U>();
result.cancellationToken = cancellationToken;
result.cancelImmediately = cancelImmediately;
if (cancelImmediately && cancellationToken.CanBeCanceled)
{
result.cancellationTokenRegistration = cancellationToken.RegisterWithoutCaptureExecutionContext(state =>
@ -586,7 +845,7 @@ namespace Cysharp.Threading.Tasks
result.equalityComparer = equalityComparer ?? UnityEqualityComparer.GetDefault<U>();
result.cancellationToken = cancellationToken;
result.cancelImmediately = cancelImmediately;
if (cancelImmediately && cancellationToken.CanBeCanceled)
{
result.cancellationTokenRegistration = cancellationToken.RegisterWithoutCaptureExecutionContext(state =>

33
src/UniTask/Assets/Tests/AsyncTest.cs
View File

@ -145,6 +145,11 @@ namespace Cysharp.Threading.TasksTests
public int MyProperty { get; set; }
}
class MyBoolenClass
{
public bool MyProperty { get; set; }
}
[UnityTest]
public IEnumerator WaitUntil() => UniTask.ToCoroutine(async () =>
{
@ -159,6 +164,20 @@ namespace Cysharp.Threading.TasksTests
diff.Should().Be(11);
});
[UnityTest]
public IEnumerator WaitUntilWithState() => UniTask.ToCoroutine(async () =>
{
var v = new MyBoolenClass { MyProperty = false };
UniTask.DelayFrame(10, PlayerLoopTiming.PostLateUpdate).ContinueWith(() => v.MyProperty = true).Forget();
var startFrame = Time.frameCount;
await UniTask.WaitUntil(v, static v => v.MyProperty, PlayerLoopTiming.EarlyUpdate);
var diff = Time.frameCount - startFrame;
diff.Should().Be(11);
});
[UnityTest]
public IEnumerator WaitWhile() => UniTask.ToCoroutine(async () =>
{
@ -173,6 +192,20 @@ namespace Cysharp.Threading.TasksTests
diff.Should().Be(11);
});
[UnityTest]
public IEnumerator WaitWhileWithState() => UniTask.ToCoroutine(async () =>
{
var v = new MyBoolenClass { MyProperty = true };
UniTask.DelayFrame(10, PlayerLoopTiming.PostLateUpdate).ContinueWith(() => v.MyProperty = false).Forget();
var startFrame = Time.frameCount;
await UniTask.WaitWhile(v, static v => v.MyProperty, PlayerLoopTiming.EarlyUpdate);
var diff = Time.frameCount - startFrame;
diff.Should().Be(11);
});
[UnityTest]
public IEnumerator WaitUntilValueChanged() => UniTask.ToCoroutine(async () =>
{