Allow modification of source generated documents

src/Features/Core/Portable/Completion/Providers/AbstractRecommendationServiceBasedCompletionProvider.cs

@@ -226,7 +226,7 @@ internal sealed override async Task<CompletionDescription> GetDescriptionWorkerA
 
         async Task<CompletionDescription?> TryGetDescriptionAsync(DocumentId documentId)
         {
-            var relatedDocument = document.Project.Solution.GetRequiredDocument(documentId);
+            var relatedDocument = await document.Project.Solution.GetRequiredDocumentAsync(documentId, includeSourceGenerated: true, cancellationToken).ConfigureAwait(false);
             var context = await Utilities.CreateSyntaxContextWithExistingSpeculativeModelAsync(relatedDocument, position, cancellationToken).ConfigureAwait(false) as TSyntaxContext;
             Contract.ThrowIfNull(context);
             var symbols = await TryGetSymbolsForContextAsync(completionContext: null, context, options, cancellationToken).ConfigureAwait(false);

src/LanguageServer/Protocol/Handler/CodeActions/CodeActionResolveHelper.cs

@@ -184,6 +184,17 @@ await AddTextDocumentEditsAsync(
                 newSolution.GetDocument,
                 solution.GetDocument).ConfigureAwait(false);
 
+            // Razor calls through our code action handlers with documents that come from the Razor source generator
+            // Those changes are not visible in project changes, because they happen in the compilation state, so we
+            // make sure to pull changes out from that too. Changed source generated documents are "frozen" because
+            // their content no longer comes from the source generator, so thats our cue to know when to handle.
+            // Changes to non-frozen documents don't need to be included, because changes to the origin document would
+            // cause the generator to re-generate the same changed content.
+            await AddTextDocumentEditsAsync(
+                changes.GetExplicitlyChangedSourceGeneratedDocuments(),
+                newSolution.GetRequiredSourceGeneratedDocumentForAlreadyGeneratedId,
+                solution.GetRequiredSourceGeneratedDocumentForAlreadyGeneratedId).ConfigureAwait(false);
+
             // Changed analyzer config documents
             await AddTextDocumentEditsAsync(
                 projectChanges.SelectMany(pc => pc.GetChangedAnalyzerConfigDocuments()),

src/Workspaces/Core/Portable/CodeActions/CodeAction_Cleanup.cs

@@ -65,8 +65,9 @@ internal static ImmutableArray<DocumentId> GetAllChangedOrAddedDocumentIds(
             .GetProjectChanges()
             .SelectMany(p => p.GetChangedDocuments(onlyGetDocumentsWithTextChanges: true).Concat(p.GetAddedDocuments()))
             .Concat(solutionChanges.GetAddedProjects().SelectMany(p => p.DocumentIds))
-            .ToImmutableArray();
-        return documentIds;
+            .Concat(solutionChanges.GetExplicitlyChangedSourceGeneratedDocuments());
+
+        return documentIds.ToImmutableArray();
     }
 
     internal static async Task<Solution> CleanSyntaxAndSemanticsAsync(
@@ -89,7 +90,10 @@ internal static async Task<Solution> CleanSyntaxAndSemanticsAsync(
             using var _ = ArrayBuilder<(DocumentId documentId, CodeCleanupOptions options)>.GetInstance(documentIds.Length, out var documentIdsAndOptions);
             foreach (var documentId in documentIds)
             {
-                var document = changedSolution.GetRequiredDocument(documentId);
+                // We include source generated documents here for Razor, which uses them. In that scenario the cleaned document is compared to the
+                // original to create a set of changes for the LSP client, and part of that will include mapping the changes back to the Razor document,
+                // so whilst it would seem like cleaning source generated documents is a waste of time, it's sometimes not.
+                var document = await changedSolution.GetRequiredDocumentAsync(documentId, includeSourceGenerated: true, cancellationToken).ConfigureAwait(false);
 
                 // Only care about documents that support syntax.  Non-C#/VB files can't be cleaned.
                 if (document.SupportsSyntaxTree)
@@ -114,7 +118,7 @@ internal static async ValueTask<Document> CleanupDocumentAsync(Document document
             CodeAnalysisProgress.None,
             cancellationToken).ConfigureAwait(false);
 
-        return cleanedSolution.GetRequiredDocument(document.Id);
+        return await cleanedSolution.GetRequiredDocumentAsync(document.Id, includeSourceGenerated: true, cancellationToken).ConfigureAwait(false);
     }
 
     private static async Task<Solution> RunAllCleanupPassesInOrderAsync(
@@ -152,7 +156,7 @@ async Task<Solution> RunParallelCleanupPassAsync(
                     var (documentId, options) = documentIdAndOptions;
 
                     // Fetch the current state of the document from this fork of the solution.
-                    var document = solution.GetRequiredDocument(documentId);
+                    var document = await solution.GetRequiredDocumentAsync(documentId, includeSourceGenerated: true, cancellationToken).ConfigureAwait(false);
                     Contract.ThrowIfFalse(document.SupportsSyntaxTree, "GetDocumentIdsAndOptionsAsync should only be returning documents that support syntax");
 
                     // Now, perform the requested cleanup pass on it.

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

@@ -388,13 +388,33 @@ public Document WithSourceCodeKind(SourceCodeKind kind)
     /// Creates a new instance of this document updated to have the text specified.
     /// </summary>
     public Document WithText(SourceText text)
-        => this.Project.Solution.WithDocumentText(this.Id, text, PreservationMode.PreserveIdentity).GetRequiredDocument(Id);
+    {
+        var solution = this.Project.Solution.WithDocumentText(this.Id, text, PreservationMode.PreserveIdentity);
+
+        if (Id.IsSourceGenerated)
+        {
+            // We just modified the text of the generated document, so it should be available synchronously, and throwing is appropriate if it isn't.
+            return solution.GetRequiredSourceGeneratedDocumentForAlreadyGeneratedId(Id);
+        }
+
+        return solution.GetRequiredDocument(Id);
+    }
 
     /// <summary>
     /// Creates a new instance of this document updated to have a syntax tree rooted by the specified syntax node.
     /// </summary>
     public Document WithSyntaxRoot(SyntaxNode root)
-        => this.Project.Solution.WithDocumentSyntaxRoot(this.Id, root, PreservationMode.PreserveIdentity).GetRequiredDocument(Id);
+    {
+        var solution = this.Project.Solution.WithDocumentSyntaxRoot(this.Id, root, PreservationMode.PreserveIdentity);
+
+        if (Id.IsSourceGenerated)
+        {
+            // We just modified the text of the generated document, so it should be available synchronously, and throwing is appropriate if it isn't.
+            return solution.GetRequiredSourceGeneratedDocumentForAlreadyGeneratedId(Id);
+        }
+
+        return solution.GetRequiredDocument(Id);
+    }
 
     /// <summary>
     /// Creates a new instance of this document updated to have the specified name.

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

@@ -549,7 +549,7 @@ internal DocumentState UpdateTree(SyntaxNode newRoot, PreservationMode mode)
         }
     }
 
-    private VersionStamp GetNewTreeVersionForUpdatedTree(SyntaxNode newRoot, VersionStamp newTextVersion, PreservationMode mode)
+    protected VersionStamp GetNewTreeVersionForUpdatedTree(SyntaxNode newRoot, VersionStamp newTextVersion, PreservationMode mode)
     {
         RoslynDebug.Assert(TreeSource != null);
 
@@ -566,7 +566,7 @@ private VersionStamp GetNewTreeVersionForUpdatedTree(SyntaxNode newRoot, Version
         return oldRoot.IsEquivalentTo(newRoot, topLevel: true) ? oldTreeAndVersion.Version : newTextVersion;
     }
 
-    private VersionStamp GetNewerVersion()
+    protected VersionStamp GetNewerVersion()
     {
         if (TextAndVersionSource.TryGetValue(LoadTextOptions, out var textAndVersion))
         {

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

@@ -332,8 +332,7 @@ internal SourceGeneratedDocument GetOrCreateSourceGeneratedDocument(SourceGenera
     /// </summary>
     /// <remarks>
     /// This is only safe to call if you already have seen the SyntaxTree or equivalent that indicates the document state has already been
-    /// generated. This method exists to implement <see cref="Solution.GetDocument(SyntaxTree?)"/> and is best avoided unless you're doing something
-    /// similarly tricky like that.
+    /// generated. This method is best avoided unless you manually ensure you realise the generated document before calling this method.
     /// </remarks>
     internal SourceGeneratedDocument? TryGetSourceGeneratedDocumentForAlreadyGeneratedId(DocumentId documentId)
     {

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

@@ -1616,7 +1616,8 @@ public Solution WithDocumentText(IEnumerable<DocumentId?> documentIds, SourceTex
     internal Document WithFrozenSourceGeneratedDocument(
         SourceGeneratedDocumentIdentity documentIdentity, DateTime generationDateTime, SourceText text)
     {
-        var newCompilationState = CompilationState.WithFrozenSourceGeneratedDocuments([(documentIdentity, generationDateTime, text)]);
+        // SyntaxNode is null here because it will be computed on demand. Other APIs, like Document.WithSyntaxRoot, specify it.
+        var newCompilationState = CompilationState.WithFrozenSourceGeneratedDocuments([(documentIdentity, generationDateTime, text, syntaxNode: null)]);
         var newSolution = WithCompilationState(newCompilationState);
 
         var newDocumentState = newCompilationState.TryGetSourceGeneratedDocumentStateForAlreadyGeneratedId(documentIdentity.DocumentId);
@@ -1627,7 +1628,7 @@ internal Document WithFrozenSourceGeneratedDocument(
     }
 
     internal Solution WithFrozenSourceGeneratedDocuments(ImmutableArray<(SourceGeneratedDocumentIdentity documentIdentity, DateTime generationDateTime, SourceText text)> documents)
-        => WithCompilationState(CompilationState.WithFrozenSourceGeneratedDocuments(documents));
+        => WithCompilationState(CompilationState.WithFrozenSourceGeneratedDocuments(documents.SelectAsArray(d => (d.documentIdentity, d.generationDateTime, d.text, (SyntaxNode?)null))));
 
     /// <inheritdoc cref="SolutionCompilationState.UpdateSpecificSourceGeneratorExecutionVersions"/>
     internal Solution UpdateSpecificSourceGeneratorExecutionVersions(SourceGeneratorExecutionVersionMap sourceGeneratorExecutionVersionMap)
@@ -1740,9 +1741,26 @@ private void CheckContainsDocument(DocumentId documentId)
             throw new ArgumentNullException(nameof(documentId));
         }
 
-        if (!ContainsDocument(documentId))
+        // For source generated documents we expect them to be already generated to use any of the APIs that call this
+        if (documentId.IsSourceGenerated && ContainsSourceGeneratedDocument(documentId))
         {
-            throw new InvalidOperationException(WorkspaceExtensionsResources.The_solution_does_not_contain_the_specified_document);
+            return;
+        }
+
+        if (ContainsDocument(documentId))
+        {
+            return;
+        }
+
+        throw new InvalidOperationException(WorkspaceExtensionsResources.The_solution_does_not_contain_the_specified_document);
+
+        bool ContainsSourceGeneratedDocument(DocumentId documentId)
+        {
+            var project = this.GetProject(documentId.ProjectId);
+            if (project is null)
+                return false;
+
+            return project.TryGetSourceGeneratedDocumentForAlreadyGeneratedId(documentId) is not null;
         }
     }
 

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

@@ -2,9 +2,12 @@
 // 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.Collections.Generic;
 using Microsoft.CodeAnalysis.Diagnostics;
+using Microsoft.CodeAnalysis.PooledObjects;
 using Microsoft.CodeAnalysis.Shared.Extensions;
+using Roslyn.Utilities;
 
 namespace Microsoft.CodeAnalysis;
 
@@ -80,4 +83,31 @@ public IEnumerable<AnalyzerReference> GetRemovedAnalyzerReferences()
             }
         }
     }
+
+    /// <summary>
+    /// Gets changed source generated document ids that were modified with <see cref="Solution.WithFrozenSourceGeneratedDocuments(System.Collections.Immutable.ImmutableArray{ValueTuple{SourceGeneratedDocumentIdentity, DateTime, Text.SourceText}})"/>
+    /// </summary>
+    /// <remarks>
+    /// It is possible for a source generated document to be "frozen" without it existing in the solution, and in that case
+    /// this method will not return that document. This only returns changes to source generated documents, hence they had
+    /// to already be observed in the old solution.
+    /// </remarks>
+    internal IEnumerable<DocumentId> GetExplicitlyChangedSourceGeneratedDocuments()
+    {
+        if (_newSolution.CompilationState.FrozenSourceGeneratedDocumentStates.IsEmpty)
+            return [];
+
+        using var _ = ArrayBuilder<SourceGeneratedDocumentState>.GetInstance(out var oldStateBuilder);
+        foreach (var (id, _) in _newSolution.CompilationState.FrozenSourceGeneratedDocumentStates.States)
+        {
+            var oldState = _oldSolution.CompilationState.TryGetSourceGeneratedDocumentStateForAlreadyGeneratedId(id);
+            oldStateBuilder.AddIfNotNull(oldState);
+        }
+
+        var oldStates = new TextDocumentStates<SourceGeneratedDocumentState>(oldStateBuilder);
+        return _newSolution.CompilationState.FrozenSourceGeneratedDocumentStates.GetChangedStateIds(
+            oldStates,
+            ignoreUnchangedContent: true,
+            ignoreUnchangeableDocuments: false);
+    }
 }

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

@@ -94,6 +94,18 @@ public ICompilationTracker WithDoNotCreateCreationPolicy()
                 : new WithFrozenSourceGeneratedDocumentsCompilationTracker(underlyingTracker, _replacementDocumentStates);
         }
 
+        /// <summary>
+        /// Updates the frozen source generated documents states being tracked
+        /// </summary>
+        /// <remarks>
+        /// NOTE: This does not merge the states currently tracked, it simply replaces them. If merging is desired, it should be done
+        /// by the caller.
+        /// </remarks>
+        public ICompilationTracker WithReplacementDocumentStates(TextDocumentStates<SourceGeneratedDocumentState> replacementDocumentStates)
+        {
+            return new WithFrozenSourceGeneratedDocumentsCompilationTracker(this.UnderlyingTracker, replacementDocumentStates);
+        }
+
         public async Task<Compilation> GetCompilationAsync(SolutionCompilationState compilationState, CancellationToken cancellationToken)
         {
             // Fast path if we've definitely already done this before