Reuse PartialCount more generally in count-methods

morelinq · Apr 9, 2018 · 7880f52 · 7880f52
1 parent 3e53e03
commit 7880f52
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 37 deletions.
diff --git a/MoreLinq/CountMethods.cs b/MoreLinq/CountMethods.cs
@@ -132,23 +132,9 @@ static bool QuantityIterator<T>(IEnumerable<T> source, int limit, int min, int m
         {
             if (source == null) throw new ArgumentNullException(nameof(source));
 
-            var count = 0;
-
-            if (source is ICollection<T> col)
-            {
-                count = col.Count;
-            }
-            else
-            {
-                using (var e = source.GetEnumerator())
-                {
-                    while (e.MoveNext())
-                    {
-                        if (++count == limit)
-                            break;
-                    }
-                }
-            }
+            var count = source is ICollection<T> col
+                         ? col.Count
+                         : source.CountUpTo(limit);
 
             return count >= min && count <= max;
         }
@@ -182,13 +168,14 @@ public static int CompareCount<TFirst, TSecond>(this IEnumerable<TFirst> first,
             {
                 return firstCol.Count.CompareTo(second is ICollection<TSecond> secondCol
                                                 ? secondCol.Count
-                                                : PartialCount(second, firstCol.Count + 1));
+                                                : second.CountUpTo(firstCol.Count + 1));
+            }
+            else if (second is ICollection<TSecond> secondCol)
+            {
+                return first.CountUpTo(secondCol.Count + 1).CompareTo(secondCol.Count);
             }
             else
             {
-                if (second is ICollection<TSecond> secondCol)
-                    return PartialCount(first, secondCol.Count + 1).CompareTo(secondCol.Count);
-
                 bool firstHasNext;
                 bool secondHasNext;
 
@@ -205,22 +192,6 @@ public static int CompareCount<TFirst, TSecond>(this IEnumerable<TFirst> first,
 
                 return firstHasNext.CompareTo(secondHasNext);
             }
-
-            int PartialCount<T>(IEnumerable<T> source, int limit)
-            {
-                var count = 0;
-
-                using (var e = source.GetEnumerator())
-                {
-                    while (e.MoveNext())
-                    {
-                        if (++count == limit)
-                            break;
-                    }
-                }
-
-                return count;
-            }
         }
     }
 }
diff --git a/MoreLinq/MoreEnumerable.cs b/MoreLinq/MoreEnumerable.cs
@@ -35,5 +35,23 @@ public static partial class MoreEnumerable
                  : source is IReadOnlyCollection<T> readOnlyCollection ? readOnlyCollection.Count
                  : (int?)null;
         }
+
+        static int CountUpTo<T>(this IEnumerable<T> source, int max)
+        {
+            if (source == null) throw new ArgumentNullException(nameof(source));
+            if (max < 0) throw new ArgumentOutOfRangeException(nameof(max), "The maximum count argument cannot be negative.");
+
+            var count = 0;
+
+            using (var e = source.GetEnumerator())
+            {
+                while (count < max && e.MoveNext())
+                {
+                    count++;
+                }
+            }
+
+            return count;
+        }
     }
 }