Merge pull request #45435 from sharwell/reference-holder

Track unrooted symbols in HashSet instead of ConditionalWeakTable

src/Workspaces/Core/Portable/Utilities/ReferenceHolder`1.cs

+// 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.
+
+#nullable enable
+
+using System;
+using System.Diagnostics.CodeAnalysis;
+
+namespace Roslyn.Utilities
+{
+    internal readonly struct ReferenceHolder<T> : IEquatable<ReferenceHolder<T>>
+        where T : class?
+    {
+        [AllowNull, MaybeNull]
+        private readonly T _strongReference;
+        private readonly WeakReference<T>? _weakReference;
+        private readonly int _hashCode;
+
+        private ReferenceHolder(T strongReference)
+        {
+            _strongReference = strongReference;
+            _weakReference = null;
+            _hashCode = 0;
+        }
+
+        private ReferenceHolder(WeakReference<T> weakReference, int hashCode)
+        {
+            _strongReference = null;
+            _weakReference = weakReference;
+            _hashCode = hashCode;
+        }
+
+        public static ReferenceHolder<T> Strong(T value)
+            => new ReferenceHolder<T>(value);
+
+        public static ReferenceHolder<T> Weak(T value)
+        {
+            if (value is null)
+            {
+                // Track this as a strong reference so we know 'Equals' should only look at values originally null.
+                return Strong(value);
+            }
+
+            return new ReferenceHolder<T>(new WeakReference<T>(value), ReferenceEqualityComparer.GetHashCode(value));
+        }
+
+        [return: MaybeNull]
+        public T TryGetTarget()
+        {
+            if (_weakReference is object)
+                return _weakReference.GetTarget();
+
+            return _strongReference;
+        }
+
+        public override bool Equals(object? obj)
+        {
+            return obj is ReferenceHolder<T> other
+                && Equals(other);
+        }
+
+        public bool Equals(ReferenceHolder<T> other)
+        {
+            var x = TryGetTarget();
+            var y = other.TryGetTarget();
+            if (x is null)
+            {
+                if (_weakReference is object)
+                {
+                    // 'x' is a weak reference that was collected. Verify 'y' is a collected weak reference with the
+                    // same runtime hash code. This code path can fail in an edge case where the references to two
+                    // different objects have both been collected, but the runtime hash codes for the objects were
+                    // equal. Callers can ensure this case is not encountered by structuring equality checks such that
+                    // at least one of the objects is alive at the time Equals is called.
+                    return y is null && other._weakReference is object && _hashCode == other._hashCode;
+                }
+                else
+                {
+                    // Null values are equal iff both were originally references to null.
+                    return y is null && other._weakReference is null;
+                }
+            }
+
+            // Intentional reference equality check
+            return x == y;
+        }
+
+        public override int GetHashCode()
+        {
+            if (_weakReference is object)
+                return _hashCode;
+
+            return ReferenceEqualityComparer.GetHashCode(_strongReference);
+        }
+
+        internal static class TestAccessor
+        {
+            /// <summary>
+            /// Creates a <see cref="ReferenceHolder{T}"/> for a weakly-held reference that has since been collected.
+            /// </summary>
+            /// <param name="hashCode">The hash code of the collected value.</param>
+            /// <returns>A weak <see cref="ReferenceHolder{T}"/> which was already collected.</returns>
+            public static ReferenceHolder<T> ReleasedWeak(int hashCode)
+                => new ReferenceHolder<T>(new WeakReference<T>(null!), hashCode);
+        }
+    }
+}

src/Workspaces/Core/Portable/Utilities/WeakSet`1.cs

+// 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.
+
+#nullable enable
+
+using System.Collections.Generic;
+
+namespace Roslyn.Utilities
+{
+    /// <summary>
+    /// A simple collection of values held as weak references. Objects in the set are compared by reference equality.
+    /// </summary>
+    /// <typeparam name="T">The type of object stored in the set.</typeparam>
+    internal sealed class WeakSet<T>
+        where T : class?
+    {
+        private readonly HashSet<ReferenceHolder<T>> _values = new HashSet<ReferenceHolder<T>>();
+
+        public WeakSet()
+        {
+        }
+
+        public bool Add(T value)
+        {
+            if (Contains(value))
+                return false;
+
+            return _values.Add(ReferenceHolder<T>.Weak(value));
+        }
+
+        public bool Contains(T value)
+        {
+            return _values.Contains(ReferenceHolder<T>.Strong(value));
+        }
+    }
+}

src/Workspaces/Core/Portable/Workspace/Solution/SolutionState.CompilationTracker.State.cs

@@ -8,7 +8,6 @@
 using System.Collections.Immutable;
 using System.Diagnostics;
 using System.Linq;
-using System.Runtime.CompilerServices;
 using Microsoft.CodeAnalysis;
 using Roslyn.Utilities;
 
@@ -48,19 +47,12 @@ private class State
                 public TrackedGeneratorDriver GeneratorDriver { get; }
 
                 /// <summary>
-                /// Weak table of the assembly, module and dynamic symbols that this compilation tracker has created.
+                /// Weak set of the assembly, module and dynamic symbols that this compilation tracker has created.
                 /// This can be used to determine which project an assembly symbol came from after the fact.  This is
                 /// needed as the compilation an assembly came from can be GC'ed and further requests to get that
                 /// compilation (or any of it's assemblies) may produce new assembly symbols.
                 /// </summary>
-                /// <remarks>
-                /// Ideally this would just be <c>ConditionalWeakSet&lt;ISymbol&gt;</c>.  Effectively we just want to
-                /// hold onto the symbols as long as someone else is keeping them alive.  And we don't actually need
-                /// them to map to anything.  We just use their existence to know if our project was the project it came
-                /// from.  However, ConditionalWeakTable is the best tool we have, so we simulate a set by just using a
-                /// table and mapping the keys to the <see langword="null"/> value.
-                /// </remarks>
-                public readonly ConditionalWeakTable<ISymbol, object?>? UnrootedSymbolSet;
+                public readonly WeakSet<ISymbol>? UnrootedSymbolSet;
 
                 /// <summary>
                 /// Specifies whether <see cref="FinalCompilation"/> and all compilations it depends on contain full information or not. This can return
@@ -77,7 +69,7 @@ private class State
                     ValueSource<Optional<Compilation>>? compilation,
                     Compilation? declarationOnlyCompilation,
                     TrackedGeneratorDriver generatorDriver,
-                    ConditionalWeakTable<ISymbol, object?>? unrootedSymbolSet)
+                    WeakSet<ISymbol>? unrootedSymbolSet)
                 {
                     // Declaration-only compilations should never have any references
                     Contract.ThrowIfTrue(declarationOnlyCompilation != null && declarationOnlyCompilation.ExternalReferences.Any());
@@ -112,18 +104,18 @@ private class State
                         : (ValueSource<Optional<Compilation>>)new ConstantValueSource<Optional<Compilation>>(compilation);
                 }
 
-                public static ConditionalWeakTable<ISymbol, object?> GetUnrootedSymbols(Compilation compilation)
+                public static WeakSet<ISymbol> GetUnrootedSymbols(Compilation compilation)
                 {
-                    var result = new ConditionalWeakTable<ISymbol, object?>();
+                    var result = new WeakSet<ISymbol>();
 
                     var compAssembly = compilation.Assembly;
-                    result.Add(compAssembly, null);
+                    result.Add(compAssembly);
 
                     // The dynamic type is also unrooted (i.e. doesn't point back at the compilation or source
                     // assembly).  So we have to keep track of it so we can get back from it to a project in case the 
                     // underlying compilation is GC'ed.
                     if (compilation.Language == LanguageNames.CSharp)
-                        result.Add(compilation.DynamicType, null);
+                        result.Add(compilation.DynamicType);
 
                     foreach (var reference in compilation.References)
                     {
@@ -131,7 +123,7 @@ private class State
                         if (symbol == null)
                             continue;
 
-                        result.Add(symbol, null);
+                        result.Add(symbol);
                     }
 
                     return result;
@@ -214,7 +206,7 @@ private sealed class FinalState : State
                     Compilation compilationWithoutGeneratedFiles,
                     TrackedGeneratorDriver generatorDriver,
                     bool hasSuccessfullyLoaded,
-                    ConditionalWeakTable<ISymbol, object?>? compilationAssemblies)
+                    WeakSet<ISymbol>? compilationAssemblies)
                     : base(compilationWithoutGeneratedFilesSource,
                            compilationWithoutGeneratedFiles.Clone().RemoveAllReferences(),
                            generatorDriver,

src/Workspaces/Core/Portable/Workspace/Solution/SolutionState.CompilationTracker.cs

@@ -107,7 +107,7 @@ public bool ContainsAssemblyOrModuleOrDynamic(ISymbol symbol)
                              symbol.Kind == SymbolKind.DynamicType);
                 var state = this.ReadState();
                 var unrootedSymbolSet = state.UnrootedSymbolSet;
-                return unrootedSymbolSet != null && unrootedSymbolSet.TryGetValue(symbol, out _);
+                return unrootedSymbolSet?.Contains(symbol) ?? false;
             }
 
             /// <summary>

src/Workspaces/CoreTest/UtilityTest/ReferenceHolderTests.cs

+// 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.
+
+#nullable enable
+
+using System.Collections.Generic;
+using Roslyn.Utilities;
+using Xunit;
+
+namespace Microsoft.CodeAnalysis.UnitTests
+{
+    public class ReferenceHolderTests
+    {
+        [Fact]
+        public void SameStrongObjectsEqual()
+        {
+            var obj = new object();
+            var first = ReferenceHolder<object?>.Strong(obj);
+            var second = ReferenceHolder<object?>.Strong(obj);
+
+            VerifyEqual(first, second);
+        }
+
+        [Fact]
+        public void SameWeakObjectsEqual()
+        {
+            var obj = new object();
+            var first = ReferenceHolder<object?>.Weak(obj);
+            var second = ReferenceHolder<object?>.Weak(obj);
+
+            // 📝 There is no need for a GC.KeepAlive(obj) here. 'VerifyEqual' will produce correct results whether
+            // or not the object is still alive. When the object is alive, the equality path is the same as
+            // SameStrongObjectsEqual. When the object is not alive, the equality path is the same as
+            // ExpiredSameValuesEqual.
+            VerifyEqual(first, second);
+        }
+
+        [Fact]
+        public void SameMixedObjectsEqual()
+        {
+            var obj = new object();
+            var first = ReferenceHolder<object?>.Strong(obj);
+            var second = ReferenceHolder<object?>.Weak(obj);
+
+            VerifyEqual(first, second);
+        }
+
+        [Fact]
+        public void NullValuesEqual()
+        {
+            var first = ReferenceHolder<object?>.Strong(null);
+            var second = ReferenceHolder<object?>.Weak(null);
+
+            VerifyEqual(first, second);
+        }
+
+        [Fact]
+        public void ExpiredValueNotEqualToNull()
+        {
+            var strongNull = ReferenceHolder<object?>.Strong(null);
+            var weakNull = ReferenceHolder<object?>.Weak(null);
+            var expired = ReferenceHolder<object?>.TestAccessor.ReleasedWeak(hashCode: EqualityComparer<object?>.Default.GetHashCode(null));
+
+            Assert.Equal(strongNull.GetHashCode(), expired.GetHashCode());
+            VerifyNotEqual(strongNull, expired);
+            VerifyNotEqual(weakNull, expired);
+        }
+
+        [Fact]
+        public void ExpiredSameValuesEqual()
+        {
+            var first = ReferenceHolder<object?>.TestAccessor.ReleasedWeak(hashCode: 1);
+            var second = ReferenceHolder<object?>.TestAccessor.ReleasedWeak(hashCode: 1);
+
+            Assert.Null(first.TryGetTarget());
+            Assert.Null(second.TryGetTarget());
+            VerifyEqual(first, second);
+        }
+
+        [Fact]
+        public void ExpiredDifferentValuesNotEqual()
+        {
+            var first = ReferenceHolder<object?>.TestAccessor.ReleasedWeak(hashCode: 1);
+            var second = ReferenceHolder<object?>.TestAccessor.ReleasedWeak(hashCode: 2);
+
+            Assert.Null(first.TryGetTarget());
+            Assert.Null(second.TryGetTarget());
+            VerifyNotEqual(first, second);
+        }
+
+        private static void VerifyEqual<T>(ReferenceHolder<T> x, ReferenceHolder<T> y)
+            where T : class?
+        {
+            Assert.Equal(x.GetHashCode(), y.GetHashCode());
+            Assert.True(x.Equals(y));
+            Assert.True(y.Equals(x));
+        }
+
+        private static void VerifyNotEqual<T>(ReferenceHolder<T> x, ReferenceHolder<T> y)
+            where T : class?
+        {
+            Assert.False(x.Equals(y));
+            Assert.False(y.Equals(x));
+        }
+    }
+}