Check count before reusing batched items

Cosifne · Mar 31, 2023 · 2984b5f · 2984b5f
1 parent 968d97f
commit 2984b5f
Showing 1 changed file with 13 additions and 1 deletion.
diff --git a/src/Compilers/Core/Portable/SourceGeneration/Nodes/BatchNode.cs b/src/Compilers/Core/Portable/SourceGeneration/Nodes/BatchNode.cs
@@ -132,7 +132,9 @@ public NodeStateTable<ImmutableArray<TInput>> UpdateStateTable(DriverStateTable.
             {
                 newTable.AddEntry(sourceValues, EntryState.Added, stopwatch.Elapsed, sourceInputs, EntryState.Added);
             }
-            else if (!sourceTable.IsCached || !newTable.TryUseCachedEntries(stopwatch.Elapsed, sourceInputs))
+            else if (!sourceTable.IsCached ||
+                GetTotalValueCount(previousTable) != sourceValues.Length ||
+                !newTable.TryUseCachedEntries(stopwatch.Elapsed, sourceInputs))
             {
                 if (!newTable.TryModifyEntry(sourceValues, _comparer, stopwatch.Elapsed, sourceInputs, EntryState.Modified))
                 {
@@ -141,6 +143,16 @@ public NodeStateTable<ImmutableArray<TInput>> UpdateStateTable(DriverStateTable.
             }
 
             return newTable.ToImmutableAndFree();
+
+            static int GetTotalValueCount(NodeStateTable<ImmutableArray<TInput>> table)
+            {
+                var count = 0;
+                foreach (var entry in table)
+                {
+                    count += entry.Item.Length;
+                }
+                return count;
+            }
         }
 
         public void RegisterOutput(IIncrementalGeneratorOutputNode output) => _sourceNode.RegisterOutput(output);