From b01c8e124771e5d902b3693927c3ea400157e388 Mon Sep 17 00:00:00 2001
From: John Call <johncall@microsoft.com>
Date: Mon, 14 Dec 2020 02:56:01 -0800
Subject: [PATCH] SortedDictionary Copy optimization (#45659)

* Sorted Dictionary Copy optimization

* Update src/libraries/System.Collections/src/System/Collections/Generic/SortedDictionary.cs

Co-authored-by: Jan Kotas <jkotas@microsoft.com>

* Feedback

Co-authored-by: Jan Kotas <jkotas@microsoft.com>
---
 .../Collections/Generic/SortedDictionary.cs   | 19 ++++++++++++++++---
 .../SortedDictionary.Generic.Tests.cs         | 13 +++++++++++--
 2 files changed, 27 insertions(+), 5 deletions(-)

diff --git a/src/libraries/System.Collections/src/System/Collections/Generic/SortedDictionary.cs b/src/libraries/System.Collections/src/System/Collections/Generic/SortedDictionary.cs
index cd352641e68a3..f9c3f1fcdfdf5 100644
--- a/src/libraries/System.Collections/src/System/Collections/Generic/SortedDictionary.cs
+++ b/src/libraries/System.Collections/src/System/Collections/Generic/SortedDictionary.cs
@@ -35,11 +35,22 @@ public SortedDictionary(IDictionary<TKey, TValue> dictionary, IComparer<TKey>? c
                 throw new ArgumentNullException(nameof(dictionary));
             }
 
-            _set = new TreeSet<KeyValuePair<TKey, TValue>>(new KeyValuePairComparer(comparer));
+            var keyValuePairComparer = new KeyValuePairComparer(comparer);
 
-            foreach (KeyValuePair<TKey, TValue> pair in dictionary)
+            if (dictionary is SortedDictionary<TKey, TValue> sortedDictionary &&
+                sortedDictionary._set.Comparer is KeyValuePairComparer kv &&
+                kv.keyComparer.Equals(keyValuePairComparer.keyComparer))
+            {
+                _set = new TreeSet<KeyValuePair<TKey, TValue>>(sortedDictionary._set, keyValuePairComparer);
+            }
+            else
             {
-                _set.Add(pair);
+                _set = new TreeSet<KeyValuePair<TKey, TValue>>(keyValuePairComparer);
+
+                foreach (KeyValuePair<TKey, TValue> pair in dictionary)
+                {
+                    _set.Add(pair);
+                }
             }
         }
 
@@ -961,6 +972,8 @@ public TreeSet()
 
         public TreeSet(IComparer<T>? comparer) : base(comparer) { }
 
+        internal TreeSet(TreeSet<T> set, IComparer<T>? comparer) : base(set, comparer) { }
+
         private TreeSet(SerializationInfo siInfo, StreamingContext context) : base(siInfo, context) { }
 
         internal override bool AddIfNotPresent(T item)
diff --git a/src/libraries/System.Collections/tests/Generic/SortedDictionary/SortedDictionary.Generic.Tests.cs b/src/libraries/System.Collections/tests/Generic/SortedDictionary/SortedDictionary.Generic.Tests.cs
index 8d8761175b897..cac00bc6723b4 100644
--- a/src/libraries/System.Collections/tests/Generic/SortedDictionary/SortedDictionary.Generic.Tests.cs
+++ b/src/libraries/System.Collections/tests/Generic/SortedDictionary/SortedDictionary.Generic.Tests.cs
@@ -57,9 +57,18 @@ public void SortedDictionary_Generic_Constructor_IDictionary_IComparer(int count
         {
             IComparer<TKey> comparer = GetKeyIComparer();
             IDictionary<TKey, TValue> source = GenericIDictionaryFactory(count);
-            SortedDictionary<TKey, TValue> copied = new SortedDictionary<TKey, TValue>(source, comparer);
-            Assert.Equal(source, copied);
+            SortedDictionary<TKey, TValue> sourceSorted = new SortedDictionary<TKey, TValue>(source, comparer);
+            Assert.Equal(source, sourceSorted);
+            Assert.Equal(comparer, sourceSorted.Comparer);
+            // Test copying a sorted dictionary.
+            SortedDictionary<TKey, TValue> copied = new SortedDictionary<TKey, TValue>(sourceSorted, comparer);
+            Assert.Equal(sourceSorted, copied);
             Assert.Equal(comparer, copied.Comparer);
+            // Test copying a sorted dictionary with a different comparer.
+            IComparer<TKey> reverseComparer = Comparer<TKey>.Create((key1, key2) => -comparer.Compare(key1, key2));
+            SortedDictionary<TKey, TValue> copiedReverse = new SortedDictionary<TKey, TValue>(sourceSorted, reverseComparer);
+            Assert.Equal(sourceSorted, copiedReverse);
+            Assert.Equal(reverseComparer, copiedReverse.Comparer);
         }
 
         #endregion