Fix issue with source generator diagnostics not refreshing on save in…

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

@@ -29,21 +29,16 @@ internal abstract partial class AbstractPullDiagnosticHandler<TDiagnosticsParams
     /// the version stamp but not the content (for example, forking LSP text).
     /// </summary>
     private sealed class DiagnosticsPullCache(IGlobalOptionService globalOptions, string uniqueKey)
-        : VersionedPullCache<(int globalStateVersion, VersionStamp? dependentVersion), (int globalStateVersion, Checksum dependentChecksum), DiagnosticsRequestState, ImmutableArray<DiagnosticData>>(uniqueKey)
+        : VersionedPullCache<(int globalStateVersion, Checksum dependentChecksum), DiagnosticsRequestState, ImmutableArray<DiagnosticData>>(uniqueKey)
     {
         private readonly IGlobalOptionService _globalOptions = globalOptions;
 
-        public override async Task<(int globalStateVersion, VersionStamp? dependentVersion)> ComputeCheapVersionAsync(DiagnosticsRequestState state, CancellationToken cancellationToken)
-        {
-            return (state.GlobalStateVersion, await state.Project.GetDependentVersionAsync(cancellationToken).ConfigureAwait(false));
-        }
-
-        public override async Task<(int globalStateVersion, Checksum dependentChecksum)> ComputeExpensiveVersionAsync(DiagnosticsRequestState state, CancellationToken cancellationToken)
+        public override async Task<(int globalStateVersion, Checksum dependentChecksum)> ComputeVersionAsync(DiagnosticsRequestState state, CancellationToken cancellationToken)
         {
             return (state.GlobalStateVersion, await state.Project.GetDiagnosticChecksumAsync(cancellationToken).ConfigureAwait(false));
         }
 
-        /// <inheritdoc cref="VersionedPullCache{TCheapVersion, TExpensiveVersion, TState, TComputedData}.ComputeDataAsync(TState, CancellationToken)"/>
+        /// <inheritdoc cref="VersionedPullCache{TVersion, TState, TComputedData}.ComputeDataAsync(TState, CancellationToken)"/>
         public override async Task<ImmutableArray<DiagnosticData>> ComputeDataAsync(DiagnosticsRequestState state, CancellationToken cancellationToken)
         {
             var diagnostics = await state.DiagnosticSource.GetDiagnosticsAsync(state.Context, cancellationToken).ConfigureAwait(false);

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

@@ -9,10 +9,10 @@
 
 namespace Microsoft.CodeAnalysis.LanguageServer.Handler;
 
-internal abstract partial class VersionedPullCache<TCheapVersion, TExpensiveVersion, TState, TComputedData>
+internal abstract partial class VersionedPullCache<TVersion, TState, TComputedData>
 {
     /// <summary>
-    /// Internal cache item that updates state for a particular <see cref="Workspace"/> and <see cref="ProjectOrDocumentId"/> in <see cref="VersionedPullCache{TCheapVersion, TExpensiveVersion, TState, TComputedData}"/>
+    /// Internal cache item that updates state for a particular <see cref="Workspace"/> and <see cref="ProjectOrDocumentId"/> in <see cref="VersionedPullCache{TVersion, TState, TComputedData}"/>
     /// This type ensures that the state for a particular key is never updated concurrently for the same key (but different key states can be concurrent).
     /// </summary>
     private sealed class CacheItem(string uniqueKey)
@@ -44,15 +44,15 @@ private sealed class CacheItem(string uniqueKey)
         /// </list>
         /// 
         /// </summary>
-        private (string resultId, TCheapVersion cheapVersion, TExpensiveVersion expensiveVersion, Checksum dataChecksum)? _lastResult;
+        private (string resultId, TVersion version, Checksum dataChecksum)? _lastResult;
 
         /// <summary>
         /// Updates the values for this cache entry.  Guarded by <see cref="_gate"/>
         /// 
         /// 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, TComputedData)?> UpdateCacheItemAsync(
-            VersionedPullCache<TCheapVersion, TExpensiveVersion, TState, TComputedData> cache,
+            VersionedPullCache<TVersion, TState, TComputedData> cache,
             PreviousPullResult? previousPullResult,
             bool isFullyLoaded,
             TState state,
@@ -63,38 +63,27 @@ private sealed class CacheItem(string uniqueKey)
             // This means that the computation of new data for this item only occurs sequentially.
             using (await _gate.DisposableWaitAsync(cancellationToken).ConfigureAwait(false))
             {
-                TCheapVersion cheapVersion;
-                TExpensiveVersion expensiveVersion;
+                TVersion version;
 
                 // Check if the version we have in the cache matches the request version.  If so we can re-use the resultId.
                 if (isFullyLoaded &&
                     _lastResult is not null &&
                     _lastResult.Value.resultId == previousPullResult?.PreviousResultId)
                 {
-                    cheapVersion = await cache.ComputeCheapVersionAsync(state, cancellationToken).ConfigureAwait(false);
-                    if (cheapVersion != null && cheapVersion.Equals(_lastResult.Value.cheapVersion))
-                    {
-                        // The client's resultId matches our cached resultId and the cheap version is an
-                        // exact match for our current cheap version. We return early here to avoid calculating
-                        // expensive versions as we know nothing is changed.
-                        return null;
-                    }
-
                     // The current cheap version does not match the last reported.  This may be because we've forked
                     // or reloaded a project, so fall back to calculating the full expensive version to determine if
                     // anything is actually changed.
-                    expensiveVersion = await cache.ComputeExpensiveVersionAsync(state, cancellationToken).ConfigureAwait(false);
-                    if (expensiveVersion != null && expensiveVersion.Equals(_lastResult.Value.expensiveVersion))
+                    version = await cache.ComputeVersionAsync(state, cancellationToken).ConfigureAwait(false);
+                    if (version != null && version.Equals(_lastResult.Value.version))
                     {
                         return null;
                     }
                 }
                 else
                 {
-                    // The versions we have in our cache (if any) do not match the ones provided by the client (if any).
+                    // The version we have in our cache does not match the one provided by the client (if any).
                     // We need to calculate new results.
-                    cheapVersion = await cache.ComputeCheapVersionAsync(state, cancellationToken).ConfigureAwait(false);
-                    expensiveVersion = await cache.ComputeExpensiveVersionAsync(state, cancellationToken).ConfigureAwait(false);
+                    version = await cache.ComputeVersionAsync(state, cancellationToken).ConfigureAwait(false);
                 }
 
                 // Compute the new result for the request.
@@ -111,7 +100,7 @@ _lastResult is not null &&
                     // subsequent requests will always fail the version comparison check (the resultId is still associated with the older version even
                     // though we reused it here for a newer version) and will trigger re-computation.
                     // By storing the updated version with the resultId we can short circuit earlier in the version checks.
-                    _lastResult = (_lastResult.Value.resultId, cheapVersion, expensiveVersion, dataChecksum);
+                    _lastResult = (_lastResult.Value.resultId, version, dataChecksum);
                     return null;
                 }
                 else
@@ -127,7 +116,7 @@ _lastResult is not null &&
                     // Note that we can safely update the map before computation as any cancellation or exception
                     // during computation means that the client will never recieve this resultId and so cannot ask us for it.
                     newResultId = $"{uniqueKey}:{cache.GetNextResultId()}";
-                    _lastResult = (newResultId, cheapVersion, expensiveVersion, dataChecksum);
+                    _lastResult = (newResultId, version, dataChecksum);
                     return (newResultId, data);
                 }
             }

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

@@ -17,7 +17,7 @@ namespace Microsoft.CodeAnalysis.LanguageServer.Handler;
 /// that existing results can be reused, or if new results need to be computed.  Multiple keys can be used,
 /// 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)
+internal abstract partial class VersionedPullCache<TVersion, TState, TComputedData>(string uniqueKey)
 {
     /// <summary>
     /// Map of workspace and diagnostic source to the data used to make the last pull report.
@@ -37,19 +37,12 @@ internal abstract partial class VersionedPullCache<TCheapVersion, TExpensiveVers
     private long _nextDocumentResultId;
 
     /// <summary>
-    /// Computes a cheap version of the current state.  This is compared to the cached version we calculated for the client's previous resultId.
+    /// Computes the version of the current state.  We compare the version of the current state against the
+    /// version we have cached for the client's previous resultId.
     /// 
     /// Note - this will run under the semaphore in <see cref="CacheItem._gate"/>.
     /// </summary>
-    public abstract Task<TCheapVersion> ComputeCheapVersionAsync(TState state, CancellationToken cancellationToken);
-
-    /// <summary>
-    /// Computes a more expensive version of the current state.  If the cheap versions are mismatched, we then compare the expensive version of the current state against the
-    /// expensive version we have cached for the client's previous resultId.
-    /// 
-    /// Note - this will run under the semaphore in <see cref="CacheItem._gate"/>.
-    /// </summary>
-    public abstract Task<TExpensiveVersion> ComputeExpensiveVersionAsync(TState state, CancellationToken cancellationToken);
+    public abstract Task<TVersion> ComputeVersionAsync(TState state, CancellationToken cancellationToken);
 
     /// <summary>
     /// Computes new data for this request.  This data must be hashable as it we store the hash with the requestId to determine if

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

@@ -14,9 +14,9 @@ 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
+internal sealed class SourceGeneratedDocumentCache(string uniqueKey) : VersionedPullCache<(SourceGeneratorExecutionVersion, VersionStamp), SourceGeneratedDocumentGetTextState, SourceText?>(uniqueKey), ILspService
 {
-    public override async Task<(SourceGeneratorExecutionVersion, VersionStamp)> ComputeCheapVersionAsync(SourceGeneratedDocumentGetTextState state, CancellationToken cancellationToken)
+    public override async Task<(SourceGeneratorExecutionVersion, VersionStamp)> ComputeVersionAsync(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).
@@ -25,11 +25,6 @@ internal sealed class SourceGeneratedDocumentCache(string uniqueKey) : Versioned
         return (executionVersion, dependentVersion);
     }
 
-    public override Task<object?> ComputeExpensiveVersionAsync(SourceGeneratedDocumentGetTextState state, CancellationToken cancellationToken)
-    {
-        return SpecializedTasks.Null<object>();
-    }
-
     public override Checksum ComputeChecksum(SourceText? data, string language)
     {
         return data is null ? Checksum.Null : Checksum.From(data.GetChecksum());

src/LanguageServer/Protocol/Handler/SpellCheck/SpellCheckPullCache.cs

@@ -13,12 +13,12 @@ namespace Microsoft.CodeAnalysis.LanguageServer.Handler.SpellCheck;
 internal record struct SpellCheckState(ISpellCheckSpanService Service, Document Document);
 
 /// <summary>
-/// Simplified version of <see cref="VersionedPullCache{TCheapVersion, TExpensiveVersion, TState, TComputedData}"/> that only uses a 
+/// Simplified version of <see cref="VersionedPullCache{TVersion, TState, TComputedData}"/> that only uses a 
 /// single cheap key to check results against.
 /// </summary>
-internal sealed class SpellCheckPullCache(string uniqueKey) : VersionedPullCache<(Checksum parseOptionsChecksum, Checksum textChecksum)?, object?, SpellCheckState, ImmutableArray<SpellCheckSpan>>(uniqueKey)
+internal sealed class SpellCheckPullCache(string uniqueKey) : VersionedPullCache<(Checksum parseOptionsChecksum, Checksum textChecksum), SpellCheckState, ImmutableArray<SpellCheckSpan>>(uniqueKey)
 {
-    public override async Task<(Checksum parseOptionsChecksum, Checksum textChecksum)?> ComputeCheapVersionAsync(SpellCheckState state, CancellationToken cancellationToken)
+    public override async Task<(Checksum parseOptionsChecksum, Checksum textChecksum)> ComputeVersionAsync(SpellCheckState state, CancellationToken cancellationToken)
     {
         var project = state.Document.Project;
         var parseOptionsChecksum = project.State.GetParseOptionsChecksum();
@@ -41,12 +41,6 @@ public override async Task<ImmutableArray<SpellCheckSpan>> ComputeDataAsync(Spel
         return spans;
     }
 
-    public override Task<object?> ComputeExpensiveVersionAsync(SpellCheckState state, CancellationToken cancellationToken)
-    {
-        // Spell check does not need an expensive version check - we return null to effectively skip this check.
-        return SpecializedTasks.Null<object>();
-    }
-
     private void SerializeSpellCheckSpan(SpellCheckSpan span, ObjectWriter writer)
     {
         writer.WriteInt32(span.TextSpan.Start);

src/LanguageServer/ProtocolUnitTests/Diagnostics/PullDiagnosticTests.cs

@@ -680,6 +680,7 @@ public partial class C
 
         // First diagnostic request should report a diagnostic since the generator does not produce any source (text does not match).
         var results = await RunGetDocumentPullDiagnosticsAsync(testLspServer, document.GetURI(), useVSDiagnostics);
+        var firstResultId = results.Single().ResultId;
         var diagnostic = AssertEx.Single(results.Single().Diagnostics);
         Assert.Equal("CS0103", diagnostic.Code);
 
@@ -703,7 +704,8 @@ public partial class C
         }
 
         await testLspServer.WaitForSourceGeneratorsAsync();
-        results = await RunGetDocumentPullDiagnosticsAsync(testLspServer, document.GetURI(), useVSDiagnostics);
+        results = await RunGetDocumentPullDiagnosticsAsync(testLspServer, document.GetURI(), useVSDiagnostics, previousResultId: firstResultId);
+        var secondResultId = results.Single().ResultId;
 
         if (executionPreference == SourceGeneratorExecutionPreference.Automatic)
         {
@@ -712,15 +714,18 @@ public partial class C
         }
         else
         {
-            // In balanced mode, the diagnostic should remain until there is a manual source generator run that updates the sg text.
-            diagnostic = AssertEx.Single(results.Single().Diagnostics);
-            Assert.Equal("CS0103", diagnostic.Code);
+            // In balanced mode, the diagnostic should be unchanged until there is a manual source generator run that updates the sg text.
+            Assert.Null(results.Single().Diagnostics);
+            Assert.Equal(firstResultId, secondResultId);
 
             testLspServer.TestWorkspace.EnqueueUpdateSourceGeneratorVersion(document.Project.Id, forceRegeneration: false);
             await testLspServer.WaitForSourceGeneratorsAsync();
 
-            results = await RunGetDocumentPullDiagnosticsAsync(testLspServer, document.GetURI(), useVSDiagnostics);
+            results = await RunGetDocumentPullDiagnosticsAsync(testLspServer, document.GetURI(), useVSDiagnostics, previousResultId: secondResultId);
+            var thirdResultId = results.Single().ResultId;
+            AssertEx.NotNull(results.Single().Diagnostics);
             Assert.Empty(results.Single().Diagnostics!);
+            Assert.NotEqual(firstResultId, thirdResultId);
         }
     }