Add DynamicCache<,> type to cache HLSL bytecode results

src/ComputeSharp.D2D1.SourceGenerators/ID2D1ShaderGenerator.CreateLoadBytecodeMethod.cs

@@ -25,6 +25,11 @@ partial class ID2D1ShaderGenerator
     /// </summary>
     internal static partial class LoadBytecode
     {
+        /// <summary>
+        /// The shared cache of <see cref="HlslBytecodeInfo"/> values.
+        /// </summary>
+        private static readonly DynamicCache<HlslBytecodeInfoKey, HlslBytecodeInfo> HlslBytecodeCache = new();
+
         /// <summary>
         /// Extracts the requested shader profile for the current shader.
         /// </summary>
@@ -141,46 +146,54 @@ public static bool IsSimpleInputShader(INamedTypeSymbol structDeclarationSymbol,
         /// <param name="key">The <see cref="HlslBytecodeInfoKey"/> instance for the shader to compile.</param>
         /// <param name="token">The <see cref="CancellationToken"/> used to cancel the operation, if needed.</param>
         /// <returns></returns>
-        public static unsafe HlslBytecodeInfo GetInfo(HlslBytecodeInfoKey key, CancellationToken token)
+        public static HlslBytecodeInfo GetInfo(ref HlslBytecodeInfoKey key, CancellationToken token)
         {
-            // No embedded shader was requested, or there were some errors earlier in the pipeline.
-            // In this case, skip the compilation, as diagnostic will be emitted for those anyway.
-            // Compiling would just add overhead and result in more errors, as the HLSL would be invalid.
-            // We also skip compilation if no shader profile has been requested (we never just assume one).
-            if (key.HasErrors || key.RequestedShaderProfile is null)
+            static unsafe HlslBytecodeInfo GetInfo(HlslBytecodeInfoKey key, CancellationToken token)
             {
-                return HlslBytecodeInfo.Missing.Instance;
-            }
+                // No embedded shader was requested, or there were some errors earlier in the pipeline.
+                // In this case, skip the compilation, as diagnostic will be emitted for those anyway.
+                // Compiling would just add overhead and result in more errors, as the HLSL would be invalid.
+                // We also skip compilation if no shader profile has been requested (we never just assume one).
+                if (key.HasErrors || key.RequestedShaderProfile is null)
+                {
+                    return HlslBytecodeInfo.Missing.Instance;
+                }
 
-            try
-            {
-                token.ThrowIfCancellationRequested();
+                try
+                {
+                    token.ThrowIfCancellationRequested();
 
-                // Compile the shader bytecode using the effective parameters
-                using ComPtr<ID3DBlob> dxcBlobBytecode = D3DCompiler.Compile(
-                    key.HlslSource.AsSpan(),
-                    key.EffectiveShaderProfile,
-                    key.EffectiveCompileOptions);
+                    // Compile the shader bytecode using the effective parameters
+                    using ComPtr<ID3DBlob> dxcBlobBytecode = D3DCompiler.Compile(
+                        key.HlslSource.AsSpan(),
+                        key.EffectiveShaderProfile,
+                        key.EffectiveCompileOptions);
 
-                token.ThrowIfCancellationRequested();
+                    token.ThrowIfCancellationRequested();
 
-                byte* buffer = (byte*)dxcBlobBytecode.Get()->GetBufferPointer();
-                int length = checked((int)dxcBlobBytecode.Get()->GetBufferSize());
+                    byte* buffer = (byte*)dxcBlobBytecode.Get()->GetBufferPointer();
+                    int length = checked((int)dxcBlobBytecode.Get()->GetBufferSize());
 
-                byte[] array = new ReadOnlySpan<byte>(buffer, length).ToArray();
+                    byte[] array = new ReadOnlySpan<byte>(buffer, length).ToArray();
 
-                ImmutableArray<byte> bytecode = Unsafe.As<byte[], ImmutableArray<byte>>(ref array);
+                    ImmutableArray<byte> bytecode = Unsafe.As<byte[], ImmutableArray<byte>>(ref array);
 
-                return new HlslBytecodeInfo.Success(bytecode);
-            }
-            catch (Win32Exception e)
-            {
-                return new HlslBytecodeInfo.Win32Error(e.NativeErrorCode, D3DCompiler.PrettifyFxcErrorMessage(e.Message));
-            }
-            catch (FxcCompilationException e)
-            {
-                return new HlslBytecodeInfo.FxcError(D3DCompiler.PrettifyFxcErrorMessage(e.Message));
+                    return new HlslBytecodeInfo.Success(bytecode);
+                }
+                catch (Win32Exception e)
+                {
+                    return new HlslBytecodeInfo.Win32Error(e.NativeErrorCode, D3DCompiler.PrettifyFxcErrorMessage(e.Message));
+                }
+                catch (FxcCompilationException e)
+                {
+                    return new HlslBytecodeInfo.FxcError(D3DCompiler.PrettifyFxcErrorMessage(e.Message));
+                }
             }
+
+            // Get or create the HLSL bytecode compilation result for the input key. The dynamic cache
+            // will take care of retrieving an existing cached value if the same shader has been compiled
+            // already with the same parameters. After this call, callers must use the updated key value.
+            return HlslBytecodeCache.GetOrCreate(ref key, GetInfo, token);
         }
 
         /// <summary>

src/ComputeSharp.D2D1.SourceGenerators/ID2D1ShaderGenerator.cs

@@ -154,7 +154,7 @@ public void Initialize(IncrementalGeneratorInitializationContext context)
                         hasErrors);
 
                     // TODO: cache this across transform runs
-                    HlslBytecodeInfo hlslInfo = LoadBytecode.GetInfo(hlslInfoKey, token);
+                    HlslBytecodeInfo hlslInfo = LoadBytecode.GetInfo(ref hlslInfoKey, token);
 
                     token.ThrowIfCancellationRequested();
 

src/ComputeSharp.SourceGeneration/ComputeSharp.SourceGeneration.projitems

@@ -17,6 +17,7 @@
     <Compile Include="$(MSBuildThisFileDirectory)Extensions\IFieldSymbolExtensions.cs" />
     <Compile Include="$(MSBuildThisFileDirectory)Extensions\ITypeSymbolExtensions.cs" />
     <Compile Include="$(MSBuildThisFileDirectory)Extensions\ISymbolExtensions.cs" />
+    <Compile Include="$(MSBuildThisFileDirectory)helpers\DynamicCache{TKey,TValue}.cs" />
     <Compile Include="$(MSBuildThisFileDirectory)Helpers\EquatableArray{T}.cs" />
     <Compile Include="$(MSBuildThisFileDirectory)Helpers\HashCode.cs" />
     <Compile Include="$(MSBuildThisFileDirectory)Helpers\ImmutableArrayBuilder{T}.cs" />