Batch overloads using an array pool

.config/dotnet-tools.json

@@ -9,7 +9,7 @@
       ]
     },
     "dotnet-reportgenerator-globaltool": {
-      "version": "4.6.4",
+      "version": "5.1.10",
       "commands": [
         "reportgenerator"
       ]

MoreLinq.Test/BatchTest.cs

@@ -23,18 +23,12 @@ namespace MoreLinq.Test
     [TestFixture]
     public class BatchTest
     {
-        [Test]
-        public void BatchZeroSize()
-        {
-            AssertThrowsArgument.OutOfRangeException("size",() =>
-                new object[0].Batch(0));
-        }
-
-        [Test]
-        public void BatchNegativeSize()
+        [TestCase(0)]
+        [TestCase(-1)]
+        public void BatchBadSize(int size)
         {
-            AssertThrowsArgument.OutOfRangeException("size",() =>
-                new object[0].Batch(-1));
+            AssertThrowsArgument.OutOfRangeException("size", () =>
+                new object[0].Batch(size));
         }
 
         [Test]
@@ -141,3 +135,259 @@ public void BatchEmptySource(SourceKind kind)
         }
     }
 }
+
+#if NETCOREAPP3_1_OR_GREATER
+
+namespace MoreLinq.Test
+{
+    using System;
+    using System.Buffers;
+    using System.Collections.Generic;
+    using MoreLinq.Experimental;
+    using NUnit.Framework;
+
+    [TestFixture]
+    public class BatchPoolTest
+    {
+        [TestCase(0)]
+        [TestCase(-1)]
+        public void BatchBadSize(int size)
+        {
+            AssertThrowsArgument.OutOfRangeException("size", () =>
+                new object[0].Batch(size, ArrayPool<object>.Shared,
+                                    BreakingFunc.Of<ICurrentBuffer<object>, IEnumerable<object>>(),
+                                    BreakingFunc.Of<IEnumerable<object>, object>()));
+        }
+
+        [Test]
+        public void BatchEvenlyDivisibleSequence()
+        {
+            using var input = TestingSequence.Of(1, 2, 3, 4, 5, 6, 7, 8, 9);
+            using var pool = new TestArrayPool<int>();
+
+            var result = input.Batch(3, pool, Enumerable.ToArray);
+
+            using var reader = result.Read();
+            reader.Read().AssertSequenceEqual(1, 2, 3);
+            reader.Read().AssertSequenceEqual(4, 5, 6);
+            reader.Read().AssertSequenceEqual(7, 8, 9);
+            reader.ReadEnd();
+        }
+
+        [Test]
+        public void BatchUnevenlyDivisibleSequence()
+        {
+            using var input = TestingSequence.Of(1, 2, 3, 4, 5, 6, 7, 8, 9);
+            using var pool = new TestArrayPool<int>();
+
+            var result = input.Batch(4, pool, Enumerable.ToArray);
+
+            using var reader = result.Read();
+            reader.Read().AssertSequenceEqual(1, 2, 3, 4);
+            reader.Read().AssertSequenceEqual(5, 6, 7, 8);
+            reader.Read().AssertSequenceEqual(9);
+            reader.ReadEnd();
+        }
+
+        [Test]
+        public void BatchIsLazy()
+        {
+            var input = new BreakingSequence<object>();
+            _ = input.Batch(1, ArrayPool<object>.Shared,
+                            BreakingFunc.Of<ICurrentBuffer<object>, IEnumerable<object>>(),
+                            BreakingFunc.Of<IEnumerable<object>, object>());
+        }
+
+        [TestCase(SourceKind.BreakingList        , 0)]
+        [TestCase(SourceKind.BreakingReadOnlyList, 0)]
+        [TestCase(SourceKind.BreakingList        , 1)]
+        [TestCase(SourceKind.BreakingReadOnlyList, 1)]
+        [TestCase(SourceKind.BreakingList        , 2)]
+        [TestCase(SourceKind.BreakingReadOnlyList, 2)]
+        public void BatchCollectionSmallerThanSize(SourceKind kind, int oversize)
+        {
+            var xs = new[] { 1, 2, 3, 4, 5 };
+            using var pool = new TestArrayPool<int>();
+
+            var result = xs.ToSourceKind(kind)
+                           .Batch(xs.Length + oversize, pool, Enumerable.ToArray);
+
+            using var reader = result.Read();
+            reader.Read().AssertSequenceEqual(1, 2, 3, 4, 5);
+            reader.ReadEnd();
+        }
+
+        [Test]
+        public void BatchReadOnlyCollectionSmallerThanSize()
+        {
+            var collection = ReadOnlyCollection.From(1, 2, 3, 4, 5);
+            using var pool = new TestArrayPool<int>();
+
+            var result = collection.Batch(collection.Count * 2, pool,
+                                          Enumerable.ToArray);
+
+            using var reader = result.Read();
+            reader.Read().AssertSequenceEqual(1, 2, 3, 4, 5);
+            reader.ReadEnd();
+        }
+
+        [TestCase(SourceKind.Sequence)]
+        [TestCase(SourceKind.BreakingList)]
+        [TestCase(SourceKind.BreakingReadOnlyList)]
+        [TestCase(SourceKind.BreakingReadOnlyCollection)]
+        [TestCase(SourceKind.BreakingCollection)]
+        public void BatchEmptySource(SourceKind kind)
+        {
+            using var pool = new TestArrayPool<int>();
+
+            var result = Enumerable.Empty<int>()
+                                   .ToSourceKind(kind)
+                                   .Batch(100, pool, Enumerable.ToArray);
+
+            Assert.That(result, Is.Empty);
+        }
+
+        [Test]
+        public void BatchFilterBucket()
+        {
+            const int scale = 2;
+            var input = TestingSequence.Of(1, 2, 3, 4, 5, 6, 7, 8, 9);
+            using var pool = new TestArrayPool<int>();
+
+            var result = input.Batch(3, pool,
+                                     current => from n in current
+                                                where n % 2 == 0
+                                                select n * scale,
+                                     Enumerable.ToArray);
+
+            using var reader = result.Read();
+            reader.Read().AssertSequenceEqual(2 * scale);
+            reader.Read().AssertSequenceEqual(4 * scale, 6 * scale);
+            reader.Read().AssertSequenceEqual(8 * scale);
+            reader.ReadEnd();
+        }
+
+        [Test]
+        public void BatchSumBucket()
+        {
+            var input = TestingSequence.Of(1, 2, 3, 4, 5, 6, 7, 8, 9);
+            using var pool = new TestArrayPool<int>();
+
+            var result = input.Batch(3, pool, Enumerable.Sum);
+
+            using var reader = result.Read();
+            Assert.That(reader.Read(), Is.EqualTo(1 + 2 + 3));
+            Assert.That(reader.Read(), Is.EqualTo(4 + 5 + 6));
+            Assert.That(reader.Read(), Is.EqualTo(7 + 8 + 9));
+            reader.ReadEnd();
+        }
+
+        /// <remarks>
+        /// This test does not exercise the intended usage!
+        /// </remarks>
+
+        [Test]
+        public void BatchUpdatesCurrentListInPlace()
+        {
+            var input = TestingSequence.Of(1, 2, 3, 4, 5, 6, 7, 8, 9);
+            using var pool = new TestArrayPool<int>();
+
+            var result = input.Batch(4, pool, current => current, current => (ICurrentBuffer<int>)current);
+
+            using var reader = result.Read();
+            var current = reader.Read();
+            current.AssertSequenceEqual(1, 2, 3, 4);
+            _ = reader.Read();
+            current.AssertSequenceEqual(5, 6, 7, 8);
+            _ = reader.Read();
+            current.AssertSequenceEqual(9);
+
+            reader.ReadEnd();
+
+            Assert.That(current, Is.Empty);
+        }
+
+        [Test]
+        public void BatchCallsBucketSelectorBeforeIteratingSource()
+        {
+            var iterations = 0;
+            IEnumerable<int> Source()
+            {
+                iterations++;
+                yield break;
+            }
+
+            var input = Source();
+            using var pool = new TestArrayPool<int>();
+            var initIterations = -1;
+
+            var result = input.Batch(4, pool,
+                                     current =>
+                                     {
+                                         initIterations = iterations;
+                                         return current;
+                                     },
+                                     _ => 0);
+
+            using var enumerator = result.GetEnumerator();
+            Assert.That(enumerator.MoveNext(), Is.False);
+            Assert.That(initIterations, Is.Zero);
+        }
+
+        [Test]
+        public void BatchBucketSelectorCurrentList()
+        {
+            var input = TestingSequence.Of(1, 2, 3, 4, 5, 6, 7, 8, 9);
+            using var pool = new TestArrayPool<int>();
+            int[] bucketSelectorItems = null;
+
+            var result = input.Batch(4, pool, current => bucketSelectorItems = current.ToArray(), _ => 0);
+
+            using var reader = result.Read();
+            _ = reader.Read();
+            Assert.That(bucketSelectorItems, Is.Not.Null);
+            Assert.That(bucketSelectorItems, Is.Empty);
+        }
+
+        /// <summary>
+        /// An <see cref="ArrayPool{T}"/> implementation for testing purposes that holds only
+        /// one array in the pool and ensures that it is returned when the pool is disposed.
+        /// </summary>
+
+        sealed class TestArrayPool<T> : ArrayPool<T>, IDisposable
+        {
+            T[] _pooledArray;
+            T[] _rentedArray;
+
+            public override T[] Rent(int minimumLength)
+            {
+                if (_pooledArray is null && _rentedArray is null)
+                    _pooledArray = new T[minimumLength * 2];
+
+                if (_pooledArray is null)
+                    throw new InvalidOperationException("The pool is exhausted.");
+
+                (_pooledArray, _rentedArray) = (null, _pooledArray);
+
+                return _rentedArray;
+            }
+
+            public override void Return(T[] array, bool clearArray = false)
+            {
+                if (_rentedArray is null)
+                    throw new InvalidOperationException("Cannot return when nothing has been rented from this pool.");
+
+                if (array != _rentedArray)
+                    throw new InvalidOperationException("Cannot return what has not been rented from this pool.");
+
+                _pooledArray = array;
+                _rentedArray = null;
+            }
+
+            public void Dispose() =>
+                Assert.That(_rentedArray, Is.Null);
+        }
+    }
+}
+
+#endif // NETCOREAPP3_1_OR_GREATER