From 8c8befe9ca32b80ed12242acb2bbecfb336d16b5 Mon Sep 17 00:00:00 2001
From: estrogen elf <87099578+BlossomiShymae@users.noreply.github.com>
Date: Mon, 16 Jun 2025 01:00:27 -0500
Subject: [PATCH] refactor: use LRU for control cache

---
 Needlework.Net/Needlework.Net.csproj |  1 +
 Needlework.Net/ViewLocator.cs        | 52 +++++++++++++++++++++++++---
 2 files changed, 49 insertions(+), 4 deletions(-)

diff --git a/Needlework.Net/Needlework.Net.csproj b/Needlework.Net/Needlework.Net.csproj
index 7cde83f..0da5007 100644
--- a/Needlework.Net/Needlework.Net.csproj
+++ b/Needlework.Net/Needlework.Net.csproj
@@ -27,6 +27,7 @@
     <PackageReference Condition="'$(Configuration)' == 'Debug'" Include="Avalonia.Diagnostics" Version="11.2.8" />
     <PackageReference Include="Avalonia.Themes.Fluent" Version="11.2.8" />
     <PackageReference Include="AvaloniaEdit.TextMate" Version="11.3.0" />
+    <PackageReference Include="BitFaster.Caching" Version="2.5.3" />
     <PackageReference Include="BlossomiShymae.Briar" Version="0.2.0" />
     <PackageReference Include="CommunityToolkit.Mvvm" Version="8.4.0" />
     <PackageReference Include="DebounceThrottle" Version="3.0.1" />
diff --git a/Needlework.Net/ViewLocator.cs b/Needlework.Net/ViewLocator.cs
index c1c6949..7ba73ea 100644
--- a/Needlework.Net/ViewLocator.cs
+++ b/Needlework.Net/ViewLocator.cs
@@ -1,8 +1,12 @@
 using Avalonia.Controls;
 using Avalonia.Controls.Templates;
+using Avalonia.VisualTree;
+using BitFaster.Caching;
+using BitFaster.Caching.Lru;
 using System;
 using System.Collections.Generic;
 using System.ComponentModel;
+using System.Diagnostics.CodeAnalysis;
 using System.Linq;
 using System.Reflection;
 
@@ -10,7 +14,46 @@ namespace Needlework.Net
 {
     public class ViewLocator : IDataTemplate
     {
-        private readonly Dictionary<object, Control> _controlCache = [];
+        private class ObjectComparer : IEqualityComparer<object>
+        {
+            public new bool Equals(object? x, object? y)
+            {
+                if (ReferenceEquals(x, y))
+                    return true;
+
+                if (x == null || y == null)
+                    return false;
+
+                return x.Equals(y);
+            }
+
+            public int GetHashCode([DisallowNull] object obj)
+            {
+                return obj.GetHashCode();
+            }
+        }
+
+        private readonly ICache<object, Control> _controlCache = new ConcurrentLruBuilder<object, Control>()
+            .WithExpireAfterAccess(TimeSpan.FromMinutes(5))
+            .WithKeyComparer(new ObjectComparer())
+            .WithCapacity(1024)
+            .WithMetrics()
+            .Build();
+
+        public ViewLocator()
+        {
+            _controlCache.Events.Value!.ItemRemoved += (source, args) =>
+            {
+                var descendants = args.Value!.GetVisualDescendants();
+                foreach (var descendant in descendants)
+                {
+                    if (descendant.DataContext is INotifyPropertyChanged key)
+                    {
+                        _controlCache.TryRemove(key, out _);
+                    }
+                }
+            };
+        }
 
         public Control Build(object? data)
         {
@@ -34,13 +77,14 @@ namespace Needlework.Net
             {
                 throw new Exception("Data type has no view.");
             }
-            if (!_controlCache.TryGetValue(data!, out var res))
+            bool isCold = !_controlCache.TryGet(data!, out var res);
+            if (isCold)
             {
                 res ??= (Control)Activator.CreateInstance(type)!;
-                _controlCache[data!] = res;
+                _controlCache.AddOrUpdate(data!, res);
             }
 
-            res.DataContext = data;
+            res!.DataContext = data;
             return res;
         }