Fix snapshot's active target framework bug

src/Microsoft.VisualStudio.ProjectSystem.Managed.VS.UnitTests/ProjectSystem/VS/Tree/Dependencies/Snapshot/DependenciesSnapshotTests.cs

@@ -251,6 +251,77 @@ public void FromChanges_ProjectPathAndTargetChange()
             Assert.Equal(@"tfm1\Xxx\dependency1", snapshot.DependenciesByTargetFramework[targetFramework].DependenciesWorld.First().Value.Id);
         }
 
+        [Fact]
+        public void SetTargets_FromEmpty()
+        {
+            const string projectPath = @"c:\somefolder\someproject\a.csproj";
+
+            ITargetFramework tfm1 = new TargetFramework("tfm1");
+            ITargetFramework tfm2 = new TargetFramework("tfm2");
+
+            var snapshot = DependenciesSnapshot.CreateEmpty(projectPath)
+                .SetTargets(new[] { tfm1, tfm2 }.ToImmutableArray(), tfm1);
+
+            Assert.Same(tfm1, snapshot.ActiveTargetFramework);
+            Assert.Equal(2, snapshot.DependenciesByTargetFramework.Count);
+            Assert.True(snapshot.DependenciesByTargetFramework.ContainsKey(tfm1));
+            Assert.True(snapshot.DependenciesByTargetFramework.ContainsKey(tfm2));
+        }
+
+        [Fact]
+        public void SetTargets_SameMembers_DifferentActive()
+        {
+            const string projectPath = @"c:\somefolder\someproject\a.csproj";
+
+            ITargetFramework tfm1 = new TargetFramework("tfm1");
+            ITargetFramework tfm2 = new TargetFramework("tfm2");
+
+            var before = DependenciesSnapshot.CreateEmpty(projectPath)
+                .SetTargets(new[] { tfm1, tfm2 }.ToImmutableArray(), tfm1);
+
+            var after = before.SetTargets(new[] { tfm1, tfm2 }.ToImmutableArray(), tfm2);
+
+            Assert.Same(tfm2, after.ActiveTargetFramework);
+            Assert.Same(before.DependenciesByTargetFramework, after.DependenciesByTargetFramework);
+        }
+
+        [Fact]
+        public void SetTargets_SameMembers_SameActive()
+        {
+            const string projectPath = @"c:\somefolder\someproject\a.csproj";
+
+            ITargetFramework tfm1 = new TargetFramework("tfm1");
+            ITargetFramework tfm2 = new TargetFramework("tfm2");
+
+            var before = DependenciesSnapshot.CreateEmpty(projectPath)
+                .SetTargets(new[] { tfm1, tfm2 }.ToImmutableArray(), tfm1);
+
+            var after = before.SetTargets(new[] { tfm1, tfm2 }.ToImmutableArray(), tfm1);
+
+            Assert.Same(before, after);
+        }
+
+        [Fact]
+        public void SetTargets_DifferentMembers_DifferentActive()
+        {
+            const string projectPath = @"c:\somefolder\someproject\a.csproj";
+
+            ITargetFramework tfm1 = new TargetFramework("tfm1");
+            ITargetFramework tfm2 = new TargetFramework("tfm2");
+            ITargetFramework tfm3 = new TargetFramework("tfm3");
+
+            var before = DependenciesSnapshot.CreateEmpty(projectPath)
+                .SetTargets(new[] { tfm1, tfm2 }.ToImmutableArray(), tfm1);
+
+            var after = before.SetTargets(new[] { tfm2, tfm3 }.ToImmutableArray(), tfm3);
+
+            Assert.Same(tfm3, after.ActiveTargetFramework);
+            Assert.Equal(2, after.DependenciesByTargetFramework.Count);
+            Assert.True(after.DependenciesByTargetFramework.ContainsKey(tfm2));
+            Assert.True(after.DependenciesByTargetFramework.ContainsKey(tfm3));
+            Assert.Same(before.DependenciesByTargetFramework[tfm2], after.DependenciesByTargetFramework[tfm2]);
+        }
+
         private static ImmutableDictionary<ITargetFramework, ITargetedDependenciesSnapshot> CreateDependenciesByTargetFramework(
             string projectPath,
             IProjectCatalogSnapshot catalogs,

src/Microsoft.VisualStudio.ProjectSystem.Managed/Tree/Dependencies/Snapshot/DependenciesSnapshot.cs

@@ -151,14 +151,29 @@ bool SyncTargetFrameworks()
             }
         }
 
-        public DependenciesSnapshot RemoveTargets(IEnumerable<ITargetFramework> targetToRemove)
+        public DependenciesSnapshot SetTargets(
+            ImmutableArray<ITargetFramework> targetFrameworks,
+            ITargetFramework activeTargetFramework)
         {
-            ImmutableDictionary<ITargetFramework, ITargetedDependenciesSnapshot> newTargets = DependenciesByTargetFramework.RemoveRange(targetToRemove);
+            bool activeChanged = !activeTargetFramework.Equals(ActiveTargetFramework);
+
+            var map = DependenciesByTargetFramework;
+
+            var diff = new SetDiff<ITargetFramework>(map.Keys, targetFrameworks);
+
+            map = map.RemoveRange(diff.Removed);
+            map = map.AddRange(
+                diff.Added.Select(
+                    added => new KeyValuePair<ITargetFramework, ITargetedDependenciesSnapshot>(
+                        added,
+                        TargetedDependenciesSnapshot.CreateEmpty(ProjectPath, added, null))));
+
+            if (activeChanged || !ReferenceEquals(map, DependenciesByTargetFramework))
+            {
+                return new DependenciesSnapshot(ProjectPath, activeTargetFramework, map);
+            }
 
-            // Return this if no targets changed
-            return ReferenceEquals(newTargets, DependenciesByTargetFramework)
-                ? this
-                : new DependenciesSnapshot(ProjectPath, ActiveTargetFramework, newTargets);
+            return this;
         }
 
         // Internal, for test use -- normal code should use the factory methods

src/Microsoft.VisualStudio.ProjectSystem.Managed/Tree/Dependencies/Subscriptions/DependenciesSnapshotProvider.cs

@@ -8,6 +8,7 @@
 using System.Threading;
 using System.Threading.Tasks;
 using System.Threading.Tasks.Dataflow;
+
 using Microsoft.Build.Execution;
 using Microsoft.VisualStudio.Build;
 using Microsoft.VisualStudio.Composition;
@@ -489,31 +490,7 @@ void OnAggregateContextChanged(
                 AggregateCrossTargetProjectContext? oldContext,
                 AggregateCrossTargetProjectContext newContext)
             {
-                if (oldContext == null)
-                {
-                    // all new rules will be sent to new context, we don't need to clean up anything
-                    return;
-                }
-
-                var targetsToClean = new HashSet<ITargetFramework>();
-
-                ImmutableArray<ITargetFramework> oldTargets = oldContext.TargetFrameworks;
-
-                if (newContext == null)
-                {
-                    targetsToClean.AddRange(oldTargets);
-                }
-                else
-                {
-                    ImmutableArray<ITargetFramework> newTargets = newContext.TargetFrameworks;
-
-                    targetsToClean.AddRange(oldTargets.Except(newTargets));
-                }
-
-                if (targetsToClean.Count != 0)
-                {
-                    TryUpdateSnapshot(snapshot => snapshot.RemoveTargets(targetsToClean));
-                }
+                TryUpdateSnapshot(snapshot => snapshot.SetTargets(newContext.TargetFrameworks, newContext.ActiveTargetFramework));
             }
         }
 

src/Microsoft.VisualStudio.ProjectSystem.Managed/Utilities/SetDiff.cs

@@ -4,8 +4,6 @@
 using System.Collections;
 using System.Collections.Generic;
 
-#nullable disable
-
 namespace Microsoft.VisualStudio.ProjectSystem
 {
     internal sealed class SetDiff<T>
@@ -70,7 +68,7 @@ public PartEnumerator(Dictionary<T, byte> dic, byte flag)
                 {
                     _flag = flag;
                     _enumerator = dic.GetEnumerator();
-                    Current = default;
+                    Current = default!;
                 }
 
                 public bool MoveNext()
@@ -89,7 +87,7 @@ public bool MoveNext()
 
                 public T Current { get; private set; }
 
-                object IEnumerator.Current => Current;
+                object IEnumerator.Current => Current!;
 
                 void IEnumerator.Reset() => throw new NotSupportedException();