feat: Use WeakReference in ObjectRegistry

b-editor · Sep 25, 2024 · 81c462d · 81c462d
1 parent 0813c0c
commit 81c462d
Showing 1 changed file with 36 additions and 11 deletions.
diff --git a/src/Beutl.Core/ObjectRegistry.cs b/src/Beutl.Core/ObjectRegistry.cs
@@ -5,7 +5,7 @@ namespace Beutl;
 
 public class ObjectRegistry
 {
-    private readonly LinkedList<(Guid, CoreObject)> _objects = new();
+    private readonly LinkedList<(Guid, WeakReference<CoreObject>)> _objects = new();
     private readonly LinkedList<(Guid, DependentHandle)> _callbacks = new();
     private readonly object _lock = new();
 
@@ -15,7 +15,7 @@ public void Register(CoreObject obj)
     {
         lock (_lock)
         {
-            _objects.AddLast((obj.Id, obj));
+            _objects.AddLast((obj.Id, new WeakReference<CoreObject>(obj)));
             SetResolved(obj.Id, obj);
             obj.PropertyChanged += OnObjectPropertyChanged;
         }
@@ -28,7 +28,15 @@ public void Unregister(CoreObject obj)
             var node = _objects.First;
             while (node != null)
             {
-                if (ReferenceEquals(node.Value.Item2, obj))
+                if (!node.Value.Item2.TryGetTarget(out var target))
+                {
+                    var next = node.Next;
+                    _objects.Remove(node);
+                    node = next;
+                    continue;
+                }
+
+                if (ReferenceEquals(target, obj))
                 {
                     obj.PropertyChanged -= OnObjectPropertyChanged;
                     _objects.Remove(node);
@@ -47,9 +55,17 @@ public void Unregister(CoreObject obj)
             var node = _objects.First;
             while (node != null)
             {
+                if (!node.Value.Item2.TryGetTarget(out var target))
+                {
+                    var next = node.Next;
+                    _objects.Remove(node);
+                    node = next;
+                    continue;
+                }
+
                 if (node.Value.Item1 == id)
                 {
-                    return node.Value.Item2;
+                    return target;
                 }
 
                 node = node.Next;
@@ -63,7 +79,10 @@ public CoreObject[] Enumerate()
     {
         lock (_lock)
         {
-            return [.._objects.Select(i => i.Item2)];
+            return _objects
+                .Select(r => r.Item2.TryGetTarget(out var t) ? t : null)
+                .Where(o => o != null)
+                .ToArray()!;
         }
     }
 
@@ -77,10 +96,8 @@ public void Resolve<TSelf>(Guid id, TSelf self, Action<TSelf, CoreObject> callba
         }
         else
         {
-            _callbacks.AddLast((id, new DependentHandle(self, new Action<object, CoreObject>((o, r) =>
-            {
-                callback((TSelf)o, r);
-            }))));
+            _callbacks.AddLast((id,
+                new DependentHandle(self, new Action<object, CoreObject>((o, r) => { callback((TSelf)o, r); }))));
         }
     }
 
@@ -117,10 +134,18 @@ private void OnObjectPropertyChanged(object? sender, PropertyChangedEventArgs e)
                 var node = _objects.First;
                 while (node != null)
                 {
-                    if (ReferenceEquals(node.Value.Item2, args.Sender))
+                    if (!node.Value.Item2.TryGetTarget(out var target))
+                    {
+                        var next = node.Next;
+                        _objects.Remove(node);
+                        node = next;
+                        continue;
+                    }
+
+                    if (ReferenceEquals(target, args.Sender))
                     {
                         node.ValueRef = (args.NewValue, node.Value.Item2);
-                        SetResolved(args.NewValue, node.Value.Item2);
+                        SetResolved(args.NewValue, target);
                         break;
                     }