EnumerableConditionalWeakTable and WeakEvent

src/Workspaces/Core/Portable/Options/GlobalOptionService.cs

@@ -267,7 +267,7 @@ public void RemoveOptionChangedHandler(object target, EventHandler<OptionChanged
     private void RaiseOptionChangedEvent(ImmutableArray<(OptionKey2, object?)> changedOptions)
     {
         Debug.Assert(!changedOptions.IsEmpty);
-        _optionChanged.RaiseEvent(this, new OptionChangedEventArgs(changedOptions));
+        _optionChanged.RaiseEvent(new OptionChangedEventArgs(changedOptions));
     }
 
     internal TestAccessor GetTestAccessor()

src/Workspaces/CoreTest/UtilityTest/WeakEventTests.cs

@@ -0,0 +1,41 @@
+using System;
+using System.Collections.Generic;
+using System.Linq;
+using System.Text;
+using System.Threading.Tasks;
+using Roslyn.Test.Utilities;
+using Roslyn.Utilities;
+using Xunit;
+
+namespace Microsoft.CodeAnalysis.UnitTests.UtilityTest;
+
+public class WeakEventTests
+{
+    [Fact]
+    public void AddAndRemove()
+    {
+        var e = new WeakEvent<int>();
+
+        var target = new List<int>();
+        var handler1 = new EventHandler<int>((sender, arg) => Assert.IsType<List<int>>(sender).Add(arg * 10));
+        var handler2 = new EventHandler<int>((sender, arg) => Assert.IsType<List<int>>(sender).Add(arg * 20));
+        var handler3 = new EventHandler<int>((sender, arg) => Assert.IsType<List<int>>(sender).Add(arg * 30));
+
+        e.AddHandler(target, handler1);
+        e.AddHandler(target, handler2);
+        e.AddHandler(target, handler3);
+
+        e.RemoveHandler(target, handler2);
+
+        e.RaiseEvent(1);
+
+        AssertEx.Equal([10, 30], target);
+        target.Clear();
+
+        e.RemoveHandler(target, handler1);
+        e.RemoveHandler(target, handler3);
+
+        e.RaiseEvent(1);
+        Assert.Empty(target);
+    }
+}

src/Workspaces/SharedUtilitiesAndExtensions/Compiler/Core/CompilerExtensions.projitems

@@ -561,6 +561,7 @@
     <Compile Include="$(MSBuildThisFileDirectory)Utilities\StringBreaker.cs" />
     <Compile Include="$(MSBuildThisFileDirectory)Utilities\ValueTaskExtensions.cs" />
     <Compile Include="$(MSBuildThisFileDirectory)Utilities\WeakEvent`1.cs" />
+    <Compile Include="$(MSBuildThisFileDirectory)Utilities\EnumerableConditionalWeakTable.cs" />
     <Compile Include="$(MSBuildThisFileDirectory)Utilities\WordSimilarityChecker.cs" />
     <Compile Include="$(MSBuildThisFileDirectory)Utilities\StringEscapeEncoder.cs" />
     <Compile Include="$(MSBuildThisFileDirectory)Utilities\SyntaxPath.cs" />

src/Workspaces/SharedUtilitiesAndExtensions/Compiler/Core/Utilities/EnumerableConditionalWeakTable.cs

@@ -0,0 +1,133 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+// See the LICENSE file in the project root for more information.
+
+using System;
+using System.Collections;
+using System.Collections.Generic;
+using System.Collections.Immutable;
+using System.Diagnostics.CodeAnalysis;
+using System.Runtime.CompilerServices;
+
+namespace Roslyn.Utilities;
+
+#if NET
+// Can't use global alias due to generic parameters. Extension types would do.
+
+internal readonly struct EnumerableConditionalWeakTable<TKey, TValue>() : IEnumerable<KeyValuePair<TKey, TValue>>
+    where TKey : class
+    where TValue : class
+{
+    private readonly ConditionalWeakTable<TKey, TValue> _table = new();
+
+    public object WriteLock => _table;
+
+    public bool TryGetValue(TKey key, [NotNullWhen(true)] out TValue? value)
+        => _table.TryGetValue(key, out value);
+
+    public void Add(TKey key, TValue value)
+        => _table.Add(key, value);
+
+    public void AddOrUpdate(TKey key, TValue value)
+        => _table.AddOrUpdate(key, value);
+
+    public bool Remove(TKey key)
+        => _table.Remove(key);
+
+    public IEnumerator<KeyValuePair<TKey, TValue>> GetEnumerator()
+        => ((IEnumerable<KeyValuePair<TKey, TValue>>)_table).GetEnumerator();
+
+    IEnumerator IEnumerable.GetEnumerator()
+        => GetEnumerator();
+}
+#else
+internal sealed class EnumerableConditionalWeakTable<TKey, TValue> : IEnumerable<KeyValuePair<TKey, TValue>>
+    where TKey : class
+    where TValue : class
+{
+    private sealed class Box(TKey key, TValue value)
+    {
+        public readonly TKey Key = key;
+        public readonly TValue Value = value;
+    }
+
+    private readonly ConditionalWeakTable<TKey, Box> _table = new();
+    private ImmutableList<WeakReference<Box>> _items = [];
+
+    public object WriteLock => _table;
+
+    public bool TryGetValue(TKey key, [NotNullWhen(true)] out TValue? value)
+    {
+        if (_table.TryGetValue(key, out var box))
+        {
+            value = box.Value;
+            return true;
+        }
+
+        value = null;
+        return false;
+    }
+
+    public void Add(TKey key, TValue value)
+    {
+        lock (WriteLock)
+        {
+            AddNoLock(key, value);
+
+            // clean up collected objects:
+            _items = _items.RemoveAll(WeakReferenceExtensions.IsNull);
+        }
+    }
+
+    public void AddOrUpdate(TKey key, TValue value)
+    {
+        lock (WriteLock)
+        {
+            _ = RemoveNoLock(key);
+            AddNoLock(key, value);
+        }
+    }
+
+    public bool Remove(TKey key)
+    {
+        lock (WriteLock)
+        {
+            return RemoveNoLock(key);
+        }
+    }
+
+    private void AddNoLock(TKey key, TValue value)
+    {
+        var box = new Box(key, value);
+        _table.Add(key, box);
+        _items = _items.Add(new WeakReference<Box>(box));
+    }
+
+    private bool RemoveNoLock(TKey key)
+    {
+        if (!_table.TryGetValue(key, out var box))
+        {
+            return false;
+        }
+
+        Contract.ThrowIfFalse(_table.Remove(key));
+        _items = _items.RemoveAll(item => !item.TryGetTarget(out var target) || ReferenceEquals(target, box));
+        return true;
+    }
+
+    public IEnumerator<KeyValuePair<TKey, TValue>> GetEnumerator()
+    {
+        foreach (var item in _items)
+        {
+            if (item.TryGetTarget(out var box))
+            {
+                yield return KeyValuePairUtil.Create(box.Key, box.Value);
+            }
+        }
+    }
+
+    IEnumerator IEnumerable.GetEnumerator()
+        => GetEnumerator();
+}
+#endif
+

src/Workspaces/SharedUtilitiesAndExtensions/Compiler/Core/Utilities/WeakEvent`1.cs

@@ -3,129 +3,63 @@
 // See the LICENSE file in the project root for more information.
 
 using System;
-using System.Collections.Immutable;
-using System.Runtime.CompilerServices;
 
 namespace Roslyn.Utilities;
 
-internal sealed class WeakEvent<TEventArgs>
+/// <summary>
+/// Implements an event that can be subscribed to without keeping the subscriber alive for the lifespan of 
+/// the object that declares <see cref="WeakEvent{TEventArgs}"/>.
+/// 
+/// Unlike <see cref="WeakEventHandler{TArgs}"/> the handlers may capture state, which makes the subscribers simpler
+/// and doesn't risk accidental leaks.
+/// </summary>
+internal readonly struct WeakEvent<TEventArgs>()
 {
     /// <summary>
     /// Each registered event handler has the lifetime of an associated owning object. This table ensures the weak
     /// references to the event handlers are not cleaned up while the owning object is still alive.
     /// </summary>
-    /// <remarks>
-    /// Access to this table is guarded by itself. For example:
-    /// <code>
-    /// <see langword="lock"/> (<see cref="_keepAliveTable"/>)
-    /// {
-    ///   // Read or write to <see cref="_keepAliveTable"/>.
-    /// }
-    /// </code>
-    /// </remarks>
-    private readonly ConditionalWeakTable<object, EventHandler<TEventArgs>> _keepAliveTable = new();
-
-    private ImmutableList<WeakReference<EventHandler<TEventArgs>>> _weakHandlers = [];
+    private readonly EnumerableConditionalWeakTable<object, EventHandler<TEventArgs>> _handlers = new();
 
     public void AddHandler(object target, EventHandler<TEventArgs> handler)
     {
-        lock (_keepAliveTable)
+        lock (_handlers.WriteLock)
         {
-            if (_keepAliveTable.TryGetValue(target, out var existingHandler))
+            if (_handlers.TryGetValue(target, out var existingHandler))
             {
-                // Combine the delegates and store the combination in the keep-alive table
-                var newHandler = existingHandler + handler;
-#if NET6_0_OR_GREATER
-                _keepAliveTable.AddOrUpdate(target, newHandler);
-#else
-                _keepAliveTable.Remove(target);
-                _keepAliveTable.Add(target, newHandler);
-#endif
+                _handlers.AddOrUpdate(target, existingHandler + handler);
             }
             else
             {
-                // This is the first handler for this target, so just store it in the table directly
-                _keepAliveTable.Add(target, handler);
+                _handlers.Add(target, handler);
             }
         }
-
-        ImmutableInterlocked.Update(
-            ref _weakHandlers,
-            static (weakHandlers, handler) =>
-            {
-                // Add a weak reference to the handler to the list of delegates to invoke for this event. During the
-                // mutation, also remove any handlers that were collected without being removed (and are now null).
-                return weakHandlers
-                    .RemoveAll(WeakReferenceExtensions.IsNull)
-                    .Add(new WeakReference<EventHandler<TEventArgs>>(handler));
-            },
-            handler);
     }
 
     public void RemoveHandler(object target, EventHandler<TEventArgs> handler)
     {
-        lock (_keepAliveTable)
+        lock (_handlers.WriteLock)
         {
-            if (_keepAliveTable.TryGetValue(target, out var existingHandler))
+            if (_handlers.TryGetValue(target, out var existingHandler))
             {
                 var newHandler = existingHandler - handler;
-                if (newHandler is null)
+                if (newHandler != null)
                 {
-                    _keepAliveTable.Remove(target);
+                    _handlers.AddOrUpdate(target, newHandler);
                 }
                 else
                 {
-#if NET6_0_OR_GREATER
-                    _keepAliveTable.AddOrUpdate(target, newHandler);
-#else
-                    _keepAliveTable.Remove(target);
-                    _keepAliveTable.Add(target, newHandler);
-#endif
+                    _handlers.Remove(target);
                 }
             }
         }
-
-        ImmutableInterlocked.Update(
-            ref _weakHandlers,
-            static (weakHandlers, handler) =>
-            {
-                // Remove any references to the handler from the list of delegates to invoke for this event. During the
-                // mutation, also remove any handlers that were collected without being removed (and are now null).
-                var builder = weakHandlers.ToBuilder();
-                for (var i = weakHandlers.Count - 1; i >= 0; i--)
-                {
-                    var weakHandler = weakHandlers[i];
-                    if (!weakHandler.TryGetTarget(out var target))
-                    {
-                        // This handler was collected
-                        builder.RemoveAt(i);
-                        continue;
-                    }
-
-                    var updatedHandler = target - handler;
-                    if (updatedHandler is null)
-                    {
-                        builder.RemoveAt(i);
-                    }
-                    else if (updatedHandler != target)
-                    {
-                        builder[i].SetTarget(updatedHandler);
-                    }
-                }
-
-                return builder.ToImmutable();
-            },
-            handler);
     }
 
-    public void RaiseEvent(object? sender, TEventArgs e)
+    public void RaiseEvent(TEventArgs e)
     {
-        foreach (var weakHandler in _weakHandlers)
+        foreach (var (target, handler) in _handlers)
         {
-            if (!weakHandler.TryGetTarget(out var handler))
-                continue;
-
-            handler(sender, e);
+            handler(target, e);
         }
     }
 }