Add server side support for refreshing source generated files (#75939)

Client side PR - dotnet/vscode-csharp#7791

Implements support for refreshing source generated files provided by
LSP. Refresh notifications are sent to the client based on workspace
changes that may affect source generated files. To handle the
overreporting of refresh notifications, I also implemented resultId
based support in the get text handler - it will now avoid re-sending the
entire text to the client if nothing has changed.


![sg_refresh](https://github.com/user-attachments/assets/065a8532-56c7-4e45-a559-0a451d37697a)

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;