Additional caching for FallbackDictionary.

Author: tillig

src/Autofac/Util/FallbackDictionary.cs

@@ -14,15 +14,14 @@ namespace Autofac.Util;
 internal class FallbackDictionary<TKey, TValue> : IDictionary<TKey, TValue>
     where TKey : notnull
 {
-    /// <summary>
-    /// Storage for local values set in the dictionary.
-    /// </summary>
     private readonly IDictionary<TKey, TValue> _localValues = new Dictionary<TKey, TValue>();
-
-    /// <summary>
-    /// The parent dictionary to which values should fall back when not present in the current dictionary.
-    /// </summary>
     private readonly IDictionary<TKey, TValue> _parent;
+    private readonly FallbackDictionary<TKey, TValue>? _fallbackParent;
+
+    private int _localVersion;
+    private List<TKey>? _cachedKeys;
+    private int _cachedKeysLocalVersion;
+    private int _cachedKeysParentVersion;
 
     /// <summary>
     /// Initializes a new instance of the <see cref="FallbackDictionary{TKey, TValue}"/> class
@@ -43,6 +42,7 @@ public FallbackDictionary()
     public FallbackDictionary(IDictionary<TKey, TValue> parent)
     {
         _parent = parent ?? throw new ArgumentNullException(nameof(parent));
+        _fallbackParent = parent as FallbackDictionary<TKey, TValue>;
     }
 
     /// <summary>
@@ -116,7 +116,7 @@ public ICollection<TValue> Values
     {
         get
         {
-            var keys = GetMergedKeysSnapshot();
+            var keys = GetMergedKeysSnapshotInternal();
             var values = new TValue[keys.Count];
             for (var i = 0; i < keys.Count; i++)
             {
@@ -154,6 +154,7 @@ public TValue this[TKey key]
         set
         {
             _localValues[key] = value;
+            IncrementVersion();
         }
     }
 
@@ -200,6 +201,7 @@ public void Add(TKey key, TValue value)
         }
 
         _localValues.Add(key, value);
+        IncrementVersion();
     }
 
     /// <summary>
@@ -208,6 +210,7 @@ public void Add(TKey key, TValue value)
     public void Clear()
     {
         _localValues.Clear();
+        IncrementVersion();
     }
 
     /// <summary>
@@ -300,7 +303,13 @@ public IEnumerator<KeyValuePair<TKey, TValue>> GetEnumerator()
     /// </remarks>
     public bool Remove(KeyValuePair<TKey, TValue> item)
     {
-        return _localValues.Remove(item);
+        var removed = _localValues.Remove(item);
+        if (removed)
+        {
+            IncrementVersion();
+        }
+
+        return removed;
     }
 
     /// <summary>
@@ -318,7 +327,13 @@ public bool Remove(KeyValuePair<TKey, TValue> item)
     /// </remarks>
     public bool Remove(TKey key)
     {
-        return _localValues.Remove(key);
+        var removed = _localValues.Remove(key);
+        if (removed)
+        {
+            IncrementVersion();
+        }
+
+        return removed;
     }
 
     /// <summary>
@@ -357,6 +372,40 @@ IEnumerator IEnumerable.GetEnumerator()
     /// </summary>
     /// <returns>A new list containing the ordered unique set of keys.</returns>
     private List<TKey> GetMergedKeysSnapshot()
+    {
+        var snapshot = GetMergedKeysSnapshotInternal();
+
+        if (snapshot.Count == 0)
+        {
+            return new List<TKey>();
+        }
+
+        return new List<TKey>(snapshot);
+    }
+
+    private List<TKey> GetMergedKeysSnapshotInternal()
+    {
+        if (_fallbackParent is null)
+        {
+            return BuildMergedKeyList();
+        }
+
+        var parentVersion = _fallbackParent._localVersion;
+        if (_cachedKeys is { } cache &&
+            _cachedKeysLocalVersion == _localVersion &&
+            _cachedKeysParentVersion == parentVersion)
+        {
+            return cache;
+        }
+
+        cache = BuildMergedKeyList();
+        _cachedKeys = cache;
+        _cachedKeysLocalVersion = _localVersion;
+        _cachedKeysParentVersion = parentVersion;
+        return cache;
+    }
+
+    private List<TKey> BuildMergedKeyList()
     {
         var keys = new List<TKey>(_localValues.Count + _parent.Count);
         keys.AddRange(_localValues.Keys);
@@ -369,24 +418,26 @@ private List<TKey> GetMergedKeysSnapshot()
             }
         }
 
+        keys.Sort();
         return keys;
     }
 
     private IEnumerable<KeyValuePair<TKey, TValue>> EnumerateKeyValuePairs()
     {
-        foreach (var kvp in _localValues)
+        foreach (var key in GetMergedKeysSnapshotInternal())
         {
-            yield return kvp;
+            yield return new KeyValuePair<TKey, TValue>(key, this[key]);
         }
+    }
 
-        foreach (var parentKvp in _parent)
+    private void IncrementVersion()
+    {
+        unchecked
         {
-            if (_localValues.ContainsKey(parentKvp.Key))
-            {
-                continue;
-            }
-
-            yield return parentKvp;
+            _localVersion++;
         }
+
+        _cachedKeys = null;
+        _cachedKeysLocalVersion = _localVersion;
     }
 }