Correctly update the list of generators for a call to WithAnalyzerRef…

…erences

When we add or remove generator refrences, we take the previous
GeneratorDriver and incrementally update it, adding or removing the
new generators. If you called WithAnalyzerReferences, however, we
completely failed to do this, so you'd be running with the incorrect
list of generators.

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

@@ -236,15 +236,17 @@ public override Task<Compilation> TransformCompilationAsync(Compilation oldCompi
                 public override bool CanUpdateCompilationWithStaleGeneratedTreesIfGeneratorsGiveSameOutput => true;
             }
 
-            internal sealed class AddAnalyzerReferencesAction : CompilationAndGeneratorDriverTranslationAction
+            internal sealed class AddOrRemoveAnalyzerReferencesAction : CompilationAndGeneratorDriverTranslationAction
             {
-                private readonly ImmutableArray<AnalyzerReference> _analyzerReferences;
                 private readonly string _language;
+                private readonly ImmutableArray<AnalyzerReference> _referencesToAdd;
+                private readonly ImmutableArray<AnalyzerReference> _referencesToRemove;
 
-                public AddAnalyzerReferencesAction(ImmutableArray<AnalyzerReference> analyzerReferences, string language)
+                public AddOrRemoveAnalyzerReferencesAction(string language, ImmutableArray<AnalyzerReference> referencesToAdd = default, ImmutableArray<AnalyzerReference> referencesToRemove = default)
                 {
-                    _analyzerReferences = analyzerReferences;
                     _language = language;
+                    _referencesToAdd = referencesToAdd;
+                    _referencesToRemove = referencesToRemove;
                 }
 
                 // Changing analyzer references doesn't change the compilation directly, so we can "apply" the
@@ -254,28 +256,17 @@ public AddAnalyzerReferencesAction(ImmutableArray<AnalyzerReference> analyzerRef
 
                 public override GeneratorDriver? TransformGeneratorDriver(GeneratorDriver generatorDriver)
                 {
-                    return generatorDriver.AddGenerators(_analyzerReferences.SelectMany(r => r.GetGenerators(_language)).ToImmutableArray());
-                }
-            }
-
-            internal sealed class RemoveAnalyzerReferencesAction : CompilationAndGeneratorDriverTranslationAction
-            {
-                private readonly ImmutableArray<AnalyzerReference> _analyzerReferences;
-                private readonly string _language;
+                    if (!_referencesToRemove.IsDefaultOrEmpty)
+                    {
+                        generatorDriver = generatorDriver.RemoveGenerators(_referencesToRemove.SelectMany(r => r.GetGenerators(_language)).ToImmutableArray());
+                    }
 
-                public RemoveAnalyzerReferencesAction(ImmutableArray<AnalyzerReference> analyzerReferences, string language)
-                {
-                    _analyzerReferences = analyzerReferences;
-                    _language = language;
-                }
+                    if (!_referencesToAdd.IsDefaultOrEmpty)
+                    {
+                        generatorDriver = generatorDriver.AddGenerators(_referencesToAdd.SelectMany(r => r.GetGenerators(_language)).ToImmutableArray());
+                    }
 
-                // Changing analyzer references doesn't change the compilation directly, so we can "apply" the
-                // translation (which is a no-op). Since we use a 'false' here to mean that it's not worth keeping
-                // the compilation with stale trees around, answering true is still important.
-                public override bool CanUpdateCompilationWithStaleGeneratedTreesIfGeneratorsGiveSameOutput => true;
-                public override GeneratorDriver? TransformGeneratorDriver(GeneratorDriver generatorDriver)
-                {
-                    return generatorDriver.RemoveGenerators(_analyzerReferences.SelectMany(r => r.GetGenerators(_language)).ToImmutableArray());
+                    return generatorDriver;
                 }
             }
 

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

@@ -1063,7 +1063,7 @@ public SolutionState AddAnalyzerReferences(ProjectId projectId, ImmutableArray<A
 
             return ForkProject(
                 oldProject.WithAnalyzerReferences(newReferences),
-                new CompilationAndGeneratorDriverTranslationAction.AddAnalyzerReferencesAction(analyzerReferences, oldProject.Language));
+                new CompilationAndGeneratorDriverTranslationAction.AddOrRemoveAnalyzerReferencesAction(oldProject.Language, referencesToAdd: analyzerReferences));
         }
 
         /// <summary>
@@ -1082,7 +1082,7 @@ public SolutionState RemoveAnalyzerReference(ProjectId projectId, AnalyzerRefere
 
             return ForkProject(
                 oldProject.WithAnalyzerReferences(newReferences),
-                new CompilationAndGeneratorDriverTranslationAction.RemoveAnalyzerReferencesAction(ImmutableArray.Create(analyzerReference), oldProject.Language));
+                new CompilationAndGeneratorDriverTranslationAction.AddOrRemoveAnalyzerReferencesAction(oldProject.Language, referencesToRemove: ImmutableArray.Create(analyzerReference)));
         }
 
         /// <summary>
@@ -1098,7 +1098,16 @@ public SolutionState WithProjectAnalyzerReferences(ProjectId projectId, IEnumera
                 return this;
             }
 
-            return ForkProject(newProject);
+            // The .Except() methods here aren't going to terribly cheap, but the assumption is adding or removing just the generators
+            // 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();
+
+            return ForkProject(
+                newProject,
+                new CompilationAndGeneratorDriverTranslationAction.AddOrRemoveAnalyzerReferencesAction(oldProject.Language, referencesToAdd: addedReferences, referencesToRemove: removedReferences));
         }
 
         /// <summary>

src/Workspaces/CoreTest/SolutionTests/SolutionWithSourceGeneratorTests.cs

@@ -69,6 +69,36 @@ public async Task SourceGeneratorBasedOnAdditionalFileGeneratesSyntaxTrees(
             Assert.Same(generatedTree, await generatedDocument.GetSyntaxTreeAsync());
         }
 
+        [Fact]
+        public async Task WithReferencesMethodCorrectlyUpdatesRunningGenerators()
+        {
+            using var workspace = CreateWorkspace();
+
+            // We always have a single generator in this test, and we add or remove a second one. This is critical
+            // to ensuring we correctly update our existing GeneratorDriver we may have from a prior run with the new
+            // generators passed to WithAnalyzerReferences. If we only swap from zero generators to one generator,
+            // we don't have a prior GeneratorDriver to update, since we don't make a GeneratorDriver if we have no generators.
+            // Similarly, once we go from one back to zero, we end up getting rid of our GeneratorDriver entirely since
+            // we have no need for it, as an optimization.
+            var generatorReferenceToKeep = new TestGeneratorReference(new SingleFileTestGenerator("// StaticContent", hintName: "generatorReferenceToKeep"));
+            var analyzerReferenceToAddAndRemove = new TestGeneratorReference(new SingleFileTestGenerator2("// More Static Content", hintName: "analyzerReferenceToAddAndRemove"));
+
+            var project = WithPreviewLanguageVersion(AddEmptyProject(workspace.CurrentSolution))
+                .AddAnalyzerReference(generatorReferenceToKeep);
+
+            Assert.Single((await project.GetRequiredCompilationAsync(CancellationToken.None)).SyntaxTrees);
+
+            // Go from one generator to two.
+            project = project.WithAnalyzerReferences(new[] { generatorReferenceToKeep, analyzerReferenceToAddAndRemove });
+
+            Assert.Equal(2, (await project.GetRequiredCompilationAsync(CancellationToken.None)).SyntaxTrees.Count());
+
+            // And go back to one
+            project = project.WithAnalyzerReferences(new[] { generatorReferenceToKeep });
+
+            Assert.Single((await project.GetRequiredCompilationAsync(CancellationToken.None)).SyntaxTrees);
+        }
+
         [Fact]
         public async Task IncrementalSourceGeneratorInvokedCorrectNumberOfTimes()
         {