EnumerableConditionalWeakTable and WeakEvent (#74160)

src/EditorFeatures/CSharpTest/Workspaces/WorkspaceTests_EditorFeatures.cs

@@ -1520,7 +1520,7 @@ public void TestOptionChangedHandlerInvokedAfterCurrentSolutionChanged()
         primaryWorkspace.GlobalOptions.RemoveOptionChangedHandler(this, OptionService_OptionChanged);
         return;
 
-        void OptionService_OptionChanged(object sender, OptionChangedEventArgs e)
+        void OptionService_OptionChanged(object sender, object target, OptionChangedEventArgs e)
         {
             // CurrentSolution has been updated when the event fires.
 

src/EditorFeatures/Core/Shared/Tagging/EventSources/TaggerEventSources.OptionChangedEventSource.cs

@@ -21,7 +21,7 @@ public override void Disconnect()
             globalOptions.RemoveOptionChangedHandler(this, OnGlobalOptionChanged);
         }
 
-        private void OnGlobalOptionChanged(object? sender, OptionChangedEventArgs e)
+        private void OnGlobalOptionChanged(object sender, object target, OptionChangedEventArgs e)
         {
             if (e.HasOption(predicate))
             {

src/EditorFeatures/Test/Options/GlobalOptionsTests.cs

@@ -92,10 +92,10 @@ public void SetGlobalOption(OptionKey2 optionKey, object? value)
         public bool SetGlobalOptions(ImmutableArray<KeyValuePair<OptionKey2, object?>> options)
             => throw new NotImplementedException();
 
-        public void AddOptionChangedHandler(object target, EventHandler<OptionChangedEventArgs> handler)
+        public void AddOptionChangedHandler(object target, WeakEventHandler<OptionChangedEventArgs> handler)
             => throw new NotImplementedException();
 
-        public void RemoveOptionChangedHandler(object target, EventHandler<OptionChangedEventArgs> handler)
+        public void RemoveOptionChangedHandler(object target, WeakEventHandler<OptionChangedEventArgs> handler)
             => throw new NotImplementedException();
 
         #endregion

src/LanguageServer/Protocol/Features/Diagnostics/DiagnosticAnalyzerService.cs

@@ -50,7 +50,7 @@ public DiagnosticAnalyzerService(
             _diagnosticsRefresher = diagnosticsRefresher;
             _createIncrementalAnalyzer = CreateIncrementalAnalyzerCallback;
 
-            globalOptions.AddOptionChangedHandler(this, (_, e) =>
+            globalOptions.AddOptionChangedHandler(this, (_, _, e) =>
             {
                 if (e.HasOption(IsGlobalOptionAffectingDiagnostics))
                 {

src/LanguageServer/Protocol/Handler/CodeLens/CodeLensRefreshQueue.cs

@@ -38,7 +38,7 @@ protected override string GetWorkspaceRefreshName()
             return Methods.WorkspaceCodeLensRefreshName;
         }
 
-        private void OnOptionChanged(object? sender, OptionChangedEventArgs e)
+        private void OnOptionChanged(object sender, object target, OptionChangedEventArgs e)
         {
             if (e.HasOption(static option => option.Equals(LspOptionsStorage.LspEnableReferencesCodeLens) || option.Equals(LspOptionsStorage.LspEnableTestsCodeLens)))
             {

src/LanguageServer/Protocol/Handler/InlayHint/InlayHintRefreshQueue.cs

@@ -31,7 +31,7 @@ public override void Dispose()
             _globalOptionService.RemoveOptionChangedHandler(this, OnOptionChanged);
         }
 
-        private void OnOptionChanged(object? sender, OptionChangedEventArgs e)
+        private void OnOptionChanged(object sender, object target, OptionChangedEventArgs e)
         {
             if (e.HasOption(static option =>
                     option.Equals(InlineHintsOptionsStorage.EnabledForParameters) ||

src/Tools/ExternalAccess/Razor/RazorGlobalOptions.cs

@@ -11,6 +11,7 @@
 using System.Collections.Immutable;
 using System.Collections.Generic;
 using System.Diagnostics.CodeAnalysis;
+using Roslyn.Utilities;
 
 namespace Microsoft.CodeAnalysis.ExternalAccess.Razor
 {
@@ -63,8 +64,8 @@ public T GetOption<T>(PerLanguageOption2<T> option, string languageName)
             public void SetGlobalOption<T>(PerLanguageOption2<T> option, string language, T value) => throw new NotImplementedException();
             public void SetGlobalOption(OptionKey2 optionKey, object? value) => throw new NotImplementedException();
             public bool SetGlobalOptions(ImmutableArray<KeyValuePair<OptionKey2, object?>> options) => throw new NotImplementedException();
-            public void AddOptionChangedHandler(object target, EventHandler<OptionChangedEventArgs> handler) => throw new NotImplementedException();
-            public void RemoveOptionChangedHandler(object target, EventHandler<OptionChangedEventArgs> handler) => throw new NotImplementedException();
+            public void AddOptionChangedHandler(object target, WeakEventHandler<OptionChangedEventArgs> handler) => throw new NotImplementedException();
+            public void RemoveOptionChangedHandler(object target, WeakEventHandler<OptionChangedEventArgs> handler) => throw new NotImplementedException();
 
             bool IOptionsReader.TryGetOption<T>(OptionKey2 optionKey, out T value)
             {

src/VisualStudio/Core/Def/InheritanceMargin/InheritanceMarginViewMargin.cs

@@ -122,7 +122,7 @@ private void OnLayoutChanged(object sender, TextViewLayoutChangedEventArgs e)
         _refreshAllGlyphs = false;
     }
 
-    private void OnGlobalOptionChanged(object sender, OptionChangedEventArgs e)
+    private void OnGlobalOptionChanged(object sender, object target, OptionChangedEventArgs e)
     {
         if (e.HasOption(option =>
                 option.Equals(InheritanceMarginOptionsStorage.ShowInheritanceMargin) ||

src/VisualStudio/Core/Def/LanguageService/AbstractLanguageService`2.VsCodeWindowManager.cs

@@ -57,7 +57,7 @@ public VsCodeWindowManager(TLanguageService languageService, IVsCodeWindow codeW
         private void SetupView(IVsTextView view)
             => _languageService.SetupNewTextView(view);
 
-        private void GlobalOptionChanged(object sender, OptionChangedEventArgs e)
+        private void GlobalOptionChanged(object sender, object target, OptionChangedEventArgs e)
         {
             if (e.ChangedOptions.Any(item => item.key.Language == _languageService.RoslynLanguageName && item.key.Option.Equals(NavigationBarViewOptionsStorage.ShowNavigationBar)))
             {

src/VisualStudio/Core/Def/StackTraceExplorer/StackTraceExplorerCommandHandler.cs

@@ -106,7 +106,7 @@ private void UnadviseBroadcastMessages()
         }
     }
 
-    private void GlobalOptionChanged(object sender, OptionChangedEventArgs e)
+    private void GlobalOptionChanged(object sender, object target, OptionChangedEventArgs e)
     {
         bool? enabled = null;
         foreach (var (key, newValue) in e.ChangedOptions)

src/VisualStudio/Core/Def/SymbolSearch/AbstractDelayStartedService.cs

@@ -79,7 +79,7 @@ public void RegisterLanguage(string language)
         _optionChangedWorkQueue.AddWork();
     }
 
-    private void OnOptionChanged(object sender, OptionChangedEventArgs e)
+    private void OnOptionChanged(object sender, object target, OptionChangedEventArgs e)
         => _optionChangedWorkQueue.AddWork();
 
     private async ValueTask ProcessOptionChangesAsync(CancellationToken arg)

src/VisualStudio/Core/Def/Telemetry/FileLogger.cs

@@ -51,7 +51,7 @@ public FileLogger(IGlobalOptionService optionService)
     private static string GetLogFileName()
         => DateTime.Now.ToString(CultureInfo.InvariantCulture).Replace(' ', '_').Replace('/', '_').Replace(':', '_') + ".log";
 
-    private void OptionService_OptionChanged(object? sender, OptionChangedEventArgs e)
+    private void OptionService_OptionChanged(object sender, object target, OptionChangedEventArgs e)
     {
         foreach (var (key, newValue) in e.ChangedOptions)
         {

src/VisualStudio/VisualStudioDiagnosticsToolWindow/OptionPages/ForceLowMemoryMode.cs

@@ -27,7 +27,7 @@ public ForceLowMemoryMode(IGlobalOptionService globalOptions)
             RefreshFromSettings();
         }
 
-        private void Options_OptionChanged(object sender, OptionChangedEventArgs e)
+        private void Options_OptionChanged(object sender, object target, OptionChangedEventArgs e)
         {
             if (e.HasOption(static option => option.Equals(Enabled) || option.Equals(SizeInMegabytes)))
             {

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

@@ -254,12 +254,12 @@ public bool RefreshOption(OptionKey2 optionKey, object? newValue)
         return true;
     }
 
-    public void AddOptionChangedHandler(object target, EventHandler<OptionChangedEventArgs> handler)
+    public void AddOptionChangedHandler(object target, WeakEventHandler<OptionChangedEventArgs> handler)
     {
         _optionChanged.AddHandler(target, handler);
     }
 
-    public void RemoveOptionChangedHandler(object target, EventHandler<OptionChangedEventArgs> handler)
+    public void RemoveOptionChangedHandler(object target, WeakEventHandler<OptionChangedEventArgs> handler)
     {
         _optionChanged.RemoveHandler(target, handler);
     }

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

@@ -5,6 +5,7 @@
 using System;
 using System.Collections.Generic;
 using System.Collections.Immutable;
+using Roslyn.Utilities;
 
 namespace Microsoft.CodeAnalysis.Options;
 
@@ -65,7 +66,7 @@ internal interface IGlobalOptionService : IOptionsReader
     /// </remarks>
     bool RefreshOption(OptionKey2 optionKey, object? newValue);
 
-    void AddOptionChangedHandler(object target, EventHandler<OptionChangedEventArgs> handler);
+    void AddOptionChangedHandler(object target, WeakEventHandler<OptionChangedEventArgs> handler);
 
-    void RemoveOptionChangedHandler(object target, EventHandler<OptionChangedEventArgs> handler);
+    void RemoveOptionChangedHandler(object target, WeakEventHandler<OptionChangedEventArgs> handler);
 }

src/Workspaces/Core/Portable/Workspace/Workspace.TextTracker.cs

@@ -35,7 +35,7 @@ internal TextTracker(
             _onChangedHandler = onChangedHandler;
 
             // use weak event so TextContainer cannot accidentally keep workspace alive.
-            _weakOnTextChanged = WeakEventHandler<TextChangeEventArgs>.Create(this, (target, sender, args) => target.OnTextChanged(sender, args));
+            _weakOnTextChanged = EventHandlerFactory<TextChangeEventArgs>.CreateWeakHandler(this, (target, sender, args) => target.OnTextChanged(sender, args));
         }
 
         public void Connect()

src/Workspaces/CoreTest/UtilityTest/WeakEventTests.cs

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

src/Workspaces/CoreTest/WorkspaceServiceTests/GlobalOptionServiceTests.cs

@@ -172,7 +172,7 @@ public void GlobalOptions()
 
             var events = new List<OptionChangedEventArgs>();
 
-            var handler = new EventHandler<OptionChangedEventArgs>((_, e) => events.Add(e));
+            var handler = new WeakEventHandler<OptionChangedEventArgs>((_, _, e) => events.Add(e));
             globalOptions.AddOptionChangedHandler(this, handler);
 
             var values = globalOptions.GetOptions([new OptionKey2(option1), new OptionKey2(option2)]);

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

@@ -561,13 +561,14 @@
     <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" />
     <Compile Include="$(MSBuildThisFileDirectory)Utilities\TaskExtensions.cs" />
     <Compile Include="$(MSBuildThisFileDirectory)Utilities\TaskFactoryExtensions.cs" />
     <Compile Include="$(MSBuildThisFileDirectory)Utilities\TopologicalSorter.cs" />
-    <Compile Include="$(MSBuildThisFileDirectory)Utilities\WeakEventHandler.cs" />
+    <Compile Include="$(MSBuildThisFileDirectory)Utilities\EventHandlerFactory.cs" />
   </ItemGroup>
   <ItemGroup>
     <Compile Include="$(MSBuildThisFileDirectory)Extensions\AccessibilityUtilities.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/WeakEventHandler.cs → src/Workspaces/SharedUtilitiesAndExtensions/Compiler/Core/Utilities/EventHandlerFactory.cs

@@ -6,15 +6,15 @@
 
 namespace Roslyn.Utilities;
 
-internal static class WeakEventHandler<TArgs>
+internal static class EventHandlerFactory<TArgs>
 {
     /// <summary>
     /// Creates an event handler that holds onto the target weakly.
     /// </summary>
     /// <param name="target">The target that is held weakly, and passed as an argument to the invoker.</param>
     /// <param name="invoker">An action that will receive the event arguments as well as the target instance. 
     /// The invoker itself must not capture any state.</param>
-    public static EventHandler<TArgs> Create<TTarget>(TTarget target, Action<TTarget, object?, TArgs> invoker)
+    public static EventHandler<TArgs> CreateWeakHandler<TTarget>(TTarget target, Action<TTarget, object?, TArgs> invoker)
         where TTarget : class
     {
         var weakTarget = new WeakReference<TTarget>(target);