Use the project cone information in Scope.FindAssetsAsync (part 2)

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

@@ -0,0 +1,38 @@
+// 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.Frozen;
+using Roslyn.Utilities;
+
+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"/>.
+/// </summary>
+internal sealed class ProjectCone : IEquatable<ProjectCone>
+{
+    public readonly ProjectId RootProjectId;
+    private readonly FrozenSet<ProjectId> _projectIds;
+
+    public ProjectCone(ProjectId rootProjectId, FrozenSet<ProjectId> projectIds)
+    {
+        Contract.ThrowIfFalse(projectIds.Contains(rootProjectId));
+        RootProjectId = rootProjectId;
+        _projectIds = projectIds;
+    }
+
+    public bool Contains(ProjectId 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);
+
+    public override int GetHashCode()
+        => throw new NotImplementedException();
+}

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

@@ -69,9 +69,12 @@ private SolutionCompilationState(
         FrozenSourceGeneratedDocumentStates = frozenSourceGeneratedDocumentStates;
 
         // when solution state is changed, we recalculate its checksum
-        _lazyChecksums = AsyncLazy.Create(static (self, c) =>
-            self.ComputeChecksumsAsync(projectId: null, c),
-            arg: this);
+        _lazyChecksums = AsyncLazy.Create(static async (self, cancellationToken) =>
+        {
+            var (checksums, projectCone) = await self.ComputeChecksumsAsync(projectId: null, cancellationToken).ConfigureAwait(false);
+            Contract.ThrowIfTrue(projectCone != null);
+            return checksums;
+        }, arg: this);
         _cachedFrozenSnapshot = cachedFrozenSnapshot ??
             AsyncLazy.Create(synchronousComputeFunction: static (self, c) =>
                 self.ComputeFrozenSnapshot(c),

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

@@ -27,25 +27,32 @@ internal partial class SolutionCompilationState
     /// Mapping from project-id to the checksums needed to synchronize it over to an OOP host.  Lock this specific
     /// field before reading/writing to it.
     /// </summary>
-    private readonly Dictionary<ProjectId, AsyncLazy<SolutionCompilationStateChecksums>> _lazyProjectChecksums = [];
+    private readonly Dictionary<ProjectId, AsyncLazy<(SolutionCompilationStateChecksums checksums, ProjectCone projectCone)>> _lazyProjectChecksums = [];
 
     public bool TryGetStateChecksums([NotNullWhen(true)] out SolutionCompilationStateChecksums? stateChecksums)
         => _lazyChecksums.TryGetValue(out stateChecksums);
 
     public bool TryGetStateChecksums(ProjectId projectId, [NotNullWhen(true)] out SolutionCompilationStateChecksums? stateChecksums)
     {
-        AsyncLazy<SolutionCompilationStateChecksums>? checksums;
+        AsyncLazy<(SolutionCompilationStateChecksums checksums, ProjectCone projectCone)>? lazyChecksums;
         lock (_lazyProjectChecksums)
         {
-            if (!_lazyProjectChecksums.TryGetValue(projectId, out checksums) ||
-                checksums == null)
+            if (!_lazyProjectChecksums.TryGetValue(projectId, out lazyChecksums) ||
+                lazyChecksums == null)
             {
                 stateChecksums = null;
                 return false;
             }
         }
 
-        return checksums.TryGetValue(out stateChecksums);
+        if (!lazyChecksums.TryGetValue(out var checksumsAndCone))
+        {
+            stateChecksums = null;
+            return false;
+        }
+
+        stateChecksums = checksumsAndCone.checksums;
+        return true;
     }
 
     public Task<SolutionCompilationStateChecksums> GetStateChecksumsAsync(CancellationToken cancellationToken)
@@ -58,20 +65,23 @@ public async Task<Checksum> GetChecksumAsync(CancellationToken cancellationToken
     }
 
     /// <summary>Gets the checksum for only the requested project (and any project it depends on)</summary>
-    public async Task<SolutionCompilationStateChecksums> GetStateChecksumsAsync(
+    public async Task<(SolutionCompilationStateChecksums checksums, ProjectCone projectCone)> GetStateChecksumsAsync(
         ProjectId projectId,
         CancellationToken cancellationToken)
     {
         Contract.ThrowIfNull(projectId);
 
-        AsyncLazy<SolutionCompilationStateChecksums>? checksums;
+        AsyncLazy<(SolutionCompilationStateChecksums checksums, ProjectCone projectCone)>? checksums;
         lock (_lazyProjectChecksums)
         {
             if (!_lazyProjectChecksums.TryGetValue(projectId, out checksums))
             {
-                checksums = AsyncLazy.Create(static (arg, c) =>
-                    arg.self.ComputeChecksumsAsync(arg.projectId, c),
-                    arg: (self: this, projectId));
+                checksums = AsyncLazy.Create(static async (arg, cancellationToken) =>
+                {
+                    var (checksum, projectCone) = await arg.self.ComputeChecksumsAsync(arg.projectId, cancellationToken).ConfigureAwait(false);
+                    Contract.ThrowIfNull(projectCone);
+                    return (checksum, projectCone);
+                }, arg: (self: this, projectId));
 
                 _lazyProjectChecksums.Add(projectId, checksums);
             }
@@ -84,21 +94,32 @@ public async Task<SolutionCompilationStateChecksums> GetStateChecksumsAsync(
     /// <summary>Gets the checksum for only the requested project (and any project it depends on)</summary>
     public async Task<Checksum> GetChecksumAsync(ProjectId projectId, CancellationToken cancellationToken)
     {
-        var checksums = await GetStateChecksumsAsync(projectId, cancellationToken).ConfigureAwait(false);
+        var (checksums, _) = await GetStateChecksumsAsync(projectId, cancellationToken).ConfigureAwait(false);
         return checksums.Checksum;
     }
 
-    private async Task<SolutionCompilationStateChecksums> ComputeChecksumsAsync(
+    private async Task<(SolutionCompilationStateChecksums checksums, ProjectCone? projectCone)> ComputeChecksumsAsync(
         ProjectId? projectId,
         CancellationToken cancellationToken)
     {
         try
         {
             using (Logger.LogBlock(FunctionId.SolutionCompilationState_ComputeChecksumsAsync, this.SolutionState.FilePath, cancellationToken))
             {
-                var solutionStateChecksum = projectId == null
-                    ? await this.SolutionState.GetChecksumAsync(cancellationToken).ConfigureAwait(false)
-                    : await this.SolutionState.GetChecksumAsync(projectId, cancellationToken).ConfigureAwait(false);
+                Checksum solutionStateChecksum;
+                ProjectCone? projectCone;
+
+                if (projectId is null)
+                {
+                    solutionStateChecksum = await this.SolutionState.GetChecksumAsync(cancellationToken).ConfigureAwait(false);
+                    projectCone = null;
+                }
+                else
+                {
+                    var stateChecksums = await this.SolutionState.GetStateChecksumsAsync(projectId, cancellationToken).ConfigureAwait(false);
+                    solutionStateChecksum = stateChecksums.Checksum;
+                    projectCone = stateChecksums.ProjectCone;
+                }
 
                 ChecksumCollection? frozenSourceGeneratedDocumentIdentities = null;
                 ChecksumsAndIds<DocumentId>? frozenSourceGeneratedDocuments = null;
@@ -112,10 +133,11 @@ private async Task<SolutionCompilationStateChecksums> ComputeChecksumsAsync(
                     frozenSourceGeneratedDocuments = await FrozenSourceGeneratedDocumentStates.Value.GetChecksumsAndIdsAsync(cancellationToken).ConfigureAwait(false);
                 }
 
-                return new SolutionCompilationStateChecksums(
+                var compilationStateChecksums = new SolutionCompilationStateChecksums(
                     solutionStateChecksum,
                     frozenSourceGeneratedDocumentIdentities,
                     frozenSourceGeneratedDocuments);
+                return (compilationStateChecksums, projectCone);
             }
         }
         catch (Exception e) when (FatalError.ReportAndPropagateUnlessCanceled(e, cancellationToken))

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

@@ -3,6 +3,7 @@
 // See the LICENSE file in the project root for more information.
 
 using System;
+using System.Collections.Frozen;
 using System.Collections.Generic;
 using System.Collections.Immutable;
 using System.Diagnostics.CodeAnalysis;
@@ -78,21 +79,15 @@ public async Task<SolutionStateChecksums> GetStateChecksumsAsync(
         {
             if (!_lazyProjectChecksums.TryGetValue(projectId, out checksums))
             {
-                checksums = Compute(projectId);
+                checksums = AsyncLazy.Create(
+                    static (arg, cancellationToken) => arg.self.ComputeChecksumsAsync(arg.projectId, cancellationToken),
+                    arg: (self: this, projectId));
                 _lazyProjectChecksums.Add(projectId, checksums);
             }
         }
 
         var collection = await checksums.GetValueAsync(cancellationToken).ConfigureAwait(false);
         return collection;
-
-        // Extracted as a local function to prevent delegate allocations when not needed.
-        AsyncLazy<SolutionStateChecksums> Compute(ProjectId projectConeId)
-        {
-            return AsyncLazy.Create(static (arg, c) =>
-                arg.self.ComputeChecksumsAsync(arg.projectConeId, c),
-                arg: (self: this, projectConeId));
-        }
     }
 
     /// <summary>Gets the checksum for only the requested project (and any project it depends on)</summary>
@@ -141,11 +136,18 @@ private async Task<SolutionStateChecksums> ComputeChecksumsAsync(
                 var analyzerReferenceChecksums = ChecksumCache.GetOrCreateChecksumCollection(
                     this.AnalyzerReferences, this.Services.GetRequiredService<ISerializerService>(), cancellationToken);
 
-                return new SolutionStateChecksums(
+                var stateChecksums = new SolutionStateChecksums(
                     projectConeId,
                     this.SolutionAttributes.Checksum,
                     new(new ChecksumCollection(projectChecksums), projectIds),
                     analyzerReferenceChecksums);
+
+#if DEBUG
+                var projectConeTemp = projectConeId is null ? null : new ProjectCone(projectConeId, projectCone.Object.ToFrozenSet());
+                RoslynDebug.Assert(Equals(projectConeTemp, stateChecksums.ProjectCone));
+#endif
+
+                return stateChecksums;
             }
         }
         catch (Exception e) when (FatalError.ReportAndPropagateUnlessCanceled(e, cancellationToken))

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

@@ -3,6 +3,7 @@
 // See the LICENSE file in the project root for more information.
 
 using System;
+using System.Collections.Frozen;
 using System.Collections.Generic;
 using System.Diagnostics;
 using System.Runtime.CompilerServices;
@@ -88,6 +89,7 @@ public static SolutionCompilationStateChecksums Deserialize(ObjectReader reader)
 
     public async Task FindAsync(
         SolutionCompilationState compilationState,
+        ProjectCone? projectCone,
         AssetHint assetHint,
         HashSet<Checksum> searchingChecksumsLeft,
         Dictionary<Checksum, object> result,
@@ -123,16 +125,18 @@ public async Task FindAsync(
         }
 
         var solutionState = compilationState.SolutionState;
-        if (solutionState.TryGetStateChecksums(out var solutionChecksums))
-            await solutionChecksums.FindAsync(solutionState, assetHint, searchingChecksumsLeft, result, cancellationToken).ConfigureAwait(false);
-
-        foreach (var projectId in solutionState.ProjectIds)
+        if (projectCone is null)
         {
-            if (searchingChecksumsLeft.Count == 0)
-                break;
-
-            if (solutionState.TryGetStateChecksums(projectId, out solutionChecksums))
-                await solutionChecksums.FindAsync(solutionState, assetHint, searchingChecksumsLeft, result, cancellationToken).ConfigureAwait(false);
+            // If we're not in a project cone, start the search at the top most state-checksum corresponding to the
+            // entire solution.
+            Contract.ThrowIfFalse(solutionState.TryGetStateChecksums(out var solutionChecksums));
+            await solutionChecksums.FindAsync(solutionState, projectCone, assetHint, searchingChecksumsLeft, result, cancellationToken).ConfigureAwait(false);
+        }
+        else
+        {
+            // Otherwise, grab the top-most state checksum for this cone and search within that.
+            Contract.ThrowIfFalse(solutionState.TryGetStateChecksums(projectCone.RootProjectId, out var solutionChecksums));
+            await solutionChecksums.FindAsync(solutionState, projectCone, assetHint, searchingChecksumsLeft, result, cancellationToken).ConfigureAwait(false);
         }
     }
 }
@@ -145,6 +149,8 @@ internal sealed class SolutionStateChecksums(
     ChecksumsAndIds<ProjectId> projects,
     ChecksumCollection analyzerReferences)
 {
+    private ProjectCone? _projectCone;
+
     public Checksum Checksum { get; } = Checksum.Create(stackalloc[]
     {
         projectConeId == null ? Checksum.Null : projectConeId.Checksum,
@@ -158,6 +164,14 @@ internal sealed class SolutionStateChecksums(
     public ChecksumsAndIds<ProjectId> Projects { get; } = projects;
     public ChecksumCollection AnalyzerReferences { get; } = analyzerReferences;
 
+    // Acceptably not threadsafe.  ProjectCone is a class, and the runtime guarantees anyone will see this field fully
+    // initialized.  It's acceptable to have multiple instances of this in a race condition as the data will be same
+    // (and our asserts don't check for reference equality, only value equality).
+    public ProjectCone? ProjectCone => _projectCone ??= ComputeProjectCone();
+
+    private ProjectCone? ComputeProjectCone()
+        => ProjectConeId == null ? null : new ProjectCone(ProjectConeId, Projects.Ids.ToFrozenSet());
+
     public void AddAllTo(HashSet<Checksum> checksums)
     {
         checksums.AddIfNotNullChecksum(this.Checksum);
@@ -193,6 +207,7 @@ public static SolutionStateChecksums Deserialize(ObjectReader reader)
 
     public async Task FindAsync(
         SolutionState solution,
+        ProjectCone? projectCone,
         AssetHint assetHint,
         HashSet<Checksum> searchingChecksumsLeft,
         Dictionary<Checksum, object> result,
@@ -216,6 +231,10 @@ public async Task FindAsync(
 
         if (assetHint.ProjectId != null)
         {
+            Contract.ThrowIfTrue(
+                projectCone != null && !projectCone.Contains(assetHint.ProjectId),
+                "Requesting an asset outside of the cone explicitly being asked for!");
+
             var projectState = solution.GetProjectState(assetHint.ProjectId);
             if (projectState != null &&
                 projectState.TryGetStateChecksums(out var projectStateChecksums))
@@ -231,11 +250,16 @@ public async Task FindAsync(
             // level. This ensures that when we are trying to sync the projects referenced by a SolutionStateChecksums'
             // instance that we don't unnecessarily walk all documents looking just for those.
 
-            foreach (var (_, projectState) in solution.ProjectStates)
+            foreach (var (projectId, projectState) in solution.ProjectStates)
             {
                 if (searchingChecksumsLeft.Count == 0)
                     break;
 
+                // If we're syncing a project cone, no point at all at looking at child projects of the solution that
+                // are not in that cone.
+                if (projectCone != null && !projectCone.Contains(projectId))
+                    continue;
+
                 if (projectState.TryGetStateChecksums(out var projectStateChecksums) &&
                     searchingChecksumsLeft.Remove(projectStateChecksums.Checksum))
                 {
@@ -245,11 +269,16 @@ public async Task FindAsync(
 
             // Now actually do the depth first search into each project.
 
-            foreach (var (_, projectState) in solution.ProjectStates)
+            foreach (var (projectId, projectState) in solution.ProjectStates)
             {
                 if (searchingChecksumsLeft.Count == 0)
                     break;
 
+                // If we're syncing a project cone, no point at all at looking at child projects of the solution that
+                // are not in that cone.
+                if (projectCone != null && !projectCone.Contains(projectId))
+                    continue;
+
                 // It's possible not all all our projects have checksums.  Specifically, we may have only been asked to
                 // compute the checksum tree for a subset of projects that were all that a feature needed.
                 if (projectState.TryGetStateChecksums(out var projectStateChecksums))

src/Workspaces/Remote/Core/SolutionAssetStorage.Scope.cs

@@ -18,13 +18,13 @@ internal partial class SolutionAssetStorage
     internal sealed partial class Scope(
         SolutionAssetStorage storage,
         Checksum solutionChecksum,
-        ProjectId? projectId,
+        ProjectCone? projectCone,
         SolutionCompilationState compilationState) : IDisposable
     {
         private readonly SolutionAssetStorage _storage = storage;
 
         public readonly Checksum SolutionChecksum = solutionChecksum;
-        public readonly ProjectId? ProjectId = projectId;
+        public readonly ProjectCone? ProjectCone = projectCone;
         public readonly SolutionCompilationState CompilationState = compilationState;
 
         /// <summary>
@@ -69,16 +69,20 @@ private async Task FindAssetsAsync(
             AssetHint assetHint, HashSet<Checksum> remainingChecksumsToFind, Dictionary<Checksum, object> result, CancellationToken cancellationToken)
         {
             var solutionState = this.CompilationState;
-            SolutionCompilationStateChecksums? stateChecksums;
 
-            if (ProjectId is null)
-                solutionState.TryGetStateChecksums(out stateChecksums);
+            if (this.ProjectCone is null)
+            {
+                // If we're not in a project cone, start the search at the top most state-checksum corresponding to the
+                // entire solution.
+                Contract.ThrowIfFalse(solutionState.TryGetStateChecksums(out var stateChecksums));
+                await stateChecksums.FindAsync(solutionState, this.ProjectCone, assetHint, remainingChecksumsToFind, result, cancellationToken).ConfigureAwait(false);
+            }
             else
-                solutionState.TryGetStateChecksums(ProjectId, out stateChecksums);
-
-            Contract.ThrowIfNull(stateChecksums);
-
-            await stateChecksums.FindAsync(solutionState, assetHint, remainingChecksumsToFind, result, cancellationToken).ConfigureAwait(false);
+            {
+                // Otherwise, grab the top-most state checksum for this cone and search within that.
+                Contract.ThrowIfFalse(solutionState.TryGetStateChecksums(this.ProjectCone.RootProjectId, out var stateChecksums));
+                await stateChecksums.FindAsync(solutionState, this.ProjectCone, assetHint, remainingChecksumsToFind, result, cancellationToken).ConfigureAwait(false);
+            }
         }
 
         public TestAccessor GetTestAccessor()