Move source-generators out of ProjectState and entirely into SolutinoCompilationState

src/Workspaces/Core/Portable/SourceGeneratorTelemetry/ISourceGeneratorTelemetryCollectorWorkspaceService.cs

@@ -2,11 +2,13 @@
 // 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 Microsoft.CodeAnalysis.Diagnostics;
 using Microsoft.CodeAnalysis.Host;
 
 namespace Microsoft.CodeAnalysis.SourceGeneratorTelemetry;
 
 internal interface ISourceGeneratorTelemetryCollectorWorkspaceService : IWorkspaceService
 {
-    void CollectRunResult(GeneratorDriverRunResult driverRunResult, GeneratorDriverTimingInfo driverTimingInfo, ProjectState project);
+    void CollectRunResult(GeneratorDriverRunResult driverRunResult, GeneratorDriverTimingInfo driverTimingInfo, Func<ISourceGenerator, AnalyzerReference> getAnalyzerReference);
 }

src/Workspaces/Core/Portable/SourceGeneratorTelemetry/SourceGeneratorTelemetryCollectorWorkspaceService.cs

@@ -42,19 +42,22 @@ public GeneratorTelemetryKey(ISourceGenerator generator, AnalyzerReference analy
     private readonly StatisticLogAggregator<GeneratorTelemetryKey> _elapsedTimeByGenerator = new();
     private readonly StatisticLogAggregator<GeneratorTelemetryKey> _producedFilesByGenerator = new();
 
-    private GeneratorTelemetryKey GetTelemetryKey(ISourceGenerator generator, ProjectState project)
-        => _generatorTelemetryKeys.GetValue(generator, g => new GeneratorTelemetryKey(g, project.GetAnalyzerReferenceForGenerator(g)));
+    private GeneratorTelemetryKey GetTelemetryKey(ISourceGenerator generator, Func<ISourceGenerator, AnalyzerReference> getAnalyzerReference)
+        => _generatorTelemetryKeys.GetValue(generator, g => new GeneratorTelemetryKey(g, getAnalyzerReference(g)));
 
-    public void CollectRunResult(GeneratorDriverRunResult driverRunResult, GeneratorDriverTimingInfo driverTimingInfo, ProjectState project)
+    public void CollectRunResult(
+        GeneratorDriverRunResult driverRunResult,
+        GeneratorDriverTimingInfo driverTimingInfo,
+        Func<ISourceGenerator, AnalyzerReference> getAnalyzerReference)
     {
         foreach (var generatorTime in driverTimingInfo.GeneratorTimes)
         {
-            _elapsedTimeByGenerator.AddDataPoint(GetTelemetryKey(generatorTime.Generator, project), generatorTime.ElapsedTime);
+            _elapsedTimeByGenerator.AddDataPoint(GetTelemetryKey(generatorTime.Generator, getAnalyzerReference), generatorTime.ElapsedTime);
         }
 
         foreach (var generatorResult in driverRunResult.Results)
         {
-            _producedFilesByGenerator.AddDataPoint(GetTelemetryKey(generatorResult.Generator, project), generatorResult.GeneratedSources.Length);
+            _producedFilesByGenerator.AddDataPoint(GetTelemetryKey(generatorResult.Generator, getAnalyzerReference), generatorResult.GeneratedSources.Length);
         }
     }
 

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

@@ -61,11 +61,6 @@ internal partial class ProjectState
 
     private AnalyzerOptions? _lazyAnalyzerOptions;
 
-    /// <summary>
-    /// The list of source generators and the analyzer reference they came from.
-    /// </summary>
-    private ImmutableDictionary<ISourceGenerator, AnalyzerReference>? _lazySourceGenerators;
-
     private ProjectState(
         ProjectInfo projectInfo,
         LanguageServices languageServices,
@@ -737,38 +732,6 @@ public ProjectState WithAnalyzerReferences(IEnumerable<AnalyzerReference> analyz
         return With(projectInfo: ProjectInfo.WithAnalyzerReferences(analyzerReferences).WithVersion(Version.GetNewerVersion()));
     }
 
-    [MemberNotNull(nameof(_lazySourceGenerators))]
-    private void EnsureSourceGeneratorsInitialized()
-    {
-        if (_lazySourceGenerators == null)
-        {
-            var builder = ImmutableDictionary.CreateBuilder<ISourceGenerator, AnalyzerReference>();
-
-            foreach (var analyzerReference in AnalyzerReferences)
-            {
-                foreach (var generator in analyzerReference.GetGenerators(Language))
-                    builder.Add(generator, analyzerReference);
-            }
-
-            Interlocked.CompareExchange(ref _lazySourceGenerators, builder.ToImmutable(), comparand: null);
-        }
-    }
-
-    public IEnumerable<ISourceGenerator> SourceGenerators
-    {
-        get
-        {
-            EnsureSourceGeneratorsInitialized();
-            return _lazySourceGenerators.Keys;
-        }
-    }
-
-    public AnalyzerReference GetAnalyzerReferenceForGenerator(ISourceGenerator generator)
-    {
-        EnsureSourceGeneratorsInitialized();
-        return _lazySourceGenerators[generator];
-    }
-
     public ProjectState AddDocuments(ImmutableArray<DocumentState> documents)
     {
         if (documents.IsEmpty)

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

@@ -422,7 +422,7 @@ async Task<InProgressState> CollapseInProgressStateAsync(InProgressState initial
                             // Also transform the compilation that has generated files; we won't do that though if the transformation either would cause problems with
                             // the generated documents, or if don't have any source generators in the first place.
                             if (translationAction.CanUpdateCompilationWithStaleGeneratedTreesIfGeneratorsGiveSameOutput &&
-                                translationAction.OldProjectState.SourceGenerators.Any())
+                                GetSourceGenerators(translationAction.OldProjectState).Any())
                             {
                                 staleCompilationWithGeneratedDocuments = await translationAction.TransformCompilationAsync(staleCompilationWithGeneratedDocuments, cancellationToken).ConfigureAwait(false);
                             }
@@ -774,7 +774,7 @@ public async ValueTask<TextDocumentStates<SourceGeneratedDocumentState>> GetSour
                 SolutionCompilationState compilationState, CancellationToken cancellationToken)
             {
                 // If we don't have any generators, then we know we have no generated files, so we can skip the computation entirely.
-                if (!this.ProjectState.SourceGenerators.Any())
+                if (!GetSourceGenerators(this.ProjectState).Any())
                 {
                     return TextDocumentStates<SourceGeneratedDocumentState>.Empty;
                 }
@@ -787,7 +787,7 @@ public async ValueTask<TextDocumentStates<SourceGeneratedDocumentState>> GetSour
             public async ValueTask<ImmutableArray<Diagnostic>> GetSourceGeneratorDiagnosticsAsync(
                 SolutionCompilationState compilationState, CancellationToken cancellationToken)
             {
-                if (!this.ProjectState.SourceGenerators.Any())
+                if (!GetSourceGenerators(this.ProjectState).Any())
                 {
                     return [];
                 }
@@ -816,7 +816,7 @@ public async ValueTask<ImmutableArray<Diagnostic>> GetSourceGeneratorDiagnostics
 
             public async ValueTask<GeneratorDriverRunResult?> GetSourceGeneratorRunResultAsync(SolutionCompilationState compilationState, CancellationToken cancellationToken)
             {
-                if (!this.ProjectState.SourceGenerators.Any())
+                if (!GetSourceGenerators(this.ProjectState).Any())
                 {
                     return null;
                 }

src/Workspaces/Core/Portable/Workspace/Solution/SolutionCompilationState.CompilationTracker_Generators.cs

@@ -233,7 +233,7 @@ await newGeneratedDocuments.States.Values.SelectAsArrayAsync(
             // the "InCurrentProcess" call so that it will normally run only in the OOP process, thus ensuring that we
             // get accurate information about what SourceGenerators we actually have (say, in case they they are rebuilt
             // by the user while VS is running).
-            if (!this.ProjectState.SourceGenerators.Any())
+            if (!GetSourceGenerators(this.ProjectState).Any())
                 return (compilationWithoutGeneratedFiles, TextDocumentStates<SourceGeneratedDocumentState>.Empty, generatorDriver);
 
             // If we don't already have an existing generator driver, create one from scratch
@@ -268,7 +268,9 @@ await newGeneratedDocuments.States.Values.SelectAsArrayAsync(
 
             var runResult = generatorDriver.GetRunResult();
 
-            telemetryCollector?.CollectRunResult(runResult, generatorDriver.GetTimingInfo(), ProjectState);
+            telemetryCollector?.CollectRunResult(
+                runResult, generatorDriver.GetTimingInfo(),
+                g => GetAnalyzerReference(this.ProjectState, g));
 
             // We may be able to reuse compilationWithStaleGeneratedTrees if the generated trees are identical. We will assign null
             // to compilationWithStaleGeneratedTrees if we at any point realize it can't be used. We'll first check the count of trees
@@ -293,7 +295,7 @@ await newGeneratedDocuments.States.Values.SelectAsArrayAsync(
                 if (IsGeneratorRunResultToIgnore(generatorResult))
                     continue;
 
-                var generatorAnalyzerReference = this.ProjectState.GetAnalyzerReferenceForGenerator(generatorResult.Generator);
+                var generatorAnalyzerReference = GetAnalyzerReference(this.ProjectState, generatorResult.Generator);
 
                 foreach (var generatedSource in generatorResult.GeneratedSources)
                 {
@@ -394,7 +396,7 @@ static GeneratorDriver CreateGeneratorDriver(ProjectState projectState)
 
                 return compilationFactory.CreateGeneratorDriver(
                     projectState.ParseOptions!,
-                    projectState.SourceGenerators.ToImmutableArray(),
+                    GetSourceGenerators(projectState),
                     projectState.AnalyzerOptions.AnalyzerConfigOptionsProvider,
                     additionalTexts);
             }

src/Workspaces/Core/Portable/Workspace/Solution/SolutionCompilationState.TranslationAction_Actions.cs

@@ -355,7 +355,7 @@ public override GeneratorDriver TransformGeneratorDriver(GeneratorDriver _)
                     .ReplaceAdditionalTexts(this.NewProjectState.AdditionalDocumentStates.SelectAsArray(static documentState => documentState.AdditionalText))
                     .WithUpdatedParseOptions(this.NewProjectState.ParseOptions!)
                     .WithUpdatedAnalyzerConfigOptions(this.NewProjectState.AnalyzerOptions.AnalyzerConfigOptionsProvider)
-                    .ReplaceGenerators(this.NewProjectState.SourceGenerators.ToImmutableArray());
+                    .ReplaceGenerators(GetSourceGenerators(this.NewProjectState));
 
                 return generatorDriver;
             }

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

@@ -0,0 +1,77 @@
+// 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.
+
+using System.Collections.Generic;
+using System.Collections.Immutable;
+using System.Runtime.CompilerServices;
+using Microsoft.CodeAnalysis.Diagnostics;
+using Microsoft.CodeAnalysis.Shared.Collections;
+using Roslyn.Utilities;
+
+namespace Microsoft.CodeAnalysis;
+
+using AnalyzerReferencesToSourceGenerators = ConditionalWeakTable<IReadOnlyList<AnalyzerReference>, SolutionCompilationState.SourceGeneratorMap>;
+
+internal partial class SolutionCompilationState
+{
+    internal sealed record SourceGeneratorMap(
+        ImmutableArray<ISourceGenerator> SourceGenerators,
+        ImmutableDictionary<ISourceGenerator, AnalyzerReference> SourceGeneratorToAnalyzerReference);
+
+    /// <summary>
+    /// Cached mapping from language (only C#/VB since those are the only languages that support analyzers) to the lists
+    /// of analyzer references (see <see cref="ProjectState.AnalyzerReferences"/>) to all the <see
+    /// cref="ISourceGenerator"/>s produced by those references.  This should only be created and cached on the OOP side
+    /// of things so that we don't cause source generators to be loaded (and fixed) within VS (which is .net framework
+    /// only).
+    /// </summary>
+    private static readonly ImmutableArray<(string language, AnalyzerReferencesToSourceGenerators referencesToGenerators, AnalyzerReferencesToSourceGenerators.CreateValueCallback callback)> s_languageToAnalyzerReferencesToSourceGeneratorsMap =
+    [
+        (LanguageNames.CSharp, new(), (static rs => ComputeSourceGenerators(rs, LanguageNames.CSharp))),
+        (LanguageNames.VisualBasic, new(), (static rs => ComputeSourceGenerators(rs, LanguageNames.VisualBasic))),
+    ];
+
+    private static SourceGeneratorMap ComputeSourceGenerators(IReadOnlyList<AnalyzerReference> analyzerReferences, string language)
+    {
+        using var generators = TemporaryArray<ISourceGenerator>.Empty;
+        var generatorToAnalyzerReference = ImmutableDictionary.CreateBuilder<ISourceGenerator, AnalyzerReference>();
+
+        foreach (var reference in analyzerReferences)
+        {
+            foreach (var generator in reference.GetGenerators(language).Distinct())
+            {
+                generators.Add(generator);
+                generatorToAnalyzerReference.Add(generator, reference);
+            }
+        }
+
+        return new(generators.ToImmutableAndClear(), generatorToAnalyzerReference.ToImmutable());
+    }
+
+    private static ImmutableArray<ISourceGenerator> GetSourceGenerators(ProjectState projectState)
+        => GetSourceGenerators(projectState.Language, projectState.AnalyzerReferences);
+
+    private static ImmutableArray<ISourceGenerator> GetSourceGenerators(string language, IReadOnlyList<AnalyzerReference> analyzerReferences)
+    {
+        var map = GetSourceGeneratorMap(language, analyzerReferences);
+        return map is null ? [] : map.SourceGenerators;
+    }
+
+    private static AnalyzerReference GetAnalyzerReference(ProjectState projectState, ISourceGenerator sourceGenerator)
+    {
+        var map = GetSourceGeneratorMap(projectState.Language, projectState.AnalyzerReferences);
+        Contract.ThrowIfNull(map);
+        return map.SourceGeneratorToAnalyzerReference[sourceGenerator];
+    }
+
+    private static SourceGeneratorMap? GetSourceGeneratorMap(string language, IReadOnlyList<AnalyzerReference> analyzerReferences)
+    {
+        var tupleOpt = s_languageToAnalyzerReferencesToSourceGeneratorsMap.FirstOrNull(static (t, language) => t.language == language, language);
+        if (tupleOpt is null)
+            return null;
+
+        var tuple = tupleOpt.Value;
+        return tuple.referencesToGenerators.GetValue(analyzerReferences, tuple.callback);
+    }
+}