Add initial tests for dependency graph state following project removal

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

@@ -202,12 +202,15 @@ internal ProjectDependencyGraph WithAdditionalProjectReferences(ProjectId projec
         /// Must be called on a non-null map.
         /// </summary>
         private static ImmutableDictionary<ProjectId, ImmutableHashSet<ProjectId>> ComputeNewReverseReferencesMapForRemovedProject(
+            ImmutableDictionary<ProjectId, ImmutableHashSet<ProjectId>> existingForwardReferencesMap,
             ImmutableDictionary<ProjectId, ImmutableHashSet<ProjectId>> existingReverseReferencesMap,
             ProjectId projectId)
         {
             var builder = existingReverseReferencesMap.ToBuilder();
 
-            if (existingReverseReferencesMap.TryGetValue(projectId, out var referencedProjectIds))
+            // Iterate over each project referenced by 'projectId', which is now being removed. Update the reverse
+            // references map for the project to no longer include 'projectId' in the list.
+            if (existingForwardReferencesMap.TryGetValue(projectId, out var referencedProjectIds))
             {
                 foreach (var referencedProjectId in referencedProjectIds)
                 {
@@ -218,6 +221,7 @@ internal ProjectDependencyGraph WithAdditionalProjectReferences(ProjectId projec
                 }
             }
 
+            // Finally, remove 'projectId' itself.
             builder.Remove(projectId);
             return builder.ToImmutable();
         }
@@ -235,6 +239,46 @@ internal ProjectDependencyGraph WithAdditionalProjectReferences(ProjectId projec
             return existingReverseReferencesMap;
         }
 
+        /// <summary>
+        /// Computes a new <see cref="_referencesMap"/> for the removal of a project.
+        /// </summary>
+        /// <param name="existingReverseReferencesMap">The <see cref="_lazyReverseReferencesMap"/> prior to the removal.
+        /// This map serves as a hint to the removal process; i.e. it is assumed correct if it contains data, but may be
+        /// omitted without impacting correctness.</param>
+        /// <param name="referencesMap">The <see cref="_referencesMap"/> prior to the removal.</param>
+        /// <param name="projectId">The ID of the project which is being removed.</param>
+        /// <returns>The <see cref="_referencesMap"/> for the project dependency graph once the project is removed.</returns>
+        private ImmutableDictionary<ProjectId, ImmutableHashSet<ProjectId>> ComputeNewReferencesMapForRemovedProject(
+            ImmutableDictionary<ProjectId, ImmutableHashSet<ProjectId>>? existingReverseReferencesMap,
+            ImmutableDictionary<ProjectId, ImmutableHashSet<ProjectId>> referencesMap,
+            ProjectId projectId)
+        {
+            var builder = referencesMap.ToBuilder();
+
+            if (existingReverseReferencesMap is object && existingReverseReferencesMap.TryGetValue(projectId, out var referencingProjects))
+            {
+                // We know all the projects directly referencing 'projectId', so remove 'projectId' from the set of
+                // references in each of those cases directly.
+                foreach (var id in referencingProjects)
+                {
+                    builder[id] = builder[id].Remove(projectId);
+                }
+            }
+            else
+            {
+                // We don't know which projects reference 'projectId', so iterate over all known projects and remove
+                // 'projectId' from the set of references if it exists.
+                foreach (var (id, references) in referencesMap)
+                {
+                    builder[id] = references.Remove(projectId);
+                }
+            }
+
+            // Finally, remove 'projectId' itself.
+            builder.Remove(projectId);
+            return builder.ToImmutable();
+        }
+
         /// <summary>
         /// Computes a new <see cref="_transitiveReferencesMap"/> for the removal of a project.
         /// Must be called on a non-null map.
@@ -245,11 +289,20 @@ internal ProjectDependencyGraph WithAdditionalProjectReferences(ProjectId projec
         {
             var builder = existingTransitiveReferencesMap.ToBuilder();
 
+            // Iterate over each project and invalidate the transitive references for the project if the project has an
+            // existing transitive reference to 'projectId'.
             foreach (var (project, references) in existingTransitiveReferencesMap)
             {
-                builder[project] = references.Remove(projectId);
+                if (references.Contains(projectId))
+                {
+                    // The project transitively referenced 'projectId', so any transitive references brought in
+                    // exclusively through this reference are no longer valid. Remove the project from the map and the
+                    // new transitive references will be recomputed the first time they are needed.
+                    builder.Remove(project);
+                }
             }
 
+            // Finally, remove 'projectId' itself.
             builder.Remove(projectId);
             return builder.ToImmutable();
         }
@@ -439,12 +492,12 @@ internal ProjectDependencyGraph WithProjectRemoved(ProjectId projectId)
             // Project ID set and direct forward references are trivially updated by removing the key corresponding to
             // the project getting removed.
             var projectIds = _projectIds.Remove(projectId);
-            var referencesMap = _referencesMap.Remove(projectId);
+            var referencesMap = ComputeNewReferencesMapForRemovedProject(_lazyReverseReferencesMap, _referencesMap, projectId);
 
             // The direct reverse references map is updated by removing the key for the project getting removed, and
             // also updating any direct references to the removed project.
             var reverseReferencesMap = _lazyReverseReferencesMap is object
-                ? ComputeNewReverseReferencesMapForRemovedProject(_lazyReverseReferencesMap, projectId)
+                ? ComputeNewReverseReferencesMapForRemovedProject(_referencesMap, _lazyReverseReferencesMap, projectId)
                 : null;
             var transitiveReferencesMap = ComputeNewTransitiveReferencesMapForRemovedProject(_transitiveReferencesMap, projectId);
             var reverseTransitiveReferencesMap = ComputeNewReverseTransitiveReferencesMapForRemovedProject(_reverseTransitiveReferencesMap, projectId);

src/Workspaces/CoreTest/SolutionTests/ProjectDependencyGraphTests.cs

@@ -399,6 +399,94 @@ public void TestReverseTransitiveReferencesForUnrelatedProjectAfterWithProjectRe
             VerifyReverseTransitiveReferences(solution, "B", new string[] { "A" });
         }
 
+        [Fact, Trait(Traits.Feature, Traits.Features.Workspace)]
+        public void TestForwardReferencesAfterProjectRemoval()
+        {
+            // We are going to create a solution with the references:
+            //
+            // A -> B -> C -> D
+            //
+            // and will then remove project B.
+
+            var solution = CreateSolutionFromReferenceMap("A:B B:C C:D D");
+            VerifyDirectReferences(solution, "A", new string[] { "B" });
+            VerifyDirectReferences(solution, "B", new string[] { "C" });
+            VerifyDirectReferences(solution, "C", new string[] { "D" });
+            VerifyDirectReferences(solution, "D", new string[] { });
+
+            solution = solution.RemoveProject(solution.GetProjectsByName("B").Single().Id);
+
+            VerifyDirectReferences(solution, "A", new string[] { });
+            VerifyDirectReferences(solution, "C", new string[] { "D" });
+            VerifyDirectReferences(solution, "D", new string[] { });
+        }
+
+        [Fact, Trait(Traits.Feature, Traits.Features.Workspace)]
+        public void TestForwardTransitiveReferencesAfterProjectRemoval()
+        {
+            // We are going to create a solution with the references:
+            //
+            // A -> B -> C -> D
+            //
+            // and will then remove project B.
+
+            var solution = CreateSolutionFromReferenceMap("A:B B:C C:D D");
+            VerifyTransitiveReferences(solution, "A", new string[] { "B", "C", "D" });
+            VerifyTransitiveReferences(solution, "B", new string[] { "C", "D" });
+            VerifyTransitiveReferences(solution, "C", new string[] { "D" });
+            VerifyTransitiveReferences(solution, "D", new string[] { });
+
+            solution = solution.RemoveProject(solution.GetProjectsByName("B").Single().Id);
+
+            VerifyTransitiveReferences(solution, "A", new string[] { });
+            VerifyTransitiveReferences(solution, "C", new string[] { "D" });
+            VerifyTransitiveReferences(solution, "D", new string[] { });
+        }
+
+        [Fact, Trait(Traits.Feature, Traits.Features.Workspace)]
+        public void TestReverseReferencesAfterProjectRemoval()
+        {
+            // We are going to create a solution with the references:
+            //
+            // A -> B -> C -> D
+            //
+            // and will then remove project B.
+
+            var solution = CreateSolutionFromReferenceMap("A:B B:C C:D D");
+            VerifyDirectReverseReferences(solution, "A", new string[] { });
+            VerifyDirectReverseReferences(solution, "B", new string[] { "A" });
+            VerifyDirectReverseReferences(solution, "C", new string[] { "B" });
+            VerifyDirectReverseReferences(solution, "D", new string[] { "C" });
+
+            solution = solution.RemoveProject(solution.GetProjectsByName("B").Single().Id);
+
+            VerifyDirectReverseReferences(solution, "A", new string[] { });
+            VerifyDirectReverseReferences(solution, "C", new string[] { });
+            VerifyDirectReverseReferences(solution, "D", new string[] { "C" });
+        }
+
+        [Fact, Trait(Traits.Feature, Traits.Features.Workspace)]
+        public void TestReverseTransitiveReferencesAfterProjectRemoval()
+        {
+            // We are going to create a solution with the references:
+            //
+            // A -> B -> C -> D
+            //
+            // and will then remove project B.
+
+            var solution = CreateSolutionFromReferenceMap("A:B B:C C:D D");
+            VerifyReverseTransitiveReferences(solution, "A", new string[] { });
+            VerifyReverseTransitiveReferences(solution, "B", new string[] { "A" });
+            VerifyReverseTransitiveReferences(solution, "C", new string[] { "A", "B" });
+            VerifyReverseTransitiveReferences(solution, "D", new string[] { "A", "B", "C" });
+
+            solution = solution.RemoveProject(solution.GetProjectsByName("B").Single().Id);
+
+            VerifyReverseTransitiveReferences(solution, "A", new string[] { });
+            VerifyReverseTransitiveReferences(solution, "C", new string[] { });
+            VerifyReverseTransitiveReferences(solution, "D", new string[] { "C" });
+        }
+
         private void VerifyDirectReverseReferences(Solution solution, string project, string[] expectedResults)
         {
             var projectDependencyGraph = solution.GetProjectDependencyGraph();