Allow modification of source generated documents (#77587)

Part of dotnet/razor#10693
Razor side of this is dotnet/razor#11619 which I
think makes us ready to internally dogfood cohosting.

This unblocks Razor completion and most of code actions in cohosting. It
allows `WithText` and `WithSyntaxRoot` to work on source generated
documents, creating a forked solution with frozen source generated
documents. Solutions can be continually frozen without issue, and
unfreezing puts them back to their original state. This also allows code
actions to run on modified source generated documents, but only if they
have been frozen.

As discussed there aren't any flags for this to only be possible
from/for Razor.

Still investigating individual code actions that aren't working in
Razor, but those will probably be a follow up. Enough of the Razor tests
pass now that it proves the system generally works (see comment below).

Author: davidwengier

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