Ensure source generated documents are up-to-date before analyzing EnC changes (#73989)

tmat · web-flow · commit c45376a4c987 · 2024-06-14T14:59:49.000-07:00
diff --git a/src/EditorFeatures/Core/EditAndContinue/ActiveStatementTrackingService.cs b/src/EditorFeatures/Core/EditAndContinue/ActiveStatementTrackingService.cs
@@ -19,6 +19,7 @@
 using Microsoft.CodeAnalysis.Host.Mef;
 using Microsoft.CodeAnalysis.PooledObjects;
 using Microsoft.CodeAnalysis.Shared.TestHooks;
+using Microsoft.CodeAnalysis.Shared.Extensions;
 using Microsoft.CodeAnalysis.Text;
 using Microsoft.CodeAnalysis.Text.Shared.Extensions;
 using Microsoft.VisualStudio.Text;
@@ -198,6 +199,9 @@ internal async Task TrackActiveSpansAsync(Solution solution)
                     return;
                 }
 
+                // Make sure the solution snapshot has all source-generated documents in the affected projects up-to-date:
+                solution = solution.WithUpToDateSourceGeneratorDocuments(openDocumentIds.Select(static d => d.ProjectId));
+
                 var baseActiveStatementSpans = await _spanProvider.GetBaseActiveStatementSpansAsync(solution, openDocumentIds, cancellationToken).ConfigureAwait(false);
                 if (baseActiveStatementSpans.IsDefault)
                 {
diff --git a/src/EditorFeatures/Core/EditAndContinue/EditAndContinueLanguageService.cs b/src/EditorFeatures/Core/EditAndContinue/EditAndContinueLanguageService.cs
@@ -220,6 +220,8 @@ public async ValueTask CommitUpdatesAsync(CancellationToken cancellationToken)
         catch (Exception e) when (FatalError.ReportAndCatchUnlessCanceled(e, cancellationToken))
         {
         }
+
+        workspaceProvider.Value.Workspace.EnqueueUpdateSourceGeneratorVersion(projectId: null, forceRegeneration: false);
     }
 
     public async ValueTask DiscardUpdatesAsync(CancellationToken cancellationToken)
diff --git a/src/Features/Core/Portable/EditAndContinue/DebuggingSession.cs b/src/Features/Core/Portable/EditAndContinue/DebuggingSession.cs
@@ -518,6 +518,9 @@ public async ValueTask<EmitSolutionUpdateResults> EmitSolutionUpdateAsync(
 
         var updateId = new UpdateId(Id, Interlocked.Increment(ref _updateOrdinal));
 
+        // Make sure the solution snapshot has all source-generated documents up-to-date.
+        solution = solution.WithUpToDateSourceGeneratorDocuments(solution.ProjectIds);
+
         var solutionUpdate = await EditSession.EmitSolutionUpdateAsync(solution, activeStatementSpanProvider, updateId, cancellationToken).ConfigureAwait(false);
 
         solutionUpdate.Log(EditAndContinueService.Log, updateId);
diff --git a/src/Features/Core/Portable/EditAndContinue/EditAndContinueService.cs b/src/Features/Core/Portable/EditAndContinue/EditAndContinueService.cs
@@ -160,6 +160,9 @@ public async ValueTask<DebuggingSessionId> StartDebuggingSessionAsync(
                 initialDocumentStates = [];
             }
 
+            // Make sure the solution snapshot has all source-generated documents up-to-date:
+            solution = solution.WithUpToDateSourceGeneratorDocuments(solution.ProjectIds);
+
             var sessionId = new DebuggingSessionId(Interlocked.Increment(ref s_debuggingSessionId));
             var session = new DebuggingSession(sessionId, solution, debuggerService, _compilationOutputsProvider, sourceTextProvider, initialDocumentStates, reportDiagnostics);
 
diff --git a/src/Features/Test/EditAndContinue/EditAndContinueWorkspaceServiceTests.cs b/src/Features/Test/EditAndContinue/EditAndContinueWorkspaceServiceTests.cs
@@ -17,6 +17,7 @@
 using Microsoft.CodeAnalysis.CSharp;
 using Microsoft.CodeAnalysis.CSharp.Test.Utilities;
 using Microsoft.CodeAnalysis.Emit;
+using Microsoft.CodeAnalysis.Host;
 using Microsoft.CodeAnalysis.Shared.Extensions;
 using Microsoft.CodeAnalysis.Test.Utilities;
 using Microsoft.CodeAnalysis.Text;
@@ -2516,8 +2517,9 @@ partial class C { int Y = 2; }
         EndDebuggingSession(debuggingSession);
     }
 
-    [Fact]
-    public async Task ValidSignificantChange_SourceGenerators_DocumentUpdate_GeneratedDocumentUpdate()
+    [Theory]
+    [CombinatorialData]
+    internal async Task ValidSignificantChange_SourceGenerators_DocumentUpdate_GeneratedDocumentUpdate(SourceGeneratorExecutionPreference executionPreference)
     {
         var sourceV1 = @"
 /* GENERATE: class G { int X => 1; } */
@@ -2532,12 +2534,21 @@ class C { int Y => 2; }
 
         var generator = new TestSourceGenerator() { ExecuteImpl = GenerateSource };
 
-        using var _ = CreateWorkspace(out var solution, out var service);
+        using var workspace = CreateWorkspace(out var solution, out var service);
+        var workspaceConfig = Assert.IsType<TestWorkspaceConfigurationService>(workspace.Services.GetRequiredService<IWorkspaceConfigurationService>());
+        workspaceConfig.Options = new WorkspaceConfigurationOptions(executionPreference);
+
         (solution, var document1) = AddDefaultTestProject(solution, sourceV1, generator);
 
         var moduleId = EmitLibrary(sourceV1, generator: generator);
         LoadLibraryToDebuggee(moduleId);
 
+        // Trigger initial source generation before debugging session starts.
+        // Causes source generator to run on the solution for the first time.
+        // Futher compilation access won't automatically trigger source generators,
+        // the EnC service has to do so.
+        _ = await solution.Projects.Single().GetCompilationAsync(CancellationToken.None);
+
         var debuggingSession = await StartDebuggingSessionAsync(service, solution);
 
         EnterBreakState(debuggingSession);
diff --git a/src/Features/TestUtilities/EditAndContinue/EditAndContinueWorkspaceTestBase.cs b/src/Features/TestUtilities/EditAndContinue/EditAndContinueWorkspaceTestBase.cs
@@ -53,7 +53,11 @@ public abstract class EditAndContinueWorkspaceTestBase : TestBase
 
     internal TestWorkspace CreateWorkspace(out Solution solution, out EditAndContinueService service, Type[]? additionalParts = null)
     {
-        var workspace = new TestWorkspace(composition: FeaturesTestCompositions.Features.AddParts(additionalParts), solutionTelemetryId: s_solutionTelemetryId);
+        var composition = FeaturesTestCompositions.Features
+            .AddParts(typeof(TestWorkspaceConfigurationService))
+            .AddParts(additionalParts);
+
+        var workspace = new TestWorkspace(composition: composition, solutionTelemetryId: s_solutionTelemetryId);
         solution = workspace.CurrentSolution;
         service = GetEditAndContinueService(workspace);
         return workspace;
diff --git a/src/Workspaces/Core/Portable/Workspace/Solution/Solution.cs b/src/Workspaces/Core/Portable/Workspace/Solution/Solution.cs
@@ -1576,8 +1576,8 @@ internal Solution WithFrozenSourceGeneratedDocuments(ImmutableArray<(SourceGener
         => WithCompilationState(_compilationState.WithFrozenSourceGeneratedDocuments(documents));
 
     /// <inheritdoc cref="SolutionCompilationState.UpdateSpecificSourceGeneratorExecutionVersions"/>
-    internal Solution UpdateSpecificSourceGeneratorExecutionVersions(SourceGeneratorExecutionVersionMap sourceGeneratorExecutionVersionMap, CancellationToken cancellationToken)
-        => WithCompilationState(_compilationState.UpdateSpecificSourceGeneratorExecutionVersions(sourceGeneratorExecutionVersionMap, cancellationToken));
+    internal Solution UpdateSpecificSourceGeneratorExecutionVersions(SourceGeneratorExecutionVersionMap sourceGeneratorExecutionVersionMap)
+        => WithCompilationState(_compilationState.UpdateSpecificSourceGeneratorExecutionVersions(sourceGeneratorExecutionVersionMap));
 
     /// <summary>
     /// Undoes the operation of <see cref="WithFrozenSourceGeneratedDocument"/>; any frozen source generated document is allowed
diff --git a/src/Workspaces/Core/Portable/Workspace/Solution/SolutionCompilationState.CompilationTracker.cs b/src/Workspaces/Core/Portable/Workspace/Solution/SolutionCompilationState.CompilationTracker.cs
@@ -679,13 +679,6 @@ private async Task<bool> HasSuccessfullyLoadedSlowAsync(
                 return finalState.HasSuccessfullyLoaded;
             }
 
-            public ICompilationTracker WithCreationPolicy(bool create, bool forceRegeneration, CancellationToken cancellationToken)
-            {
-                return create
-                    ? WithCreateCreationPolicy(forceRegeneration)
-                    : WithDoNotCreateCreationPolicy(forceRegeneration, cancellationToken);
-            }
-
             public ICompilationTracker WithCreateCreationPolicy(bool forceRegeneration)
             {
                 var state = this.ReadState();
@@ -742,13 +735,8 @@ public ICompilationTracker WithCreateCreationPolicy(bool forceRegeneration)
                     skeletonReferenceCacheToClone: _skeletonReferenceCache);
             }
 
-            public ICompilationTracker WithDoNotCreateCreationPolicy(
-                bool forceRegeneration, CancellationToken cancellationToken)
+            public ICompilationTracker WithDoNotCreateCreationPolicy(CancellationToken cancellationToken)
             {
-                // We do not expect this to ever be passed true.  This is for freezing generators, and no callers
-                // (currently) will ask to drop drivers when they do that.
-                Contract.ThrowIfTrue(forceRegeneration);
-
                 var state = this.ReadState();
 
                 // We're freezing the solution for features where latency performance is paramount.  Do not run SGs or
diff --git a/src/Workspaces/Core/Portable/Workspace/Solution/SolutionCompilationState.GeneratedFileReplacingCompilationTracker.cs b/src/Workspaces/Core/Portable/Workspace/Solution/SolutionCompilationState.GeneratedFileReplacingCompilationTracker.cs
@@ -75,9 +75,17 @@ public ICompilationTracker Fork(ProjectState newProject, TranslationAction? tran
             throw new NotImplementedException();
         }
 
-        public ICompilationTracker WithCreationPolicy(bool create, bool forceRegeneration, CancellationToken cancellationToken)
+        public ICompilationTracker WithCreateCreationPolicy(bool forceRegeneration)
         {
-            var underlyingTracker = this.UnderlyingTracker.WithCreationPolicy(create, forceRegeneration, cancellationToken);
+            var underlyingTracker = this.UnderlyingTracker.WithCreateCreationPolicy(forceRegeneration);
+            return underlyingTracker == this.UnderlyingTracker
+                ? this
+                : new GeneratedFileReplacingCompilationTracker(underlyingTracker, _replacementDocumentStates);
+        }
+
+        public ICompilationTracker WithDoNotCreateCreationPolicy(CancellationToken cancellationToken)
+        {
+            var underlyingTracker = this.UnderlyingTracker.WithDoNotCreateCreationPolicy(cancellationToken);
             return underlyingTracker == this.UnderlyingTracker
                 ? this
                 : new GeneratedFileReplacingCompilationTracker(underlyingTracker, _replacementDocumentStates);
diff --git a/src/Workspaces/Core/Portable/Workspace/Solution/SolutionCompilationState.ICompilationTracker.cs b/src/Workspaces/Core/Portable/Workspace/Solution/SolutionCompilationState.ICompilationTracker.cs
@@ -40,13 +40,16 @@ bool ContainsAssemblyOrModuleOrDynamic(
         Task<Compilation> GetCompilationAsync(SolutionCompilationState compilationState, CancellationToken cancellationToken);
 
         /// <summary>
-        /// Updates the creation policy for this tracker.  A value of <see langword="true"/> will set this to <see
-        /// cref="CreationPolicy.Create"/> and <see langword="false"/> will set it to <see
-        /// cref="CreationPolicy.DoNotCreate"/>.
+        /// Updates the creation policy for this tracker.  Setting it to <see cref="CreationPolicy.Create"/>.
         /// </summary>
-        /// <param name="forceRegeneration">When switching to <paramref name="create"/>, this will force source
-        /// generated documents to be created.</param>
-        ICompilationTracker WithCreationPolicy(bool create, bool forceRegeneration, CancellationToken cancellationToken);
+        /// <param name="forceRegeneration">Forces source generated documents to be created by dumping any existing <see
+        /// cref="GeneratorDriver"/> and rerunning generators from scratch for this tracker.</param>
+        ICompilationTracker WithCreateCreationPolicy(bool forceRegeneration);
+
+        /// <summary>
+        /// Updates the creation policy for this tracker.  Setting it to <see cref="CreationPolicy.DoNotCreate"/>.
+        /// </summary>
+        ICompilationTracker WithDoNotCreateCreationPolicy(CancellationToken cancellationToken);
 
         Task<VersionStamp> GetDependentVersionAsync(SolutionCompilationState compilationState, CancellationToken cancellationToken);
         Task<VersionStamp> GetDependentSemanticVersionAsync(SolutionCompilationState compilationState, CancellationToken cancellationToken);
diff --git a/src/Workspaces/Core/Portable/Workspace/Solution/SolutionCompilationState.cs b/src/Workspaces/Core/Portable/Workspace/Solution/SolutionCompilationState.cs
@@ -1262,16 +1262,14 @@ public SolutionCompilationState WithOptions(SolutionOptionSet options)
     /// <paramref name="sourceGeneratorExecutionVersions"/> will not be touched (and they will stay in the map).
     /// </summary>
     public SolutionCompilationState UpdateSpecificSourceGeneratorExecutionVersions(
-        SourceGeneratorExecutionVersionMap sourceGeneratorExecutionVersions, CancellationToken cancellationToken)
+        SourceGeneratorExecutionVersionMap sourceGeneratorExecutionVersions)
     {
         var versionMapBuilder = _sourceGeneratorExecutionVersionMap.Map.ToBuilder();
         var newIdToTrackerMapBuilder = _projectIdToTrackerMap.ToBuilder();
         var changed = false;
 
         foreach (var (projectId, sourceGeneratorExecutionVersion) in sourceGeneratorExecutionVersions.Map)
         {
-            cancellationToken.ThrowIfCancellationRequested();
-
             var currentExecutionVersion = versionMapBuilder[projectId];
 
             // Nothing to do if already at this version.
@@ -1289,7 +1287,7 @@ public SolutionCompilationState UpdateSpecificSourceGeneratorExecutionVersions(
                 // if the major version has changed then we also want to drop the generator driver so that we're rerun
                 // generators from scratch.
                 var forceRegeneration = currentExecutionVersion.MajorVersion != sourceGeneratorExecutionVersion.MajorVersion;
-                var newTracker = existingTracker.WithCreationPolicy(create: true, forceRegeneration, cancellationToken);
+                var newTracker = existingTracker.WithCreateCreationPolicy(forceRegeneration);
                 if (newTracker != existingTracker)
                     newIdToTrackerMapBuilder[projectId] = newTracker;
             }
@@ -1325,7 +1323,7 @@ private SolutionCompilationState ComputeFrozenSnapshot(CancellationToken cancell
 
             // Since we're freezing, set both generators and skeletons to not be created.  We don't want to take any
             // perf hit on either of those at all for our clients.
-            var newTracker = oldTracker.WithCreationPolicy(create: false, forceRegeneration: false, cancellationToken);
+            var newTracker = oldTracker.WithDoNotCreateCreationPolicy(cancellationToken);
             if (oldTracker == newTracker)
                 continue;
 
diff --git a/src/Workspaces/Core/Portable/Workspace/Workspace_SourceGeneration.cs b/src/Workspaces/Core/Portable/Workspace/Workspace_SourceGeneration.cs
@@ -48,7 +48,7 @@ await this.SetCurrentSolutionAsync(
             oldSolution =>
             {
                 var updates = GetUpdatedSourceGeneratorVersions(oldSolution, projectIds);
-                return oldSolution.UpdateSpecificSourceGeneratorExecutionVersions(updates, cancellationToken);
+                return oldSolution.UpdateSpecificSourceGeneratorExecutionVersions(updates);
             },
             static (_, _) => (WorkspaceChangeKind.SolutionChanged, projectId: null, documentId: null),
             onBeforeUpdate: null,
diff --git a/src/Workspaces/Remote/ServiceHub/Host/RemoteWorkspace.SolutionCreator.cs b/src/Workspaces/Remote/ServiceHub/Host/RemoteWorkspace.SolutionCreator.cs
@@ -137,7 +137,7 @@ public async Task<Solution> CreateSolutionAsync(Checksum newSolutionChecksum, Ca
                         }
 #endif
 
-                        solution = solution.UpdateSpecificSourceGeneratorExecutionVersions(newVersions, cancellationToken);
+                        solution = solution.UpdateSpecificSourceGeneratorExecutionVersions(newVersions);
                     }
 
 #if DEBUG
diff --git a/src/Workspaces/Remote/ServiceHub/Host/RemoteWorkspace.cs b/src/Workspaces/Remote/ServiceHub/Host/RemoteWorkspace.cs
@@ -256,7 +256,7 @@ private async Task<Solution> ComputeDisconnectedSolutionAsync(
                         var newVersions = await assetProvider.GetAssetAsync<SourceGeneratorExecutionVersionMap>(
                             AssetPathKind.SolutionSourceGeneratorExecutionVersionMap, newSolutionCompilationChecksums.SourceGeneratorExecutionVersionMap, cancellationToken).ConfigureAwait(false);
 
-                        solution = solution.UpdateSpecificSourceGeneratorExecutionVersions(newVersions, cancellationToken);
+                        solution = solution.UpdateSpecificSourceGeneratorExecutionVersions(newVersions);
                     }
 
                     return solution;
diff --git a/src/Workspaces/SharedUtilitiesAndExtensions/Workspace/Core/Extensions/ISolutionExtensions.cs b/src/Workspaces/SharedUtilitiesAndExtensions/Workspace/Core/Extensions/ISolutionExtensions.cs
@@ -4,8 +4,10 @@
 
 using System;
 using System.Collections.Generic;
+using System.Collections.Immutable;
 using System.Threading;
 using System.Threading.Tasks;
+using Microsoft.CodeAnalysis.Host;
 
 namespace Microsoft.CodeAnalysis.Shared.Extensions;
 
@@ -79,4 +81,28 @@ public static TextDocument GetRequiredTextDocument(this Solution solution, Docum
 
     private static Exception CreateDocumentNotFoundException()
         => new InvalidOperationException(WorkspaceExtensionsResources.The_solution_does_not_contain_the_specified_document);
+
+#if !CODE_STYLE
+    public static Solution WithUpToDateSourceGeneratorDocuments(this Solution solution, IEnumerable<ProjectId> projectIds)
+    {
+        // If the solution is already in automatic mode, then SG documents are already always up to date.
+        var configuration = solution.Services.GetRequiredService<IWorkspaceConfigurationService>().Options;
+        if (configuration.SourceGeneratorExecution is SourceGeneratorExecutionPreference.Automatic)
+            return solution;
+
+        var projectIdToSourceGenerationVersion = ImmutableSortedDictionary.CreateBuilder<ProjectId, SourceGeneratorExecutionVersion>();
+
+        foreach (var projectId in projectIds)
+        {
+            if (!projectIdToSourceGenerationVersion.ContainsKey(projectId))
+            {
+                var currentVersion = solution.GetSourceGeneratorExecutionVersion(projectId);
+                projectIdToSourceGenerationVersion.Add(projectId, currentVersion.IncrementMinorVersion());
+            }
+        }
+
+        return solution.UpdateSpecificSourceGeneratorExecutionVersions(
+            new SourceGeneratorExecutionVersionMap(projectIdToSourceGenerationVersion.ToImmutable()));
+    }
+#endif
 }

Original file line number	Diff line number	Diff line change
`@@ -220,6 +220,8 @@ public async ValueTask CommitUpdatesAsync(CancellationToken cancellationToken)`
`220`	`220`	`catch (Exception e) when (FatalError.ReportAndCatchUnlessCanceled(e, cancellationToken))`
`221`	`221`	`{`
`222`	`222`	`}`
	`223`	`+`
	`224`	`+ workspaceProvider.Value.Workspace.EnqueueUpdateSourceGeneratorVersion(projectId: null, forceRegeneration: false);`
`223`	`225`	`}`
`224`	`226`
`225`	`227`	`public async ValueTask DiscardUpdatesAsync(CancellationToken cancellationToken)`
Original file line number	Diff line number	Diff line change
`@@ -160,6 +160,9 @@ public async ValueTask<DebuggingSessionId> StartDebuggingSessionAsync(`
`160`	`160`	`initialDocumentStates = [];`
`161`	`161`	`}`
`162`	`162`
	`163`	`+ // Make sure the solution snapshot has all source-generated documents up-to-date:`
	`164`	`+ solution = solution.WithUpToDateSourceGeneratorDocuments(solution.ProjectIds);`
	`165`	`+`
`163`	`166`	`var sessionId = new DebuggingSessionId(Interlocked.Increment(ref s_debuggingSessionId));`
`164`	`167`	`var session = new DebuggingSession(sessionId, solution, debuggerService, _compilationOutputsProvider, sourceTextProvider, initialDocumentStates, reportDiagnostics);`
`165`	`168`