Fix CountBy to handle null keys

Fixes morelinq#536 This is a squashed merge of PR morelinq#537.
leandromoh · Oct 9, 2018 · 1cb6494 · 1cb6494
1 parent dbbef18
commit 1cb6494
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 4 deletions.
diff --git a/MoreLinq.Test/CountByTest.cs b/MoreLinq.Test/CountByTest.cs
@@ -87,5 +87,21 @@ public void CountByIsLazy()
         {
             new BreakingSequence<string>().CountBy(x => x.Length);
         }
+
+        [Test]
+        public void CountByWithSomeNullKeys()
+        {
+            var ss = new[]
+            {
+                "foo", null, "bar", "baz", null, null, "baz", "bar", null, "foo"
+            };
+            var result = ss.CountBy(s => s);
+
+            result.AssertSequenceEqual(
+                KeyValuePair.Create("foo", 2),
+                KeyValuePair.Create((string) null, 4),
+                KeyValuePair.Create("bar", 2),
+                KeyValuePair.Create("baz", 2));
+        }
     }
 }
diff --git a/MoreLinq/CountBy.cs b/MoreLinq/CountBy.cs
@@ -79,6 +79,19 @@ public static IEnumerable<KeyValuePair<TKey, int>> CountBy<TSource, TKey>(this I
                 void Loop(IEqualityComparer<TKey> cmp)
                 {
                     var dic = new Dictionary<TKey, int>(cmp);
+                    var nullIndex = (int?) null;
+
+                    bool TryGetIndex(TKey key, out int i)
+                    {
+                        if (key == null)
+                        {
+                            i = nullIndex.GetValueOrDefault();
+                            return nullIndex.HasValue;
+                        }
+
+                        return dic.TryGetValue(key, out i);
+                    }
+
                     keys = new List<TKey>();
                     counts = new List<int>();
                     var havePrevKey = false;
@@ -90,17 +103,20 @@ void Loop(IEqualityComparer<TKey> cmp)
                         var key = keySelector(item);
 
                         if (// key same as the previous? then re-use the index
-                            (havePrevKey && cmp.GetHashCode(prevKey) == cmp.GetHashCode(key)
-                                          && cmp.Equals(prevKey, key))
+                            havePrevKey && cmp.GetHashCode(prevKey) == cmp.GetHashCode(key)
+                                         && cmp.Equals(prevKey, key)
                             // otherwise try & find index of the key
-                            || dic.TryGetValue(key, out index))
+                            || TryGetIndex(key, out index))
                         {
                             counts[index]++;
                         }
                         else
                         {
-                            dic[key] = keys.Count;
                             index = keys.Count;
+                            if (key != null)
+                                dic[key] = index;
+                            else
+                                nullIndex = index;
                             keys.Add(key);
                             counts.Add(1);
                         }