Merge pull request #72215 from CyrusNajmabadi/nullToInProgress

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

@@ -281,13 +281,14 @@ async Task<FinalCompilationTrackerState> BuildFinalStateAsync()
                     if (state is FinalCompilationTrackerState finalState)
                         return finalState;
 
-                    // Transition from wherever we're currently at to the 'all trees parsed' state.
+                    // Transition from wherever we're currently at to an in-progress-state.
                     var expandedInProgressState = state switch
                     {
+                        // We're already there, so no transition needed.
                         InProgressState inProgressState => inProgressState,
 
                         // We've got nothing.  Build it from scratch :(
-                        null => await BuildInProgressStateFromNoCompilationStateAsync().ConfigureAwait(false),
+                        null => BuildInProgressStateFromNoCompilationState(),
 
                         _ => throw ExceptionUtilities.UnexpectedValue(state.GetType())
                     };
@@ -301,26 +302,43 @@ async Task<FinalCompilationTrackerState> BuildFinalStateAsync()
                     "https://github.com/dotnet/roslyn/issues/23582",
                     Constraint = "Avoid calling " + nameof(Compilation.AddSyntaxTrees) + " in a loop due to allocation overhead.")]
 
-                async Task<InProgressState> BuildInProgressStateFromNoCompilationStateAsync()
+                InProgressState BuildInProgressStateFromNoCompilationState()
                 {
                     try
                     {
-                        var compilation = CreateEmptyCompilation();
+                        // Create a chain of translation steps where we add each document one at a time to an initially
+                        // empty compilation.  This allows us to then process that chain of actions like we would do any
+                        // other.  It also means that if we're in the process of parsing documents in that chain, that
+                        // we'll see the results of how far we've gotten if someone asks for a frozen snapshot midway
+                        // through.
+                        var initialProjectState = this.ProjectState.RemoveAllDocuments();
+                        var initialCompilation = this.CreateEmptyCompilation();
+
+                        var translationActionsBuilder = ImmutableList.CreateBuilder<TranslationAction>();
+
+                        var oldProjectState = initialProjectState;
+                        foreach (var documentState in this.ProjectState.DocumentStates.GetStatesInCompilationOrder())
+                        {
+                            var documentStates = ImmutableArray.Create(documentState);
+                            var newProjectState = oldProjectState.AddDocuments(documentStates);
+                            translationActionsBuilder.Add(new TranslationAction.AddDocumentsAction(
+                                oldProjectState, newProjectState, documentStates));
 
-                        var trees = await GetAllSyntaxTreesAsync(
-                            this.ProjectState.DocumentStates.GetStatesInCompilationOrder(),
-                            this.ProjectState.DocumentStates.Count,
-                            cancellationToken).ConfigureAwait(false);
+                            oldProjectState = newProjectState;
+                        }
 
-                        compilation = compilation.AddSyntaxTrees(trees);
+                        var compilationWithoutGeneratedDocuments = CreateEmptyCompilation();
 
                         // We only got here when we had no compilation state at all.  So we couldn't have gotten
                         // here from a frozen state (as a frozen state always ensures we have a
                         // WithCompilationTrackerState).  As such, we can safely still preserve that we're not
                         // frozen here.
                         var allSyntaxTreesParsedState = InProgressState.Create(
-                            isFrozen: false, compilation, CompilationTrackerGeneratorInfo.Empty, staleCompilationWithGeneratedDocuments: null,
-                            pendingTranslationActions: []);
+                            isFrozen: false,
+                            compilationWithoutGeneratedDocuments,
+                            CompilationTrackerGeneratorInfo.Empty,
+                            staleCompilationWithGeneratedDocuments: null,
+                            pendingTranslationActions: translationActionsBuilder.ToImmutable());
 
                         WriteState(allSyntaxTreesParsedState);
                         return allSyntaxTreesParsedState;
@@ -336,6 +354,28 @@ async Task<InProgressState> CollapseInProgressStateAsync(InProgressState initial
                     try
                     {
                         var currentState = initialState;
+
+                        // To speed things up, we look for all the added documents in the chain and we preemptively kick
+                        // off work to parse the documents for it in parallel.  This has the added benefit that if
+                        // someone asks for a frozen partial snapshot while we're in the middle of doing this, they can
+                        // use however many document states have successfully parsed their syntax trees.  For example,
+                        // if you had one extremely large file that took a long time to parse, and dozens of tiny ones,
+                        // it's more likely that the frozen tree would have many more documents in it.
+                        //
+                        // Note: we intentionally kick these off in a fire-and-forget fashion.  If we get canceled, all
+                        // the tasks will attempt to cancel.  If we complete, that's only because these tasks would
+                        // complete as well.  There's no need to track this with any sort of listener as this work just
+                        // acts to speed up getting to a compilation, but is otherwise unobservable.
+                        foreach (var action in currentState.PendingTranslationActions)
+                        {
+                            if (action is TranslationAction.AddDocumentsAction { Documents: var documents })
+                            {
+                                foreach (var document in documents)
+                                    _ = Task.Run(async () => await document.GetSyntaxTreeAsync(cancellationToken).ConfigureAwait(false), cancellationToken);
+                            }
+                        }
+
+                        // Then, we serially process the chain while that parsing is happening concurrently.
                         while (currentState.PendingTranslationActions.Count > 0)
                         {
                             cancellationToken.ThrowIfCancellationRequested();

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

@@ -126,9 +126,16 @@ internal sealed class AddDocumentsAction(
                 ImmutableArray<DocumentState> documents)
                 : TranslationAction(oldProjectState, newProjectState)
             {
+                public readonly ImmutableArray<DocumentState> Documents = documents;
+
                 public override async Task<Compilation> TransformCompilationAsync(Compilation oldCompilation, CancellationToken cancellationToken)
                 {
-                    var trees = await GetAllSyntaxTreesAsync(documents, documents.Length, cancellationToken).ConfigureAwait(false);
+                    // Parse all the documents in parallel.
+                    using var _ = ArrayBuilder<Task<SyntaxTree>>.GetInstance(this.Documents.Length, out var tasks);
+                    foreach (var document in this.Documents)
+                        tasks.Add(Task.Run(async () => await document.GetSyntaxTreeAsync(cancellationToken).ConfigureAwait(false), cancellationToken));
+
+                    var trees = await Task.WhenAll(tasks).ConfigureAwait(false);
                     return oldCompilation.AddSyntaxTrees(trees);
                 }
 
@@ -209,14 +216,11 @@ public override GeneratorDriver TransformGeneratorDriver(GeneratorDriver generat
 
             internal sealed class ProjectAssemblyNameAction(
                 ProjectState oldProjectState,
-                ProjectState newProjectState,
-                string assemblyName)
+                ProjectState newProjectState)
                 : TranslationAction(oldProjectState, newProjectState)
             {
                 public override Task<Compilation> TransformCompilationAsync(Compilation oldCompilation, CancellationToken cancellationToken)
-                {
-                    return Task.FromResult(oldCompilation.WithAssemblyName(assemblyName));
-                }
+                    => Task.FromResult(oldCompilation.WithAssemblyName(NewProjectState.AssemblyName));
 
                 // Updating the options of a compilation doesn't require us to reparse trees, so we can use this to update
                 // compilations with stale generated trees.

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

@@ -269,7 +269,7 @@ public SolutionCompilationState WithProjectAssemblyName(
         return ForkProject(
             this.SolutionState.WithProjectAssemblyName(projectId, assemblyName),
             static (stateChange, assemblyName) => new TranslationAction.ProjectAssemblyNameAction(
-                stateChange.OldProjectState, stateChange.NewProjectState, assemblyName),
+                stateChange.OldProjectState, stateChange.NewProjectState),
             forkTracker: true,
             arg: assemblyName);
     }
@@ -1479,16 +1479,6 @@ public SolutionCompilationState WithDocumentText(IEnumerable<DocumentId?> docume
         return result;
     }
 
-    private static async Task<SyntaxTree[]> GetAllSyntaxTreesAsync(IEnumerable<DocumentState> documents, int documentCount, CancellationToken cancellationToken)
-    {
-        // Parse all the documents in parallel.
-        using var _ = ArrayBuilder<Task<SyntaxTree>>.GetInstance(documentCount, out var tasks);
-        foreach (var document in documents)
-            tasks.Add(Task.Run(async () => await document.GetSyntaxTreeAsync(cancellationToken).ConfigureAwait(false), cancellationToken));
-
-        return await Task.WhenAll(tasks).ConfigureAwait(false);
-    }
-
     internal TestAccessor GetTestAccessor()
         => new(this);