From 7e1d9dbc72eb9c6371f8bf33800caab19916cf20 Mon Sep 17 00:00:00 2001
From: Nikita Tsukanov <keks9n@gmail.com>
Date: Sun, 16 Jan 2022 19:28:49 +0300
Subject: [PATCH] Optimized WeakEvent (9332ms to 5ms using #6660 bench)

---
 src/Avalonia.Base/Utilities/WeakEvent.cs | 85 ++++++++++++++++++++----
 1 file changed, 72 insertions(+), 13 deletions(-)

diff --git a/src/Avalonia.Base/Utilities/WeakEvent.cs b/src/Avalonia.Base/Utilities/WeakEvent.cs
index 0b32015a8a..21c165afae 100644
--- a/src/Avalonia.Base/Utilities/WeakEvent.cs
+++ b/src/Avalonia.Base/Utilities/WeakEvent.cs
@@ -1,5 +1,6 @@
 using System;
 using System.Collections.Generic;
+using System.Diagnostics.CodeAnalysis;
 using System.Linq;
 using System.Reflection;
 using System.Runtime.CompilerServices;
@@ -36,7 +37,7 @@ public class WeakEvent<TSender, TEventArgs> : WeakEvent where TEventArgs : Event
     {
         if (!_subscriptions.TryGetValue(target, out var subscription))
             _subscriptions.Add(target, subscription = new Subscription(this, target));
-        subscription.Add(new WeakReference<IWeakEventSubscriber<TEventArgs>>(subscriber));
+        subscription.Add(subscriber);
     }
 
     public void Unsubscribe(TSender target, IWeakEventSubscriber<TEventArgs> subscriber)
@@ -51,11 +52,61 @@ public class WeakEvent<TSender, TEventArgs> : WeakEvent where TEventArgs : Event
         private readonly TSender _target;
         private readonly Action _compact;
 
-        private WeakReference<IWeakEventSubscriber<TEventArgs>>?[] _data =
-            new WeakReference<IWeakEventSubscriber<TEventArgs>>[16];
+        struct Entry
+        {
+            WeakReference<IWeakEventSubscriber<TEventArgs>>? _reference;
+            int _hashCode;
+
+            public Entry(IWeakEventSubscriber<TEventArgs> r)
+            {
+                if (r == null)
+                {
+                    _reference = null;
+                    _hashCode = 0;
+                    return;
+                }
+
+                _hashCode = r.GetHashCode();
+                _reference = new WeakReference<IWeakEventSubscriber<TEventArgs>>(r);
+            }
+
+            public bool IsEmpty
+            {
+                get
+                {
+                    if (_reference == null)
+                        return false;
+                    if (_reference.TryGetTarget(out var target))
+                        return true;
+                    _reference = null;
+                    return false;
+                }
+            }
+
+            public bool TryGetTarget([MaybeNullWhen(false)]out IWeakEventSubscriber<TEventArgs> target)
+            {
+                if (_reference == null)
+                {
+                    target = null!;
+                    return false;
+                }
+                return _reference.TryGetTarget(out target);
+            }
+
+            public bool Equals(IWeakEventSubscriber<TEventArgs> r)
+            {
+                if (_reference == null || r.GetHashCode() != _hashCode)
+                    return false;
+                return _reference.TryGetTarget(out var target) && target == r;
+            }
+        }
+
+        private Entry[] _data =
+            new Entry[16];
         private int _count;
         private readonly Action _unsubscribe;
         private bool _compactScheduled;
+        private int _removedSinceLastCompact;
 
         public Subscription(WeakEvent<TSender, TEventArgs> ev, TSender target)
         {
@@ -71,17 +122,17 @@ public class WeakEvent<TSender, TEventArgs> : WeakEvent where TEventArgs : Event
             _ev._subscriptions.Remove(_target);
         }
 
-        public void Add(WeakReference<IWeakEventSubscriber<TEventArgs>> s)
+        public void Add(IWeakEventSubscriber<TEventArgs> s)
         {
             if (_count == _data.Length)
             {
                 //Extend capacity
-                var extendedData = new WeakReference<IWeakEventSubscriber<TEventArgs>>?[_data.Length * 2];
+                var extendedData = new Entry[_data.Length * 2];
                 Array.Copy(_data, extendedData, _data.Length);
                 _data = extendedData;
             }
 
-            _data[_count] = s;
+            _data[_count] = new(s);
             _count++;
         }
 
@@ -93,16 +144,21 @@ public class WeakEvent<TSender, TEventArgs> : WeakEvent where TEventArgs : Event
             {
                 var reference = _data[c];
 
-                if (reference != null && reference.TryGetTarget(out var instance) && instance == s)
+                if (reference.Equals(s))
                 {
-                    _data[c] = null;
+                    _data[c] = default;
                     removed = true;
+                    break;
                 }
             }
 
             if (removed)
             {
+                _removedSinceLastCompact++;
                 ScheduleCompact();
+                
+                if (_removedSinceLastCompact > 500)
+                    Compact();
             }
         }
 
@@ -116,18 +172,21 @@ public class WeakEvent<TSender, TEventArgs> : WeakEvent where TEventArgs : Event
         
         void Compact()
         {
+            if(!_compactScheduled || _removedSinceLastCompact == 0)
+                return;
             _compactScheduled = false;
+            _removedSinceLastCompact = 0;
             int empty = -1;
             for (var c = 0; c < _count; c++)
             {
-                var r = _data[c];
+                ref var r = ref _data[c];
                 //Mark current index as first empty
-                if (r == null && empty == -1)
+                if (r.IsEmpty && empty == -1)
                     empty = c;
                 //If current element isn't null and we have an empty one
-                if (r != null && empty != -1)
+                if (!r.IsEmpty && empty != -1)
                 {
-                    _data[c] = null;
+                    _data[c] = default;
                     _data[empty] = r;
                     empty++;
                 }
@@ -145,7 +204,7 @@ public class WeakEvent<TSender, TEventArgs> : WeakEvent where TEventArgs : Event
             for (var c = 0; c < _count; c++)
             {
                 var r = _data[c];
-                if (r?.TryGetTarget(out var sub) == true)
+                if (r.TryGetTarget(out var sub))
                     sub!.OnEvent(_target, _ev, eventArgs);
                 else
                     needCompact = true;