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

Buffer

pull/73/head
neuecc 2020-05-12 03:33:11 +09:00
parent 72efadd0a2
commit 61a3744fdd
2 changed files with 387 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

340
src/UniTask.NetCore/Linq/Buffer.cs Normal file
View File

@ -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;
}
}
}
}

47
src/UniTask.NetCoreTests/Linq/Projection.cs
View File

@ -272,5 +272,52 @@ namespace NetCoreTests.Linq
await Assert.ThrowsAsync<UniTaskTestException>(async () => await xs); 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);
}
}
} }
} }