Allow BatchNodes to reuse prior computed arrays if they would generate hte same value

src/Compilers/Core/Portable/CodeAnalysisEventSource.cs

@@ -22,6 +22,9 @@ public static class Tasks
         {
             public const EventTask GeneratorDriverRunTime = (EventTask)1;
             public const EventTask SingleGeneratorRunTime = (EventTask)2;
+
+            public const EventTask PooledWhenSparse = (EventTask)3;
+            public const EventTask HalvedPooledCapacity = (EventTask)4;
         }
 
         private CodeAnalysisEventSource() { }
@@ -37,5 +40,13 @@ private CodeAnalysisEventSource() { }
 
         [Event(4, Message = "Generator {0} ran for {2} ticks", Keywords = Keywords.Performance, Level = EventLevel.Informational, Opcode = EventOpcode.Stop, Task = Tasks.SingleGeneratorRunTime)]
         internal void StopSingleGeneratorRunTime(string generatorName, string assemblyPath, long elapsedTicks, string id) => WriteEvent(4, generatorName, assemblyPath, elapsedTicks, id);
+
+        [Event(5, Message = "Pooled when sparse: {0} {1}/{2}", Keywords = Keywords.Performance, Level = EventLevel.Informational, Task = Tasks.PooledWhenSparse)]
+        internal void PooledWhenSparseImpl(Type type, int count, int capacity, string id)
+            => WriteEvent(5, type.FullName, count, capacity, id);
+
+        [Event(6, Message = "Halved capacity: {0} {1}/{2}", Keywords = Keywords.Performance, Level = EventLevel.Informational, Task = Tasks.HalvedPooledCapacity)]
+        internal void HalvedCapacityImpl(Type type, int numberOfTimesPooledWhenSparse, int numberOfTimesPooled, string id)
+            => WriteEvent(6, type.FullName, numberOfTimesPooledWhenSparse, numberOfTimesPooled, id);
     }
 }

src/Compilers/Core/Portable/SourceGeneration/GeneratorTimerExtensions.cs

@@ -42,6 +42,22 @@ public static RunTimer CreateSingleGeneratorRunTimer(this CodeAnalysisEventSourc
             }
         }
 
+        public static void PooledWhenSparse(this CodeAnalysisEventSource eventSource, Type type, int count, int capacity)
+        {
+            if (eventSource.IsEnabled(EventLevel.Informational, Keywords.Performance))
+            {
+                eventSource.PooledWhenSparseImpl(type, count, capacity, Guid.NewGuid().ToString());
+            }
+        }
+
+        public static void HalvedCapacity(this CodeAnalysisEventSource eventSource, Type type, int numberOfTimesPooledWhenSparse, int numberOfTimesPooled)
+        {
+            if (eventSource.IsEnabled(EventLevel.Informational, Keywords.Performance))
+            {
+                eventSource.HalvedCapacityImpl(type, numberOfTimesPooledWhenSparse, numberOfTimesPooled, Guid.NewGuid().ToString());
+            }
+        }
+
         internal readonly struct RunTimer : IDisposable
         {
             private readonly SharedStopwatch _timer;

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

@@ -3,10 +3,11 @@
 // See the LICENSE file in the project root for more information.
 
 using System;
-using System.Collections;
+using System.Collections.Concurrent;
 using System.Collections.Generic;
 using System.Collections.Immutable;
 using System.Diagnostics;
+using System.Linq;
 using System.Threading;
 using Microsoft.CodeAnalysis.PooledObjects;
 using Roslyn.Utilities;
@@ -15,6 +16,8 @@ namespace Microsoft.CodeAnalysis
 {
     internal sealed class BatchNode<TInput> : IIncrementalGeneratorNode<ImmutableArray<TInput>>
     {
+        private static readonly ConcurrentQueue<BuilderAndStatistics<TInput>> s_builderPool = new();
+
         private readonly IIncrementalGeneratorNode<TInput> _sourceNode;
         private readonly IEqualityComparer<ImmutableArray<TInput>> _comparer;
         private readonly string? _name;
@@ -50,22 +53,7 @@ public NodeStateTable<ImmutableArray<TInput>> UpdateStateTable(DriverStateTable.
 
             var stopwatch = SharedStopwatch.StartNew();
 
-            var sourceValuesBuilder = ArrayBuilder<TInput>.GetInstance();
-            var sourceInputsBuilder = newTable.TrackIncrementalSteps ? ArrayBuilder<(IncrementalGeneratorRunStep InputStep, int OutputIndex)>.GetInstance() : null;
-
-            foreach (var entry in sourceTable)
-            {
-                // At this point, we can remove any 'Removed' items and ensure they're not in our list of states.
-                if (entry.State != EntryState.Removed)
-                    sourceValuesBuilder.Add(entry.Item);
-
-                // However, regardless of if the entry was removed or not, we still keep track of its step information
-                // so we can accurately report how long it took and what actually happened (for testing validation).
-                sourceInputsBuilder?.Add((entry.Step!, entry.OutputIndex));
-            }
-
-            var sourceValues = sourceValuesBuilder.ToImmutableAndFree();
-            var sourceInputs = newTable.TrackIncrementalSteps ? sourceInputsBuilder!.ToImmutableAndFree() : default;
+            var (sourceValues, sourceInputs) = GetValuesAndInputs(sourceTable, previousTable, newTable);
 
             if (previousTable.IsEmpty)
             {
@@ -82,6 +70,45 @@ public NodeStateTable<ImmutableArray<TInput>> UpdateStateTable(DriverStateTable.
             return newTable.ToImmutableAndFree();
         }
 
+        private (ImmutableArray<TInput>, ImmutableArray<(IncrementalGeneratorRunStep InputStep, int OutputIndex)>) GetValuesAndInputs(
+            NodeStateTable<TInput> sourceTable,
+            NodeStateTable<ImmutableArray<TInput>> previousTable,
+            NodeStateTable<ImmutableArray<TInput>>.Builder newTable)
+        {
+            var sourceValuesBuilder = s_builderPool.DequeuePooledItem();
+            try
+            {
+                var sourceInputsBuilder = newTable.TrackIncrementalSteps ? ArrayBuilder<(IncrementalGeneratorRunStep InputStep, int OutputIndex)>.GetInstance() : null;
+
+                foreach (var entry in sourceTable)
+                {
+                    // At this point, we can remove any 'Removed' items and ensure they're not in our list of states.
+                    if (entry.State != EntryState.Removed)
+                        sourceValuesBuilder.Add(entry.Item);
+
+                    // However, regardless of if the entry was removed or not, we still keep track of its step information
+                    // so we can accurately report how long it took and what actually happened (for testing validation).
+                    sourceInputsBuilder?.Add((entry.Step!, entry.OutputIndex));
+                }
+
+                var sourceInputs = newTable.TrackIncrementalSteps ? sourceInputsBuilder!.ToImmutableAndFree() : default;
+
+                // If we're producing the exact same values as the last time, then just point at the prior array
+                // wholesale to avoid a costly copy.
+                if (previousTable.Count == 1 &&
+                    sourceValuesBuilder.Builder.SequenceEqual(previousTable.Single().item, EqualityComparer<TInput>.Default))
+                {
+                    return (previousTable.Single().item, sourceInputs);
+                }
+
+                return (sourceValuesBuilder.ToImmutable(), sourceInputs);
+            }
+            finally
+            {
+                s_builderPool.ReturnPooledItem(sourceValuesBuilder);
+            }
+        }
+
         public void RegisterOutput(IIncrementalGeneratorOutputNode output) => _sourceNode.RegisterOutput(output);
     }
 }

src/Compilers/Core/Portable/SourceGeneration/Pooling.cs

@@ -0,0 +1,118 @@
+// 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.Concurrent;
+using System.Collections.Generic;
+using System.Collections.Immutable;
+using System.Diagnostics;
+using System.Diagnostics.CodeAnalysis;
+using System.Diagnostics.Contracts;
+using System.Linq;
+using Microsoft.CodeAnalysis.PooledObjects;
+using Roslyn.Utilities;
+
+namespace Microsoft.CodeAnalysis
+{
+    internal struct PoolingStatistics
+    {
+        /// <summary>
+        /// The number of times this item has been added back to the pool.  Once this goes past some threshold
+        /// we will start checking if we're continually returning a large array that is mostly empty.  If so, we
+        /// will try to lower the capacity of the array to prevent wastage.
+        /// </summary>
+        public int NumberOfTimesPooled;
+
+        /// <summary>
+        /// The number of times we returned a large array to the pool that was barely filled.  If this is a
+        /// significant number of the total times pooled, then we will attempt to lower the capacity of the
+        /// array.
+        /// </summary>
+        public int NumberOfTimesPooledWhenSparse;
+    }
+
+    internal readonly record struct BuilderAndStatistics<TValue>(
+        ImmutableArray<TValue>.Builder Builder,
+        PoolingStatistics Statistics)
+    {
+        public int Count => Builder.Count;
+
+        public void Add(TValue value)
+            => Builder.Add(value);
+
+        public ImmutableArray<TValue> ToImmutable()
+            => Builder.ToImmutable();
+
+        public TValue this[int index]
+            => Builder[index];
+
+        public bool Any(Func<TValue, bool> predicate)
+            => Builder.Any(predicate);
+    }
+
+    internal static class PoolingExtensions
+    {
+        public static BuilderAndStatistics<TValue> DequeuePooledItem<TValue>(
+            this ConcurrentQueue<BuilderAndStatistics<TValue>> queue)
+            => queue.TryDequeue(out var item) ? item : new BuilderAndStatistics<TValue> { Builder = ImmutableArray.CreateBuilder<TValue>() };
+
+        public static void ReturnPooledItem<TValue>(
+            this ConcurrentQueue<BuilderAndStatistics<TValue>> queue,
+            BuilderAndStatistics<TValue> item)
+        {
+            // Don't bother shrinking the array for arrays less than this capacity.  They're not going to be a
+            // huge waste of space so we can just pool them forever.
+            const int MinCapacityToConsiderThreshold = 1000;
+
+            // The number of times something is added/removed to the pool before we start considering
+            // statistics. This is so that we have enough data to reasonably see if something is consistently
+            // sparse.
+            const int MinTimesPooledToConsiderStatistics = 100;
+
+            // The ratio of Count/Capacity to be at to be considered sparse.  under this, there is a lot of
+            // wasted space and we would prefer to just throw the array away.  Above this and we're reasonably
+            // filling the array and should keep it around.
+            const double SparseThresholdRatio = 0.25;
+
+            // The ratio of times we pooled something sparse.  Once above this, we will jettison the array as
+            // being not worth keeping.
+            const double ConsistentlySparseRatio = 0.75;
+
+            // Note: the values 0.25 and 0.75 were picked as they reflect the common array practicing of growing
+            // by doubling.  So once we've grown so much that we're consistently under 25% of the array, then we
+            // want to shrink down.  To prevent shrinking and inflating over and over again, we only shrink when
+            // we're highly confident we're going to stay small.
+
+            var (builder, statistics) = item;
+            statistics.NumberOfTimesPooled++;
+
+            // See if we're pooling something both large and sparse.
+            if (builder.Capacity > MinCapacityToConsiderThreshold &&
+                ((double)builder.Count / builder.Capacity) < SparseThresholdRatio)
+            {
+                CodeAnalysisEventSource.Log.PooledWhenSparse(builder.GetType(), builder.Count, builder.Capacity);
+                statistics.NumberOfTimesPooledWhenSparse++;
+            }
+
+            builder.Clear();
+
+            // See if this builder has been consistently sparse. If so then time to lower its capacity.
+            if (statistics.NumberOfTimesPooled > MinTimesPooledToConsiderStatistics &&
+                ((double)statistics.NumberOfTimesPooledWhenSparse / statistics.NumberOfTimesPooled) > ConsistentlySparseRatio)
+            {
+                CodeAnalysisEventSource.Log.HalvedCapacity(builder.GetType(), builder.Count, builder.Capacity);
+                builder.Capacity /= 2;
+
+                // Reset our statistics.  We'll wait another 100 pooling attempts to reassess if we need to
+                // adjust the capacity here.
+                statistics = new PoolingStatistics
+                {
+                    NumberOfTimesPooled = 1,
+                };
+            }
+
+            queue.Enqueue(new BuilderAndStatistics<TValue>(builder, statistics));
+        }
+    }
+}