Merge pull request #72835 from CyrusNajmabadi/sgOnlyOop3

src/Workspaces/Core/Portable/Remote/RemoteHostClient.cs

@@ -151,6 +151,17 @@ public async ValueTask<bool> TryInvokeAsync<TService>(
         return await connection.TryInvokeAsync(project, invocation, cancellationToken).ConfigureAwait(false);
     }
 
+    public async ValueTask<Optional<TResult>> TryInvokeAsync<TService, TResult>(
+        SolutionCompilationState compilationState,
+        ProjectId projectId,
+        Func<TService, Checksum, CancellationToken, ValueTask<TResult>> invocation,
+        CancellationToken cancellationToken)
+        where TService : class
+    {
+        using var connection = CreateConnection<TService>(callbackTarget: null);
+        return await connection.TryInvokeAsync(compilationState, projectId, invocation, cancellationToken).ConfigureAwait(false);
+    }
+
     /// <summary>
     /// Equivalent to <see cref="TryInvokeAsync{TService}(Solution, Func{TService, Checksum, CancellationToken, ValueTask}, CancellationToken)"/>
     /// except that only the project (and its dependent projects) will be sync'ed to the remote host before executing.

src/Workspaces/Core/Portable/SourceGeneration/IRemoteSourceGenerationService.cs

@@ -31,6 +31,12 @@ internal interface IRemoteSourceGenerationService
     /// </summary>
     ValueTask<ImmutableArray<string>> GetContentsAsync(
         Checksum solutionChecksum, ProjectId projectId, ImmutableArray<DocumentId> documentIds, CancellationToken cancellationToken);
+
+    /// <summary>
+    /// Whether or not the specified <paramref name="projectId"/> has source generators or not.
+    /// </summary>
+    ValueTask<bool> HasGeneratorsAsync(
+        Checksum solutionChecksum, ProjectId projectId, CancellationToken cancellationToken);
 }
 
 /// <summary>

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

@@ -359,6 +359,9 @@ async Task<InProgressState> CollapseInProgressStateAsync(InProgressState initial
                 {
                     try
                     {
+                        // Only bother keeping track of staleCompilationWithGeneratedDocuments for projects that
+                        // actually have generators in them.
+                        var hasSourceGenerators = await compilationState.HasSourceGeneratorsAsync(this.ProjectState.Id, cancellationToken).ConfigureAwait(false);
                         var currentState = initialState;
 
                         // Then, we serially process the chain while that parsing is happening concurrently.
@@ -368,7 +371,7 @@ async Task<InProgressState> CollapseInProgressStateAsync(InProgressState initial
 
                             // We have a list of transformations to get to our final compilation; take the first transformation and apply it.
                             var (compilationWithoutGeneratedDocuments, staleCompilationWithGeneratedDocuments, generatorInfo) =
-                                await ApplyFirstTransformationAsync(currentState).ConfigureAwait(false);
+                                await ApplyFirstTransformationAsync(currentState, hasSourceGenerators).ConfigureAwait(false);
 
                             // We have updated state, so store this new result; this allows us to drop the intermediate state we already processed
                             // even if we were to get cancelled at a later point.
@@ -397,7 +400,7 @@ async Task<InProgressState> CollapseInProgressStateAsync(InProgressState initial
                     }
 
                     async Task<(Compilation compilationWithoutGeneratedDocuments, Compilation? staleCompilationWithGeneratedDocuments, CompilationTrackerGeneratorInfo generatorInfo)>
-                        ApplyFirstTransformationAsync(InProgressState inProgressState)
+                        ApplyFirstTransformationAsync(InProgressState inProgressState, bool hasSourceGenerators)
                     {
                         Contract.ThrowIfTrue(inProgressState.PendingTranslationActions.IsEmpty);
                         var translationAction = inProgressState.PendingTranslationActions[0];
@@ -422,7 +425,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 &&
-                                GetSourceGenerators(translationAction.OldProjectState).Any())
+                                hasSourceGenerators)
                             {
                                 staleCompilationWithGeneratedDocuments = await translationAction.TransformCompilationAsync(staleCompilationWithGeneratedDocuments, cancellationToken).ConfigureAwait(false);
                             }
@@ -774,10 +777,8 @@ 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 (!GetSourceGenerators(this.ProjectState).Any())
-                {
+                if (!await compilationState.HasSourceGeneratorsAsync(this.ProjectState.Id, cancellationToken).ConfigureAwait(false))
                     return TextDocumentStates<SourceGeneratedDocumentState>.Empty;
-                }
 
                 var finalState = await GetOrBuildFinalStateAsync(
                     compilationState, cancellationToken: cancellationToken).ConfigureAwait(false);
@@ -787,10 +788,8 @@ public async ValueTask<TextDocumentStates<SourceGeneratedDocumentState>> GetSour
             public async ValueTask<ImmutableArray<Diagnostic>> GetSourceGeneratorDiagnosticsAsync(
                 SolutionCompilationState compilationState, CancellationToken cancellationToken)
             {
-                if (!GetSourceGenerators(this.ProjectState).Any())
-                {
+                if (!await compilationState.HasSourceGeneratorsAsync(this.ProjectState.Id, cancellationToken).ConfigureAwait(false))
                     return [];
-                }
 
                 var finalState = await GetOrBuildFinalStateAsync(
                     compilationState, cancellationToken: cancellationToken).ConfigureAwait(false);
@@ -816,10 +815,8 @@ public async ValueTask<ImmutableArray<Diagnostic>> GetSourceGeneratorDiagnostics
 
             public async ValueTask<GeneratorDriverRunResult?> GetSourceGeneratorRunResultAsync(SolutionCompilationState compilationState, CancellationToken cancellationToken)
             {
-                if (!GetSourceGenerators(this.ProjectState).Any())
-                {
+                if (!await compilationState.HasSourceGeneratorsAsync(this.ProjectState.Id, cancellationToken).ConfigureAwait(false))
                     return null;
-                }
 
                 var finalState = await GetOrBuildFinalStateAsync(
                     compilationState, cancellationToken).ConfigureAwait(false);

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

@@ -57,14 +57,13 @@ private partial class CompilationTracker : ICompilationTracker
                 if (result.HasValue)
                 {
                     // Since we ran the SG work out of process, we could not have created or modified the driver passed in.
-                    // So just pass what we got in right back out.
-                    return (result.Value.compilationWithGeneratedFiles, new(result.Value.generatedDocuments, generatorInfo.Driver));
+                    // Just return `null` for the driver as there's nothing to track for it on the host side.
+                    return (result.Value.compilationWithGeneratedFiles, new(result.Value.generatedDocuments, Driver: null));
                 }
 
                 // If that failed (OOP crash, or we are the OOP process ourselves), then generate the SG docs locally.
-                var telemetryCollector = compilationState.SolutionState.Services.GetService<ISourceGeneratorTelemetryCollectorWorkspaceService>();
                 var (compilationWithGeneratedFiles, nextGeneratedDocuments, nextGeneratorDriver) = await ComputeNewGeneratorInfoInCurrentProcessAsync(
-                    telemetryCollector,
+                    compilationState,
                     compilationWithoutGeneratedFiles,
                     generatorInfo.Documents,
                     generatorInfo.Driver,
@@ -222,18 +221,15 @@ await newGeneratedDocuments.States.Values.SelectAsArrayAsync(
         }
 
         private async Task<(Compilation compilationWithGeneratedFiles, TextDocumentStates<SourceGeneratedDocumentState> generatedDocuments, GeneratorDriver? generatorDriver)> ComputeNewGeneratorInfoInCurrentProcessAsync(
-            ISourceGeneratorTelemetryCollectorWorkspaceService? telemetryCollector,
+            SolutionCompilationState compilationState,
             Compilation compilationWithoutGeneratedFiles,
             TextDocumentStates<SourceGeneratedDocumentState> oldGeneratedDocuments,
             GeneratorDriver? generatorDriver,
             Compilation? compilationWithStaleGeneratedTrees,
             CancellationToken cancellationToken)
         {
-            // If we don't have any source generators.  Trivially bail out.  Note: this check is intentionally don't in
-            // 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 (!GetSourceGenerators(this.ProjectState).Any())
+            // If we don't have any source generators.  Trivially bail out.
+            if (!await compilationState.HasSourceGeneratorsAsync(this.ProjectState.Id, cancellationToken).ConfigureAwait(false))
                 return (compilationWithoutGeneratedFiles, TextDocumentStates<SourceGeneratedDocumentState>.Empty, generatorDriver);
 
             // If we don't already have an existing generator driver, create one from scratch
@@ -268,6 +264,7 @@ await newGeneratedDocuments.States.Values.SelectAsArrayAsync(
 
             var runResult = generatorDriver.GetRunResult();
 
+            var telemetryCollector = compilationState.SolutionState.Services.GetService<ISourceGeneratorTelemetryCollectorWorkspaceService>();
             telemetryCollector?.CollectRunResult(
                 runResult, generatorDriver.GetTimingInfo(),
                 g => GetAnalyzerReference(this.ProjectState, g));

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

@@ -2,22 +2,26 @@
 // 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.Frozen;
 using System.Collections.Immutable;
+using System.Linq;
 using System.Runtime.CompilerServices;
+using System.Threading;
+using System.Threading.Tasks;
 using Microsoft.CodeAnalysis.Diagnostics;
+using Microsoft.CodeAnalysis.PooledObjects;
+using Microsoft.CodeAnalysis.Remote;
 using Microsoft.CodeAnalysis.Shared.Collections;
+using Microsoft.CodeAnalysis.SourceGeneration;
 using Roslyn.Utilities;
 
 namespace Microsoft.CodeAnalysis;
 
-using AnalyzerReferencesToSourceGenerators = ConditionalWeakTable<IReadOnlyList<AnalyzerReference>, SolutionCompilationState.SourceGeneratorMap>;
-
 internal partial class SolutionCompilationState
 {
-    internal sealed record SourceGeneratorMap(
+    private sealed record SourceGeneratorMap(
         ImmutableArray<ISourceGenerator> SourceGenerators,
-        ImmutableDictionary<ISourceGenerator, AnalyzerReference> SourceGeneratorToAnalyzerReference);
+        FrozenDictionary<ISourceGenerator, AnalyzerReference> SourceGeneratorToAnalyzerReference);
 
     /// <summary>
     /// Cached mapping from language (only C#/VB since those are the only languages that support analyzers) to the lists
@@ -26,52 +30,102 @@ internal sealed record SourceGeneratorMap(
     /// 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);
-            }
-        }
+    private static readonly ConditionalWeakTable<ProjectState, SourceGeneratorMap> s_projectStateToSourceGeneratorsMap = new();
 
-        return new(generators.ToImmutableAndClear(), generatorToAnalyzerReference.ToImmutable());
-    }
+    /// <summary>
+    /// Cached information about if a project has source generators or not.  Note: this is distinct from <see
+    /// cref="s_projectStateToSourceGeneratorsMap"/> as we want to be able to compute it by calling over to our OOP
+    /// process (if present) and having it make the determination, without the host necessarily loading generators
+    /// itself.
+    /// </summary>
+    private static readonly ConditionalWeakTable<ProjectState, AsyncLazy<bool>> s_hasSourceGeneratorsMap = new();
 
+    /// <summary>
+    /// This method should only be called in a .net core host like our out of process server.
+    /// </summary>
     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);
+        var map = GetSourceGeneratorMap(projectState);
         return map is null ? [] : map.SourceGenerators;
     }
 
+    /// <summary>
+    /// This method should only be called in a .net core host like our out of process server.
+    /// </summary>
     private static AnalyzerReference GetAnalyzerReference(ProjectState projectState, ISourceGenerator sourceGenerator)
     {
-        var map = GetSourceGeneratorMap(projectState.Language, projectState.AnalyzerReferences);
+        // We must be talking about a project that supports compilations, since we already got a source generator from it.
+        Contract.ThrowIfFalse(projectState.SupportsCompilation);
+
+        var map = GetSourceGeneratorMap(projectState);
+
+        // It should not be possible for this to be null. We have the source generator, as such we must have mapped from
+        // the project state to the SG info for it.
         Contract.ThrowIfNull(map);
+
         return map.SourceGeneratorToAnalyzerReference[sourceGenerator];
     }
 
-    private static SourceGeneratorMap? GetSourceGeneratorMap(string language, IReadOnlyList<AnalyzerReference> analyzerReferences)
+    private static SourceGeneratorMap? GetSourceGeneratorMap(ProjectState projectState)
     {
-        var tupleOpt = s_languageToAnalyzerReferencesToSourceGeneratorsMap.FirstOrNull(static (t, language) => t.language == language, language);
-        if (tupleOpt is null)
+        if (!projectState.SupportsCompilation)
             return null;
 
-        var tuple = tupleOpt.Value;
-        return tuple.referencesToGenerators.GetValue(analyzerReferences, tuple.callback);
+        return s_projectStateToSourceGeneratorsMap.GetValue(projectState, ComputeSourceGenerators);
+
+        static SourceGeneratorMap ComputeSourceGenerators(ProjectState projectState)
+        {
+            using var generators = TemporaryArray<ISourceGenerator>.Empty;
+            using var _ = PooledDictionary<ISourceGenerator, AnalyzerReference>.GetInstance(out var generatorToAnalyzerReference);
+
+            foreach (var reference in projectState.AnalyzerReferences)
+            {
+                foreach (var generator in reference.GetGenerators(projectState.Language).Distinct())
+                {
+                    generators.Add(generator);
+                    generatorToAnalyzerReference.Add(generator, reference);
+                }
+            }
+
+            return new(generators.ToImmutableAndClear(), generatorToAnalyzerReference.ToFrozenDictionary());
+        }
+    }
+
+    public async Task<bool> HasSourceGeneratorsAsync(ProjectId projectId, CancellationToken cancellationToken)
+    {
+        var projectState = this.SolutionState.GetRequiredProjectState(projectId);
+
+        if (!s_hasSourceGeneratorsMap.TryGetValue(projectState, out var lazy))
+        {
+            // Extracted into local function to prevent allocations in the case where we find a value in the cache.
+            lazy = GetLazy(projectState);
+        }
+
+        return await lazy.GetValueAsync(cancellationToken).ConfigureAwait(false);
+
+        AsyncLazy<bool> GetLazy(ProjectState projectState)
+            => s_hasSourceGeneratorsMap.GetValue(
+                projectState,
+                projectState => AsyncLazy.Create(
+                    static (tuple, cancellationToken) => ComputeHasSourceGeneratorsAsync(tuple.@this, tuple.projectState, cancellationToken),
+                    (@this: this, projectState)));
+
+        static async Task<bool> ComputeHasSourceGeneratorsAsync(
+            SolutionCompilationState solution, ProjectState projectState, CancellationToken cancellationToken)
+        {
+            var client = await RemoteHostClient.TryGetClientAsync(solution.Services, cancellationToken).ConfigureAwait(false);
+            // If in proc, just load the generators and see if we have any.
+            if (client is null)
+                return GetSourceGenerators(projectState).Any();
+
+            // Out of process, call to the remote to figure this out.
+            var projectId = projectState.Id;
+            var result = await client.TryInvokeAsync<IRemoteSourceGenerationService, bool>(
+                solution,
+                projectId,
+                (service, solution, cancellationToken) => service.HasGeneratorsAsync(solution, projectId, cancellationToken),
+                cancellationToken).ConfigureAwait(false);
+            return result.HasValue && result.Value;
+        }
     }
 }

src/Workspaces/Remote/ServiceHub/Services/SourceGeneration/RemoteSourceGenerationService.cs

@@ -63,4 +63,12 @@ public ValueTask<ImmutableArray<string>> GetContentsAsync(
             return result.ToImmutableAndClear();
         }, cancellationToken);
     }
+
+    public ValueTask<bool> HasGeneratorsAsync(Checksum solutionChecksum, ProjectId projectId, CancellationToken cancellationToken)
+    {
+        return RunServiceAsync(solutionChecksum, async solution =>
+        {
+            return await solution.CompilationState.HasSourceGeneratorsAsync(projectId, cancellationToken).ConfigureAwait(false);
+        }, cancellationToken);
+    }
 }