Cleanup Partition Nullability

morelinq · Nov 23, 2022 · 26b9656 · 26b9656
1 parent c3a7094
commit 26b9656
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 7 deletions.
diff --git a/MoreLinq.Test/PartitionTest.cs b/MoreLinq.Test/PartitionTest.cs
@@ -15,6 +15,8 @@
 // limitations under the License.
 #endregion
 
+#nullable enable
+
 namespace MoreLinq.Test
 {
     using System;

diff --git a/MoreLinq/Partition.cs b/MoreLinq/Partition.cs
@@ -177,7 +177,8 @@ public static TResult Partition<TKey, TElement, TResult>(this IEnumerable<IGroup
             Func<IEnumerable<TElement>, IEnumerable<IGrouping<TKey, TElement>>, TResult> resultSelector)
         {
             if (resultSelector == null) throw new ArgumentNullException(nameof(resultSelector));
-            return PartitionImpl(source, 1, key, key2: default!, key3: default!, comparer,
+
+            return PartitionImpl(source, 1, key, key2: default, key3: default, comparer,
                                  (a, _, _, rest) => resultSelector(a, rest));
         }
 
@@ -235,8 +236,9 @@ public static TResult Partition<TKey, TElement, TResult>(this IEnumerable<IGroup
             Func<IEnumerable<TElement>, IEnumerable<TElement>, IEnumerable<IGrouping<TKey, TElement>>, TResult> resultSelector)
         {
             if (resultSelector == null) throw new ArgumentNullException(nameof(resultSelector));
-            return PartitionImpl(source, 2, key1, key2, key3: default!, comparer,
-                                 (a, b, _, rest) => resultSelector(a, b, rest));
+
+            return PartitionImpl(source, 2, key1, key2, key3: default, comparer,
+                                 (a, b, c, rest) => resultSelector(a, b, rest));
         }
 
         /// <summary>
@@ -296,7 +298,7 @@ public static TResult Partition<TKey, TElement, TResult>(this IEnumerable<IGroup
             PartitionImpl(source, 3, key1, key2, key3, comparer, resultSelector);
 
         static TResult PartitionImpl<TKey, TElement, TResult>(IEnumerable<IGrouping<TKey, TElement>> source,
-            int count, TKey key1, TKey key2, TKey key3, IEqualityComparer<TKey>? comparer,
+            int count, TKey? key1, TKey? key2, TKey? key3, IEqualityComparer<TKey>? comparer,
             Func<IEnumerable<TElement>, IEnumerable<TElement>, IEnumerable<TElement>, IEnumerable<IGrouping<TKey, TElement>>, TResult> resultSelector)
         {
             Debug.Assert(count is > 0 and <= 3);
@@ -317,9 +319,9 @@ static TResult PartitionImpl<TKey, TElement, TResult>(IEnumerable<IGrouping<TKey
 
             foreach (var e in source)
             {
-                var i = count > 0 && comparer.Equals(e.Key, key1) ? 0
-                      : count > 1 && comparer.Equals(e.Key, key2) ? 1
-                      : count > 2 && comparer.Equals(e.Key, key3) ? 2
+                var i = count > 0 && comparer.Equals(e.Key, key1!) ? 0
+                      : count > 1 && comparer.Equals(e.Key, key2!) ? 1
+                      : count > 2 && comparer.Equals(e.Key, key3!) ? 2
                       : -1;
 
                 if (i < 0)