From 4a02e6c6cc6453d3d13fc587d64e0f29626bea5a Mon Sep 17 00:00:00 2001
From: RoosterDragon <RoosterDragon@outlook.com>
Date: Fri, 7 Jul 2023 19:43:59 +0100
Subject: [PATCH] Improve Exts.GetOrAdd method to avoid multiple lookups.

Use CollectionsMarshal to hold a ref to the dictionary entry. When the value needs to be added this allows us to set the value directly into it without having to locate the entry a second time.
---
 OpenRA.Game/Exts.cs | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/OpenRA.Game/Exts.cs b/OpenRA.Game/Exts.cs
index f9cca3b774..06a9f70a38 100644
--- a/OpenRA.Game/Exts.cs
+++ b/OpenRA.Game/Exts.cs
@@ -109,13 +109,23 @@ namespace OpenRA
 
 		public static V GetOrAdd<K, V>(this Dictionary<K, V> d, K k, V v)
 		{
+#if NET5_0_OR_GREATER
+			// SAFETY: Dictionary cannot be modified whilst the ref is alive.
+			ref var value = ref System.Runtime.InteropServices.CollectionsMarshal.GetValueRefOrAddDefault(d, k, out var exists);
+			if (!exists)
+				value = v;
+			return value;
+#else
 			if (!d.TryGetValue(k, out var ret))
 				d.Add(k, ret = v);
 			return ret;
+#endif
 		}
 
 		public static V GetOrAdd<K, V>(this Dictionary<K, V> d, K k, Func<K, V> createFn)
 		{
+			// Cannot use CollectionsMarshal.GetValueRefOrAddDefault here,
+			// the creation function could mutate the dictionary which would invalidate the ref.
 			if (!d.TryGetValue(k, out var ret))
 				d.Add(k, ret = createFn(k));
 			return ret;