[17.6] Cache source generator node tables only if their state counts match

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

@@ -2053,6 +2053,130 @@ class C { }
                 });
         }
 
+        [Fact, WorkItem(61162, "https://github.com/dotnet/roslyn/issues/61162")]
+        public void IncrementalGenerator_Collect_SyntaxProvider_01()
+        {
+            var generator = new IncrementalGeneratorWrapper(new PipelineCallbackGenerator(static ctx =>
+            {
+                var invokedMethodsProvider = ctx.SyntaxProvider
+                    .CreateSyntaxProvider(
+                        static (node, _) => node is InvocationExpressionSyntax,
+                        static (ctx, ct) => ctx.SemanticModel.GetSymbolInfo(ctx.Node, ct).Symbol?.Name ?? "(method not found)")
+                    .Collect();
+
+                ctx.RegisterSourceOutput(invokedMethodsProvider, static (spc, invokedMethods) =>
+                {
+                    spc.AddSource("InvokedMethods.g.cs", string.Join(Environment.NewLine,
+                        invokedMethods.Select(m => $"// {m}")));
+                });
+            }));
+
+            var source = """
+                System.Console.WriteLine();
+                System.Console.WriteLine();
+                System.Console.WriteLine();
+                System.Console.WriteLine();
+                """;
+            var parseOptions = TestOptions.RegularPreview;
+            Compilation compilation = CreateCompilation(source, options: TestOptions.DebugExeThrowing, parseOptions: parseOptions);
+
+            GeneratorDriver driver = CSharpGeneratorDriver.Create(new[] { generator }, parseOptions: parseOptions);
+            verify(ref driver, compilation, """
+                // WriteLine
+                // WriteLine
+                // WriteLine
+                // WriteLine
+                """);
+
+            replace(ref compilation, parseOptions, """
+                System.Console.WriteLine();
+                System.Console.WriteLine();
+                """);
+            verify(ref driver, compilation, """
+                // WriteLine
+                // WriteLine
+                """);
+
+            replace(ref compilation, parseOptions, "_ = 0;");
+            verify(ref driver, compilation, "");
+
+            static void verify(ref GeneratorDriver driver, Compilation compilation, string generatedContent)
+            {
+                driver = driver.RunGeneratorsAndUpdateCompilation(compilation, out var outputCompilation, out var generatorDiagnostics);
+                outputCompilation.VerifyDiagnostics();
+                generatorDiagnostics.Verify();
+                var generatedTree = driver.GetRunResult().GeneratedTrees.Single();
+                AssertEx.EqualOrDiff(generatedContent, generatedTree.ToString());
+            }
+
+            static void replace(ref Compilation compilation, CSharpParseOptions parseOptions, string source)
+            {
+                compilation = compilation.ReplaceSyntaxTree(compilation.SyntaxTrees.Single(), CSharpSyntaxTree.ParseText(source, parseOptions));
+            }
+        }
+
+        [Fact, WorkItem(61162, "https://github.com/dotnet/roslyn/issues/61162")]
+        public void IncrementalGenerator_Collect_SyntaxProvider_02()
+        {
+            var generator = new IncrementalGeneratorWrapper(new PipelineCallbackGenerator(static ctx =>
+            {
+                var invokedMethodsProvider = ctx.SyntaxProvider
+                    .CreateSyntaxProvider(
+                        static (node, _) => node is InvocationExpressionSyntax,
+                        static (ctx, ct) => ctx.SemanticModel.GetSymbolInfo(ctx.Node, ct).Symbol?.Name ?? "(method not found)")
+                    .Select((n, _) => n);
+
+                ctx.RegisterSourceOutput(invokedMethodsProvider, static (spc, invokedMethod) =>
+                {
+                    spc.AddSource(invokedMethod, "// " + invokedMethod);
+                });
+            }));
+
+            var source = """
+                System.Console.WriteLine();
+                System.Console.ReadLine();
+                """;
+            var parseOptions = TestOptions.RegularPreview;
+            Compilation compilation = CreateCompilation(source, options: TestOptions.DebugExeThrowing, parseOptions: parseOptions);
+
+            GeneratorDriver driver = CSharpGeneratorDriver.Create(new[] { generator }, parseOptions: parseOptions);
+            verify(ref driver, compilation, new[]
+            {
+                "// WriteLine",
+                "// ReadLine"
+            });
+
+            replace(ref compilation, parseOptions, """
+                System.Console.WriteLine();
+                """);
+
+            verify(ref driver, compilation, new[]
+            {
+                "// WriteLine"
+            });
+
+            replace(ref compilation, parseOptions, "_ = 0;");
+            verify(ref driver, compilation, Array.Empty<string>());
+
+            static void verify(ref GeneratorDriver driver, Compilation compilation, string[] generatedContent)
+            {
+                driver = driver.RunGeneratorsAndUpdateCompilation(compilation, out var outputCompilation, out var generatorDiagnostics);
+                outputCompilation.VerifyDiagnostics();
+                generatorDiagnostics.Verify();
+                var trees = driver.GetRunResult().GeneratedTrees;
+                Assert.Equal(generatedContent.Length, trees.Length);
+                for (int i = 0; i < generatedContent.Length; i++)
+                {
+                    AssertEx.EqualOrDiff(generatedContent[i], trees[i].ToString());
+                }
+            }
+
+            static void replace(ref Compilation compilation, CSharpParseOptions parseOptions, string source)
+            {
+                compilation = compilation.ReplaceSyntaxTree(compilation.SyntaxTrees.Single(), CSharpSyntaxTree.ParseText(source, parseOptions));
+            }
+        }
+
         [Fact]
         public void IncrementalGenerator_Register_End_Node_Only_Once_Through_Combines()
         {

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

@@ -585,6 +585,45 @@ public void Batch_Node_Records_InputModified_Step_When_Inputs_Are_Changed()
                 });
         }
 
+        [Fact, WorkItem(61162, "https://github.com/dotnet/roslyn/issues/61162")]
+        public void Batch_Node_Remove_From_Beginning()
+        {
+            // [A], [B]
+            var input = new[] { ("A", EntryState.Added), ("B", EntryState.Added) };
+            var inputNode = new CallbackNode<string>((_, _) =>
+            {
+                // Simulate syntax node.
+                var builder = NodeStateTable<string>.Empty.ToBuilder(null, false);
+                foreach (var (value, state) in input)
+                {
+                    builder.AddEntry(value, state, TimeSpan.Zero, default, state);
+                }
+                return builder.ToImmutableAndFree();
+            });
+            var dstBuilder = GetBuilder(DriverStateTable.Empty);
+            var table1 = dstBuilder.GetLatestStateTableForNode(inputNode);
+            AssertTableEntries(table1, new[] { ("A", EntryState.Added, 0), ("B", EntryState.Added, 0) });
+            AssertTableEntries(table1.AsCached(), new[] { ("A", EntryState.Cached, 0), ("B", EntryState.Cached, 0) });
+
+            // batch => [[A], [B]]
+            var batchNode = new BatchNode<string>(inputNode);
+            var table2 = dstBuilder.GetLatestStateTableForNode(batchNode);
+            AssertTableEntries(table2, new[] { (ImmutableArray.Create("A", "B"), EntryState.Added, 0) });
+            AssertTableEntries(table2.AsCached(), new[] { (ImmutableArray.Create("A", "B"), EntryState.Cached, 0) });
+
+            // [B]
+            input = new[] { ("B", EntryState.Cached) };
+            dstBuilder = GetBuilder(dstBuilder.ToImmutable());
+            table1 = dstBuilder.GetLatestStateTableForNode(inputNode);
+            AssertTableEntries(table1, new[] { ("B", EntryState.Cached, 0) });
+            AssertTableEntries(table1.AsCached(), new[] { ("B", EntryState.Cached, 0) });
+
+            // batch => [[B]]
+            table2 = dstBuilder.GetLatestStateTableForNode(batchNode);
+            AssertTableEntries(table2, new[] { (ImmutableArray.Create("B"), EntryState.Modified, 0) });
+            AssertTableEntries(table2.AsCached(), new[] { (ImmutableArray.Create("B"), EntryState.Cached, 0) });
+        }
+
         [Fact]
         [WorkItem(54832, "https://github.com/dotnet/roslyn/issues/54832")]
         public void Transform_Node_Records_NewInput_OnFirst_Run()

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

@@ -60,24 +60,24 @@ internal interface IStateTable
     /// <typeparam name="T">The type of the items tracked by this table</typeparam>
     internal sealed class NodeStateTable<T> : IStateTable
     {
-        internal static NodeStateTable<T> Empty { get; } = new NodeStateTable<T>(ImmutableArray<TableEntry>.Empty, ImmutableArray<IncrementalGeneratorRunStep>.Empty, hasTrackedSteps: true);
+        internal static NodeStateTable<T> Empty { get; } = new NodeStateTable<T>(ImmutableArray<TableEntry>.Empty, ImmutableArray<IncrementalGeneratorRunStep>.Empty, hasTrackedSteps: true, isCached: false);
 
         private readonly ImmutableArray<TableEntry> _states;
 
-        private NodeStateTable(ImmutableArray<TableEntry> states, ImmutableArray<IncrementalGeneratorRunStep> steps, bool hasTrackedSteps)
+        private NodeStateTable(ImmutableArray<TableEntry> states, ImmutableArray<IncrementalGeneratorRunStep> steps, bool hasTrackedSteps, bool isCached)
         {
             Debug.Assert(!hasTrackedSteps || steps.Length == states.Length);
 
             _states = states;
             Steps = steps;
-            IsCached = !states.IsEmpty && states.All(s => s.IsCached);
+            IsCached = isCached;
             HasTrackedSteps = hasTrackedSteps;
         }
 
         public int Count => _states.Length;
 
         /// <summary>
-        /// Indicates if every entry in this table has a state of <see cref="EntryState.Cached"/>
+        /// Indicates that this table is unchanged from the previous version.
         /// </summary>
         public bool IsCached { get; }
 
@@ -126,7 +126,7 @@ public NodeStateTable<T> AsCached()
 
             // Ensure we are completely full so that ToImmutable translates to a MoveToImmutable
             Debug.Assert(compacted.Count == nonRemovedCount);
-            return new NodeStateTable<T>(compacted.ToImmutableAndFree(), ImmutableArray<IncrementalGeneratorRunStep>.Empty, hasTrackedSteps: false);
+            return new NodeStateTable<T>(compacted.ToImmutableAndFree(), ImmutableArray<IncrementalGeneratorRunStep>.Empty, hasTrackedSteps: false, isCached: true);
         }
 
         IStateTable IStateTable.AsCached() => AsCached();
@@ -451,7 +451,8 @@ public NodeStateTable<T> ToImmutableAndFree()
                 return new NodeStateTable<T>(
                     finalStates,
                     TrackIncrementalSteps ? _steps.ToImmutableAndFree() : default,
-                    hasTrackedSteps: TrackIncrementalSteps);
+                    hasTrackedSteps: TrackIncrementalSteps,
+                    isCached: finalStates.All(static s => s.IsCached) && _previous.GetTotalEntryItemCount() == finalStates.Sum(static s => s.Count));
             }
 
             private static (T chosen, EntryState state, bool chosePrevious) GetModifiedItemAndState(T previous, T replacement, IEqualityComparer<T> comparer)