Fix a couple of bugs in SG (#55032)

* Fixes a couple of bugs in sg:
- Ensure that .Collect() reports as cached if all the inputs are cached
- Don't use user comparer for matching inputs

src/Compilers/CSharp/Test/Semantic/SourceGeneration/StateTableTests.cs

@@ -8,6 +8,8 @@
 using System.Threading;
 using Microsoft.CodeAnalysis.Diagnostics;
 using Xunit;
+using Roslyn.Test.Utilities.TestGenerators;
+using Roslyn.Test.Utilities;
 
 namespace Microsoft.CodeAnalysis.CSharp.Semantic.UnitTests.SourceGeneration
 {
@@ -320,6 +322,59 @@ public void Driver_Table_Builder_Doesnt_Build_Twice()
             Assert.Equal(2, callCount);
         }
 
+        [Fact]
+        public void Batch_Node_Is_Cached_If_All_Inputs_Are_Cached()
+        {
+            var inputNode = new InputNode<int>((_) => ImmutableArray.Create(1, 2, 3));
+            BatchNode<int> batchNode = new BatchNode<int>(inputNode);
+
+            // first time through will always be added (because it's not been run before)
+            DriverStateTable.Builder dstBuilder = GetBuilder(DriverStateTable.Empty);
+            _ = dstBuilder.GetLatestStateTableForNode(batchNode);
+
+            // second time through should show as cached
+            dstBuilder = GetBuilder(dstBuilder.ToImmutable());
+            var table = dstBuilder.GetLatestStateTableForNode(batchNode);
+
+            AssertTableEntries(table, new[] { (ImmutableArray.Create(1, 2, 3), EntryState.Cached) });
+        }
+
+        [Fact]
+        public void Batch_Node_Is_Not_Cached_When_Inputs_Are_Changed()
+        {
+            int third = 3;
+
+            var inputNode = new InputNode<int>((_) => ImmutableArray.Create(1, 2, third++));
+            BatchNode<int> batchNode = new BatchNode<int>(inputNode);
+
+            // first time through will always be added (because it's not been run before)
+            DriverStateTable.Builder dstBuilder = GetBuilder(DriverStateTable.Empty);
+            _ = dstBuilder.GetLatestStateTableForNode(batchNode);
+
+            // second time through should show as modified
+            dstBuilder = GetBuilder(dstBuilder.ToImmutable());
+            var table = dstBuilder.GetLatestStateTableForNode(batchNode);
+
+            AssertTableEntries(table, new[] { (ImmutableArray.Create(1, 2, 4), EntryState.Modified) });
+        }
+
+        [Fact]
+        public void User_Comparer_Is_Not_Used_To_Determine_Inputs()
+        {
+            var inputNode = new InputNode<int>((_) => ImmutableArray.Create(1, 2, 3))
+                                .WithComparer(new LambdaComparer<int>((a, b) => false));
+
+            // first time through will always be added (because it's not been run before)
+            DriverStateTable.Builder dstBuilder = GetBuilder(DriverStateTable.Empty);
+            _ = dstBuilder.GetLatestStateTableForNode(inputNode);
+
+            // second time through should show as cached, even though we supplied a comparer (comparer should only used to turn modified => cached)
+            dstBuilder = GetBuilder(dstBuilder.ToImmutable());
+            var table = dstBuilder.GetLatestStateTableForNode(inputNode);
+
+            AssertTableEntries(table, new[] { (1, EntryState.Cached), (2, EntryState.Cached), (3, EntryState.Cached) });
+        }
+
         private void AssertTableEntries<T>(NodeStateTable<T> table, IList<(T item, EntryState state)> expected)
         {
             int index = 0;
@@ -331,6 +386,17 @@ private void AssertTableEntries<T>(NodeStateTable<T> table, IList<(T item, Entry
             }
         }
 
+        private void AssertTableEntries<T>(NodeStateTable<ImmutableArray<T>> table, IList<(ImmutableArray<T> item, EntryState state)> expected)
+        {
+            int index = 0;
+            foreach (var entry in table)
+            {
+                AssertEx.Equal(expected[index].item, entry.item);
+                Assert.Equal(expected[index].state, entry.state);
+                index++;
+            }
+        }
+
         private DriverStateTable.Builder GetBuilder(DriverStateTable previous)
         {
             var options = CSharpParseOptions.Default.WithLanguageVersion(LanguageVersion.CSharp10);

src/Compilers/Core/Portable/SourceGeneration/Nodes/BatchNode.cs

@@ -34,13 +34,16 @@ public NodeStateTable<ImmutableArray<TInput>> UpdateStateTable(DriverStateTable.
             // - Added when the previous table was empty
             // - Modified otherwise
 
-            var source = ImmutableArray.Create(sourceTable.Batch());
+            var source = sourceTable.Batch();
 
             // update the table 
             var newTable = previousTable.ToBuilder();
-            if (!newTable.TryModifyEntries(source, _comparer))
+            if (!sourceTable.IsCached || !newTable.TryUseCachedEntries())
             {
-                newTable.AddEntries(source, EntryState.Added);
+                if (!newTable.TryModifyEntry(source, _comparer))
+                {
+                    newTable.AddEntry(source, EntryState.Added);
+                }
             }
 
             return newTable.ToImmutableAndFree();

src/Compilers/Core/Portable/SourceGeneration/Nodes/InputNode.cs

@@ -38,7 +38,7 @@ public NodeStateTable<T> UpdateStateTable(DriverStateTable.Builder graphState, N
             var inputItems = _getInput(graphState);
 
             // create a mutable hashset of the new items we can check against
-            HashSet<T> itemsSet = new HashSet<T>(_comparer);
+            HashSet<T> itemsSet = new HashSet<T>();
             foreach (var item in inputItems)
             {
                 var added = itemsSet.Add(item);