Don't break if we fork a snapshot after freezing partial semantics

src/EditorFeatures/CSharpTest/Completion/CompletionServiceTests.cs

@@ -52,7 +52,7 @@ public class C1
             var generatorRanCount = 0;
             var generator = new CallbackGenerator(onInit: _ => { }, onExecute: _ => Interlocked.Increment(ref generatorRanCount));
 
-            using var workspace = WorkspaceTestUtilities.CreateWorkspaceWithPartalSemantics();
+            using var workspace = WorkspaceTestUtilities.CreateWorkspaceWithPartialSemantics();
             var analyzerReference = new TestGeneratorReference(generator);
             var project = SolutionUtilities.AddEmptyProject(workspace.CurrentSolution)
                 .AddAnalyzerReference(analyzerReference)

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

@@ -0,0 +1,68 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// 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;
+
+namespace Microsoft.CodeAnalysis
+{
+    internal partial class SolutionState
+    {
+        private partial class CompilationTracker
+        {
+            /// <summary>
+            /// When we're working with compilations, we often have two: a compilation that does not contain generated files
+            /// (which we might need later to run generators again), and one that has the stale generated files that we might
+            /// be able to reuse as well. In those cases we have to do the same transformations to both, and this gives us
+            /// a handy way to do precisely that while not forking compilations twice if there are no generated files anywhere.
+            /// </summary>
+            internal readonly struct CompilationPair
+            {
+                public CompilationPair(Compilation withoutGeneratedDocuments, Compilation withGeneratedDocuments) : this()
+                {
+                    CompilationWithoutGeneratedDocuments = withoutGeneratedDocuments;
+                    CompilationWithGeneratedDocuments = withGeneratedDocuments;
+                }
+
+                public Compilation CompilationWithoutGeneratedDocuments { get; }
+                public Compilation CompilationWithGeneratedDocuments { get; }
+
+                public CompilationPair ReplaceSyntaxTree(SyntaxTree oldTree, SyntaxTree newTree)
+                {
+                    return WithChange(static (compilation, trees) => compilation.ReplaceSyntaxTree(trees.oldTree, trees.newTree), (oldTree, newTree));
+                }
+
+                public CompilationPair AddSyntaxTree(SyntaxTree newTree)
+                {
+                    return WithChange(static (compilation, t) => compilation.AddSyntaxTrees(t), newTree);
+                }
+
+                public CompilationPair WithPreviousScriptCompilation(Compilation previousScriptCompilation)
+                {
+                    return WithChange(static (compilation, priorCompilation) => compilation.WithScriptCompilationInfo(compilation.ScriptCompilationInfo!.WithPreviousScriptCompilation(priorCompilation)), previousScriptCompilation);
+                }
+
+                public CompilationPair WithReferences(IReadOnlyCollection<MetadataReference> metadataReferences)
+                {
+                    return WithChange(static (c, r) => c.WithReferences(r), metadataReferences);
+                }
+
+                private CompilationPair WithChange<TArg>(Func<Compilation, TArg, Compilation> change, TArg arg)
+                {
+                    var changedWithoutGeneratedDocuments = change(CompilationWithoutGeneratedDocuments, arg);
+
+                    if (CompilationWithoutGeneratedDocuments == CompilationWithGeneratedDocuments)
+                    {
+                        // If we didn't have any generated files, then no reason to transform twice
+                        return new CompilationPair(changedWithoutGeneratedDocuments, changedWithoutGeneratedDocuments);
+                    }
+
+                    var changedWithGeneratedDocuments = change(CompilationWithGeneratedDocuments, arg);
+
+                    return new CompilationPair(changedWithoutGeneratedDocuments, changedWithGeneratedDocuments);
+                }
+            }
+        }
+    }
+}

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

@@ -6,6 +6,8 @@
 using System.Collections.Generic;
 using System.Collections.Immutable;
 using System.Diagnostics;
+using System.Linq;
+using System.Threading;
 using Roslyn.Utilities;
 
 namespace Microsoft.CodeAnalysis
@@ -39,21 +41,55 @@ private readonly struct CompilationTrackerGeneratorInfo
                 /// </summary>
                 public readonly bool DocumentsAreFinal;
 
+                /// <summary>
+                /// Whether the generated documents are frozen and generators should never be ran again, ever, even if a document
+                /// is later changed. This is used to ensure that when we produce a frozen solution for partial semantics,
+                /// further downstream forking of that solution won't rerun generators. This is because of two reasons:
+                /// <list type="number">
+                /// <item>Generally once we've produced a frozen solution with partial semantics, we now want speed rather
+                /// than accuracy; a generator running in a later path will still cause issues there.</item>
+                /// <item>The frozen solution with partial semantics makes no guarantee that other syntax trees exist or
+                /// whether we even have references -- it's pretty likely that running a generator might produce worse results
+                /// than what we originally had.</item>
+                /// </list>
+                /// </summary>
+                public readonly bool DocumentsAreFinalAndFrozen;
+
                 public CompilationTrackerGeneratorInfo(
                     TextDocumentStates<SourceGeneratedDocumentState> documents,
                     GeneratorDriver? driver,
-                    bool documentsAreFinal)
+                    bool documentsAreFinal,
+                    bool documentsAreFinalAndFrozen = false)
                 {
                     Documents = documents;
                     Driver = driver;
                     DocumentsAreFinal = documentsAreFinal;
+                    DocumentsAreFinalAndFrozen = documentsAreFinalAndFrozen;
+
+                    // If we're frozen, that implies final as well
+                    Contract.ThrowIfTrue(documentsAreFinalAndFrozen && !documentsAreFinal);
                 }
 
                 public CompilationTrackerGeneratorInfo WithDocumentsAreFinal(bool documentsAreFinal)
-                    => DocumentsAreFinal == documentsAreFinal ? this : new(Documents, Driver, documentsAreFinal);
+                {
+                    // If we're already frozen, then we won't do anything even if somebody calls WithDocumentsAreFinal(false);
+                    // this for example would happen if we had a frozen snapshot, and then we fork it further with additional changes.
+                    // In that case we would be calling WithDocumentsAreFinal(false) to force generators to run again, but if we've
+                    // frozen in partial semantics, we're done running them period. So we'll just keep treating them as final,
+                    // no matter the wishes of the caller.
+                    if (DocumentsAreFinalAndFrozen || DocumentsAreFinal == documentsAreFinal)
+                        return this;
+                    else
+                        return new(Documents, Driver, documentsAreFinal);
+                }
+
+                public CompilationTrackerGeneratorInfo WithDocumentsAreFinalAndFrozen()
+                {
+                    return DocumentsAreFinalAndFrozen ? this : new(Documents, Driver, documentsAreFinal: true, documentsAreFinalAndFrozen: true);
+                }
 
                 public CompilationTrackerGeneratorInfo WithDriver(GeneratorDriver? driver)
-                    => Driver == driver ? this : new(Documents, driver, DocumentsAreFinal);
+                    => Driver == driver ? this : new(Documents, driver, DocumentsAreFinal, DocumentsAreFinalAndFrozen);
             }
 
             /// <summary>
@@ -101,6 +137,21 @@ protected CompilationTrackerState(
                 {
                     CompilationWithoutGeneratedDocuments = compilationWithoutGeneratedDocuments;
                     GeneratorInfo = generatorInfo;
+
+#if DEBUG
+
+                    // As a sanity check, we should never see the generated trees inside of the compilation that should not
+                    // have generated trees.
+                    var compilation = compilationWithoutGeneratedDocuments?.GetValueOrNull();
+
+                    if (compilation != null)
+                    {
+                        foreach (var generatedDocument in generatorInfo.Documents.States.Values)
+                        {
+                            Contract.ThrowIfTrue(compilation.SyntaxTrees.Contains(generatedDocument.GetSyntaxTree(CancellationToken.None)));
+                        }
+                    }
+#endif
                 }
 
                 public static CompilationTrackerState Create(