Fix for distinct calculator with duplicate keys

src/DynamicData.Tests/Cache/DistinctFixture.cs

@@ -1,6 +1,7 @@
 
 using System;
 using System.Linq;
+using System.Reactive.Linq;
 using DynamicData.Tests.Domain;
 using FluentAssertions;
 using Xunit;
@@ -100,5 +101,29 @@ public void BreakWithLoadsOfUpdates()
 
             _results.Data.Items.ShouldAllBeEquivalentTo(new[] {1, 12, 14});
         }
+
+        [Fact]
+        public void DuplicateKeysRefreshAfterRemove()
+        {
+            var source1 = new SourceCache<Person, string>(p => p.Name);
+            var source2 = new SourceCache<Person, string>(p => p.Name);
+
+            var person = new Person("Person2", 12);
+
+            var results = source1.Connect().Merge(source2.Connect()).DistinctValues(p => p.Age).AsAggregator();
+
+            source1.AddOrUpdate(person);
+            source2.AddOrUpdate(person);
+            source2.Remove(person);
+            source1.Refresh(person); // would previously throw KeyNotFoundException here
+
+            results.Messages.Should().HaveCount(1);
+            results.Data.Items.ShouldAllBeEquivalentTo(new[] {12});
+
+            source1.Remove(person);
+
+            results.Messages.Should().HaveCount(2);
+            results.Data.Items.Should().BeEmpty();
+        }
     }
 }

src/DynamicData/Cache/Internal/DistinctCalculator.cs

@@ -14,6 +14,7 @@ internal sealed class DistinctCalculator<TObject, TKey, TValue>
         private readonly IObservable<IChangeSet<TObject, TKey>> _source;
         private readonly Func<TObject, TValue> _valueSelector;
         private readonly IDictionary<TValue, int> _valueCounters = new Dictionary<TValue, int>();
+        private readonly IDictionary<TKey, int> _keyCounters = new Dictionary<TKey, int>();
         private readonly IDictionary<TKey, TValue> _itemCache = new Dictionary<TKey, TValue>();
 
         public DistinctCalculator(IObservable<IChangeSet<TObject, TKey>> source, Func<TObject, TValue> valueSelector)
@@ -31,15 +32,44 @@ private DistinctChangeSet<TValue> Calculate(IChangeSet<TObject, TKey> changes)
         {
             var result = new DistinctChangeSet<TValue>();
 
-            void AddAction(TValue value) => _valueCounters.Lookup(value)
+            void AddKeyAction( TKey key, TValue value) => _keyCounters.Lookup(key)
+                .IfHasValue(count => _keyCounters[key] = count + 1)
+                .Else(() =>
+                {
+                    _keyCounters[key] = 1;
+                    _itemCache[key] = value; // add to cache
+                });
+
+            void AddValueAction( TValue value) => _valueCounters.Lookup(value)
                 .IfHasValue(count => _valueCounters[value] = count + 1)
                 .Else(() =>
                 {
                     _valueCounters[value] = 1;
                     result.Add(new Change<TValue, TValue>(ChangeReason.Add, value, value));
                 });
 
-            void RemoveAction(TValue value)
+            void RemoveKeyAction(TKey key)
+            {
+                var counter = _keyCounters.Lookup(key);
+                if (!counter.HasValue)
+                {
+                    return;
+                }
+
+                //decrement counter
+                var newCount = counter.Value - 1;
+                _keyCounters[key] = newCount;
+                if (newCount != 0)
+                {
+                    return;
+                }
+
+                //if there are none, then remove from cache
+                _keyCounters.Remove(key);
+                _itemCache.Remove(key);
+            }
+
+            void RemoveValueAction(TValue value)
             {
                 var counter = _valueCounters.Lookup(value);
                 if (!counter.HasValue)
@@ -69,11 +99,10 @@ void RemoveAction(TValue value)
                     case ChangeReason.Add:
                         {
                             var value = _valueSelector(change.Current);
-                            AddAction(value);
-                            _itemCache[key] = value;
+                            AddKeyAction(key, value);
+                            AddValueAction(value);
                             break;
                         }
-
                     case ChangeReason.Refresh:
                     case ChangeReason.Update:
                         {
@@ -84,22 +113,20 @@ void RemoveAction(TValue value)
                                 continue;
                             }
 
-                            RemoveAction(previous);
-                            AddAction(value);
+                            RemoveValueAction(previous);
+                            AddValueAction(value);
                             _itemCache[key] = value;
                             break;
                         }
-
                     case ChangeReason.Remove:
                         {
                             var previous = _itemCache[key];
-                            RemoveAction(previous);
-                            _itemCache.Remove(key);
+                            RemoveKeyAction(key);
+                            RemoveValueAction(previous);
                             break;
                         }
                 }
             }
-
             return result;
         }
     }