PR feedback on recent work to sync solutions consistently

src/Workspaces/Core/Portable/SourceGeneration/IRemoteSourceGenerationService.cs

@@ -27,24 +27,22 @@ internal interface IRemoteSourceGenerationService
     /// compare that to the prior generated documents it has to see if it can reuse those directly, or if it needs to
     /// remove any documents no longer around, add any new documents, or change the contents of any existing documents.
     /// </summary>
-    /// <remarks>
-    /// Should only be called by the "RegularCompilationTracker", and should only return data from its view of the
-    /// world.  Not from the view of a "GeneratedFileReplacingCompilationTracker".
-    /// </remarks>
-    ValueTask<ImmutableArray<RegularCompilationTrackerSourceGenerationInfo>> GetRegularCompilationTrackerSourceGenerationInfoAsync(
-        Checksum solutionChecksum, ProjectId projectId, CancellationToken cancellationToken);
+    /// <param name="withFrozenSourceGeneratedDocuments">Controls if the caller wants frozen source generator documents
+    /// included in the result, or if only the most underlying generated documents (produced by the real compiler <see
+    /// cref="GeneratorDriver"/> should be included.</param>
+    ValueTask<ImmutableArray<RegularCompilationTrackerSourceGenerationInfo>> GetSourceGenerationInfoAsync(
+        Checksum solutionChecksum, ProjectId projectId, bool withFrozenSourceGeneratedDocuments, CancellationToken cancellationToken);
 
     /// <summary>
     /// Given a particular set of generated document ids, returns the fully generated content for those documents.
     /// Should only be called by the host for documents it does not know about, or documents whose checksum contents are
     /// different than the last time the document was queried.
     /// </summary>
-    /// <remarks>
-    /// Should only be called by the "RegularCompilationTracker", and should only return data from its view of the
-    /// world.  Not from the view of a "GeneratedFileReplacingCompilationTracker".
-    /// </remarks>
-    ValueTask<ImmutableArray<string>> GetRegularCompilationTrackerContentsAsync(
-        Checksum solutionChecksum, ProjectId projectId, ImmutableArray<DocumentId> documentIds, CancellationToken cancellationToken);
+    /// <param name="withFrozenSourceGeneratedDocuments">Controls if the caller wants frozen source generator documents
+    /// included in the result, or if only the most underlying generated documents (produced by the real compiler <see
+    /// cref="GeneratorDriver"/> should be included.</param>
+    ValueTask<ImmutableArray<string>> GetContentsAsync(
+        Checksum solutionChecksum, ProjectId projectId, ImmutableArray<DocumentId> documentIds, bool withFrozenSourceGeneratedDocuments, CancellationToken cancellationToken);
 
     /// <summary>
     /// Whether or not the specified analyzer references have source generators or not.

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

@@ -56,21 +56,15 @@ bool ContainsAssemblyOrModuleOrDynamic(
         Task<Checksum> GetDependentChecksumAsync(SolutionCompilationState compilationState, CancellationToken cancellationToken);
 
         /// <summary>
-        /// Gets the *final* view of the generated documents for this tracker.  If this is a <see
-        /// cref="RegularCompilationTracker"/> this will be the true generated documents for this tracker (generated by
-        /// our underlying <see cref="GeneratorDriver"/>).  If this is a <see
-        /// cref="GeneratedFileReplacingCompilationTracker"/>, then this will be the generated documents of its <see
-        /// cref="GeneratedFileReplacingCompilationTracker.UnderlyingTracker"/>, along with all of its replacement
-        /// frozen documents overlaid on top.
+        /// Gets the source generator files generated by this <see cref="ICompilationTracker"/>. <paramref
+        /// name="withFrozenSourceGeneratedDocuments"/>Controls whether frozen source generated documents are included
+        /// in the result.  If <see langword="false"/> this will call all the way through to the most underlying <see
+        /// cref="RegularCompilationTracker"/> to get its generated documents.  If this is <see langword="true"/> then
+        /// this will be those same generated documents, along with all the generated documents from all wrapping <see
+        /// cref="GeneratedFileReplacingCompilationTracker"/>'s frozen documents overlaid on top.
         /// </summary>
-        ValueTask<TextDocumentStates<SourceGeneratedDocumentState>> GetSourceGeneratedDocumentStatesAsync(SolutionCompilationState compilationState, CancellationToken cancellationToken);
-
-        /// <summary>
-        /// Equivalent to <see cref="GetSourceGeneratedDocumentStatesAsync"/>, but only returning from the innermost
-        /// underlying <see cref="RegularCompilationTracker"/>.  Any frozen generated documents in a <see
-        /// cref="GeneratedFileReplacingCompilationTracker"/> will *not* be overlaid on top of this.
-        /// </summary>
-        ValueTask<TextDocumentStates<SourceGeneratedDocumentState>> GetRegularCompilationTrackerSourceGeneratedDocumentStatesAsync(SolutionCompilationState compilationState, CancellationToken cancellationToken);
+        ValueTask<TextDocumentStates<SourceGeneratedDocumentState>> GetSourceGeneratedDocumentStatesAsync(
+            SolutionCompilationState compilationState, bool withFrozenSourceGeneratedDocuments, CancellationToken cancellationToken);
 
         ValueTask<ImmutableArray<Diagnostic>> GetSourceGeneratorDiagnosticsAsync(SolutionCompilationState compilationState, CancellationToken cancellationToken);
         ValueTask<GeneratorDriverRunResult?> GetSourceGeneratorRunResultAsync(SolutionCompilationState solution, CancellationToken cancellationToken);

src/Workspaces/Core/Portable/Workspace/Solution/SolutionCompilationState.CompilationTracker.cs → src/Workspaces/Core/Portable/Workspace/Solution/SolutionCompilationState.RegularCompilationTracker.cs

@@ -837,13 +837,12 @@ public ICompilationTracker WithDoNotCreateCreationPolicy(CancellationToken cance
                 }
             }
 
-            public ValueTask<TextDocumentStates<SourceGeneratedDocumentState>> GetSourceGeneratedDocumentStatesAsync(SolutionCompilationState compilationState, CancellationToken cancellationToken)
-                // Just defer to the core function that creates these.  They are the right values for both of these calls.
-                => GetRegularCompilationTrackerSourceGeneratedDocumentStatesAsync(compilationState, cancellationToken);
-
-            public async ValueTask<TextDocumentStates<SourceGeneratedDocumentState>> GetRegularCompilationTrackerSourceGeneratedDocumentStatesAsync(
-                SolutionCompilationState compilationState, CancellationToken cancellationToken)
+            public async ValueTask<TextDocumentStates<SourceGeneratedDocumentState>> GetSourceGeneratedDocumentStatesAsync(
+                SolutionCompilationState compilationState, bool withFrozenSourceGeneratedDocuments, CancellationToken cancellationToken)
             {
+                // Note: withFrozenSourceGeneratedDocuments has no impact on is.  We're always returning real generated
+                // docs, not frozen docs.  Frozen docs are only involved with a GeneratedFileReplacingCompilationTracker
+
                 // If we don't have any generators, then we know we have no generated files, so we can skip the computation entirely.
                 if (!await compilationState.HasSourceGeneratorsAsync(this.ProjectState.Id, cancellationToken).ConfigureAwait(false))
                     return TextDocumentStates<SourceGeneratedDocumentState>.Empty;

src/Workspaces/Core/Portable/Workspace/Solution/SolutionCompilationState.CompilationTracker_Generators.cs → src/Workspaces/Core/Portable/Workspace/Solution/SolutionCompilationState.RegularCompilationTracker_Generators.cs

@@ -93,14 +93,14 @@ private partial class RegularCompilationTracker : ICompilationTracker
 
             // First, grab the info from our external host about the generated documents it has for this project.  Note:
             // we ourselves are the innermost "RegularCompilationTracker" responsible for actually running generators.
-            // As such, our call to the oop side reflects that by asking for only *its* innermost
-            // RegularCompilationTracker to do the same.
+            // As such, our call to the oop side reflects that by asking for the real source generated docs, and *not*
+            // any overlaid 'frozen' source generated documents.
             var projectId = this.ProjectState.Id;
             var infosOpt = await connection.TryInvokeAsync(
                 compilationState,
                 projectId,
-                (service, solutionChecksum, cancellationToken) => service.GetRegularCompilationTrackerSourceGenerationInfoAsync(
-                    solutionChecksum, projectId, cancellationToken),
+                (service, solutionChecksum, cancellationToken) => service.GetSourceGenerationInfoAsync(
+                    solutionChecksum, projectId, withFrozenSourceGeneratedDocuments: false, cancellationToken),
                 cancellationToken).ConfigureAwait(false);
 
             if (!infosOpt.HasValue)
@@ -160,13 +160,13 @@ private partial class RegularCompilationTracker : ICompilationTracker
             // Either we generated a different number of files, and/or we had contents of files that changed. Ensure we
             // know the contents of any new/changed files.  Note: we ourselves are the innermost
             // "RegularCompilationTracker" responsible for actually running generators. As such, our call to the oop
-            // side reflects that by asking for the source gen contents produced by *its* innermost
-            // RegularCompilationTracker.
+            // side reflects that by asking for the real source generated docs, and *not* any overlaid 'frozen' source
+            // generated documents.
             var generatedSourcesOpt = await connection.TryInvokeAsync(
                 compilationState,
                 projectId,
-                (service, solutionChecksum, cancellationToken) => service.GetRegularCompilationTrackerContentsAsync(
-                    solutionChecksum, projectId, documentsToAddOrUpdate.ToImmutable(), cancellationToken),
+                (service, solutionChecksum, cancellationToken) => service.GetContentsAsync(
+                    solutionChecksum, projectId, documentsToAddOrUpdate.ToImmutable(), withFrozenSourceGeneratedDocuments: false, cancellationToken),
                 cancellationToken).ConfigureAwait(false);
 
             if (!generatedSourcesOpt.HasValue)

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

@@ -14,11 +14,17 @@ namespace Microsoft.CodeAnalysis;
 internal partial class SolutionCompilationState
 {
     /// <summary>
-    /// An implementation of <see cref="ICompilationTracker"/> that takes a compilation from another compilation tracker and updates it
-    /// to return a generated document with a specific content, regardless of what the generator actually produces. In other words, it says
-    /// "take the compilation this other thing produced, and pretend the generator gave this content, even if it wouldn't."
+    /// An implementation of <see cref="ICompilationTracker"/> that takes a compilation from another compilation tracker
+    /// and updates it to return a generated document with a specific content, regardless of what the generator actually
+    /// produces. In other words, it says "take the compilation this other thing produced, and pretend the generator
+    /// gave this content, even if it wouldn't."  This is used by <see
+    /// cref="Solution.WithFrozenSourceGeneratedDocuments"/> to ensure that a particular solution snapshot contains a
+    /// pre-existing generated document from a prior run that the user is interacting with in the host.  The current
+    /// snapshot might not produce the same content from before (or may not even produce that document anymore).  But we
+    /// want to still let the user work with that doc effectively up until the point that new generated documents are
+    /// produced and replace it in the host view.
     /// </summary>
-    private sealed class GeneratedFileReplacingCompilationTracker : ICompilationTracker
+    private sealed class WithFrozenSourceGeneratedDocumentsCompilationTracker : ICompilationTracker
     {
         private readonly TextDocumentStates<SourceGeneratedDocumentState> _replacementDocumentStates;
 
@@ -41,7 +47,7 @@ private sealed class GeneratedFileReplacingCompilationTracker : ICompilationTrac
         /// </summary>
         private SkeletonReferenceCache _skeletonReferenceCache;
 
-        public GeneratedFileReplacingCompilationTracker(
+        public WithFrozenSourceGeneratedDocumentsCompilationTracker(
             ICompilationTracker underlyingTracker,
             TextDocumentStates<SourceGeneratedDocumentState> replacementDocumentStates)
         {
@@ -80,15 +86,15 @@ public ICompilationTracker WithCreateCreationPolicy(bool forceRegeneration)
             var underlyingTracker = this.UnderlyingTracker.WithCreateCreationPolicy(forceRegeneration);
             return underlyingTracker == this.UnderlyingTracker
                 ? this
-                : new GeneratedFileReplacingCompilationTracker(underlyingTracker, _replacementDocumentStates);
+                : new WithFrozenSourceGeneratedDocumentsCompilationTracker(underlyingTracker, _replacementDocumentStates);
         }
 
         public ICompilationTracker WithDoNotCreateCreationPolicy(CancellationToken cancellationToken)
         {
             var underlyingTracker = this.UnderlyingTracker.WithDoNotCreateCreationPolicy(cancellationToken);
             return underlyingTracker == this.UnderlyingTracker
                 ? this
-                : new GeneratedFileReplacingCompilationTracker(underlyingTracker, _replacementDocumentStates);
+                : new WithFrozenSourceGeneratedDocumentsCompilationTracker(underlyingTracker, _replacementDocumentStates);
         }
 
         public async Task<Compilation> GetCompilationAsync(SolutionCompilationState compilationState, CancellationToken cancellationToken)
@@ -99,7 +105,8 @@ public async Task<Compilation> GetCompilationAsync(SolutionCompilationState comp
                 return _compilationWithReplacements;
             }
 
-            var underlyingSourceGeneratedDocuments = await UnderlyingTracker.GetSourceGeneratedDocumentStatesAsync(compilationState, cancellationToken).ConfigureAwait(false);
+            var underlyingSourceGeneratedDocuments = await UnderlyingTracker.GetSourceGeneratedDocumentStatesAsync(
+                compilationState, withFrozenSourceGeneratedDocuments: true, cancellationToken).ConfigureAwait(false);
             var newCompilation = await UnderlyingTracker.GetCompilationAsync(compilationState, cancellationToken).ConfigureAwait(false);
 
             foreach (var (id, replacementState) in _replacementDocumentStates.States)
@@ -159,41 +166,37 @@ await UnderlyingTracker.GetDependentChecksumAsync(compilationState, cancellation
                 (await _replacementDocumentStates.GetDocumentChecksumsAndIdsAsync(cancellationToken).ConfigureAwait(false)).Checksum);
 
         public async ValueTask<TextDocumentStates<SourceGeneratedDocumentState>> GetSourceGeneratedDocumentStatesAsync(
-            SolutionCompilationState compilationState, CancellationToken cancellationToken)
+            SolutionCompilationState compilationState, bool withFrozenSourceGeneratedDocuments, CancellationToken cancellationToken)
         {
             var newStates = await UnderlyingTracker.GetSourceGeneratedDocumentStatesAsync(
-                compilationState, cancellationToken).ConfigureAwait(false);
+                compilationState, withFrozenSourceGeneratedDocuments, cancellationToken).ConfigureAwait(false);
 
-            foreach (var (id, replacementState) in _replacementDocumentStates.States)
+            // Only if the caller *wants* frozen source generated documents, then we will overlay the real underlying
+            // generated docs with the frozen ones we're pointing at.
+            if (withFrozenSourceGeneratedDocuments)
             {
-                if (newStates.Contains(id))
-                {
-                    // The generated file still exists in the underlying compilation, but the contents may not match the open file if the open file
-                    // is stale. Replace the syntax tree so we have a tree that matches the text.
-                    newStates = newStates.SetState(replacementState);
-                }
-                else
+                foreach (var (id, replacementState) in _replacementDocumentStates.States)
                 {
-                    // The generated output no longer exists in the underlying compilation. This could happen if the user made
-                    // an edit which would cause this file to no longer exist, but they're still operating on an open representation
-                    // of that file. To ensure that this snapshot is still usable, we'll just add this document back in. This is not a
-                    // semantically correct operation, but working on stale snapshots never has that guarantee.
-                    newStates = newStates.AddRange([replacementState]);
+                    if (newStates.Contains(id))
+                    {
+                        // The generated file still exists in the underlying compilation, but the contents may not match the open file if the open file
+                        // is stale. Replace the syntax tree so we have a tree that matches the text.
+                        newStates = newStates.SetState(replacementState);
+                    }
+                    else
+                    {
+                        // The generated output no longer exists in the underlying compilation. This could happen if the user made
+                        // an edit which would cause this file to no longer exist, but they're still operating on an open representation
+                        // of that file. To ensure that this snapshot is still usable, we'll just add this document back in. This is not a
+                        // semantically correct operation, but working on stale snapshots never has that guarantee.
+                        newStates = newStates.AddRange([replacementState]);
+                    }
                 }
             }
 
             return newStates;
         }
 
-        public ValueTask<TextDocumentStates<SourceGeneratedDocumentState>> GetRegularCompilationTrackerSourceGeneratedDocumentStatesAsync(
-            SolutionCompilationState compilationState, CancellationToken cancellationToken)
-        {
-            // Just defer to the underlying tracker.  The caller only wants the innermost generated documents, not any
-            // frozen docs we'll overlay on top of it.  The caller will already know about those frozen documents and
-            // will do its own overlay on these results.
-            return this.UnderlyingTracker.GetRegularCompilationTrackerSourceGeneratedDocumentStatesAsync(compilationState, cancellationToken);
-        }
-
         public Task<bool> HasSuccessfullyLoadedAsync(
             SolutionCompilationState compilationState, CancellationToken cancellationToken)
         {