WithFrozenPartialSemantics freezing linked documents

src/Workspaces/Core/Portable/Shared/Extensions/ISolutionExtensions.cs

@@ -59,10 +59,5 @@ public static Solution WithTextDocumentText(this Solution solution, DocumentId d
                     throw ExceptionUtilities.UnexpectedValue(documentKind);
             }
         }
-
-        public static ImmutableArray<DocumentId> FilterDocumentIdsByLanguage(this Solution solution, ImmutableArray<DocumentId> documentIds, string language)
-            => documentIds.WhereAsArray(
-                (documentId, args) => args.solution.GetDocument(documentId)?.Project.Language == args.language,
-                (solution, language));
     }
 }

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

@@ -430,8 +430,7 @@ public async Task<IEnumerable<TextChange>> GetTextChangesAsync(Document oldDocum
         /// </summary>
         public ImmutableArray<DocumentId> GetLinkedDocumentIds()
         {
-            var documentIdsWithPath = this.Project.Solution.GetDocumentIdsWithFilePath(this.FilePath);
-            var filteredDocumentIds = this.Project.Solution.FilterDocumentIdsByLanguage(documentIdsWithPath, this.Project.Language);
+            var filteredDocumentIds = this.Project.Solution.GetRelatedDocumentIds(this.Id);
             return filteredDocumentIds.Remove(this.Id);
         }
 

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

@@ -1696,29 +1696,7 @@ internal async Task<Solution> WithMergedLinkedFileChangesAsync(
 
         internal ImmutableArray<DocumentId> GetRelatedDocumentIds(DocumentId documentId)
         {
-            var projectState = _state.GetProjectState(documentId.ProjectId);
-            if (projectState == null)
-            {
-                // this document no longer exist
-                return ImmutableArray<DocumentId>.Empty;
-            }
-
-            var documentState = projectState.DocumentStates.GetState(documentId);
-            if (documentState == null)
-            {
-                // this document no longer exist
-                return ImmutableArray<DocumentId>.Empty;
-            }
-
-            var filePath = documentState.FilePath;
-            if (string.IsNullOrEmpty(filePath))
-            {
-                // this document can't have any related document. only related document is itself.
-                return ImmutableArray.Create(documentId);
-            }
-
-            var documentIds = GetDocumentIdsWithFilePath(filePath);
-            return this.FilterDocumentIdsByLanguage(documentIds, projectState.ProjectInfo.Language);
+            return _state.GetRelatedDocumentIds(documentId);
         }
 
         internal Solution WithNewWorkspace(string? workspaceKind, int workspaceVersion, SolutionServices services)

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

@@ -1633,8 +1633,14 @@ public SolutionState WithFrozenPartialCompilationIncludingSpecificDocument(Docum
         {
             try
             {
-                var doc = this.GetRequiredDocumentState(documentId);
-                var tree = doc.GetSyntaxTree(cancellationToken);
+                var allDocumentIds = GetRelatedDocumentIds(documentId);
+                using var _ = ArrayBuilder<(DocumentState, SyntaxTree)>.GetInstance(allDocumentIds.Length, out var builder);
+
+                foreach (var currentDocumentId in allDocumentIds)
+                {
+                    var document = this.GetRequiredDocumentState(currentDocumentId);
+                    builder.Add((document, document.GetSyntaxTree(cancellationToken)));
+                }
 
                 using (this.StateLock.DisposableWait(cancellationToken))
                 {
@@ -1658,13 +1664,19 @@ public SolutionState WithFrozenPartialCompilationIncludingSpecificDocument(Docum
                         return currentPartialSolution!;
                     }
 
-                    // if we don't have one or it is stale, create a new partial solution
-                    var tracker = this.GetCompilationTracker(documentId.ProjectId);
-                    var newTracker = tracker.FreezePartialStateWithTree(this, doc, tree, cancellationToken);
+                    var newIdToProjectStateMap = _projectIdToProjectStateMap;
+                    var newIdToTrackerMap = _projectIdToTrackerMap;
 
-                    Contract.ThrowIfFalse(_projectIdToProjectStateMap.ContainsKey(documentId.ProjectId));
-                    var newIdToProjectStateMap = _projectIdToProjectStateMap.SetItem(documentId.ProjectId, newTracker.ProjectState);
-                    var newIdToTrackerMap = _projectIdToTrackerMap.SetItem(documentId.ProjectId, newTracker);
+                    foreach (var (doc, tree) in builder)
+                    {
+                        // if we don't have one or it is stale, create a new partial solution
+                        var tracker = this.GetCompilationTracker(doc.Id.ProjectId);
+                        var newTracker = tracker.FreezePartialStateWithTree(this, doc, tree, cancellationToken);
+
+                        Contract.ThrowIfFalse(newIdToProjectStateMap.ContainsKey(doc.Id.ProjectId));
+                        newIdToProjectStateMap = newIdToProjectStateMap.SetItem(doc.Id.ProjectId, newTracker.ProjectState);
+                        newIdToTrackerMap = newIdToTrackerMap.SetItem(doc.Id.ProjectId, newTracker);
+                    }
 
                     currentPartialSolution = this.Branch(
                         idToProjectStateMap: newIdToProjectStateMap,
@@ -1685,6 +1697,48 @@ public SolutionState WithFrozenPartialCompilationIncludingSpecificDocument(Docum
             }
         }
 
+        public ImmutableArray<DocumentId> GetRelatedDocumentIds(DocumentId documentId)
+        {
+            var projectState = this.GetProjectState(documentId.ProjectId);
+            if (projectState == null)
+            {
+                // this document no longer exist
+                return ImmutableArray<DocumentId>.Empty;
+            }
+
+            var documentState = projectState.DocumentStates.GetState(documentId);
+            if (documentState == null)
+            {
+                // this document no longer exist
+                return ImmutableArray<DocumentId>.Empty;
+            }
+
+            var filePath = documentState.FilePath;
+            if (string.IsNullOrEmpty(filePath))
+            {
+                // this document can't have any related document. only related document is itself.
+                return ImmutableArray.Create(documentId);
+            }
+
+            var documentIds = GetDocumentIdsWithFilePath(filePath);
+            return FilterDocumentIdsByLanguage(this, documentIds, projectState.ProjectInfo.Language);
+        }
+
+        private static ImmutableArray<DocumentId> FilterDocumentIdsByLanguage(SolutionState solution, ImmutableArray<DocumentId> documentIds, string language)
+            => documentIds.WhereAsArray(
+                static (documentId, args) =>
+                {
+                    var projectState = args.solution.GetProjectState(documentId.ProjectId);
+                    if (projectState == null)
+                    {
+                        // this document no longer exist
+                        return false;
+                    }
+
+                    return projectState.ProjectInfo.Language == args.language;
+                },
+                (solution, language));
+
         /// <summary>
         /// Creates a new solution instance with all the documents specified updated to have the same specified text.
         /// </summary>

src/Workspaces/CoreTest/SolutionTests/SolutionWithSourceGeneratorTests.cs

@@ -752,6 +752,41 @@ public async Task ForkAfterFreezeNoLongerRunsGenerators()
             Assert.Equal("// Something else", (await document.GetRequiredSyntaxRootAsync(CancellationToken.None)).ToFullString());
         }
 
+        [Fact]
+        public async Task LinkedDocumentOfFrozenShouldNotRunSourceGenerator()
+        {
+            using var workspace = CreateWorkspaceWithPartialSemantics();
+            var generatorRan = false;
+            var analyzerReference = new TestGeneratorReference(new CallbackGenerator(_ => { }, onExecute: _ => { generatorRan = true; }, source: "// Hello World!"));
+
+            var originalDocument1 = AddEmptyProject(workspace.CurrentSolution, name: "Project1")
+                .AddAnalyzerReference(analyzerReference)
+                .AddDocument("RegularDocument.cs", "// Source File", filePath: "RegularDocument.cs");
+
+            // this is a linked document of document1 above
+            var originalDocument2 = AddEmptyProject(originalDocument1.Project.Solution, name: "Project2")
+                .AddAnalyzerReference(analyzerReference)
+                .AddDocument(originalDocument1.Name, await originalDocument1.GetTextAsync().ConfigureAwait(false), filePath: originalDocument1.FilePath);
+
+            var project2 = originalDocument2.Project;
+            Assert.True(workspace.SetCurrentSolution(_ => project2.Solution, WorkspaceChangeKind.SolutionChanged));
+            originalDocument2 = project2.GetRequiredDocument(originalDocument2.Id).WithFrozenPartialSemantics(CancellationToken.None);
+            var frozenSolution = originalDocument2.Project.Solution;
+            var documentIdsToTest = new[] { originalDocument1.Id, originalDocument2.Id };
+
+            foreach (var documentIdToTest in documentIdsToTest)
+            {
+                var document = frozenSolution.GetRequiredDocument(documentIdToTest);
+                Assert.Equal(document.GetLinkedDocumentIds().Single(), documentIdsToTest.Except(new[] { documentIdToTest }).Single());
+                document = document.WithText(SourceText.From("// Something else"));
+                var project = document.Project;
+
+                var compilation = await project.GetRequiredCompilationAsync(CancellationToken.None);
+                Assert.Single(compilation.SyntaxTrees);
+                Assert.False(generatorRan);
+            }
+        }
+
         [Fact]
         public async Task DynamicFilesNotPassedToSourceGenerators()
         {