Ensure we sync source-generator versions over properly when doing a cone-sync

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

@@ -860,6 +860,51 @@ public async Task TestPartialProjectSync_ReferenceToNonExistentProject()
         var project1Checksum = await solution.CompilationState.GetChecksumAsync(project1.Id, CancellationToken.None);
     }
 
+    [Fact]
+    public async Task TestPartialProjectSync_SourceGeneratorExecutionVersion_1()
+    {
+        var code = @"class Test { void Method() { } }";
+
+        using var workspace = TestWorkspace.CreateCSharp(code);
+        using var remoteWorkspace = CreateRemoteWorkspace();
+
+        var solution = workspace.CurrentSolution;
+
+        var project1 = solution.Projects.Single();
+        var project2 = solution.AddProject("P2", "P2", LanguageNames.CSharp);
+
+        solution = project2.Solution;
+
+        var map = new Dictionary<Checksum, object>();
+        var assetProvider = new AssetProvider(
+            Checksum.Create(ImmutableArray.CreateRange(Guid.NewGuid().ToByteArray())), new SolutionAssetCache(), new SimpleAssetSource(workspace.Services.GetService<ISerializerService>(), map), remoteWorkspace.Services.GetService<ISerializerService>());
+
+        // Do the initial full sync
+        await solution.AppendAssetMapAsync(map, CancellationToken.None);
+        var solutionChecksum = await solution.CompilationState.GetChecksumAsync(CancellationToken.None);
+        var fullSyncedSolution = await remoteWorkspace.GetTestAccessor().GetSolutionAsync(assetProvider, solutionChecksum, updatePrimaryBranch: true, CancellationToken.None);
+        Assert.Equal(2, fullSyncedSolution.Projects.Count());
+
+        // Update the source generator versions for all projects for the local workspace.
+        workspace.EnqueueUpdateSourceGeneratorVersion(projectId: null, forceRegeneration: true);
+        await GetWorkspaceWaiter(workspace).ExpeditedWaitAsync();
+        solution = workspace.CurrentSolution;
+
+        // Now just sync project1's cone over.  This will validate that that we get the right checksums, even with a
+        // partial cone sync.
+        {
+            await solution.AppendAssetMapAsync(map, project1.Id, CancellationToken.None);
+            var project1Checksum = await solution.CompilationState.GetChecksumAsync(project1.Id, CancellationToken.None);
+            var project1SyncedSolution = await remoteWorkspace.GetTestAccessor().GetSolutionAsync(assetProvider, project1Checksum, updatePrimaryBranch: false, CancellationToken.None);
+        }
+    }
+
+    private static IAsynchronousOperationWaiter GetWorkspaceWaiter(TestWorkspace workspace)
+    {
+        var operations = workspace.ExportProvider.GetExportedValue<AsynchronousOperationListenerProvider>();
+        return operations.GetWaiter(FeatureAttribute.Workspace);
+    }
+
     [Fact]
     public void TestNoActiveDocumentSemanticModelNotCached()
     {

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

@@ -10,28 +10,28 @@ namespace Microsoft.CodeAnalysis;
 
 /// <summary>
 /// Represents a 'cone' of projects that is being sync'ed between the local and remote hosts.  A project cone starts
-/// with a <see cref="RootProjectId"/>, and contains both it and all dependent projects within <see cref="_projectIds"/>.
+/// with a <see cref="RootProjectId"/>, and contains both it and all dependent projects within <see cref="ProjectIds"/>.
 /// </summary>
 internal sealed class ProjectCone : IEquatable<ProjectCone>
 {
     public readonly ProjectId RootProjectId;
-    private readonly FrozenSet<ProjectId> _projectIds;
+    public readonly FrozenSet<ProjectId> ProjectIds;
 
     public ProjectCone(ProjectId rootProjectId, FrozenSet<ProjectId> projectIds)
     {
         Contract.ThrowIfFalse(projectIds.Contains(rootProjectId));
         RootProjectId = rootProjectId;
-        _projectIds = projectIds;
+        ProjectIds = projectIds;
     }
 
     public bool Contains(ProjectId projectId)
-        => _projectIds.Contains(projectId);
+        => ProjectIds.Contains(projectId);
 
     public override bool Equals(object? obj)
         => obj is ProjectCone cone && Equals(cone);
 
     public bool Equals(ProjectCone? other)
-        => other is not null && this.RootProjectId == other.RootProjectId && this._projectIds.SetEquals(other._projectIds);
+        => other is not null && this.RootProjectId == other.RootProjectId && this.ProjectIds.SetEquals(other.ProjectIds);
 
     public override int GetHashCode()
         => throw new NotImplementedException();

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

@@ -1579,8 +1579,9 @@ internal Document WithFrozenSourceGeneratedDocument(
     internal Solution WithFrozenSourceGeneratedDocuments(ImmutableArray<(SourceGeneratedDocumentIdentity documentIdentity, DateTime generationDateTime, SourceText text)> documents)
         => WithCompilationState(_compilationState.WithFrozenSourceGeneratedDocuments(documents));
 
-    internal Solution WithSourceGeneratorExecutionVersions(SourceGeneratorExecutionVersionMap sourceGeneratorExecutionVersionMap, CancellationToken cancellationToken)
-        => WithCompilationState(_compilationState.WithSourceGeneratorExecutionVersions(sourceGeneratorExecutionVersionMap, cancellationToken));
+    /// <inheritdoc cref="SolutionCompilationState.UpdateSpecificSourceGeneratorExecutionVersions"/>
+    internal Solution UpdateSpecificSourceGeneratorExecutionVersions(SourceGeneratorExecutionVersionMap sourceGeneratorExecutionVersionMap, CancellationToken cancellationToken)
+        => WithCompilationState(_compilationState.UpdateSpecificSourceGeneratorExecutionVersions(sourceGeneratorExecutionVersionMap, cancellationToken));
 
     /// <summary>
     /// Undoes the operation of <see cref="WithFrozenSourceGeneratedDocument"/>; any frozen source generated document is allowed

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

@@ -1247,7 +1247,13 @@ public SolutionCompilationState WithOptions(SolutionOptionSet options)
             this.SolutionState.WithOptions(options));
     }
 
-    public SolutionCompilationState WithSourceGeneratorExecutionVersions(
+    /// <summary>
+    /// Updates entries in our <see cref="_sourceGeneratorExecutionVersionMap"/> to the corresponding values in the
+    /// given <paramref name="sourceGeneratorExecutionVersions"/>.  Importantly, <paramref
+    /// name="sourceGeneratorExecutionVersions"/> must refer to projects in this solution.  Projects not mentioned in
+    /// <paramref name="sourceGeneratorExecutionVersions"/> will not be touched (and they will stay in the map).
+    /// </summary>
+    public SolutionCompilationState UpdateSpecificSourceGeneratorExecutionVersions(
         SourceGeneratorExecutionVersionMap sourceGeneratorExecutionVersions, CancellationToken cancellationToken)
     {
         var versionMapBuilder = _sourceGeneratorExecutionVersionMap.Map.ToBuilder();

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

@@ -140,10 +140,8 @@ public async Task<Checksum> GetChecksumAsync(ProjectId projectId, CancellationTo
                     frozenSourceGeneratedDocumentGenerationDateTimes = FrozenSourceGeneratedDocumentStates.SelectAsArray(d => d.GenerationDateTime);
                 }
 
-                var versionMapChecksum = ChecksumCache.GetOrCreate(
-                    this.SourceGeneratorExecutionVersionMap,
-                    static (map, @this) => GetVersionMapChecksum(@this),
-                    this);
+                // Ensure we only send the execution map over for projects in the project cone.
+                var versionMapChecksum = this.GetFilteredSourceGenerationExecutionMap(projectCone).GetChecksum();
 
                 var compilationStateChecksums = new SolutionCompilationStateChecksums(
                     solutionStateChecksum,
@@ -158,30 +156,27 @@ public async Task<Checksum> GetChecksumAsync(ProjectId projectId, CancellationTo
         {
             throw ExceptionUtilities.Unreachable();
         }
+    }
 
-        static Checksum GetVersionMapChecksum(SolutionCompilationState @this)
-        {
-            // We want the projects in sorted order so we can generate the checksum for the
-            // source-generation-execution-map consistently.
-            var sortedProjectIds = SolutionState.GetOrCreateSortedProjectIds(@this.SolutionState.ProjectIds);
-            var supportedCount = sortedProjectIds.Count(
-                static (projectId, @this) => RemoteSupportedLanguages.IsSupported(@this.SolutionState.GetRequiredProjectState(projectId).Language),
-                @this);
-
-            // For each project, we'll add one checksum for the project id and one for the version map.
-            using var _ = ArrayBuilder<Checksum>.GetInstance(2 * supportedCount, out var checksums);
+    public SourceGeneratorExecutionVersionMap GetFilteredSourceGenerationExecutionMap(ProjectCone? projectCone)
+    {
+        var builder = this.SourceGeneratorExecutionVersionMap.Map.ToBuilder();
 
-            foreach (var projectId in sortedProjectIds)
+        foreach (var (projectId, projectState) in this.SolutionState.ProjectStates)
+        {
+            if (!RemoteSupportedLanguages.IsSupported(projectState.Language))
             {
-                var projectState = @this.SolutionState.GetRequiredProjectState(projectId);
-                if (!RemoteSupportedLanguages.IsSupported(projectState.Language))
-                    continue;
-
-                checksums.Add(projectId.Checksum);
-                checksums.Add(Checksum.Create(@this.SourceGeneratorExecutionVersionMap[projectId], static (v, w) => v.WriteTo(w)));
+                builder.Remove(projectId);
+            }
+            else if (projectCone != null && !projectCone.Contains(projectId))
+            {
+                builder.Remove(projectId);
             }
-
-            return Checksum.Create(checksums);
         }
+
+        if (builder.Count == this.SourceGeneratorExecutionVersionMap.Map.Count)
+            return this.SourceGeneratorExecutionVersionMap;
+
+        return new SourceGeneratorExecutionVersionMap(builder.ToImmutable());
     }
 }

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

@@ -103,6 +103,19 @@ public static SourceGeneratorExecutionVersionMap Deserialize(ObjectReader reader
         return new(builder.ToImmutable());
     }
 
+    public Checksum GetChecksum()
+    {
+        using var _ = ArrayBuilder<Checksum>.GetInstance(this.Map.Count * 2, out var checksums);
+
+        foreach (var (projectId, version) in this.Map)
+        {
+            checksums.Add(projectId.Checksum);
+            checksums.Add(Checksum.Create(version, static (v, w) => v.WriteTo(w)));
+        }
+
+        return Checksum.Create(checksums);
+    }
+
     public override string ToString()
     {
         using var _ = PooledStringBuilder.GetInstance(out var builder);

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

@@ -131,7 +131,11 @@ public async Task FindAsync<TArg>(
                 onAssetFound(this.Checksum, this, arg);
 
             if (assetPath.IncludeSolutionSourceGeneratorExecutionVersionMap && searchingChecksumsLeft.Remove(this.SourceGeneratorExecutionVersionMap))
-                onAssetFound(this.SourceGeneratorExecutionVersionMap, compilationState.SourceGeneratorExecutionVersionMap, arg);
+            {
+                // Only send over the part of the execution map corresponding to the project cone.
+                var filteredExecutionMap = compilationState.GetFilteredSourceGenerationExecutionMap(projectCone);
+                onAssetFound(this.SourceGeneratorExecutionVersionMap, filteredExecutionMap, arg);
+            }
 
             if (compilationState.FrozenSourceGeneratedDocumentStates != null)
             {

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

@@ -52,7 +52,7 @@ await this.SetCurrentSolutionAsync(
             oldSolution =>
             {
                 var updates = GetUpdatedSourceGeneratorVersions(oldSolution, projectIds);
-                return oldSolution.WithSourceGeneratorExecutionVersions(updates, cancellationToken);
+                return oldSolution.UpdateSpecificSourceGeneratorExecutionVersions(updates, cancellationToken);
             },
             static (_, _) => (WorkspaceChangeKind.SolutionChanged, projectId: null, documentId: null),
             onBeforeUpdate: null,

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

@@ -5,6 +5,7 @@
 using System;
 using System.Collections.Generic;
 using System.Collections.Immutable;
+using System.Diagnostics;
 using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
@@ -122,11 +123,21 @@ public async Task<Solution> CreateSolutionAsync(Checksum newSolutionChecksum, Ca
                         var newVersions = await _assetProvider.GetAssetAsync<SourceGeneratorExecutionVersionMap>(
                             AssetPathKind.SolutionSourceGeneratorExecutionVersionMap, newSolutionCompilationChecksums.SourceGeneratorExecutionVersionMap, cancellationToken).ConfigureAwait(false);
 
-                        // The execution version map will be for the entire solution on the host side.  However, we may
-                        // only be syncing over a partial cone.  In that case, filter down the version map we apply to
-                        // the local solution to only be for that cone as well.
-                        newVersions = FilterToProjectCone(newVersions, newSolutionChecksums.ProjectCone);
-                        solution = solution.WithSourceGeneratorExecutionVersions(newVersions, cancellationToken);
+#if DEBUG
+                        var projectCone = newSolutionChecksums.ProjectCone;
+                        if (projectCone != null)
+                        {
+                            Debug.Assert(projectCone.ProjectIds.Count == newVersions.Map.Count);
+                            Debug.Assert(projectCone.ProjectIds.All(id => newVersions.Map.ContainsKey(id)));
+                        }
+                        else
+                        {
+                            Debug.Assert(solution.ProjectIds.Count == newVersions.Map.Count);
+                            Debug.Assert(solution.ProjectIds.All(id => newVersions.Map.ContainsKey(id)));
+                        }
+#endif
+
+                        solution = solution.UpdateSpecificSourceGeneratorExecutionVersions(newVersions, cancellationToken);
                     }
 
 #if DEBUG
@@ -140,21 +151,6 @@ public async Task<Solution> CreateSolutionAsync(Checksum newSolutionChecksum, Ca
                 {
                     throw ExceptionUtilities.Unreachable();
                 }
-
-                static SourceGeneratorExecutionVersionMap FilterToProjectCone(SourceGeneratorExecutionVersionMap map, ProjectCone? projectCone)
-                {
-                    if (projectCone is null)
-                        return map;
-
-                    var builder = map.Map.ToBuilder();
-                    foreach (var (projectId, _) in map.Map)
-                    {
-                        if (!projectCone.Contains(projectId))
-                            builder.Remove(projectId);
-                    }
-
-                    return new(builder.ToImmutable());
-                }
             }
 
             private async Task<Solution> UpdateProjectsAsync(

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

@@ -10,6 +10,7 @@
 using Microsoft.CodeAnalysis.Host;
 using Microsoft.CodeAnalysis.PooledObjects;
 using Microsoft.CodeAnalysis.Serialization;
+using Microsoft.VisualStudio.Telemetry;
 using Microsoft.VisualStudio.Threading;
 using Roslyn.Utilities;
 using static Microsoft.VisualStudio.Threading.ThreadingTools;
@@ -207,7 +208,7 @@ async ValueTask DecrementInFlightCountAsync(InFlightSolution inFlightSolution)
 
         /// <summary>
         /// Create an appropriate <see cref="Solution"/> instance corresponding to the <paramref
-        /// name="solutionChecksum"/> passed in.  Note: this method changes no Workspace state and exists purely to
+        /// name="newSolutionChecksum"/> passed in.  Note: this method changes no Workspace state and exists purely to
         /// compute the corresponding solution.  Updating of our caches, or storing this solution as the <see
         /// cref="Workspace.CurrentSolution"/> of this <see cref="RemoteWorkspace"/> is the responsibility of any
         /// callers.
@@ -224,26 +225,41 @@ async ValueTask DecrementInFlightCountAsync(InFlightSolution inFlightSolution)
         /// </summary>
         private async Task<Solution> ComputeDisconnectedSolutionAsync(
             AssetProvider assetProvider,
-            Checksum solutionChecksum,
+            Checksum newSolutionChecksum,
             CancellationToken cancellationToken)
         {
             try
             {
                 // Try to create the solution snapshot incrementally off of the workspaces CurrentSolution first.
                 var updater = new SolutionCreator(Services.HostServices, assetProvider, this.CurrentSolution);
-                if (await updater.IsIncrementalUpdateAsync(solutionChecksum, cancellationToken).ConfigureAwait(false))
+                if (await updater.IsIncrementalUpdateAsync(newSolutionChecksum, cancellationToken).ConfigureAwait(false))
                 {
-                    return await updater.CreateSolutionAsync(solutionChecksum, cancellationToken).ConfigureAwait(false);
+                    return await updater.CreateSolutionAsync(newSolutionChecksum, cancellationToken).ConfigureAwait(false);
                 }
                 else
                 {
                     // Otherwise, this is a different solution, or the first time we're creating this solution.  Bulk
                     // sync over all assets for it.
-                    await assetProvider.SynchronizeSolutionAssetsAsync(solutionChecksum, cancellationToken).ConfigureAwait(false);
+                    await assetProvider.SynchronizeSolutionAssetsAsync(newSolutionChecksum, cancellationToken).ConfigureAwait(false);
 
                     // get new solution info and options
-                    var solutionInfo = await assetProvider.CreateSolutionInfoAsync(solutionChecksum, cancellationToken).ConfigureAwait(false);
-                    return CreateSolutionFromInfo(solutionInfo);
+                    var solutionInfo = await assetProvider.CreateSolutionInfoAsync(newSolutionChecksum, cancellationToken).ConfigureAwait(false);
+                    var solution = CreateSolutionFromInfo(solutionInfo);
+
+                    // ensure that the solution has the correct source generator execution versions. note we should do
+                    // this all in a unified fashion with CreateSolutionAsync above.  However, that is blocked in
+                    // https://github.com/dotnet/roslyn/pull/72860
+                    {
+                        var newSolutionCompilationChecksums = await assetProvider.GetAssetAsync<SolutionCompilationStateChecksums>(
+                            AssetPathKind.SolutionCompilationStateChecksums, newSolutionChecksum, cancellationToken).ConfigureAwait(false);
+
+                        var newVersions = await assetProvider.GetAssetAsync<SourceGeneratorExecutionVersionMap>(
+                            AssetPathKind.SolutionSourceGeneratorExecutionVersionMap, newSolutionCompilationChecksums.SourceGeneratorExecutionVersionMap, cancellationToken).ConfigureAwait(false);
+
+                        solution = solution.UpdateSpecificSourceGeneratorExecutionVersions(newVersions, cancellationToken);
+                    }
+
+                    return solution;
                 }
             }
             catch (Exception e) when (FatalError.ReportAndPropagateUnlessCanceled(e, cancellationToken))