Fix morelinq#626. SkipLast is not optimal in memory

Added a test for Collection
Orace · Oct 31, 2019 · 1dc3a88 · 1dc3a88
1 parent b9a5348
commit 1dc3a88
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 7 deletions.
diff --git a/MoreLinq.Test/SkipLastTest.cs b/MoreLinq.Test/SkipLastTest.cs
@@ -32,7 +32,7 @@ public void SkipLastWithCountLesserThanOne(int skip)
         }
 
         [Test]
-        public void SkipLast()
+        public void SkipLastEnumerable()
         {
             const int take = 100;
             const int skip = 20;
@@ -44,6 +44,19 @@ public void SkipLast()
             Assert.That(expectations, Is.EqualTo(sequence.SkipLast(skip)));
         }
 
+        [Test]
+        public void SkipLastCollection()
+        {
+            const int take = 100;
+            const int skip = 20;
+
+            var sequence = Enumerable.Range(1, take).ToArray();
+
+            var expectations = sequence.Take(take - skip);
+
+            Assert.That(expectations, Is.EqualTo(sequence.SkipLast(skip)));
+        }
+
         [TestCase(5)]
         [TestCase(6)]
         public void SkipLastWithSequenceShorterThanCount(int skip)

diff --git a/MoreLinq/SkipLast.cs b/MoreLinq/SkipLast.cs
@@ -40,12 +40,26 @@ public static IEnumerable<T> SkipLast<T>(this IEnumerable<T> source, int count)
             if (count < 1)
                 return source;
 
-            return
-                source.TryGetCollectionCount() is int collectionCount
-                ? source.Take(collectionCount - count)
-                : source.CountDown(count, (e, cd) => (Element: e, Countdown: cd ))
-                        .TakeWhile(e => e.Countdown == null)
-                        .Select(e => e.Element);
+            var collectionCount = source.TryGetCollectionCount();
+            if (collectionCount.HasValue)
+                return source.Take(collectionCount.Value - count);
+
+            return _(); IEnumerable<T> _()
+            {
+                var queue = new Queue<T>(count);
+                using var enumerator = source.GetEnumerator();
+
+                while (count-- > 0 && enumerator.MoveNext())
+                {
+                    queue.Enqueue(enumerator.Current);
+                }
+
+                while (enumerator.MoveNext())
+                {
+                    yield return queue.Dequeue();
+                    queue.Enqueue(enumerator.Current);
+                }
+            }
         }
     }
 }