improves QuotaCache

pardeike · Apr 7, 2024 · 774e7a0 · 774e7a0
1 parent c15d9bb
commit 774e7a0
Show file tree

Hide file tree

Showing 3 changed files with 43 additions and 12 deletions.
diff --git a/1.5/Assemblies/CameraPlus.dll b/1.5/Assemblies/CameraPlus.dll
diff --git a/Source/QuotaCache.cs b/Source/QuotaCache.cs
@@ -3,28 +3,59 @@
 
 namespace CameraPlus
 {
-	public class QuotaCache<S, T>(int maxRetrievals, Func<S, T> fetchCallback)
+	public class WeakQuotaCache<KEY, KEYID, VALUE>(int maxRetrievals, Func<KEY, KEYID> keyConverter, Func<KEY, VALUE> fetchCallback) where VALUE : class
 	{
-		private readonly Dictionary<S, (T value, int count)> cache = [];
+		private readonly Dictionary<KEYID, (WeakReference<VALUE> value, int count)> cache = [];
 		private readonly int maxRetrievals = maxRetrievals;
-		private readonly Func<S, T> fetchCallback = fetchCallback;
+		private readonly Func<KEY, KEYID> keyConverter = keyConverter;
+		private readonly Func<KEY, VALUE> fetchCallback = fetchCallback;
 
-		public T Get(S key)
+		public VALUE Get(KEY key)
 		{
 			if (key == null)
 				return default;
 
-			if (!cache.ContainsKey(key) || cache[key].count >= maxRetrievals)
+			var rawKey = keyConverter(key);
+
+			if (!cache.ContainsKey(rawKey) || cache[rawKey].count >= maxRetrievals)
 			{
-				T newValue = fetchCallback(key);
-				cache[key] = (newValue, 0);
+				VALUE newValue = fetchCallback(key);
+				cache[rawKey] = (new WeakReference<VALUE>(newValue), 0);
 			}
 
-			var (value, count) = cache[key];
-			cache[key] = (value, count + 1);
+			var (value, count) = cache[rawKey];
+			cache[rawKey] = (value, count + 1);
 
-			return value;
+			if (value.TryGetTarget(out VALUE returnValue))
+				return returnValue;
+			return default;
 		}
 	}
 
+	public class QuotaCache<KEY, KEYID, VALUE>(int maxRetrievals, Func<KEY, KEYID> keyConverter, Func<KEY, VALUE> fetchCallback)
+	{
+		private readonly Dictionary<KEYID, (VALUE value, int count)> cache = [];
+		private readonly int maxRetrievals = maxRetrievals;
+		private readonly Func<KEY, KEYID> keyConverter = keyConverter;
+		private readonly Func<KEY, VALUE> fetchCallback = fetchCallback;
+
+		public VALUE Get(KEY key)
+		{
+			if (key == null)
+				return default;
+
+			var rawKey = keyConverter(key);
+
+			if (!cache.ContainsKey(rawKey) || cache[rawKey].count >= maxRetrievals)
+			{
+				VALUE newValue = fetchCallback(key);
+				cache[rawKey] = (newValue, 0);
+			}
+
+			var (value, count) = cache[rawKey];
+			cache[rawKey] = (value, count + 1);
+
+			return value;
+		}
+	}
 }
diff --git a/Source/Tools.cs b/Source/Tools.cs
@@ -29,7 +29,7 @@ class Tools : CameraPlusSettings
 		static readonly Color downedColor = new(0.9f, 0f, 0f);
 		static readonly Color draftedColor = new(0f, 0.5f, 0f);
 
-		static readonly QuotaCache<Pawn, bool> shouldShowDotCache = new(60, pawn =>
+		static readonly QuotaCache<Pawn, int, bool> shouldShowDotCache = new(60, pawn => pawn.thingIDNumber, pawn =>
 		{
 			if (CameraPlusMain.Settings.customNameStyle == LabelStyle.HideAnimals && pawn.RaceProps.Animal)
 				return false;
@@ -46,7 +46,7 @@ class Tools : CameraPlusSettings
 			return isSmall && (CameraPlusMain.Settings.includeNotTamedAnimals || pawn.RaceProps.Animal == false || tamedAnimal);
 		});
 
-		static readonly QuotaCache<Thing, bool> shouldShowLabelCache = new(60, thing =>
+		static readonly QuotaCache<Thing, int, bool> shouldShowLabelCache = new(60, thing => thing.thingIDNumber, thing =>
 		{
 			var len = FastUI.CurUICellSize;
 			var isPawn = thing is Pawn;