// MIT License - Copyright (c) 2025 wallstop
// Full license text: https://github.com/wallstop/unity-helpers/blob/main/LICENSE
namespace WallstopStudios.UnityHelpers.Core.Extension
{
using System;
using System.Collections;
using System.Collections.Concurrent;
using System.Collections.Generic;
using System.Runtime.CompilerServices;
using System.Threading.Tasks;
using UnityEngine;
///
/// Extension methods for Unity AsyncOperation and Task/ValueTask/IEnumerator interoperability.
///
public static class AsyncOperationExtensions
{
private static readonly ConcurrentDictionary<
AsyncOperation,
Action
> Handlers = new();
private static readonly ConcurrentDictionary Continuations = new();
///
/// Provides an awaiter for Unity AsyncOperation objects, enabling async/await syntax.
///
///
/// This struct is used internally to enable async/await on AsyncOperation.
/// Thread safety: Thread-safe using concurrent dictionaries for handler storage. Must complete on Unity main thread.
/// Performance: O(1) for completion checks. Allocations occur for continuation storage in dictionaries.
/// Allocations: Allocates dictionary entries for tracking completions. Cleaned up on completion.
///
public readonly struct AsyncOperationAwaiter : INotifyCompletion
{
private readonly AsyncOperation _operation;
///
/// Initializes a new instance of the AsyncOperationAwaiter struct.
///
/// The AsyncOperation to await.
/// Thrown when operation is null.
public AsyncOperationAwaiter(AsyncOperation operation)
{
_operation = operation ?? throw new ArgumentNullException(nameof(operation));
}
///
/// Gets a value indicating whether the async operation has completed.
///
public bool IsCompleted => _operation.isDone;
///
/// Schedules the continuation action to be invoked when the operation completes.
///
/// The action to invoke when the operation completes.
public void OnCompleted(Action continuation)
{
Continuations[_operation] = continuation;
Action handler = CachedHandler;
if (!Handlers.TryAdd(_operation, handler))
{
return;
}
Handlers[_operation] = handler;
_operation.completed += handler;
}
///
/// Gets the result of the async operation. Since AsyncOperation has no return value, this is a no-op.
///
public void GetResult() { }
}
private static readonly Action CachedHandler = OnOperationCompleted;
private static void OnOperationCompleted(AsyncOperation operation)
{
Handlers.Remove(operation, out Action _);
if (!Continuations.Remove(operation, out Action completionCondition))
{
return;
}
completionCondition?.Invoke();
}
///
/// Converts a Unity AsyncOperation to a Task.
///
/// The AsyncOperation to convert.
/// A Task that completes when the AsyncOperation completes.
///
/// Null handling: Throws NullReferenceException if asyncOp is null when awaiting.
/// Thread safety: Must complete on Unity main thread. No Unity main thread requirement for initial call.
/// Performance: O(1). Returns immediately if already complete.
/// Allocations: Allocates Task state machine if operation not complete.
/// Edge cases: Returns immediately if operation is already done.
///
public static async Task AsTask(this AsyncOperation asyncOp)
{
if (asyncOp.isDone)
{
return;
}
await asyncOp;
}
///
/// Converts a Unity AsyncOperation to a ValueTask.
///
/// The AsyncOperation to convert.
/// A ValueTask that completes when the AsyncOperation completes.
///
/// Null handling: Throws NullReferenceException if asyncOp is null when awaiting.
/// Thread safety: Must complete on Unity main thread. No Unity main thread requirement for initial call.
/// Performance: O(1). Returns immediately if already complete.
/// Allocations: No allocations if operation is already done, otherwise allocates ValueTask state machine.
/// Edge cases: Returns immediately if operation is already done. Prefer over AsTask for completed operations.
///
public static async ValueTask AsValueTask(this AsyncOperation asyncOp)
{
if (asyncOp.isDone)
{
return;
}
await asyncOp;
}
#if !UNITY_2023_1_OR_NEWER
///
/// Gets an awaiter for the AsyncOperation, enabling async/await syntax.
/// Only available in Unity versions before 2023.1 (Unity 2023.1+ provides this natively).
///
/// The AsyncOperation to get an awaiter for.
/// An AsyncOperationAwaiter for the operation.
///
/// Null handling: Throws ArgumentNullException if op is null (thrown by AsyncOperationAwaiter constructor).
/// Thread safety: Thread-safe. Must complete on Unity main thread.
/// Performance: O(1).
/// Allocations: Allocates AsyncOperationAwaiter struct (stack allocation).
/// Edge cases: Not available in Unity 2023.1+, where AsyncOperation implements INotifyCompletion natively.
///
/// Thrown when op is null.
public static AsyncOperationAwaiter GetAwaiter(this AsyncOperation op)
{
return new AsyncOperationAwaiter(op);
}
#endif
///
/// Executes a continuation action after a ValueTask completes.
///
/// The task to await.
/// The action to execute after the task completes. Can be null.
/// A ValueTask that completes after the continuation executes.
///
/// Null handling: If continuation is null, no action is taken after the task completes.
/// Thread safety: Continuation executes on the same context as the task completion. No Unity main thread requirement unless task requires it.
/// Performance: O(1) overhead for continuation invocation.
/// Allocations: Allocates async state machine.
/// Edge cases: Null continuation is allowed and does nothing.
///
public static async ValueTask WithContinuation(this ValueTask task, Action continuation)
{
await task;
continuation?.Invoke();
}
///
/// Executes a continuation function that transforms the result after a ValueTask completes.
///
/// The type of the task result.
/// The task to await.
/// The function to execute on the result. Can be null, in which case the original result is returned.
/// A ValueTask containing the transformed result, or the original result if continuation is null.
///
/// Null handling: If continuation is null, returns the original task result unchanged.
/// Thread safety: Continuation executes on the same context as the task completion. No Unity main thread requirement unless task requires it.
/// Performance: O(1) overhead for continuation invocation.
/// Allocations: Allocates async state machine.
/// Edge cases: Null continuation is allowed and returns original result.
///
public static async ValueTask WithContinuation(
this ValueTask task,
Func continuation
)
{
TResult result = await task;
return continuation != null ? continuation(result) : result;
}
///
/// Executes a continuation action with the result after a ValueTask completes.
///
/// The type of the task result.
/// The task to await.
/// The action to execute with the result. Can be null.
/// A ValueTask that completes after the continuation executes.
///
/// Null handling: If continuation is null, no action is taken after the task completes.
/// Thread safety: Continuation executes on the same context as the task completion. No Unity main thread requirement unless task requires it.
/// Performance: O(1) overhead for continuation invocation.
/// Allocations: Allocates async state machine.
/// Edge cases: Null continuation is allowed and does nothing.
///
public static async ValueTask WithContinuation(
this ValueTask task,
Action continuation
)
{
TResult result = await task;
continuation?.Invoke(result);
}
// Task/ValueTask to IEnumerator conversions
///
/// Converts a Task to a Unity coroutine (IEnumerator).
///
/// The task to convert.
/// An IEnumerator that can be used with StartCoroutine.
///
/// Null handling: Throws NullReferenceException if task is null when checking IsCompleted.
/// Thread safety: Must be iterated on Unity main thread. No Unity main thread requirement for task execution.
/// Performance: Yields every frame until task completes. O(1) per iteration.
/// Allocations: Allocates iterator state machine.
/// Edge cases: Throws task.Exception if task is faulted. Blocks coroutine execution until task completes.
///
/// Throws the task's exception if the task is faulted.
public static IEnumerator AsCoroutine(this Task task)
{
while (!task.IsCompleted)
{
yield return null;
}
if (task.IsFaulted)
{
throw task.Exception;
}
}
///
/// Converts a Task with a result to a Unity coroutine (IEnumerator), optionally invoking a callback with the result.
///
/// The type of the task result.
/// The task to convert.
/// Optional callback to receive the task result. Can be null.
/// An IEnumerator that can be used with StartCoroutine.
///
/// Null handling: Throws NullReferenceException if task is null. onResult can be null.
/// Thread safety: Must be iterated on Unity main thread. No Unity main thread requirement for task execution.
/// Performance: Yields every frame until task completes. O(1) per iteration.
/// Allocations: Allocates iterator state machine.
/// Edge cases: Throws task.Exception if task is faulted. onResult is invoked with result after successful completion.
///
/// Throws the task's exception if the task is faulted.
public static IEnumerator AsCoroutine(this Task task, Action onResult = null)
{
while (!task.IsCompleted)
{
yield return null;
}
if (task.IsFaulted)
{
throw task.Exception;
}
onResult?.Invoke(task.Result);
}
///
/// Converts a Task returning a tuple with two elements to a Unity coroutine (IEnumerator), optionally invoking a callback with the tuple elements.
///
/// The type of the first tuple element.
/// The type of the second tuple element.
/// The task to convert.
/// Optional callback invoked with the tuple elements. Can be null.
/// An IEnumerator that can be used with StartCoroutine.
///
/// Null handling: Task cannot be null. onResult can be null.
/// Thread safety: Must be iterated on Unity main thread. No Unity main thread requirement for task execution.
/// Performance: Delegates work to the generic Task overload. O(1) per iteration.
/// Allocations: Iterator allocation only.
/// Edge cases: Returns immediately if task is complete. Propagates task exceptions.
///
/// Throws the task's exception if the task is faulted.
public static IEnumerator AsCoroutine(
this Task<(T1 First, T2 Second)> task,
Action onResult = null
)
{
return task.AsCoroutine(tuple =>
{
if (onResult != null)
{
onResult(tuple.First, tuple.Second);
}
});
}
///
/// Converts a Task returning a tuple with three elements to a Unity coroutine (IEnumerator), optionally invoking a callback with the tuple elements.
///
/// The type of the first tuple element.
/// The type of the second tuple element.
/// The type of the third tuple element.
/// The task to convert.
/// Optional callback invoked with the tuple elements. Can be null.
/// An IEnumerator that can be used with StartCoroutine.
///
/// Null handling: Task cannot be null. onResult can be null.
/// Thread safety: Must be iterated on Unity main thread. No Unity main thread requirement for task execution.
/// Performance: Delegates work to the generic Task overload. O(1) per iteration.
/// Allocations: Iterator allocation only.
/// Edge cases: Returns immediately if task is complete. Propagates task exceptions.
///
/// Throws the task's exception if the task is faulted.
public static IEnumerator AsCoroutine(
this Task<(T1 First, T2 Second, T3 Third)> task,
Action onResult = null
)
{
return task.AsCoroutine(tuple =>
{
if (onResult != null)
{
onResult(tuple.First, tuple.Second, tuple.Third);
}
});
}
///
/// Converts a ValueTask to a Unity coroutine (IEnumerator).
///
/// The ValueTask to convert.
/// An IEnumerator that can be used with StartCoroutine.
///
/// Null handling: ValueTask is a value type and cannot be null.
/// Thread safety: Must be iterated on Unity main thread. No Unity main thread requirement for task execution.
/// Performance: No yielding if already complete. Otherwise yields every frame. O(1) per iteration.
/// Allocations: No allocations if already complete. Otherwise allocates iterator and converts to Task internally.
/// Edge cases: Returns immediately if task is already completed. Throws task exception if faulted.
///
/// Throws the task's exception if the task is faulted.
public static IEnumerator AsCoroutine(this ValueTask task)
{
if (task.IsCompleted)
{
if (task.IsFaulted)
{
throw task.AsTask().Exception;
}
yield break;
}
Task innerTask = task.AsTask();
while (!innerTask.IsCompleted)
{
yield return null;
}
if (innerTask.IsFaulted)
{
throw innerTask.Exception;
}
}
///
/// Converts a ValueTask with a result to a Unity coroutine (IEnumerator), optionally invoking a callback with the result.
///
/// The type of the task result.
/// The ValueTask to convert.
/// Optional callback to receive the task result. Can be null.
/// An IEnumerator that can be used with StartCoroutine.
///
/// Null handling: ValueTask is a value type and cannot be null. onResult can be null.
/// Thread safety: Must be iterated on Unity main thread. No Unity main thread requirement for task execution.
/// Performance: No yielding if already complete. Otherwise yields every frame. O(1) per iteration.
/// Allocations: No allocations if already complete. Otherwise allocates iterator and converts to Task internally.
/// Edge cases: Returns immediately if task is already completed. onResult invoked with result after successful completion.
///
/// Throws the task's exception if the task is faulted.
public static IEnumerator AsCoroutine(this ValueTask task, Action onResult = null)
{
if (task.IsCompleted)
{
if (task.IsFaulted)
{
throw task.AsTask().Exception;
}
onResult?.Invoke(task.Result);
yield break;
}
Task innerTask = task.AsTask();
while (!innerTask.IsCompleted)
{
yield return null;
}
if (innerTask.IsFaulted)
{
throw innerTask.Exception;
}
onResult?.Invoke(innerTask.Result);
}
///
/// Converts a ValueTask returning a tuple with two elements to a Unity coroutine (IEnumerator), optionally invoking a callback with the tuple elements.
///
/// The type of the first tuple element.
/// The type of the second tuple element.
/// The ValueTask to convert.
/// Optional callback invoked with the tuple elements. Can be null.
/// An IEnumerator that can be used with StartCoroutine.
///
/// Null handling: ValueTask is a value type and cannot be null. onResult can be null.
/// Thread safety: Must be iterated on Unity main thread. No Unity main thread requirement for task execution.
/// Performance: Delegates work to the generic ValueTask overload. O(1) per iteration.
/// Allocations: Iterator allocation only.
/// Edge cases: Returns immediately if task is complete. Propagates task exceptions.
///
/// Throws the task's exception if the task is faulted.
public static IEnumerator AsCoroutine(
this ValueTask<(T1 First, T2 Second)> task,
Action onResult = null
)
{
return task.AsCoroutine(tuple =>
{
if (onResult != null)
{
onResult(tuple.First, tuple.Second);
}
});
}
///
/// Converts a ValueTask returning a tuple with three elements to a Unity coroutine (IEnumerator), optionally invoking a callback with the tuple elements.
///
/// The type of the first tuple element.
/// The type of the second tuple element.
/// The type of the third tuple element.
/// The ValueTask to convert.
/// Optional callback invoked with the tuple elements. Can be null.
/// An IEnumerator that can be used with StartCoroutine.
///
/// Null handling: ValueTask is a value type and cannot be null. onResult can be null.
/// Thread safety: Must be iterated on Unity main thread. No Unity main thread requirement for task execution.
/// Performance: Delegates work to the generic ValueTask overload. O(1) per iteration.
/// Allocations: Iterator allocation only.
/// Edge cases: Returns immediately if task is complete. Propagates task exceptions.
///
/// Throws the task's exception if the task is faulted.
public static IEnumerator AsCoroutine(
this ValueTask<(T1 First, T2 Second, T3 Third)> task,
Action onResult = null
)
{
return task.AsCoroutine(tuple =>
{
if (onResult != null)
{
onResult(tuple.First, tuple.Second, tuple.Third);
}
});
}
// IEnumerator to Task/ValueTask conversions
///
/// Converts a Unity coroutine (IEnumerator) to a Task.
///
/// The coroutine to convert.
/// A Task that completes when the coroutine finishes.
///
/// Null handling: Throws ArgumentNullException if coroutine is null.
/// Thread safety: Coroutine must be iterated on the thread where it's executed. Typically requires Unity main thread.
/// Performance: O(n) where n is the number of iterations. Yields control between iterations.
/// Allocations: Allocates async state machine.
/// Edge cases: Does not use Unity's StartCoroutine - manually iterates the enumerator. Task.Yield() returns control to caller between iterations.
///
/// Thrown when coroutine is null.
public static async Task AsTask(this IEnumerator coroutine)
{
if (coroutine == null)
{
throw new ArgumentNullException(nameof(coroutine));
}
while (coroutine.MoveNext())
{
await Task.Yield();
}
}
///
/// Converts a Unity coroutine (IEnumerator) to a ValueTask.
///
/// The coroutine to convert.
/// A ValueTask that completes when the coroutine finishes.
///
/// Null handling: Throws ArgumentNullException if coroutine is null.
/// Thread safety: Coroutine must be iterated on the thread where it's executed. Typically requires Unity main thread.
/// Performance: O(n) where n is the number of iterations. Yields control between iterations.
/// Allocations: Allocates async state machine.
/// Edge cases: Does not use Unity's StartCoroutine - manually iterates the enumerator. Task.Yield() returns control to caller between iterations.
///
/// Thrown when coroutine is null.
public static async ValueTask AsValueTask(this IEnumerator coroutine)
{
if (coroutine == null)
{
throw new ArgumentNullException(nameof(coroutine));
}
while (coroutine.MoveNext())
{
await Task.Yield();
}
}
}
}