Merge pull request #59009 from dotnet/merges/release/dev17.1-to-relea…

…se/dev17.1-vs-deps

Merge release/dev17.1 to release/dev17.1-vs-deps

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

@@ -35,7 +35,7 @@ private static Project WithPreviewLanguageVersion(Project project)
         [Theory]
         [CombinatorialData]
         public async Task SourceGeneratorBasedOnAdditionalFileGeneratesSyntaxTrees(
-            bool fetchCompilationBeforeAddingGenerator,
+            bool fetchCompilationBeforeAddingAdditionalFile,
             bool useRecoverableTrees)
         {
             // This test is just the sanity test to make sure generators work at all. There's not a special scenario being
@@ -50,7 +50,7 @@ private static Project WithPreviewLanguageVersion(Project project)
             // when the file already exists, and when it is added incrementally.
             Compilation? originalCompilation = null;
 
-            if (fetchCompilationBeforeAddingGenerator)
+            if (fetchCompilationBeforeAddingAdditionalFile)
             {
                 originalCompilation = await project.GetRequiredCompilationAsync(CancellationToken.None);
             }
@@ -69,6 +69,36 @@ private static Project WithPreviewLanguageVersion(Project project)
             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()
         {