Merge pull request #3167 from smoogipoo/improve-weaklist-performance

Improve performance/allocations of WeakList
ppy · Jan 10, 2020 · 500ca96 · 500ca96
2 parents 4839a9d + 0a2a466
commit 500ca96
Show file tree

Hide file tree

Showing 2 changed files with 71 additions and 43 deletions.
diff --git a/osu.Framework.Benchmarks/BenchmarkWeakList.cs b/osu.Framework.Benchmarks/BenchmarkWeakList.cs
@@ -3,59 +3,82 @@
 
 using System.Linq;
 using BenchmarkDotNet.Attributes;
-using NUnit.Framework;
 using osu.Framework.Lists;
 
 namespace osu.Framework.Benchmarks
 {
     public class BenchmarkWeakList : BenchmarkTest
     {
+        [Params(1, 10, 100, 1000)]
+        public int ItemCount { get; set; }
+
         private readonly object[] objects = new object[1000];
         private WeakList<object> weakList;
 
         public override void SetUp()
         {
-            for (int i = 0; i < 1000; i++)
+            for (int i = 0; i < ItemCount; i++)
                 objects[i] = new object();
         }
 
-        [IterationSetup]
-        [SetUp]
-        public void IterationSetup()
+        [Benchmark(Baseline = true)]
+        public void Add() => init();
+
+        [Benchmark]
+        public void RemoveOne()
         {
-            weakList = new WeakList<object>();
+            init();
 
-            foreach (var obj in objects)
-                weakList.Add(obj);
+            weakList.Remove(objects[^1]);
         }
 
         [Benchmark]
-        public void Add() => weakList.Add(new object());
+        public void RemoveAllIteratively()
+        {
+            init();
 
-        [Benchmark]
-        public void Remove() => weakList.Remove(objects[0]);
+            for (int i = 0; i < ItemCount; i++)
+                weakList.Remove(objects[i]);
+        }
 
         [Benchmark]
-        public void RemoveEach()
+        public void Clear()
         {
-            foreach (var obj in objects)
-                weakList.Remove(obj);
+            init();
+
+            weakList.Clear();
         }
 
         [Benchmark]
-        public void Clear() => weakList.Clear();
+        public bool Contains()
+        {
+            init();
 
-        [Benchmark]
-        public bool Contains() => weakList.Contains(objects[0]);
+            return weakList.Contains(objects[^1]);
+        }
 
         [Benchmark]
-        public object[] Enumerate() => weakList.ToArray();
+        public object[] AddAndEnumerate()
+        {
+            init();
+
+            return weakList.ToArray();
+        }
 
         [Benchmark]
         public object[] ClearAndEnumerate()
         {
+            init();
+
             weakList.Clear();
             return weakList.ToArray();
         }
+
+        private void init()
+        {
+            weakList = new WeakList<object>();
+            for (int i = 0; i < ItemCount; i++)
+                weakList.Add(objects[i]);
+        }
     }
 }
diff --git a/osu.Framework/Lists/WeakList.cs b/osu.Framework/Lists/WeakList.cs
@@ -4,6 +4,7 @@
 using System;
 using System.Collections;
 using System.Collections.Generic;
+using JetBrains.Annotations;
 
 namespace osu.Framework.Lists
 {
@@ -22,27 +23,31 @@ public class WeakList<T> : IWeakList<T>, IEnumerable<T>
 
         public void Remove(T item)
         {
-            foreach (var i in list)
+            for (int i = 0; i < list.Count; i++)
             {
-                if (i.Reference.TryGetTarget(out var obj) && obj == item)
-                {
-                    i.Invalidate();
-                    return;
-                }
+                if (list[i].Reference == null)
+                    continue;
+
+                if (!list[i].Reference.TryGetTarget(out var obj) || obj != item)
+                    continue;
+
+                list[i] = default;
+                break;
             }
         }
 
         public bool Remove(WeakReference<T> weakReference)
         {
             bool found = false;
 
-            foreach (var item in list)
+            for (int i = 0; i < list.Count; i++)
             {
-                if (item.Reference == weakReference)
-                {
-                    item.Invalidate();
-                    found = true;
-                }
+                if (list[i].Reference != weakReference)
+                    continue;
+
+                list[i] = default;
+                found = true;
+                break;
             }
 
             return found;
@@ -52,7 +57,10 @@ public bool Contains(T item)
         {
             foreach (var t in list)
             {
-                if (!t.Invalid && t.Reference.TryGetTarget(out var obj) && obj == item)
+                if (t.Reference == null)
+                    continue;
+
+                if (t.Reference.TryGetTarget(out var obj) && obj == item)
                     return true;
             }
 
@@ -63,7 +71,7 @@ public bool Contains(WeakReference<T> weakReference)
         {
             foreach (var t in list)
             {
-                if (!t.Invalid && t.Reference == weakReference)
+                if (t.Reference != null && t.Reference == weakReference)
                     return true;
             }
 
@@ -72,13 +80,13 @@ public bool Contains(WeakReference<T> weakReference)
 
         public void Clear()
         {
-            foreach (var item in list)
-                item.Invalidate();
+            for (int i = 0; i < list.Count; i++)
+                list[i] = default;
         }
 
         public Enumerator GetEnumerator()
         {
-            list.RemoveAll(item => item.Invalid || !item.Reference.TryGetTarget(out _));
+            list.RemoveAll(item => item.Reference == null || !item.Reference.TryGetTarget(out _));
 
             return new Enumerator(list);
         }
@@ -106,7 +114,7 @@ public bool MoveNext()
             {
                 while (++currentIndex < list.Count)
                 {
-                    if (list[currentIndex].Invalid || !list[currentIndex].Reference.TryGetTarget(out currentObject))
+                    if (list[currentIndex].Reference == null || !list[currentIndex].Reference.TryGetTarget(out currentObject))
                         continue;
 
                     return true;
@@ -132,23 +140,20 @@ public void Dispose()
             }
         }
 
-        internal class InvalidatableWeakReference
+        internal readonly struct InvalidatableWeakReference
         {
+            [CanBeNull]
             public readonly WeakReference<T> Reference;
 
-            public bool Invalid { get; private set; }
-
-            public InvalidatableWeakReference(T reference)
+            public InvalidatableWeakReference([CanBeNull] T reference)
                 : this(new WeakReference<T>(reference))
             {
             }
 
-            public InvalidatableWeakReference(WeakReference<T> weakReference)
+            public InvalidatableWeakReference([CanBeNull] WeakReference<T> weakReference)
             {
                 Reference = weakReference;
             }
-
-            public void Invalidate() => Invalid = true;
         }
     }
 }