Hold onto and reuse GeneratorDrivers

src/Workspaces/CSharp/Portable/Workspace/LanguageServices/CSharpCompilationFactoryService.cs

@@ -41,7 +41,7 @@ Compilation ICompilationFactoryService.CreateSubmissionCompilation(string assemb
         CompilationOptions ICompilationFactoryService.GetDefaultCompilationOptions()
             => s_defaultOptions;
 
-        GeneratorDriver? ICompilationFactoryService.CreateGeneratorDriver(ParseOptions parseOptions, ImmutableArray<ISourceGenerator> generators, AnalyzerConfigOptionsProvider optionsProvider, ImmutableArray<AdditionalText> additionalTexts)
+        GeneratorDriver ICompilationFactoryService.CreateGeneratorDriver(ParseOptions parseOptions, ImmutableArray<ISourceGenerator> generators, AnalyzerConfigOptionsProvider optionsProvider, ImmutableArray<AdditionalText> additionalTexts)
         {
             return CSharpGeneratorDriver.Create(generators, additionalTexts, (CSharpParseOptions)parseOptions, optionsProvider);
         }

src/Workspaces/Core/Portable/Workspace/Host/CompilationFactory/ICompilationFactoryService.cs

@@ -13,6 +13,6 @@ internal interface ICompilationFactoryService : ILanguageService
         Compilation CreateCompilation(string assemblyName, CompilationOptions options);
         Compilation CreateSubmissionCompilation(string assemblyName, CompilationOptions options, Type? hostObjectType);
         CompilationOptions GetDefaultCompilationOptions();
-        GeneratorDriver? CreateGeneratorDriver(ParseOptions parseOptions, ImmutableArray<ISourceGenerator> generators, AnalyzerConfigOptionsProvider optionsProvider, ImmutableArray<AdditionalText> additionalTexts);
+        GeneratorDriver CreateGeneratorDriver(ParseOptions parseOptions, ImmutableArray<ISourceGenerator> generators, AnalyzerConfigOptionsProvider optionsProvider, ImmutableArray<AdditionalText> additionalTexts);
     }
 }

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

@@ -3,6 +3,7 @@
 // See the LICENSE file in the project root for more information.
 
 using System;
+using System.Runtime.CompilerServices;
 using System.Threading;
 using Microsoft.CodeAnalysis.Text;
 
@@ -14,11 +15,18 @@ namespace Microsoft.CodeAnalysis.Diagnostics
     internal sealed class AdditionalTextWithState : AdditionalText
     {
         private readonly TextDocumentState _documentState;
+        private static readonly ConditionalWeakTable<TextDocumentState, AdditionalText> _additionalTextsForDocumentStates = new();
+        private static readonly ConditionalWeakTable<TextDocumentState, AdditionalText>.CreateValueCallback s_createAdditionalText = static ts => new AdditionalTextWithState(ts);
+
+        public static AdditionalText FromState(TextDocumentState state)
+        {
+            return _additionalTextsForDocumentStates.GetValue(state, s_createAdditionalText);
+        }
 
         /// <summary>
         /// Create a <see cref="SourceText"/> from a <see cref="TextDocumentState"/>.
         /// </summary>
-        public AdditionalTextWithState(TextDocumentState documentState)
+        private AdditionalTextWithState(TextDocumentState documentState)
             => _documentState = documentState ?? throw new ArgumentNullException(nameof(documentState));
 
         /// <summary>

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

@@ -245,7 +245,7 @@ internal DocumentState CreateDocument(DocumentInfo documentInfo, ParseOptions? p
 
         public AnalyzerOptions AnalyzerOptions
             => _lazyAnalyzerOptions ??= new AnalyzerOptions(
-                additionalFiles: AdditionalDocumentStates.SelectAsArray<AdditionalText>(static state => new AdditionalTextWithState(state)),
+                additionalFiles: AdditionalDocumentStates.SelectAsArray(AdditionalTextWithState.FromState),
                 optionsProvider: new WorkspaceAnalyzerConfigOptionsProvider(this));
 
         public async Task<ImmutableDictionary<string, string>> GetAnalyzerOptionsForPathAsync(

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

@@ -9,6 +9,10 @@ namespace Microsoft.CodeAnalysis
 {
     internal partial class SolutionState
     {
+        /// <summary>
+        /// Represents a change that needs to be made to a <see cref="Compilation"/>, <see cref="GeneratorDriver"/>, or both in response to
+        /// some user edit.
+        /// </summary>
         private abstract partial class CompilationAndGeneratorDriverTranslationAction
         {
             public virtual Task<Compilation> TransformCompilationAsync(Compilation oldCompilation, CancellationToken cancellationToken)
@@ -28,6 +32,8 @@ public virtual Task<Compilation> TransformCompilationAsync(Compilation oldCompil
             /// </remarks>
             public abstract bool CanUpdateCompilationWithStaleGeneratedTreesIfGeneratorsGiveSameOutput { get; }
 
+            public virtual GeneratorDriver? TransformGeneratorDriver(GeneratorDriver generatorDriver) => generatorDriver;
+
             /// <summary>
             /// When changes are made to a solution, we make a list of translation actions. If multiple similar changes happen in rapid
             /// succession, we may be able to merge them without holding onto intermediate state.

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

@@ -4,6 +4,7 @@
 
 using System.Collections.Generic;
 using System.Collections.Immutable;
+using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.CodeAnalysis.Diagnostics;
@@ -50,11 +51,8 @@ public override Task<Compilation> TransformCompilationAsync(Compilation oldCompi
 
             internal sealed class TouchAdditionalDocumentAction : CompilationAndGeneratorDriverTranslationAction
             {
-#pragma warning disable IDE0052 // Remove unread private members
-                // https://github.com/dotnet/roslyn/issues/44161: right now there is no way to tell a GeneratorDriver that an additional document changed
                 private readonly TextDocumentState _oldState;
                 private readonly TextDocumentState _newState;
-#pragma warning restore IDE0052 // Remove unread private members
 
                 public TouchAdditionalDocumentAction(TextDocumentState oldState, TextDocumentState newState)
                 {
@@ -77,6 +75,16 @@ public TouchAdditionalDocumentAction(TextDocumentState oldState, TextDocumentSta
 
                     return null;
                 }
+
+                public override GeneratorDriver? TransformGeneratorDriver(GeneratorDriver generatorDriver)
+                {
+                    var oldText = AdditionalTextWithState.FromState(_oldState);
+                    var newText = AdditionalTextWithState.FromState(_newState);
+
+                    // TODO: have the compiler add an API for replacing an additional text
+                    return generatorDriver.RemoveAdditionalTexts(ImmutableArray.Create(oldText))
+                                          .AddAdditionalTexts(ImmutableArray.Create(newText));
+                }
             }
 
             internal sealed class RemoveDocumentsAction : CompilationAndGeneratorDriverTranslationAction
@@ -132,10 +140,12 @@ public override async Task<Compilation> TransformCompilationAsync(Compilation ol
             internal sealed class ReplaceAllSyntaxTreesAction : CompilationAndGeneratorDriverTranslationAction
             {
                 private readonly ProjectState _state;
+                private readonly bool _isParseOptionChange;
 
-                public ReplaceAllSyntaxTreesAction(ProjectState state)
+                public ReplaceAllSyntaxTreesAction(ProjectState state, bool isParseOptionChange)
                 {
                     _state = state;
+                    _isParseOptionChange = isParseOptionChange;
                 }
 
                 public override async Task<Compilation> TransformCompilationAsync(Compilation oldCompilation, CancellationToken cancellationToken)
@@ -153,15 +163,33 @@ public override async Task<Compilation> TransformCompilationAsync(Compilation ol
 
                 // Because this removes all trees, it'd also remove the generated trees.
                 public override bool CanUpdateCompilationWithStaleGeneratedTreesIfGeneratorsGiveSameOutput => false;
+
+                public override GeneratorDriver? TransformGeneratorDriver(GeneratorDriver generatorDriver)
+                {
+                    if (_isParseOptionChange)
+                    {
+                        // TODO: update the existing generator driver; the compiler needs to add an API for that.
+                        // In the mean time, drop it and we'll recreate it from scratch.
+                        return null;
+                    }
+                    else
+                    {
+                        // We are using this as a way to reorder syntax trees -- we don't need to do anything as the driver
+                        // will get the new compilation once we pass it to it.
+                        return generatorDriver;
+                    }
+                }
             }
 
             internal sealed class ProjectCompilationOptionsAction : CompilationAndGeneratorDriverTranslationAction
             {
                 private readonly CompilationOptions _options;
+                private readonly bool _isAnalyzerConfigChange;
 
-                public ProjectCompilationOptionsAction(CompilationOptions options)
+                public ProjectCompilationOptionsAction(CompilationOptions options, bool isAnalyzerConfigChange)
                 {
                     _options = options;
+                    _isAnalyzerConfigChange = isAnalyzerConfigChange;
                 }
 
                 public override Task<Compilation> TransformCompilationAsync(Compilation oldCompilation, CancellationToken cancellationToken)
@@ -172,6 +200,22 @@ public override Task<Compilation> TransformCompilationAsync(Compilation oldCompi
                 // Updating the options of a compilation doesn't require us to reparse trees, so we can use this to update
                 // compilations with stale generated trees.
                 public override bool CanUpdateCompilationWithStaleGeneratedTreesIfGeneratorsGiveSameOutput => true;
+
+                public override GeneratorDriver? TransformGeneratorDriver(GeneratorDriver generatorDriver)
+                {
+                    if (_isAnalyzerConfigChange)
+                    {
+                        // TODO: update the existing generator driver; the compiler needs to add an API for that.
+                        // In the mean time, drop it and we'll recreate it from scratch.
+                        return null;
+                    }
+                    else
+                    {
+                        // Changing any other option is fine and the driver can be reused. The driver
+                        // will get the new compilation once we pass it to it.
+                        return generatorDriver;
+                    }
+                }
             }
 
             internal sealed class ProjectAssemblyNameAction : CompilationAndGeneratorDriverTranslationAction
@@ -195,11 +239,8 @@ public override Task<Compilation> TransformCompilationAsync(Compilation oldCompi
 
             internal sealed class AddAnalyzerReferencesAction : CompilationAndGeneratorDriverTranslationAction
             {
-#pragma warning disable IDE0052 // Remove unread private members
-                // https://github.com/dotnet/roslyn/issues/44161: right now there is no way to tell a GeneratorDriver that an analyzer reference has been added
                 private readonly ImmutableArray<AnalyzerReference> _analyzerReferences;
                 private readonly string _language;
-#pragma warning restore IDE0052 // Remove unread private members
 
                 public AddAnalyzerReferencesAction(ImmutableArray<AnalyzerReference> analyzerReferences, string language)
                 {
@@ -211,15 +252,17 @@ public AddAnalyzerReferencesAction(ImmutableArray<AnalyzerReference> analyzerRef
                 // 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.AddGenerators(_analyzerReferences.SelectMany(r => r.GetGenerators(_language)).ToImmutableArray());
+                }
             }
 
             internal sealed class RemoveAnalyzerReferencesAction : CompilationAndGeneratorDriverTranslationAction
             {
-#pragma warning disable IDE0052 // Remove unread private members
-                // https://github.com/dotnet/roslyn/issues/44161: right now there is no way to tell a GeneratorDriver that an analyzer reference has been removed
                 private readonly ImmutableArray<AnalyzerReference> _analyzerReferences;
                 private readonly string _language;
-#pragma warning restore IDE0052 // Remove unread private members
 
                 public RemoveAnalyzerReferencesAction(ImmutableArray<AnalyzerReference> analyzerReferences, string language)
                 {
@@ -231,14 +274,15 @@ public RemoveAnalyzerReferencesAction(ImmutableArray<AnalyzerReference> analyzer
                 // 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());
+                }
             }
 
             internal sealed class AddAdditionalDocumentsAction : CompilationAndGeneratorDriverTranslationAction
             {
-#pragma warning disable IDE0052 // Remove unread private members
-                // https://github.com/dotnet/roslyn/issues/44161: right now there is no way to tell a GeneratorDriver that an additional file has been added
                 private readonly ImmutableArray<TextDocumentState> _additionalDocuments;
-#pragma warning restore IDE0052 // Remove unread private members
 
                 public AddAdditionalDocumentsAction(ImmutableArray<TextDocumentState> additionalDocuments)
                 {
@@ -249,14 +293,16 @@ public AddAdditionalDocumentsAction(ImmutableArray<TextDocumentState> additional
                 // 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.AddAdditionalTexts(_additionalDocuments.SelectAsArray(AdditionalTextWithState.FromState));
+                }
             }
 
             internal sealed class RemoveAdditionalDocumentsAction : CompilationAndGeneratorDriverTranslationAction
             {
-#pragma warning disable IDE0052 // Remove unread private members
-                // https://github.com/dotnet/roslyn/issues/44161: right now there is no way to tell a GeneratorDriver that an additional file has been added
                 private readonly ImmutableArray<TextDocumentState> _additionalDocuments;
-#pragma warning restore IDE0052 // Remove unread private members
 
                 public RemoveAdditionalDocumentsAction(ImmutableArray<TextDocumentState> additionalDocuments)
                 {
@@ -267,6 +313,11 @@ public RemoveAdditionalDocumentsAction(ImmutableArray<TextDocumentState> additio
                 // 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.AddAdditionalTexts(_additionalDocuments.SelectAsArray(AdditionalTextWithState.FromState));
+                }
             }
         }
     }