Add generator driver cache to VBCSCompiler

src/Compilers/CSharp/Portable/CommandLine/CSharpCompiler.cs

@@ -25,8 +25,8 @@ internal abstract class CSharpCompiler : CommonCompiler
         private readonly CommandLineDiagnosticFormatter _diagnosticFormatter;
         private readonly string? _tempDirectory;
 
-        protected CSharpCompiler(CSharpCommandLineParser parser, string? responseFile, string[] args, BuildPaths buildPaths, string? additionalReferenceDirectories, IAnalyzerAssemblyLoader assemblyLoader)
-            : base(parser, responseFile, args, buildPaths, additionalReferenceDirectories, assemblyLoader)
+        protected CSharpCompiler(CSharpCommandLineParser parser, string? responseFile, string[] args, BuildPaths buildPaths, string? additionalReferenceDirectories, IAnalyzerAssemblyLoader assemblyLoader, GeneratorDriverCache? driverCache = null)
+            : base(parser, responseFile, args, buildPaths, additionalReferenceDirectories, assemblyLoader, driverCache)
         {
             _diagnosticFormatter = new CommandLineDiagnosticFormatter(buildPaths.WorkingDirectory, Arguments.PrintFullPaths, Arguments.ShouldIncludeErrorEndLocation);
             _tempDirectory = buildPaths.TempDirectory;
@@ -372,12 +372,9 @@ protected override void ResolveEmbeddedFilesFromExternalSourceDirectives(
             }
         }
 
-        private protected override Compilation RunGenerators(Compilation input, ParseOptions parseOptions, ImmutableArray<ISourceGenerator> generators, AnalyzerConfigOptionsProvider analyzerConfigProvider, ImmutableArray<AdditionalText> additionalTexts, DiagnosticBag diagnostics)
+        private protected override GeneratorDriver CreateGeneratorDriver(ParseOptions parseOptions, ImmutableArray<ISourceGenerator> generators, AnalyzerConfigOptionsProvider analyzerConfigOptionsProvider, ImmutableArray<AdditionalText> additionalTexts)
         {
-            var driver = CSharpGeneratorDriver.Create(generators, additionalTexts, (CSharpParseOptions)parseOptions, analyzerConfigProvider);
-            driver.RunGeneratorsAndUpdateCompilation(input, out var compilationOut, out var generatorDiagnostics);
-            diagnostics.AddRange(generatorDiagnostics);
-            return compilationOut;
+            return CSharpGeneratorDriver.Create(generators, additionalTexts, (CSharpParseOptions)parseOptions, analyzerConfigOptionsProvider);
         }
     }
 }

src/Compilers/CSharp/Test/CommandLine/CommandLineTestBase.cs

@@ -54,15 +54,15 @@ internal CSharpCommandLineArguments DefaultParse(IEnumerable<string> args, strin
             return CSharpCommandLineParser.Default.Parse(args, baseDirectory, sdkDirectory, additionalReferenceDirectories);
         }
 
-        internal MockCSharpCompiler CreateCSharpCompiler(string[] args, ImmutableArray<DiagnosticAnalyzer> analyzers = default, ImmutableArray<ISourceGenerator> generators = default, AnalyzerAssemblyLoader loader = null)
+        internal MockCSharpCompiler CreateCSharpCompiler(string[] args, ImmutableArray<DiagnosticAnalyzer> analyzers = default, ImmutableArray<ISourceGenerator> generators = default, AnalyzerAssemblyLoader loader = null, GeneratorDriverCache driverCache = null)
         {
-            return CreateCSharpCompiler(null, WorkingDirectory, args, analyzers, generators, loader);
+            return CreateCSharpCompiler(null, WorkingDirectory, args, analyzers, generators, loader, driverCache);
         }
 
-        internal MockCSharpCompiler CreateCSharpCompiler(string responseFile, string workingDirectory, string[] args, ImmutableArray<DiagnosticAnalyzer> analyzers = default, ImmutableArray<ISourceGenerator> generators = default, AnalyzerAssemblyLoader loader = null)
+        internal MockCSharpCompiler CreateCSharpCompiler(string responseFile, string workingDirectory, string[] args, ImmutableArray<DiagnosticAnalyzer> analyzers = default, ImmutableArray<ISourceGenerator> generators = default, AnalyzerAssemblyLoader loader = null, GeneratorDriverCache driverCache = null)
         {
             var buildPaths = RuntimeUtilities.CreateBuildPaths(workingDirectory, sdkDirectory: SdkDirectory);
-            return new MockCSharpCompiler(responseFile, buildPaths, args, analyzers, generators, loader);
+            return new MockCSharpCompiler(responseFile, buildPaths, args, analyzers, generators, loader, driverCache);
         }
     }
 }

src/Compilers/CSharp/Test/CommandLine/CommandLineTests.cs

@@ -9793,6 +9793,111 @@ string[] compileAndRun(string featureOpt)
             };
         }
 
+        [Fact]
+        public void Compiler_Uses_DriverCache()
+        {
+            var dir = Temp.CreateDirectory();
+            var src = dir.CreateFile("temp.cs").WriteAllText(@"
+class C
+{
+}");
+            int sourceCallbackCount = 0;
+            var generator = new PipelineCallbackGenerator((ctx) =>
+            {
+                ctx.RegisterSourceOutput(ctx.ParseOptionsProvider, (spc, po) =>
+                {
+                    sourceCallbackCount++;
+                });
+            });
+
+            // with no cache, we'll see the callback execute multiple times
+            RunWithNoCache();
+            Assert.Equal(1, sourceCallbackCount);
+
+            RunWithNoCache();
+            Assert.Equal(2, sourceCallbackCount);
+
+            RunWithNoCache();
+            Assert.Equal(3, sourceCallbackCount);
+
+            // now re-run with a cache
+            GeneratorDriverCache cache = new GeneratorDriverCache();
+            sourceCallbackCount = 0;
+
+            RunWithCache();
+            Assert.Equal(1, sourceCallbackCount);
+
+            RunWithCache();
+            Assert.Equal(1, sourceCallbackCount);
+
+            RunWithCache();
+            Assert.Equal(1, sourceCallbackCount);
+
+            // Clean up temp files
+            CleanupAllGeneratedFiles(src.Path);
+            Directory.Delete(dir.Path, true);
+
+            void RunWithNoCache() => VerifyOutput(dir, src, includeCurrentAssemblyAsAnalyzerReference: false, additionalFlags: new[] { "/langversion:preview" }, generators: new[] { generator.AsSourceGenerator() }, analyzers: null);
+
+            void RunWithCache() => VerifyOutput(dir, src, includeCurrentAssemblyAsAnalyzerReference: false, additionalFlags: new[] { "/langversion:preview" }, generators: new[] { generator.AsSourceGenerator() }, driverCache: cache, analyzers: null);
+        }
+
+        [Fact]
+        public void Compiler_Doesnt_Use_Cache_From_Other_Compilation()
+        {
+            var dir = Temp.CreateDirectory();
+            var src = dir.CreateFile("temp.cs").WriteAllText(@"
+class C
+{
+}");
+            int sourceCallbackCount = 0;
+            var generator = new PipelineCallbackGenerator((ctx) =>
+            {
+                ctx.RegisterSourceOutput(ctx.ParseOptionsProvider, (spc, po) =>
+                {
+                    sourceCallbackCount++;
+                });
+            });
+
+            // now re-run with a cache
+            GeneratorDriverCache cache = new GeneratorDriverCache();
+            sourceCallbackCount = 0;
+
+            RunWithCache("1.dll");
+            Assert.Equal(1, sourceCallbackCount);
+
+            RunWithCache("1.dll");
+            Assert.Equal(1, sourceCallbackCount);
+
+            // now emulate a new compilation, and check we were invoked, but only once
+            RunWithCache("2.dll");
+            Assert.Equal(2, sourceCallbackCount);
+
+            RunWithCache("2.dll");
+            Assert.Equal(2, sourceCallbackCount);
+
+            // now re-run our first compilation
+            RunWithCache("1.dll");
+            Assert.Equal(2, sourceCallbackCount);
+
+            // a new one
+            RunWithCache("3.dll");
+            Assert.Equal(3, sourceCallbackCount);
+
+            // and another old one
+            RunWithCache("2.dll");
+            Assert.Equal(3, sourceCallbackCount);
+
+            RunWithCache("1.dll");
+            Assert.Equal(3, sourceCallbackCount);
+
+            // Clean up temp files
+            CleanupAllGeneratedFiles(src.Path);
+            Directory.Delete(dir.Path, true);
+
+            void RunWithCache(string outputPath) => VerifyOutput(dir, src, includeCurrentAssemblyAsAnalyzerReference: false, additionalFlags: new[] { "/langversion:preview", "/out:" + outputPath }, generators: new[] { generator.AsSourceGenerator() }, driverCache: cache, analyzers: null);
+        }
+
         private static int OccurrenceCount(string source, string word)
         {
             var n = 0;
@@ -9814,6 +9919,7 @@ private string VerifyOutput(TempDirectory sourceDir, TempFile sourceFile,
                                            int? expectedExitCode = null,
                                            bool errorlog = false,
                                            IEnumerable<ISourceGenerator> generators = null,
+                                           GeneratorDriverCache driverCache = null,
                                            params DiagnosticAnalyzer[] analyzers)
         {
             var args = new[] {
@@ -9835,7 +9941,7 @@ private string VerifyOutput(TempDirectory sourceDir, TempFile sourceFile,
                 args = args.Append(additionalFlags);
             }
 
-            var csc = CreateCSharpCompiler(null, sourceDir.Path, args, analyzers: analyzers.ToImmutableArrayOrEmpty(), generators: generators.ToImmutableArrayOrEmpty());
+            var csc = CreateCSharpCompiler(null, sourceDir.Path, args, analyzers: analyzers.ToImmutableArrayOrEmpty(), generators: generators.ToImmutableArrayOrEmpty(), driverCache: driverCache);
             var outWriter = new StringWriter(CultureInfo.InvariantCulture);
             var exitCode = csc.Run(outWriter);
             var output = outWriter.ToString();

src/Compilers/CSharp/Test/CommandLine/DriverCacheTests.cs

@@ -0,0 +1,109 @@
+// 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.Generic;
+using System.Linq;
+using System.Text;
+using System.Threading.Tasks;
+using System.Collections.Immutable;
+using Xunit;
+using Microsoft.CodeAnalysis.PooledObjects;
+using Roslyn.Test.Utilities.TestGenerators;
+using System.IO;
+using static Roslyn.Test.Utilities.SharedResourceHelpers;
+
+namespace Microsoft.CodeAnalysis.CSharp.CommandLine.UnitTests
+{
+    public class DriverCacheTests : CommandLineTestBase
+    {
+
+        [Fact]
+        public void DriverCache_Returns_Null_For_No_Match()
+        {
+            var driverCache = new GeneratorDriverCache();
+            var driver = driverCache.TryGetDriver("0");
+
+            Assert.Null(driver);
+        }
+
+        [Fact]
+        public void DriverCache_Returns_Cached_Driver()
+        {
+            var drivers = GetDrivers(1);
+            var driverCache = new GeneratorDriverCache();
+            driverCache.CacheGenerator("0", drivers[0]);
+
+            var driver = driverCache.TryGetDriver("0");
+            Assert.Same(driver, drivers[0]);
+        }
+
+        [Fact]
+        public void DriverCache_Can_Cache_Multiple_Drivers()
+        {
+            var drivers = GetDrivers(3);
+
+            var driverCache = new GeneratorDriverCache();
+            driverCache.CacheGenerator("0", drivers[0]);
+            driverCache.CacheGenerator("1", drivers[1]);
+            driverCache.CacheGenerator("2", drivers[2]);
+
+            var driver = driverCache.TryGetDriver("0");
+            Assert.Same(driver, drivers[0]);
+
+            driver = driverCache.TryGetDriver("1");
+            Assert.Same(driver, drivers[1]);
+
+            driver = driverCache.TryGetDriver("2");
+            Assert.Same(driver, drivers[2]);
+        }
+
+        [Fact]
+        public void DriverCache_Evicts_Least_Recently_Used()
+        {
+            var drivers = GetDrivers(GeneratorDriverCache.MaxCacheSize + 2);
+            var driverCache = new GeneratorDriverCache();
+
+            // put n+1 drivers into the cache
+            for (int i = 0; i < GeneratorDriverCache.MaxCacheSize + 1; i++)
+            {
+                driverCache.CacheGenerator(i.ToString(), drivers[i]);
+            }
+            // current cache state is 
+            // (10, 9, 8, 7, 6, 5, 4, 3, 2, 1)
+
+            // now try and retreive the first driver which should no longer be in the cache
+            var driver = driverCache.TryGetDriver("0");
+            Assert.Null(driver);
+
+            // add it back
+            driverCache.CacheGenerator("0", drivers[0]);
+
+            // current cache state is 
+            // (0, 10, 9, 8, 7, 6, 5, 4, 3, 2)
+
+            // access some drivers in the middle
+            driver = driverCache.TryGetDriver("7");
+            driver = driverCache.TryGetDriver("4");
+            driver = driverCache.TryGetDriver("2");
+
+            // current cache state is 
+            // (2, 4, 7, 0, 10, 9, 8, 6, 5, 3)
+
+            // try and get a new driver that was never in the cache
+            driver = driverCache.TryGetDriver("11");
+            Assert.Null(driver);
+            driverCache.CacheGenerator("11", drivers[11]);
+
+            // current cache state is 
+            // (11, 2, 4, 7, 0, 10, 9, 8, 6, 5)
+
+            // get a driver that has been evicted
+            driver = driverCache.TryGetDriver("3");
+            Assert.Null(driver);
+        }
+
+        private GeneratorDriver[] GetDrivers(int count) =>  Enumerable.Range(0, count).Select(i => CSharpGeneratorDriver.Create(Array.Empty<ISourceGenerator>())).ToArray();
+    }
+}

src/Compilers/Core/Portable/CommandLine/CommonCompiler.cs

@@ -72,6 +72,7 @@ internal abstract partial class CommonCompiler
         public CommonMessageProvider MessageProvider { get; }
         public CommandLineArguments Arguments { get; }
         public IAnalyzerAssemblyLoader AssemblyLoader { get; private set; }
+        public GeneratorDriverCache? DriverCache { get; }
         public abstract DiagnosticFormatter DiagnosticFormatter { get; }
 
         /// <summary>
@@ -116,7 +117,7 @@ protected abstract void ResolveAnalyzersFromArguments(
             out ImmutableArray<DiagnosticAnalyzer> analyzers,
             out ImmutableArray<ISourceGenerator> generators);
 
-        public CommonCompiler(CommandLineParser parser, string? responseFile, string[] args, BuildPaths buildPaths, string? additionalReferenceDirectories, IAnalyzerAssemblyLoader assemblyLoader)
+        public CommonCompiler(CommandLineParser parser, string? responseFile, string[] args, BuildPaths buildPaths, string? additionalReferenceDirectories, IAnalyzerAssemblyLoader assemblyLoader, GeneratorDriverCache? driverCache)
         {
             IEnumerable<string> allArgs = args;
 
@@ -129,6 +130,7 @@ public CommonCompiler(CommandLineParser parser, string? responseFile, string[] a
             this.Arguments = parser.Parse(allArgs, buildPaths.WorkingDirectory, buildPaths.SdkDirectory, additionalReferenceDirectories);
             this.MessageProvider = parser.MessageProvider;
             this.AssemblyLoader = assemblyLoader;
+            this.DriverCache = driverCache;
             this.EmbeddedSourcePaths = GetEmbeddedSourcePaths(Arguments);
 
             if (Arguments.ParseOptions.Features.ContainsKey("debug-determinism"))
@@ -734,7 +736,18 @@ public virtual int Run(TextWriter consoleOutput, CancellationToken cancellationT
         /// <param name="additionalTexts">Any additional texts that should be passed to the generators when run.</param>
         /// <param name="generatorDiagnostics">Any diagnostics that were produced during generation</param>
         /// <returns>A compilation that represents the original compilation with any additional, generated texts added to it.</returns>
-        private protected virtual Compilation RunGenerators(Compilation input, ParseOptions parseOptions, ImmutableArray<ISourceGenerator> generators, AnalyzerConfigOptionsProvider analyzerConfigOptionsProvider, ImmutableArray<AdditionalText> additionalTexts, DiagnosticBag generatorDiagnostics) { return input; }
+        private protected Compilation RunGenerators(Compilation input, ParseOptions parseOptions, ImmutableArray<ISourceGenerator> generators, AnalyzerConfigOptionsProvider analyzerConfigOptionsProvider, ImmutableArray<AdditionalText> additionalTexts, DiagnosticBag generatorDiagnostics)
+        {
+            var cacheKey = Arguments.GetOutputFilePath(GetOutputFileName(input, cancellationToken: default));
+
+            var driver = this.DriverCache?.TryGetDriver(cacheKey) ?? CreateGeneratorDriver(parseOptions, generators, analyzerConfigOptionsProvider, additionalTexts);
+            driver = driver.RunGeneratorsAndUpdateCompilation(input, out var compilationOut, out var diagnostics);
+            this.DriverCache?.CacheGenerator(cacheKey, driver);
+            generatorDiagnostics.AddRange(diagnostics);
+            return compilationOut;
+        }
+
+        private protected abstract GeneratorDriver CreateGeneratorDriver(ParseOptions parseOptions, ImmutableArray<ISourceGenerator> generators, AnalyzerConfigOptionsProvider analyzerConfigOptionsProvider, ImmutableArray<AdditionalText> additionalTexts);
 
         private int RunCore(TextWriter consoleOutput, ErrorLogger? errorLogger, CancellationToken cancellationToken)
         {