Do not fork when source generator document is unchanged and add tests

src/EditorFeatures/TestUtilities/LanguageServer/AbstractLanguageServerProtocolTests.cs

@@ -427,7 +427,7 @@ protected static void AddMappedDocument(Workspace workspace, string markup)
             workspace.TryApplyChanges(newSolution);
         }
 
-        protected static async Task AddGeneratorAsync(ISourceGenerator generator, EditorTestWorkspace workspace)
+        protected static async Task<AnalyzerReference> AddGeneratorAsync(ISourceGenerator generator, EditorTestWorkspace workspace)
         {
             var analyzerReference = new TestGeneratorReference(generator);
 
@@ -438,7 +438,18 @@ protected static async Task AddGeneratorAsync(ISourceGenerator generator, Editor
 
             await workspace.ChangeSolutionAsync(solution);
             await WaitForWorkspaceOperationsAsync(workspace);
+            return analyzerReference;
+        }
+
+        protected static async Task RemoveGeneratorAsync(AnalyzerReference reference, EditorTestWorkspace workspace)
+        {
+            var solution = workspace.CurrentSolution
+                .Projects.Single()
+                .RemoveAnalyzerReference(reference)
+                .Solution;
 
+            await workspace.ChangeSolutionAsync(solution);
+            await WaitForWorkspaceOperationsAsync(workspace);
         }
 
         internal static async Task<Dictionary<string, IList<LSP.Location>>> GetAnnotatedLocationsAsync(EditorTestWorkspace workspace, Solution solution)

src/LanguageServer/Protocol/Workspaces/LspWorkspaceManager.cs

@@ -389,17 +389,23 @@ public void UpdateTrackedDocument(Uri uri, SourceText newSourceText)
             workspaceCurrentSolution = workspace.CurrentSolution;
 
             // Step 3: Check to see if the LSP text matches the workspace text.
+
             var documentsInWorkspace = GetDocumentsForUris(_trackedDocuments.Keys.ToImmutableArray(), workspaceCurrentSolution);
             var sourceGeneratedDocuments =
                 _trackedDocuments.Keys.Where(static uri => uri.Scheme == SourceGeneratedDocumentUri.Scheme)
                     .Select(uri => (identity: SourceGeneratedDocumentUri.DeserializeIdentity(workspaceCurrentSolution, uri), _trackedDocuments[uri].Text))
                     .Where(tuple => tuple.identity.HasValue)
                     .SelectAsArray(tuple => (tuple.identity!.Value, DateTime.Now, tuple.Text));
 
-            // We don't want to check if the source generated document text matches the workspace text because that
-            // will trigger generators to run.  We don't want that to happen in the queue dispatch as it could be quite slow.
+            // First we check if normal document text matches the workspace solution.
+            // This does not look at source generated documents.
             var doesAllTextMatch = await DoesAllTextMatchWorkspaceSolutionAsync(documentsInWorkspace, cancellationToken).ConfigureAwait(false);
-            if (doesAllTextMatch && !sourceGeneratedDocuments.Any())
+
+            // Then we check if source generated document text matches the workspace solution.
+            // This is intentionally done differently from normal documents because the normal method will cause
+            // source generators to run which we do not want to do in queue dispatch.
+            var doesAllSourceGeneratedTextMatch = DoesAllSourceGeneratedTextMatchWorkspaceSolution(sourceGeneratedDocuments, workspaceCurrentSolution);
+            if (doesAllTextMatch && doesAllSourceGeneratedTextMatch)
             {
                 // Remember that the current LSP text matches the text in this workspace solution.
                 _cachedLspSolutions[workspace] = (forkedFromVersion: null, workspaceCurrentSolution);
@@ -412,30 +418,22 @@ public void UpdateTrackedDocument(Uri uri, SourceText newSourceText)
 
             // Step 5: Fork a new solution from the workspace with the LSP text applied.
             var lspSolution = workspaceCurrentSolution;
-            // If the workspace text matched but we have source generated documents open, we can
-            // leave the normal documents as-is (the text matched) and just fork with the frozen sg docs.
+            // If the workspace text matched we can leave the normal documents as-is
             if (!doesAllTextMatch)
             {
                 foreach (var (uri, workspaceDocuments) in documentsInWorkspace)
                     lspSolution = lspSolution.WithDocumentText(workspaceDocuments.Select(d => d.Id), _trackedDocuments[uri].Text);
             }
 
-            lspSolution = lspSolution.WithFrozenSourceGeneratedDocuments(sourceGeneratedDocuments);
-
-            // Did we actually have to fork anything? WithFrozenSourceGeneratedDocuments will return the same instance if we were able to
-            // immediately determine we already had the same generated contents
-            if (lspSolution == workspaceCurrentSolution)
+            // If the source generated documents matched we can leave the source generated documents as-is
+            if (!doesAllSourceGeneratedTextMatch)
             {
-                // Remember that the current LSP text matches the text in this workspace solution.
-                _cachedLspSolutions[workspace] = (forkedFromVersion: null, workspaceCurrentSolution);
-                return (workspaceCurrentSolution, IsForked: false);
-            }
-            else
-            {
-                // Remember this forked solution and the workspace version it was forked from.
-                _cachedLspSolutions[workspace] = (workspaceCurrentSolution.WorkspaceVersion, lspSolution);
-                return (lspSolution, IsForked: true);
+                lspSolution = lspSolution.WithFrozenSourceGeneratedDocuments(sourceGeneratedDocuments);
             }
+
+            // Remember this forked solution and the workspace version it was forked from.
+            _cachedLspSolutions[workspace] = (workspaceCurrentSolution.WorkspaceVersion, lspSolution);
+            return (lspSolution, IsForked: true);
         }
 
         async ValueTask TryOpenAndEditDocumentsInMutatingWorkspaceAsync(Workspace workspace)
@@ -471,6 +469,34 @@ await workspace.TryOnDocumentOpenedAsync(
         }
     }
 
+    /// <summary>
+    /// Checks if the open source generator document contents matches the contents of the workspace solution.
+    /// This looks at the source generator state explicitly to avoid actually running source generators
+    /// </summary>
+    private static bool DoesAllSourceGeneratedTextMatchWorkspaceSolution(
+        ImmutableArray<(SourceGeneratedDocumentIdentity Identity, DateTime Generated, SourceText Text)> sourceGenereatedDocuments,
+        Solution workspaceSolution)
+    {
+        var compilationState = workspaceSolution.CompilationState;
+        foreach (var (identity, _, text) in sourceGenereatedDocuments)
+        {
+            var existingState = compilationState.TryGetSourceGeneratedDocumentStateForAlreadyGeneratedId(identity.DocumentId);
+            if (existingState is null)
+            {
+                // We don't have existing state for at least one of the documents, so the text does cannot match.
+                return false;
+            }
+
+            var newState = existingState.WithText(text);
+            if (newState != existingState)
+            {
+                return false;
+            }
+        }
+
+        return true;
+    }
+
     /// <summary>
     /// Given a set of documents from the workspace current solution, verify that the LSP text is the same as the document contents.
     /// </summary>

src/LanguageServer/ProtocolUnitTests/Workspaces/LspWorkspaceManagerTests.cs

@@ -693,7 +693,7 @@ await testLspServer.ReplaceTextAsync(documentUri,
     }
 
     [Theory, CombinatorialData]
-    public async Task TestUsesForkForUnchangedGeneratedFileAsync(bool mutatingLspWorkspace)
+    public async Task TestDoesNotUseForkForUnchangedGeneratedFileAsync(bool mutatingLspWorkspace)
     {
         var generatorText = "// Hello World!";
         await using var testLspServer = await CreateTestLspServerAsync(string.Empty, mutatingLspWorkspace);
@@ -710,6 +710,47 @@ public async Task TestUsesForkForUnchangedGeneratedFileAsync(bool mutatingLspWor
 
         var sourceGeneratedDocument = await OpenDocumentAndVerifyLspTextAsync(sourceGeneratorDocumentUri, testLspServer, generatorText) as SourceGeneratedDocument;
         AssertEx.NotNull(sourceGeneratedDocument);
+        Assert.Same(testLspServer.TestWorkspace.CurrentSolution, sourceGeneratedDocument.Project.Solution);
+    }
+
+    [Theory, CombinatorialData]
+    public async Task TestForksWithDifferentGeneratedContentsAsync(bool mutatingLspWorkspace)
+    {
+        var workspaceGeneratedText = "// Hello World!";
+        var lspGeneratedText = "// Hello LSP!";
+        await using var testLspServer = await CreateTestLspServerAsync(string.Empty, mutatingLspWorkspace);
+        await AddGeneratorAsync(new SingleFileTestGenerator(workspaceGeneratedText), testLspServer.TestWorkspace);
+
+        var sourceGeneratedDocuments = await testLspServer.GetCurrentSolution().Projects.Single().GetSourceGeneratedDocumentsAsync();
+        var sourceGeneratedDocumentIdentity = sourceGeneratedDocuments.Single().Identity;
+        var sourceGeneratorDocumentUri = SourceGeneratedDocumentUri.Create(sourceGeneratedDocumentIdentity);
+
+        var sourceGeneratedDocument = await OpenDocumentAndVerifyLspTextAsync(sourceGeneratorDocumentUri, testLspServer, lspGeneratedText) as SourceGeneratedDocument;
+        AssertEx.NotNull(sourceGeneratedDocument);
+        Assert.NotSame(testLspServer.TestWorkspace.CurrentSolution, sourceGeneratedDocument.Project.Solution);
+    }
+
+    [Theory, CombinatorialData]
+    public async Task TestForksWithRemovedGeneratorAsync(bool mutatingLspWorkspace)
+    {
+        var generatorText = "// Hello World!";
+        await using var testLspServer = await CreateTestLspServerAsync(string.Empty, mutatingLspWorkspace);
+        var generatorReference = await AddGeneratorAsync(new SingleFileTestGenerator(generatorText), testLspServer.TestWorkspace);
+
+        var sourceGeneratedDocuments = await testLspServer.GetCurrentSolution().Projects.Single().GetSourceGeneratedDocumentsAsync();
+        var sourceGeneratedDocumentIdentity = sourceGeneratedDocuments.Single().Identity;
+        var sourceGeneratorDocumentUri = SourceGeneratedDocumentUri.Create(sourceGeneratedDocumentIdentity);
+
+        var sourceGeneratedDocument = await OpenDocumentAndVerifyLspTextAsync(sourceGeneratorDocumentUri, testLspServer, generatorText) as SourceGeneratedDocument;
+        AssertEx.NotNull(sourceGeneratedDocument);
+        Assert.Same(testLspServer.TestWorkspace.CurrentSolution, sourceGeneratedDocument.Project.Solution);
+
+        // Remove the generator and verify the document is forked.
+        await RemoveGeneratorAsync(generatorReference, testLspServer.TestWorkspace);
+
+        var (_, removedSourceGeneratorDocument) = await GetLspWorkspaceAndDocumentAsync(sourceGeneratorDocumentUri, testLspServer).ConfigureAwait(false);
+        AssertEx.NotNull(sourceGeneratedDocument as SourceGeneratedDocument);
+        Assert.Equal(generatorText, (await sourceGeneratedDocument.GetTextAsync(CancellationToken.None)).ToString());
         Assert.NotSame(testLspServer.TestWorkspace.CurrentSolution, sourceGeneratedDocument.Project.Solution);
     }