Fix disposal of enumerators in EnumerablePartition

src/System.Linq/src/System/Linq/Partition.cs

@@ -415,6 +415,17 @@ public override Iterator<TSource> Clone()
                 return new EnumerablePartition<TSource>(_source, _minIndexInclusive, _maxIndexInclusive);
             }
 
+            public override void Dispose()
+            {
+                if (_enumerator != null)
+                {
+                    _enumerator.Dispose();
+                    _enumerator = null;
+                }
+
+                base.Dispose();
+            }
+
             public int GetCount(bool onlyIfCheap)
             {
                 if (onlyIfCheap)

src/System.Linq/tests/EnumerableTests.cs

@@ -286,28 +286,47 @@ protected static List<Func<IEnumerable<T>, IEnumerable<T>>> IdentityTransforms<T
             };
         }
 
-        protected class DelegateBasedEnumerator<T> : IEnumerator<T>
+        protected sealed class DelegateIterator<TSource> : IEnumerable<TSource>, IEnumerator<TSource>
         {
-            public Func<bool> MoveNextWorker { get; set; }
-            public Func<T> CurrentWorker { get; set; }
-            public Action DisposeWorker { get; set; }
-            public Func<object> NonGenericCurrentWorker { get; set; }
-            public Action ResetWorker { get; set; }
-
-            public T Current => CurrentWorker();
-            public bool MoveNext() => MoveNextWorker();
-            public void Dispose() => DisposeWorker();
-            void IEnumerator.Reset() => ResetWorker();
-            object IEnumerator.Current => NonGenericCurrentWorker();
-        }
+            private readonly Func<IEnumerator<TSource>> _getEnumerator;
+            private readonly Func<bool> _moveNext;
+            private readonly Func<TSource> _current;
+            private readonly Func<IEnumerator> _explicitGetEnumerator;
+            private readonly Func<object> _explicitCurrent;
+            private readonly Action _reset;
+            private readonly Action _dispose;
+
+            public DelegateIterator(
+                Func<IEnumerator<TSource>> getEnumerator = null,
+                Func<bool> moveNext = null,
+                Func<TSource> current = null,
+                Func<IEnumerator> explicitGetEnumerator = null,
+                Func<object> explicitCurrent = null,
+                Action reset = null,
+                Action dispose = null)
+            {
+                _getEnumerator = getEnumerator ?? (() => this);
+                _moveNext = moveNext ?? (() => { throw new NotImplementedException(); });
+                _current = current ?? (() => { throw new NotImplementedException(); });
+                _explicitGetEnumerator = explicitGetEnumerator ?? (() => { throw new NotImplementedException(); });
+                _explicitCurrent = explicitCurrent ?? (() => { throw new NotImplementedException(); });
+                _reset = reset ?? (() => { throw new NotImplementedException(); });
+                _dispose = dispose ?? (() => { throw new NotImplementedException(); });
+            }
 
-        protected class DelegateBasedEnumerable<T> : IEnumerable<T>
-        {
-            public Func<IEnumerator<T>> GetEnumeratorWorker { get; set; }
-            public Func<IEnumerator> NonGenericGetEnumeratorWorker { get; set; }
+            public IEnumerator<TSource> GetEnumerator() => _getEnumerator();
 
-            public IEnumerator<T> GetEnumerator() => GetEnumeratorWorker();
-            IEnumerator IEnumerable.GetEnumerator() => NonGenericGetEnumeratorWorker();
+            public bool MoveNext() => _moveNext();
+
+            public TSource Current => _current();
+
+            IEnumerator IEnumerable.GetEnumerator() => _explicitGetEnumerator();
+
+            object IEnumerator.Current => _explicitCurrent();
+
+            void IEnumerator.Reset() => _reset();
+
+            void IDisposable.Dispose() => _dispose();
         }
     }
 }

src/System.Linq/tests/SelectManyTests.cs

@@ -409,31 +409,15 @@ public void DisposeAfterEnumeration(int sourceLength, int subLength)
             bool sourceDisposed = false;
             bool[] subCollectionDisposed = new bool[sourceLength];
 
-            var sourceEnumerator = new DelegateBasedEnumerator<int>
-            {
-                MoveNextWorker = () => ++sourceState <= sourceLength,
-                CurrentWorker = () => 0,
-                DisposeWorker = () => sourceDisposed = true
-            };
-
-            var source = new DelegateBasedEnumerable<int>
-            {
-                GetEnumeratorWorker = () => sourceEnumerator
-            };
-
-            var subEnumerator = new DelegateBasedEnumerator<int>
-            {
-                // MoveNext: Return true subLength times.
-                // Dispose: Record that Dispose was called & move to the next index.
-                MoveNextWorker = () => ++subState[subIndex] <= subLength,
-                CurrentWorker = () => subState[subIndex],
-                DisposeWorker = () => subCollectionDisposed[subIndex++] = true
-            };
-
-            var subCollection = new DelegateBasedEnumerable<int>
-            {
-                GetEnumeratorWorker = () => subEnumerator
-            };
+            var source = new DelegateIterator<int>(
+                moveNext: () => ++sourceState <= sourceLength,
+                current: () => 0,
+                dispose: () => sourceDisposed = true);
+
+            var subCollection = new DelegateIterator<int>(
+                moveNext: () => ++subState[subIndex] <= subLength, // Return true `subLength` times.
+                current: () => subState[subIndex],
+                dispose: () => subCollectionDisposed[subIndex++] = true); // Record that Dispose was called, and move on to the next index.
 
             var iterator = source.SelectMany(_ => subCollection);
 

src/System.Linq/tests/SkipTests.cs

@@ -508,5 +508,29 @@ public void IteratorStateShouldNotChangeIfNumberOfElementsIsUnbounded()
                 }
             }
         }
+
+        [Theory]
+        [InlineData(0, -1)]
+        [InlineData(0, 0)]
+        [InlineData(1, 0)]
+        [InlineData(2, 1)]
+        [InlineData(2, 2)]
+        [InlineData(2, 3)]
+        public void DisposeSource(int sourceCount, int count)
+        {
+            int state = 0;
+
+            var source = new DelegateIterator<int>(
+                moveNext: () => ++state <= sourceCount,
+                current: () => 0,
+                dispose: () => state = -1);
+
+            IEnumerator<int> iterator = source.Skip(count).GetEnumerator();
+            int iteratorCount = Math.Max(0, sourceCount - Math.Max(0, count));
+            Assert.All(Enumerable.Range(0, iteratorCount), _ => Assert.True(iterator.MoveNext()));
+
+            Assert.False(iterator.MoveNext());
+            Assert.Equal(-1, state);
+        }
     }
 }

src/System.Linq/tests/TakeTests.cs

@@ -534,5 +534,34 @@ public void ElementAtOfLazySkipTakeChain(IEnumerable<int> source, int skip, int
                 Assert.Equal(expectedValues[i], partition.ElementAtOrDefault(indices[i]));
             }
         }
+
+        [Theory]
+        [InlineData(0, -1)]
+        [InlineData(0, 0)]
+        [InlineData(1, 0)]
+        [InlineData(2, 1)]
+        [InlineData(2, 2)]
+        [InlineData(2, 3)]
+        public void DisposeSource(int sourceCount, int count)
+        {
+            int state = 0;
+
+            var source = new DelegateIterator<int>(
+                moveNext: () => ++state <= sourceCount,
+                current: () => 0,
+                dispose: () => state = -1);
+
+            IEnumerator<int> iterator = source.Take(count).GetEnumerator();
+            int iteratorCount = Math.Min(sourceCount, Math.Max(0, count));
+            Assert.All(Enumerable.Range(0, iteratorCount), _ => Assert.True(iterator.MoveNext()));
+
+            Assert.False(iterator.MoveNext());
+
+            // Unlike Skip, Take can tell straightaway that it can return a sequence with no elements if count <= 0.
+            // The enumerable it returns is a specialized empty iterator that has no connections to the source. Hence,
+            // after MoveNext returns false under those circumstances, it won't invoke Dispose on our enumerator.
+            int expected = count <= 0 ? 0 : -1;
+            Assert.Equal(expected, state);
+        }
     }
 }