Fix IUniTaskAsyncEnumerable.Take
parent
6f4d1183cc
commit
49ba57f20a
|
@ -30,39 +30,92 @@ namespace Cysharp.Threading.Tasks.Linq
|
||||||
return new _Take(source, count, cancellationToken);
|
return new _Take(source, count, cancellationToken);
|
||||||
}
|
}
|
||||||
|
|
||||||
sealed class _Take : AsyncEnumeratorBase<TSource, TSource>
|
sealed class _Take : MoveNextSource, IUniTaskAsyncEnumerator<TSource>
|
||||||
{
|
{
|
||||||
readonly int count;
|
static readonly Action<object> MoveNextCoreDelegate = MoveNextCore;
|
||||||
|
|
||||||
|
readonly IUniTaskAsyncEnumerable<TSource> source;
|
||||||
|
readonly int count;
|
||||||
|
CancellationToken cancellationToken;
|
||||||
|
|
||||||
|
IUniTaskAsyncEnumerator<TSource> enumerator;
|
||||||
|
UniTask<bool>.Awaiter awaiter;
|
||||||
int index;
|
int index;
|
||||||
|
|
||||||
public _Take(IUniTaskAsyncEnumerable<TSource> source, int count, CancellationToken cancellationToken)
|
public _Take(IUniTaskAsyncEnumerable<TSource> source, int count, CancellationToken cancellationToken)
|
||||||
: base(source, cancellationToken)
|
|
||||||
{
|
{
|
||||||
|
this.source = source;
|
||||||
this.count = count;
|
this.count = count;
|
||||||
|
this.cancellationToken = cancellationToken;
|
||||||
}
|
}
|
||||||
|
|
||||||
protected override bool TryMoveNextCore(bool sourceHasCurrent, out bool result)
|
public TSource Current { get; private set; }
|
||||||
|
|
||||||
|
public UniTask<bool> MoveNextAsync()
|
||||||
{
|
{
|
||||||
if (sourceHasCurrent)
|
cancellationToken.ThrowIfCancellationRequested();
|
||||||
|
|
||||||
|
if (enumerator == null)
|
||||||
{
|
{
|
||||||
if (checked(index++) < count)
|
enumerator = source.GetAsyncEnumerator(cancellationToken);
|
||||||
|
}
|
||||||
|
|
||||||
|
if (checked(index) >= count)
|
||||||
{
|
{
|
||||||
Current = SourceCurrent;
|
return CompletedTasks.False;
|
||||||
result = true;
|
}
|
||||||
return true;
|
|
||||||
|
completionSource.Reset();
|
||||||
|
SourceMoveNext();
|
||||||
|
return new UniTask<bool>(this, completionSource.Version);
|
||||||
|
}
|
||||||
|
|
||||||
|
void SourceMoveNext()
|
||||||
|
{
|
||||||
|
try
|
||||||
|
{
|
||||||
|
awaiter = enumerator.MoveNextAsync().GetAwaiter();
|
||||||
|
if (awaiter.IsCompleted)
|
||||||
|
{
|
||||||
|
MoveNextCore(this);
|
||||||
}
|
}
|
||||||
else
|
else
|
||||||
{
|
{
|
||||||
result = false;
|
awaiter.SourceOnCompleted(MoveNextCoreDelegate, this);
|
||||||
return true;
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
catch (Exception ex)
|
||||||
|
{
|
||||||
|
completionSource.TrySetException(ex);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
static void MoveNextCore(object state)
|
||||||
|
{
|
||||||
|
var self = (_Take)state;
|
||||||
|
|
||||||
|
if (self.TryGetResult(self.awaiter, out var result))
|
||||||
|
{
|
||||||
|
if (result)
|
||||||
|
{
|
||||||
|
self.index++;
|
||||||
|
self.Current = self.enumerator.Current;
|
||||||
|
self.completionSource.TrySetResult(true);
|
||||||
|
}
|
||||||
else
|
else
|
||||||
{
|
{
|
||||||
result = false;
|
self.completionSource.TrySetResult(false);
|
||||||
return true;
|
}
|
||||||
}
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
public UniTask DisposeAsync()
|
||||||
|
{
|
||||||
|
if (enumerator != null)
|
||||||
|
{
|
||||||
|
return enumerator.DisposeAsync();
|
||||||
|
}
|
||||||
|
return default;
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
Loading…
Reference in New Issue