Correctly handle a cached solution having frozen generated documents

src/VisualStudio/Core/Test.Next/Services/SolutionServiceTests.cs

@@ -486,18 +486,29 @@ public async Task TestFrozenSourceGeneratedDocument()
                 .AddAnalyzerReference(new AnalyzerFileReference(typeof(Microsoft.CodeAnalysis.TestSourceGenerator.HelloWorldGenerator).Assembly.Location, new TestAnalyzerAssemblyLoader()))
                 .Solution;
 
+            // First sync the solution over that has a generator
             var assetProvider = await GetAssetProviderAsync(workspace, remoteWorkspace, solution);
             var solutionChecksum = await solution.State.GetChecksumAsync(CancellationToken.None);
             var synched = await remoteWorkspace.GetSolutionAsync(assetProvider, solutionChecksum, fromPrimaryBranch: true, workspaceVersion: 0, projectId: null, CancellationToken.None);
             Assert.Equal(solutionChecksum, await synched.State.GetChecksumAsync(CancellationToken.None));
 
+            // Now freeze with some content
             var documentIdentity = (await solution.Projects.Single().GetSourceGeneratedDocumentsAsync()).First().Identity;
-            var frozenText = SourceText.From("// Hello, World!");
-            solution = solution.WithFrozenSourceGeneratedDocument(documentIdentity, frozenText).Project.Solution;
+            var frozenText1 = SourceText.From("// Hello, World!");
+            var frozenSolution1 = solution.WithFrozenSourceGeneratedDocument(documentIdentity, frozenText1).Project.Solution;
 
-            assetProvider = await GetAssetProviderAsync(workspace, remoteWorkspace, solution);
-            solutionChecksum = await solution.State.GetChecksumAsync(CancellationToken.None);
-            synched = await remoteWorkspace.GetSolutionAsync(assetProvider, solutionChecksum, fromPrimaryBranch: false, workspaceVersion: 1, projectId: null, CancellationToken.None);
+            assetProvider = await GetAssetProviderAsync(workspace, remoteWorkspace, frozenSolution1);
+            solutionChecksum = await frozenSolution1.State.GetChecksumAsync(CancellationToken.None);
+            synched = await remoteWorkspace.GetSolutionAsync(assetProvider, solutionChecksum, fromPrimaryBranch: true, workspaceVersion: 1, projectId: null, CancellationToken.None);
+            Assert.Equal(solutionChecksum, await synched.State.GetChecksumAsync(CancellationToken.None));
+
+            // Try freezing with some different content from the original solution
+            var frozenText2 = SourceText.From("// Hello, World! A second time!");
+            var frozenSolution2 = solution.WithFrozenSourceGeneratedDocument(documentIdentity, frozenText2).Project.Solution;
+
+            assetProvider = await GetAssetProviderAsync(workspace, remoteWorkspace, frozenSolution2);
+            solutionChecksum = await frozenSolution2.State.GetChecksumAsync(CancellationToken.None);
+            synched = await remoteWorkspace.GetSolutionAsync(assetProvider, solutionChecksum, fromPrimaryBranch: true, workspaceVersion: 2, projectId: null, CancellationToken.None);
             Assert.Equal(solutionChecksum, await synched.State.GetChecksumAsync(CancellationToken.None));
         }
 

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

@@ -1758,6 +1758,21 @@ internal Document WithFrozenSourceGeneratedDocument(SourceGeneratedDocumentIdent
             return newProject.GetOrCreateSourceGeneratedDocument(newDocumentState);
         }
 
+        /// <summary>
+        /// Undoes the operation of <see cref="WithFrozenSourceGeneratedDocument"/>; any frozen source generated document is allowed
+        /// to have it's real output again.
+        /// </summary>
+        internal Solution WithoutFrozenSourceGeneratedDocuments()
+        {
+            var newState = _state.WithoutFrozenSourceGeneratedDocuments();
+            if (newState == _state)
+            {
+                return this;
+            }
+
+            return new Solution(newState);
+        }
+
         /// <summary>
         /// Gets an objects that lists the added, changed and removed projects between
         /// this solution and the specified solution.

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

@@ -20,7 +20,6 @@ internal partial class SolutionState
         /// </summary>
         private class GeneratedFileReplacingCompilationTracker : ICompilationTracker
         {
-            private readonly ICompilationTracker _underlyingTracker;
             private readonly SourceGeneratedDocumentState _replacedGeneratedDocumentState;
 
             private AsyncLazy<Checksum>? _lazyDependentChecksum;
@@ -32,17 +31,18 @@ private class GeneratedFileReplacingCompilationTracker : ICompilationTracker
             [DisallowNull]
             private Compilation? _compilationWithReplacement;
 
+            public ICompilationTracker UnderlyingTracker { get; }
             public SkeletonReferenceCache SkeletonReferenceCache { get; }
-            public ProjectState ProjectState => _underlyingTracker.ProjectState;
+            public ProjectState ProjectState => UnderlyingTracker.ProjectState;
 
             public GeneratedFileReplacingCompilationTracker(ICompilationTracker underlyingTracker, SourceGeneratedDocumentState replacementDocumentState)
             {
-                _underlyingTracker = underlyingTracker;
+                UnderlyingTracker = underlyingTracker;
                 _replacedGeneratedDocumentState = replacementDocumentState;
                 SkeletonReferenceCache = underlyingTracker.SkeletonReferenceCache.Clone();
             }
 
-            public GeneratorDriver? GeneratorDriver => _underlyingTracker.GeneratorDriver;
+            public GeneratorDriver? GeneratorDriver => UnderlyingTracker.GeneratorDriver;
 
             public bool ContainsAssemblyOrModuleOrDynamic(ISymbol symbol, bool primary)
             {
@@ -79,8 +79,8 @@ public async Task<Compilation> GetCompilationAsync(SolutionState solution, Cance
                     return _compilationWithReplacement;
                 }
 
-                var underlyingCompilation = await _underlyingTracker.GetCompilationAsync(solution, cancellationToken).ConfigureAwait(false);
-                var underlyingSourceGeneratedDocuments = await _underlyingTracker.GetSourceGeneratedDocumentStatesAsync(solution, cancellationToken).ConfigureAwait(false);
+                var underlyingCompilation = await UnderlyingTracker.GetCompilationAsync(solution, cancellationToken).ConfigureAwait(false);
+                var underlyingSourceGeneratedDocuments = await UnderlyingTracker.GetSourceGeneratedDocumentStatesAsync(solution, cancellationToken).ConfigureAwait(false);
 
                 underlyingSourceGeneratedDocuments.TryGetState(_replacedGeneratedDocumentState.Id, out var existingState);
 
@@ -110,10 +110,10 @@ public async Task<Compilation> GetCompilationAsync(SolutionState solution, Cance
             }
 
             public Task<VersionStamp> GetDependentVersionAsync(SolutionState solution, CancellationToken cancellationToken)
-                => _underlyingTracker.GetDependentVersionAsync(solution, cancellationToken);
+                => UnderlyingTracker.GetDependentVersionAsync(solution, cancellationToken);
 
             public Task<VersionStamp> GetDependentSemanticVersionAsync(SolutionState solution, CancellationToken cancellationToken)
-                => _underlyingTracker.GetDependentSemanticVersionAsync(solution, cancellationToken);
+                => UnderlyingTracker.GetDependentSemanticVersionAsync(solution, cancellationToken);
 
             public Task<Checksum> GetDependentChecksumAsync(SolutionState solution, CancellationToken cancellationToken)
             {
@@ -129,7 +129,7 @@ public Task<Checksum> GetDependentChecksumAsync(SolutionState solution, Cancella
 
             private async Task<Checksum> ComputeDependentChecksumAsync(SolutionState solution, CancellationToken cancellationToken)
                 => Checksum.Create(
-                    await _underlyingTracker.GetDependentChecksumAsync(solution, cancellationToken).ConfigureAwait(false),
+                    await UnderlyingTracker.GetDependentChecksumAsync(solution, cancellationToken).ConfigureAwait(false),
                     await _replacedGeneratedDocumentState.GetChecksumAsync(cancellationToken).ConfigureAwait(false));
 
             public CompilationReference? GetPartialMetadataReference(ProjectState fromProject, ProjectReference projectReference)
@@ -146,7 +146,7 @@ await _underlyingTracker.GetDependentChecksumAsync(solution, cancellationToken).
 
             public async ValueTask<TextDocumentStates<SourceGeneratedDocumentState>> GetSourceGeneratedDocumentStatesAsync(SolutionState solution, CancellationToken cancellationToken)
             {
-                var underlyingGeneratedDocumentStates = await _underlyingTracker.GetSourceGeneratedDocumentStatesAsync(solution, cancellationToken).ConfigureAwait(false);
+                var underlyingGeneratedDocumentStates = await UnderlyingTracker.GetSourceGeneratedDocumentStatesAsync(solution, cancellationToken).ConfigureAwait(false);
 
                 if (underlyingGeneratedDocumentStates.Contains(_replacedGeneratedDocumentState.Id))
                 {
@@ -166,7 +166,7 @@ public async ValueTask<TextDocumentStates<SourceGeneratedDocumentState>> GetSour
 
             public Task<bool> HasSuccessfullyLoadedAsync(SolutionState solution, CancellationToken cancellationToken)
             {
-                return _underlyingTracker.HasSuccessfullyLoadedAsync(solution, cancellationToken);
+                return UnderlyingTracker.HasSuccessfullyLoadedAsync(solution, cancellationToken);
             }
 
             public bool TryGetCompilation([NotNullWhen(true)] out Compilation? compilation)
@@ -183,7 +183,7 @@ public bool TryGetCompilation([NotNullWhen(true)] out Compilation? compilation)
                 }
                 else
                 {
-                    return _underlyingTracker.TryGetSourceGeneratedDocumentStateForAlreadyGeneratedId(documentId);
+                    return UnderlyingTracker.TryGetSourceGeneratedDocumentStateForAlreadyGeneratedId(documentId);
                 }
             }
         }

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

@@ -314,7 +314,7 @@ private SolutionState CreatePrimarySolution(
                 _filePathToDocumentIdsMap,
                 _dependencyGraph,
                 _lazyAnalyzers,
-                frozenSourceGeneratedDocument: null);
+                _frozenSourceGeneratedDocumentState);
         }
 
         private BranchId GetBranchId()
@@ -1908,6 +1908,31 @@ public SolutionState WithFrozenSourceGeneratedDocument(SourceGeneratedDocumentId
                 frozenSourceGeneratedDocument: newGeneratedState);
         }
 
+        /// <summary>
+        /// Undoes the operation of <see cref="WithFrozenSourceGeneratedDocument"/>; any frozen source generated document is allowed
+        /// to have it's real output again.
+        /// </summary>
+        public SolutionState WithoutFrozenSourceGeneratedDocuments()
+        {
+            // If there's nothing frozen, there's nothing to do.
+            if (_frozenSourceGeneratedDocumentState == null)
+                return this;
+
+            var projectId = _frozenSourceGeneratedDocumentState.Id.ProjectId;
+
+            // Since we previously froze this document, we should have a CompilationTracker entry for it, and it should be a
+            // GeneratedFileReplacingCompilationTracker. To undo the operation, we'll just restore the original CompilationTracker.
+            var newTrackerMap = CreateCompilationTrackerMap(projectId, _dependencyGraph);
+            Contract.ThrowIfFalse(newTrackerMap.TryGetValue(projectId, out var existingTracker));
+            var replacingItemTracker = existingTracker as GeneratedFileReplacingCompilationTracker;
+            Contract.ThrowIfNull(replacingItemTracker);
+            newTrackerMap = newTrackerMap.SetItem(projectId, replacingItemTracker.UnderlyingTracker);
+
+            return this.Branch(
+                projectIdToTrackerMap: newTrackerMap,
+                frozenSourceGeneratedDocument: null);
+        }
+
         /// <summary>
         /// Symbols need to be either <see cref="IAssemblySymbol"/> or <see cref="IModuleSymbol"/>.
         /// </summary>

src/Workspaces/Remote/ServiceHub/Host/RemoteWorkspace.SolutionCreator.cs

@@ -57,6 +57,10 @@ public async Task<Solution> CreateSolutionAsync(Checksum newSolutionChecksum)
                 {
                     var solution = _baseSolution;
 
+                    // If we previously froze a source generated document and then held onto that, unfreeze it now. We'll re-freeze the new document
+                    // if needed again later.
+                    solution = solution.WithoutFrozenSourceGeneratedDocuments();
+
                     var oldSolutionChecksums = await solution.State.GetStateChecksumsAsync(_cancellationToken).ConfigureAwait(false);
                     var newSolutionChecksums = await _assetProvider.GetAssetAsync<SolutionStateChecksums>(newSolutionChecksum, _cancellationToken).ConfigureAwait(false);
 
@@ -85,10 +89,6 @@ public async Task<Solution> CreateSolutionAsync(Checksum newSolutionChecksum)
                             newSolutionChecksums.AnalyzerReferences, _cancellationToken).ConfigureAwait(false));
                     }
 
-                    // The old solution should never have any frozen source generated documents -- those are only created and forked off of
-                    // a workspaces's CurrentSolution
-                    Contract.ThrowIfFalse(solution.State.FrozenSourceGeneratedDocumentState == null);
-
                     if (newSolutionChecksums.FrozenSourceGeneratedDocumentIdentity != Checksum.Null && newSolutionChecksums.FrozenSourceGeneratedDocumentText != Checksum.Null)
                     {
                         var identity = await _assetProvider.GetAssetAsync<SourceGeneratedDocumentIdentity>(newSolutionChecksums.FrozenSourceGeneratedDocumentIdentity, _cancellationToken).ConfigureAwait(false);