Add server side support for refreshing source generated files

src/LanguageServer/Protocol/Handler/Diagnostics/DiagnosticsPullCache.cs

@@ -24,7 +24,7 @@ internal record struct DiagnosticsRequestState(Project Project, int GlobalStateV
     /// and works well for us in the normal case.  The latter still allows us to reuse diagnostics when changes happen that
     /// update the version stamp but not the content (for example, forking LSP text).
     /// </summary>
-    private sealed class DiagnosticsPullCache(string uniqueKey) : VersionedPullCache<(int globalStateVersion, VersionStamp? dependentVersion), (int globalStateVersion, Checksum dependentChecksum), DiagnosticsRequestState, DiagnosticData>(uniqueKey)
+    private sealed class DiagnosticsPullCache(string uniqueKey) : VersionedPullCache<(int globalStateVersion, VersionStamp? dependentVersion), (int globalStateVersion, Checksum dependentChecksum), DiagnosticsRequestState, ImmutableArray<DiagnosticData>>(uniqueKey)
     {
         public override async Task<(int globalStateVersion, VersionStamp? dependentVersion)> ComputeCheapVersionAsync(DiagnosticsRequestState state, CancellationToken cancellationToken)
         {

src/LanguageServer/Protocol/Handler/PullHandlers/VersionedPullCache.CacheItem.cs

@@ -52,7 +52,7 @@ private sealed class CacheItem(string uniqueKey)
         /// 
         /// Returns <see langword="null"/> if the previousPullResult can be re-used, otherwise returns a new resultId and the new data associated with it.
         /// </summary>
-        public async Task<(string, ImmutableArray<TComputedData>)?> UpdateCacheItemAsync(
+        public async Task<(string, TComputedData)?> UpdateCacheItemAsync(
             VersionedPullCache<TCheapVersion, TExpensiveVersion, TState, TComputedData> cache,
             PreviousPullResult? previousPullResult,
             bool isFullyLoaded,

src/LanguageServer/Protocol/Handler/PullHandlers/VersionedPullCache.cs

@@ -4,7 +4,6 @@
 
 using System.Collections.Concurrent;
 using System.Collections.Generic;
-using System.Collections.Immutable;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.CodeAnalysis.Host;
@@ -19,7 +18,6 @@ namespace Microsoft.CodeAnalysis.LanguageServer.Handler;
 /// with different computation costs to determine if the previous cached data is still valid.
 /// </summary>
 internal abstract partial class VersionedPullCache<TCheapVersion, TExpensiveVersion, TState, TComputedData>(string uniqueKey)
-    where TComputedData : notnull
 {
     /// <summary>
     /// Map of workspace and diagnostic source to the data used to make the last pull report.
@@ -59,9 +57,9 @@ internal abstract partial class VersionedPullCache<TCheapVersion, TExpensiveVers
     /// 
     /// Note - this will run under the semaphore in <see cref="CacheItem._gate"/>.
     /// </summary>
-    public abstract Task<ImmutableArray<TComputedData>> ComputeDataAsync(TState state, CancellationToken cancellationToken);
+    public abstract Task<TComputedData> ComputeDataAsync(TState state, CancellationToken cancellationToken);
 
-    public abstract Checksum ComputeChecksum(ImmutableArray<TComputedData> data);
+    public abstract Checksum ComputeChecksum(TComputedData data);
 
     /// <summary>
     /// If results have changed since the last request this calculates and returns a new
@@ -70,7 +68,7 @@ internal abstract partial class VersionedPullCache<TCheapVersion, TExpensiveVers
     /// <param name="idToClientLastResult">a map of roslyn document or project id to the previous result the client sent us for that doc.</param>
     /// <param name="projectOrDocumentId">the id of the project or document that we are checking to see if it has changed.</param>
     /// <returns>Null when results are unchanged, otherwise returns a non-null new resultId.</returns>
-    public async Task<(string ResultId, ImmutableArray<TComputedData> Data)?> GetOrComputeNewDataAsync(
+    public async Task<(string ResultId, TComputedData Data)?> GetOrComputeNewDataAsync(
         Dictionary<ProjectOrDocumentId, PreviousPullResult> idToClientLastResult,
         ProjectOrDocumentId projectOrDocumentId,
         Project project,

src/LanguageServer/Protocol/Handler/SourceGenerators/SourceGeneratedDocumentCache.cs

@@ -0,0 +1,61 @@
+// 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;
+using System.Composition;
+using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.CodeAnalysis.Host.Mef;
+using Microsoft.CodeAnalysis.Text;
+using Roslyn.Utilities;
+
+namespace Microsoft.CodeAnalysis.LanguageServer.Handler.SourceGenerators;
+
+internal record struct SourceGeneratedDocumentGetTextState(Document Document);
+
+internal sealed class SourceGeneratedDocumentCache(string uniqueKey) : VersionedPullCache<(SourceGeneratorExecutionVersion, VersionStamp), object?, SourceGeneratedDocumentGetTextState, SourceText?>(uniqueKey), ILspService
+{
+    public override async Task<(SourceGeneratorExecutionVersion, VersionStamp)> ComputeCheapVersionAsync(SourceGeneratedDocumentGetTextState state, CancellationToken cancellationToken)
+    {
+        // The execution version and the dependent version must be considered as one version cached together -
+        // it is not correct to say that if the execution version is the same then we can re-use results (as in automatic mode the execution version never changes).
+        var executionVersion = state.Document.Project.Solution.GetSourceGeneratorExecutionVersion(state.Document.Project.Id);
+        var dependentVersion = await state.Document.Project.GetDependentVersionAsync(cancellationToken).ConfigureAwait(false);
+        return (executionVersion, dependentVersion);
+    }
+
+    public override Task<object?> ComputeExpensiveVersionAsync(SourceGeneratedDocumentGetTextState state, CancellationToken cancellationToken)
+    {
+        return SpecializedTasks.Null<object>();
+    }
+
+    public override Checksum ComputeChecksum(SourceText? data)
+    {
+        return data is null ? Checksum.Null : Checksum.From(data.GetChecksum());
+    }
+
+    public override async Task<SourceText?> ComputeDataAsync(SourceGeneratedDocumentGetTextState state, CancellationToken cancellationToken)
+    {
+        // When a user has a open source-generated file, we ensure that the contents in the LSP snapshot match the contents that we
+        // get through didOpen/didChanges, like any other file. That way operations in LSP file are in sync with the
+        // contents the user has. However in this case, we don't want to look at that frozen text, but look at what the
+        // generator would generate if we ran it again. Otherwise, we'll get "stuck" and never update the file with something new.
+        // This can return null when the source generated file has been removed (but the queue itself is using the frozen non-null document).
+        var unfrozenDocument = await state.Document.Project.Solution.WithoutFrozenSourceGeneratedDocuments().GetDocumentAsync(state.Document.Id, includeSourceGenerated: true, cancellationToken).ConfigureAwait(false);
+        return unfrozenDocument == null
+            ? null
+            : await unfrozenDocument.GetTextAsync(cancellationToken).ConfigureAwait(false);
+    }
+}
+
+[ExportCSharpVisualBasicLspServiceFactory(typeof(SourceGeneratedDocumentCache)), Shared]
+[method: ImportingConstructor]
+[method: Obsolete(MefConstruction.ImportingConstructorMessage, error: true)]
+internal sealed class SourceGeneratedDocumentCacheFactory() : ILspServiceFactory
+{
+    public ILspService CreateILspService(LspServices lspServices, WellKnownLspServerKinds serverKind)
+    {
+        return new SourceGeneratedDocumentCache(this.GetType().Name);
+    }
+}

src/LanguageServer/Protocol/Handler/SourceGenerators/SourceGeneratedDocumentGetTextHandler.cs

@@ -3,14 +3,17 @@
 // See the LICENSE file in the project root for more information.
 
 using System;
+using System.Collections.Generic;
 using System.Composition;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.CodeAnalysis.Host.Mef;
+using Microsoft.CodeAnalysis.LanguageServer.Handler.Diagnostics;
+using Microsoft.CodeAnalysis.PooledObjects;
 using Roslyn.Utilities;
 using LSP = Roslyn.LanguageServer.Protocol;
 
-namespace Microsoft.CodeAnalysis.LanguageServer.Handler;
+namespace Microsoft.CodeAnalysis.LanguageServer.Handler.SourceGenerators;
 
 [ExportCSharpVisualBasicStatelessLspService(typeof(SourceGeneratedDocumentGetTextHandler)), Shared]
 [Method(MethodName)]
@@ -29,18 +32,46 @@ public async Task<SourceGeneratedDocumentText> HandleRequestAsync(SourceGenerato
     {
         var document = context.Document;
 
+        if (document is null)
+        {
+            // The source generated file being asked about is not present.
+            // This is a rare case the request queue always gives us a frozen, non-null document for any opened sg document,
+            // even if the generator itself was removed and the document no longer exists in the host solution.
+            //
+            // We can only get a null document here if the sg document has not been opened and
+            // the source generated document does not exist in the workspace.
+            //
+            // Return a value indicating that the document is removed.
+            return new SourceGeneratedDocumentText(ResultId: null, Text: null);
+        }
+
         // Nothing here strictly prevents this from working on any other document, but we'll assert we got a source-generated file, since
         // it wouldn't really make sense for the server to be asked for the contents of a regular file. Since this endpoint is intended for
         // source-generated files only, this would indicate that something else has gone wrong.
         Contract.ThrowIfFalse(document is SourceGeneratedDocument);
 
-        // When a user has a open source-generated file, we ensure that the contents in the LSP snapshot match the contents that we
-        // get through didOpen/didChanges, like any other file. That way operations in LSP file are in sync with the
-        // contents the user has. However in this case, we don't want to look at that frozen text, but look at what the
-        // generator would generate if we ran it again. Otherwise, we'll get "stuck" and never update the file with something new.
-        document = await document.Project.Solution.WithoutFrozenSourceGeneratedDocuments().GetDocumentAsync(document.Id, includeSourceGenerated: true, cancellationToken).ConfigureAwait(false);
+        var cache = context.GetRequiredLspService<SourceGeneratedDocumentCache>();
+        var projectOrDocument = new ProjectOrDocumentId(document.Id);
+
+        using var _ = PooledDictionary<ProjectOrDocumentId, PreviousPullResult>.GetInstance(out var previousPullResults);
+        if (request.ResultId is not null)
+        {
+            previousPullResults.Add(projectOrDocument, new PreviousPullResult(request.ResultId, request.TextDocument));
+        }
+
+        var newResult = await cache.GetOrComputeNewDataAsync(previousPullResults, projectOrDocument, document.Project, new SourceGeneratedDocumentGetTextState(document), cancellationToken).ConfigureAwait(false);
 
-        var text = document != null ? await document.GetTextAsync(cancellationToken).ConfigureAwait(false) : null;
-        return new SourceGeneratedDocumentText(text?.ToString());
+        if (newResult is null)
+        {
+            Contract.ThrowIfNull(request.ResultId, "Attempted to reuse cache entry but given no resultId");
+            // The generated document is the same, we can return the same resultId.
+            return new SourceGeneratedDocumentText(request.ResultId, Text: null);
+        }
+        else
+        {
+            // We may get no text back if the unfrozen source generated file no longer exists.
+            var data = newResult.Value.Data?.ToString();
+            return new SourceGeneratedDocumentText(newResult.Value.ResultId, data);
+        }
     }
 }

src/LanguageServer/Protocol/Handler/SourceGenerators/SourceGeneratedDocumentText.cs

@@ -1,9 +1,21 @@
-// Licensed to the .NET Foundation under one or more agreements.
+// 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.Text.Json.Serialization;
 
 namespace Microsoft.CodeAnalysis.LanguageServer.Handler;
 
-internal sealed record SourceGeneratedDocumentText([property: JsonPropertyName("text")] string? Text);
+/// <summary>
+/// Source generated file text result.  The client uses the resultId to inform what the text value is.
+/// 
+/// An unchanged result has a non-null resultId (same as client request resultId) + null text.
+/// 
+/// A changed result has a new non-null resultId + possibly null text (if the sg document no longer exists).
+/// 
+/// In rare circumstances it is possible to get a null resultId + null text - this happens when
+/// the source generated document is not open AND the source generated document no longer exists
+/// </summary>
+internal sealed record SourceGeneratedDocumentText(
+    [property: JsonPropertyName("resultId")] string? ResultId,
+    [property: JsonPropertyName("text")] string? Text);

src/LanguageServer/Protocol/Handler/SourceGenerators/SourceGeneratorGetTextParams.cs

@@ -1,10 +1,12 @@
-// Licensed to the .NET Foundation under one or more agreements.
+// 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.Text.Json.Serialization;
 using Roslyn.LanguageServer.Protocol;
 
-namespace Microsoft.CodeAnalysis.LanguageServer.Handler;
+namespace Microsoft.CodeAnalysis.LanguageServer.Handler.SourceGenerators;
 
-internal sealed record SourceGeneratorGetTextParams([property: JsonPropertyName("textDocument")] TextDocumentIdentifier TextDocument) : ITextDocumentParams;
+internal sealed record SourceGeneratorGetTextParams(
+    [property: JsonPropertyName("textDocument")] TextDocumentIdentifier TextDocument,
+    [property: JsonPropertyName("resultId")] string? ResultId) : ITextDocumentParams;