Buffer
parent
72efadd0a2
commit
61a3744fdd
|
@ -0,0 +1,340 @@
|
|||
using Cysharp.Threading.Tasks.Internal;
|
||||
using System;
|
||||
using System.Collections.Generic;
|
||||
using System.Threading;
|
||||
|
||||
namespace Cysharp.Threading.Tasks.Linq
|
||||
{
|
||||
public static partial class UniTaskAsyncEnumerable
|
||||
{
|
||||
public static IUniTaskAsyncEnumerable<IList<TSource>> Buffer<TSource>(this IUniTaskAsyncEnumerable<TSource> source, Int32 count)
|
||||
{
|
||||
Error.ThrowArgumentNullException(source, nameof(source));
|
||||
if (count <= 0) throw Error.ArgumentOutOfRange(nameof(count));
|
||||
|
||||
return new Buffer<TSource>(source, count);
|
||||
}
|
||||
|
||||
public static IUniTaskAsyncEnumerable<IList<TSource>> Buffer<TSource>(this IUniTaskAsyncEnumerable<TSource> source, Int32 count, Int32 skip)
|
||||
{
|
||||
Error.ThrowArgumentNullException(source, nameof(source));
|
||||
if (count <= 0) throw Error.ArgumentOutOfRange(nameof(count));
|
||||
if (skip <= 0) throw Error.ArgumentOutOfRange(nameof(skip));
|
||||
|
||||
return new BufferSkip<TSource>(source, count, skip);
|
||||
}
|
||||
}
|
||||
|
||||
internal sealed class Buffer<TSource> : IUniTaskAsyncEnumerable<IList<TSource>>
|
||||
{
|
||||
readonly IUniTaskAsyncEnumerable<TSource> source;
|
||||
readonly int count;
|
||||
|
||||
public Buffer(IUniTaskAsyncEnumerable<TSource> source, int count)
|
||||
{
|
||||
this.source = source;
|
||||
this.count = count;
|
||||
}
|
||||
|
||||
public IUniTaskAsyncEnumerator<IList<TSource>> GetAsyncEnumerator(CancellationToken cancellationToken = default)
|
||||
{
|
||||
return new Enumerator(source, count, cancellationToken);
|
||||
}
|
||||
|
||||
sealed class Enumerator : MoveNextSource, IUniTaskAsyncEnumerator<IList<TSource>>
|
||||
{
|
||||
static readonly Action<object> MoveNextCoreDelegate = MoveNextCore;
|
||||
|
||||
readonly IUniTaskAsyncEnumerable<TSource> source;
|
||||
readonly int count;
|
||||
CancellationToken cancellationToken;
|
||||
|
||||
IUniTaskAsyncEnumerator<TSource> enumerator;
|
||||
UniTask<bool>.Awaiter awaiter;
|
||||
bool continueNext;
|
||||
|
||||
bool completed;
|
||||
List<TSource> buffer;
|
||||
|
||||
public Enumerator(IUniTaskAsyncEnumerable<TSource> source, int count, CancellationToken cancellationToken)
|
||||
{
|
||||
this.source = source;
|
||||
this.count = count;
|
||||
this.cancellationToken = cancellationToken;
|
||||
}
|
||||
|
||||
public IList<TSource> Current { get; private set; }
|
||||
|
||||
public UniTask<bool> MoveNextAsync()
|
||||
{
|
||||
cancellationToken.ThrowIfCancellationRequested();
|
||||
|
||||
if (enumerator == null)
|
||||
{
|
||||
enumerator = source.GetAsyncEnumerator(cancellationToken);
|
||||
buffer = new List<TSource>(count);
|
||||
}
|
||||
|
||||
completionSource.Reset();
|
||||
SourceMoveNext();
|
||||
return new UniTask<bool>(this, completionSource.Version);
|
||||
}
|
||||
|
||||
void SourceMoveNext()
|
||||
{
|
||||
if (completed)
|
||||
{
|
||||
if (buffer != null && buffer.Count > 0)
|
||||
{
|
||||
var ret = buffer;
|
||||
buffer = null;
|
||||
Current = ret;
|
||||
completionSource.TrySetResult(true);
|
||||
return;
|
||||
}
|
||||
else
|
||||
{
|
||||
completionSource.TrySetResult(false);
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
try
|
||||
{
|
||||
|
||||
LOOP:
|
||||
awaiter = enumerator.MoveNextAsync().GetAwaiter();
|
||||
if (awaiter.IsCompleted)
|
||||
{
|
||||
continueNext = true;
|
||||
MoveNextCore(this);
|
||||
if (continueNext)
|
||||
{
|
||||
continueNext = false;
|
||||
goto LOOP; // avoid recursive
|
||||
}
|
||||
}
|
||||
else
|
||||
{
|
||||
awaiter.SourceOnCompleted(MoveNextCoreDelegate, this);
|
||||
}
|
||||
}
|
||||
catch (Exception ex)
|
||||
{
|
||||
completionSource.TrySetException(ex);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
static void MoveNextCore(object state)
|
||||
{
|
||||
var self = (Enumerator)state;
|
||||
|
||||
if (self.TryGetResult(self.awaiter, out var result))
|
||||
{
|
||||
if (result)
|
||||
{
|
||||
self.buffer.Add(self.enumerator.Current);
|
||||
|
||||
if (self.buffer.Count == self.count)
|
||||
{
|
||||
self.Current = self.buffer;
|
||||
self.buffer = new List<TSource>(self.count);
|
||||
self.continueNext = false;
|
||||
self.completionSource.TrySetResult(true);
|
||||
return;
|
||||
}
|
||||
else
|
||||
{
|
||||
if (!self.continueNext)
|
||||
{
|
||||
self.SourceMoveNext();
|
||||
}
|
||||
}
|
||||
}
|
||||
else
|
||||
{
|
||||
self.continueNext = false;
|
||||
self.completed = true;
|
||||
self.SourceMoveNext();
|
||||
}
|
||||
}
|
||||
else
|
||||
{
|
||||
self.continueNext = false;
|
||||
}
|
||||
}
|
||||
|
||||
public UniTask DisposeAsync()
|
||||
{
|
||||
if (enumerator != null)
|
||||
{
|
||||
return enumerator.DisposeAsync();
|
||||
}
|
||||
return default;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
internal sealed class BufferSkip<TSource> : IUniTaskAsyncEnumerable<IList<TSource>>
|
||||
{
|
||||
readonly IUniTaskAsyncEnumerable<TSource> source;
|
||||
readonly int count;
|
||||
readonly int skip;
|
||||
|
||||
public BufferSkip(IUniTaskAsyncEnumerable<TSource> source, int count, int skip)
|
||||
{
|
||||
this.source = source;
|
||||
this.count = count;
|
||||
this.skip = skip;
|
||||
}
|
||||
|
||||
public IUniTaskAsyncEnumerator<IList<TSource>> GetAsyncEnumerator(CancellationToken cancellationToken = default)
|
||||
{
|
||||
return new Enumerator(source, count, skip, cancellationToken);
|
||||
}
|
||||
|
||||
sealed class Enumerator : MoveNextSource, IUniTaskAsyncEnumerator<IList<TSource>>
|
||||
{
|
||||
static readonly Action<object> MoveNextCoreDelegate = MoveNextCore;
|
||||
|
||||
readonly IUniTaskAsyncEnumerable<TSource> source;
|
||||
readonly int count;
|
||||
readonly int skip;
|
||||
CancellationToken cancellationToken;
|
||||
|
||||
IUniTaskAsyncEnumerator<TSource> enumerator;
|
||||
UniTask<bool>.Awaiter awaiter;
|
||||
bool continueNext;
|
||||
|
||||
bool completed;
|
||||
Queue<List<TSource>> buffers;
|
||||
int index = 0;
|
||||
|
||||
public Enumerator(IUniTaskAsyncEnumerable<TSource> source, int count, int skip, CancellationToken cancellationToken)
|
||||
{
|
||||
this.source = source;
|
||||
this.count = count;
|
||||
this.skip = skip;
|
||||
this.cancellationToken = cancellationToken;
|
||||
}
|
||||
|
||||
public IList<TSource> Current { get; private set; }
|
||||
|
||||
public UniTask<bool> MoveNextAsync()
|
||||
{
|
||||
cancellationToken.ThrowIfCancellationRequested();
|
||||
|
||||
if (enumerator == null)
|
||||
{
|
||||
enumerator = source.GetAsyncEnumerator(cancellationToken);
|
||||
buffers = new Queue<List<TSource>>();
|
||||
}
|
||||
|
||||
completionSource.Reset();
|
||||
SourceMoveNext();
|
||||
return new UniTask<bool>(this, completionSource.Version);
|
||||
}
|
||||
|
||||
void SourceMoveNext()
|
||||
{
|
||||
if (completed)
|
||||
{
|
||||
if (buffers.Count > 0)
|
||||
{
|
||||
Current = buffers.Dequeue();
|
||||
completionSource.TrySetResult(true);
|
||||
return;
|
||||
}
|
||||
else
|
||||
{
|
||||
completionSource.TrySetResult(false);
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
try
|
||||
{
|
||||
|
||||
LOOP:
|
||||
awaiter = enumerator.MoveNextAsync().GetAwaiter();
|
||||
if (awaiter.IsCompleted)
|
||||
{
|
||||
continueNext = true;
|
||||
MoveNextCore(this);
|
||||
if (continueNext)
|
||||
{
|
||||
continueNext = false;
|
||||
goto LOOP; // avoid recursive
|
||||
}
|
||||
}
|
||||
else
|
||||
{
|
||||
awaiter.SourceOnCompleted(MoveNextCoreDelegate, this);
|
||||
}
|
||||
}
|
||||
catch (Exception ex)
|
||||
{
|
||||
completionSource.TrySetException(ex);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
static void MoveNextCore(object state)
|
||||
{
|
||||
var self = (Enumerator)state;
|
||||
|
||||
if (self.TryGetResult(self.awaiter, out var result))
|
||||
{
|
||||
if (result)
|
||||
{
|
||||
if (self.index++ % self.skip == 0)
|
||||
{
|
||||
self.buffers.Enqueue(new List<TSource>(self.count));
|
||||
}
|
||||
|
||||
var item = self.enumerator.Current;
|
||||
foreach (var buffer in self.buffers)
|
||||
{
|
||||
buffer.Add(item);
|
||||
}
|
||||
|
||||
if (self.buffers.Count > 0 && self.buffers.Peek().Count == self.count)
|
||||
{
|
||||
self.Current = self.buffers.Dequeue();
|
||||
self.continueNext = false;
|
||||
self.completionSource.TrySetResult(true);
|
||||
return;
|
||||
}
|
||||
else
|
||||
{
|
||||
if (!self.continueNext)
|
||||
{
|
||||
self.SourceMoveNext();
|
||||
}
|
||||
}
|
||||
}
|
||||
else
|
||||
{
|
||||
self.continueNext = false;
|
||||
self.completed = true;
|
||||
self.SourceMoveNext();
|
||||
}
|
||||
}
|
||||
else
|
||||
{
|
||||
self.continueNext = false;
|
||||
}
|
||||
}
|
||||
|
||||
public UniTask DisposeAsync()
|
||||
{
|
||||
if (enumerator != null)
|
||||
{
|
||||
return enumerator.DisposeAsync();
|
||||
}
|
||||
return default;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
|
@ -272,5 +272,52 @@ namespace NetCoreTests.Linq
|
|||
await Assert.ThrowsAsync<UniTaskTestException>(async () => await xs);
|
||||
}
|
||||
}
|
||||
|
||||
[Theory]
|
||||
// [InlineData(0, 0)]
|
||||
[InlineData(0, 3)]
|
||||
[InlineData(9, 1)]
|
||||
[InlineData(9, 2)]
|
||||
[InlineData(9, 3)]
|
||||
[InlineData(17, 3)]
|
||||
[InlineData(17, 16)]
|
||||
[InlineData(17, 17)]
|
||||
[InlineData(17, 27)]
|
||||
public async Task Buffer(int rangeCount, int bufferCount)
|
||||
{
|
||||
var xs = await UniTaskAsyncEnumerable.Range(0, rangeCount).Buffer(bufferCount).Select(x => string.Join(",", x)).ToArrayAsync();
|
||||
var ys = await AsyncEnumerable.Range(0, rangeCount).Buffer(bufferCount).Select(x => string.Join(",", x)).ToArrayAsync();
|
||||
|
||||
xs.Should().BeEquivalentTo(ys);
|
||||
}
|
||||
|
||||
[Theory]
|
||||
// [InlineData(0, 0)]
|
||||
[InlineData(0, 3, 2)]
|
||||
[InlineData(9, 1, 1)]
|
||||
[InlineData(9, 2, 3)]
|
||||
[InlineData(9, 3, 4)]
|
||||
[InlineData(17, 3, 3)]
|
||||
[InlineData(17, 16, 5)]
|
||||
[InlineData(17, 17, 19)]
|
||||
public async Task BufferSkip(int rangeCount, int bufferCount, int skipCount)
|
||||
{
|
||||
var xs = await UniTaskAsyncEnumerable.Range(0, rangeCount).Buffer(bufferCount, skipCount).Select(x => string.Join(",", x)).ToArrayAsync();
|
||||
var ys = await AsyncEnumerable.Range(0, rangeCount).Buffer(bufferCount, skipCount).Select(x => string.Join(",", x)).ToArrayAsync();
|
||||
|
||||
xs.Should().BeEquivalentTo(ys);
|
||||
}
|
||||
|
||||
[Fact]
|
||||
public async Task BufferError()
|
||||
{
|
||||
foreach (var item in UniTaskTestException.Throws())
|
||||
{
|
||||
var xs = item.Buffer(3).ToArrayAsync();
|
||||
var ys = item.Buffer(3, 2).ToArrayAsync();
|
||||
await Assert.ThrowsAsync<UniTaskTestException>(async () => await xs);
|
||||
await Assert.ThrowsAsync<UniTaskTestException>(async () => await ys);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
Loading…
Reference in New Issue