Fix generators not being added or removed correctly

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

@@ -824,7 +824,12 @@ private async Task<CompilationInfo> FinalizeCompilationAsync(
                     {
                         using var generatedDocumentsBuilder = new TemporaryArray<SourceGeneratedDocumentState>();
 
-                        if (ProjectState.SourceGenerators.Any())
+                        if (!ProjectState.SourceGenerators.Any())
+                        {
+                            // We don't have any generators, so if we have a compilation from a previous run with generated files, we definitely can't use it anymore
+                            compilationWithStaleGeneratedTrees = null;
+                        }
+                        else // we have a generator
                         {
                             // If we don't already have a generator driver, we'll have to create one from scratch
                             if (generatorInfo.Driver == null)

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

@@ -1031,8 +1031,15 @@ public SolutionState WithProjectAnalyzerReferences(ProjectId projectId, IEnumera
             // we changed, rather than creating an entire new generator driver from scratch and rerunning all generators, is cheaper
             // in the end. This was written without data backing up that assumption, so if a profile indicates to the contrary,
             // this could be changed.
-            var addedReferences = newProject.AnalyzerReferences.Except(oldProject.AnalyzerReferences).ToImmutableArray();
-            var removedReferences = oldProject.AnalyzerReferences.Except(newProject.AnalyzerReferences).ToImmutableArray();
+            //
+            // When we're comparing AnalyzerReferences, we'll compare with reference equality; AnalyzerReferences like AnalyzerFileReference
+            // may implement their own equality, but that can result in things getting out of sync: two references that are value equal can still
+            // have their own generator instances; it's important that as we're adding and removing references that are value equal that we
+            // still update with the correct generator instances that are coming from the new reference that is actually held in the project state from above.
+            // An alternative approach would be to call oldProject.WithAnalyzerReferences keeping all the references in there that are value equal the same,
+            // but this avoids any surprises where other components calling WithAnalyzerReferences might not expect that.
+            var addedReferences = newProject.AnalyzerReferences.Except<AnalyzerReference>(oldProject.AnalyzerReferences, ReferenceEqualityComparer.Instance).ToImmutableArray();
+            var removedReferences = oldProject.AnalyzerReferences.Except<AnalyzerReference>(newProject.AnalyzerReferences, ReferenceEqualityComparer.Instance).ToImmutableArray();
 
             return ForkProject(
                 newProject,

src/Workspaces/CoreTest/SolutionTests/SolutionWithSourceGeneratorTests.cs

@@ -2,12 +2,14 @@
 // 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.Immutable;
 using System.Linq;
 using System.Text;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.CodeAnalysis.CSharp;
+using Microsoft.CodeAnalysis.Diagnostics;
 using Microsoft.CodeAnalysis.Shared.Extensions;
 using Microsoft.CodeAnalysis.Test.Utilities;
 using Microsoft.CodeAnalysis.Text;
@@ -62,7 +64,56 @@ public async Task SourceGeneratorBasedOnAdditionalFileGeneratesSyntaxTrees(
         }
 
         [Fact]
-        public async Task WithReferencesMethodCorrectlyUpdatesRunningGenerators()
+        [WorkItem(1655835, "https://devdiv.visualstudio.com/DevDiv/_workitems/edit/1655835")]
+        public async Task WithReferencesMethodCorrectlyUpdatesWithEqualReferences()
+        {
+            using var workspace = CreateWorkspace();
+
+            // AnalyzerReferences may implement equality (AnalyezrFileReference does), and we want to make sure if we substitute out one
+            // reference with another reference that's equal, we correctly update generators. We'll have the underlying generators
+            // be different since two AnalyzerFileReferences that are value equal but different instances would have their own generators as well.
+            const string SharedPath = "Z:\\Generator.dll";
+            ISourceGenerator CreateGenerator() => new SingleFileTestGenerator("// StaticContent", hintName: "generated");
+
+            var analyzerReference1 = new TestGeneratorReferenceWithFilePathEquality(CreateGenerator(), SharedPath);
+            var analyzerReference2 = new TestGeneratorReferenceWithFilePathEquality(CreateGenerator(), SharedPath);
+
+            var project = AddEmptyProject(workspace.CurrentSolution)
+                .AddAnalyzerReference(analyzerReference1);
+
+            Assert.Single((await project.GetRequiredCompilationAsync(CancellationToken.None)).SyntaxTrees);
+
+            // Go from one analyzer reference to the other
+            project = project.WithAnalyzerReferences(new[] { analyzerReference2 });
+
+            Assert.Single((await project.GetRequiredCompilationAsync(CancellationToken.None)).SyntaxTrees);
+
+            // Now remove and confirm that we don't have any files
+            project = project.WithAnalyzerReferences(SpecializedCollections.EmptyEnumerable<AnalyzerReference>());
+
+            Assert.Empty((await project.GetRequiredCompilationAsync(CancellationToken.None)).SyntaxTrees);
+        }
+
+        private class TestGeneratorReferenceWithFilePathEquality : TestGeneratorReference, IEquatable<AnalyzerReference>
+        {
+            public TestGeneratorReferenceWithFilePathEquality(ISourceGenerator generator, string analyzerFilePath) : base(generator)
+            {
+                FullPath = analyzerFilePath;
+            }
+
+            public override bool Equals(object? obj) => Equals(obj as AnalyzerReference);
+            public override string FullPath { get; }
+            public override int GetHashCode() => this.FullPath.GetHashCode();
+
+            public bool Equals(AnalyzerReference? other)
+            {
+                return other is TestGeneratorReferenceWithFilePathEquality otherReference &&
+                    this.FullPath == otherReference.FullPath;
+            }
+        }
+
+        [Fact]
+        public async Task WithReferencesMethodCorrectlyAddsAndRemovesRunningGenerators()
         {
             using var workspace = CreateWorkspace();
 

src/Workspaces/CoreTestUtilities/TestGeneratorReference.cs

@@ -13,7 +13,7 @@ namespace Roslyn.Test.Utilities
     /// A simple deriviation of <see cref="AnalyzerReference"/> that returns the source generator
     /// passed, for ease in unit tests.
     /// </summary>
-    public sealed class TestGeneratorReference : AnalyzerReference, IChecksummedObject
+    public class TestGeneratorReference : AnalyzerReference, IChecksummedObject
     {
         private readonly ISourceGenerator _generator;
         private readonly Checksum _checksum;