[graphite] Consolidate functionality in the SkShaderCodeDictionary

At some point we'll need to go through the dictionary for user provided SkSL

Bug: skia:12701
Change-Id: I484ae30626dad64f2bce1e0948071380d9f8282e
Reviewed-on: https://skia-review.googlesource.com/c/skia/+/504596
Reviewed-by: Michael Ludwig <michaelludwig@google.com>
Commit-Queue: Robert Phillips <robertphillips@google.com>
diff --git a/experimental/graphite/src/Context.cpp b/experimental/graphite/src/Context.cpp
index e217058..f1899a5 100644
--- a/experimental/graphite/src/Context.cpp
+++ b/experimental/graphite/src/Context.cpp
@@ -78,7 +78,8 @@
         for (auto& shaderCombo: paintCombo.fShaders) {
             for (auto shaderType: shaderCombo.fTypes) {
                 for (auto tm: shaderCombo.fTileModes) {
-                    SkPaintParamsKey key = CreateKey(SkBackend::kGraphite, shaderType, tm, bm);
+                    SkPaintParamsKey key = CreateKey(fGlobalCache->shaderCodeDictionary(),
+                                                     SkBackend::kGraphite, shaderType, tm, bm);
 
                     GraphicsPipelineDesc desc;
 
diff --git a/experimental/graphite/src/ContextUtils.cpp b/experimental/graphite/src/ContextUtils.cpp
index 21fe935..8e5b7bb 100644
--- a/experimental/graphite/src/ContextUtils.cpp
+++ b/experimental/graphite/src/ContextUtils.cpp
@@ -21,79 +21,6 @@
 
 namespace skgpu {
 
-using GradientData = GradientShaderBlocks::GradientData;
-
-namespace {
-
-//--------------------------------------------------------------------------------------------------
-// TODO: For the sprint we unify all the gradient uniforms into a standard set of 6:
-//   kMaxStops colors
-//   kMaxStops offsets
-//   2 points
-//   2 radii
-static constexpr int kNumGradientUniforms = 6;
-static constexpr SkUniform kGradientUniforms[kNumGradientUniforms] = {
-        { "colors",  SkSLType::kHalf4, GradientData::kMaxStops },
-        { "offsets", SkSLType::kFloat, GradientData::kMaxStops },
-        { "point0",  SkSLType::kFloat2 },
-        { "point1",  SkSLType::kFloat2 },
-        { "radius0", SkSLType::kFloat },
-        { "radius1", SkSLType::kFloat },
-};
-
-static const char *kLinearGradient4Name = "linear_grad_4_shader";
-static const char *kLinearGradient4SkSL =
-        // TODO: This should use local coords
-        "half4 linear_grad_4_shader() {\n"
-        "    float2 pos = sk_FragCoord.xy;\n"
-        "    float2 delta = point1 - point0;\n"
-        "    float2 pt = pos - point0;\n"
-        "    float t = dot(pt, delta) / dot(delta, delta);\n"
-        "    float4 result = colors[0];\n"
-        "    result = mix(result, colors[1],\n"
-        "                 clamp((t-offsets[0])/(offsets[1]-offsets[0]),\n"
-        "                       0, 1));\n"
-        "    result = mix(result, colors[2],\n"
-        "                 clamp((t-offsets[1])/(offsets[2]-offsets[1]),\n"
-        "                       0, 1));\n"
-        "    result = mix(result, colors[3],\n"
-        "                 clamp((t-offsets[2])/(offsets[3]-offsets[2]),\n"
-        "                 0, 1));\n"
-        "    return half4(result);\n"
-        "}\n";
-
-//--------------------------------------------------------------------------------------------------
-static constexpr int kNumSolidUniforms = 1;
-static constexpr SkUniform kSolidUniforms[kNumSolidUniforms] = {
-        { "color", SkSLType::kFloat4 }
-};
-
-static const char* kSolidColorName = "solid_shader";
-static const char* kSolidColorSkSL =
-        "half4 solid_shader() {\n"
-        "    return half4(color);\n"
-        "}\n";
-
-//--------------------------------------------------------------------------------------------------
-static constexpr int kNumImageUniforms = 0;
-static constexpr SkUniform kImageUniforms[kNumImageUniforms] = {
-};
-
-static const char* kImageName = "image_shader";
-static const char* kImageSkSL =
-        "half4 image_shader() {\n"
-        "    float r = fract(abs(sk_FragCoord.x/10.0));\n"
-        "    return half4(r, 0.0, 0.0, 1.0);\n"
-        "}\n";
-
-//--------------------------------------------------------------------------------------------------
-// TODO: kNone is for depth-only draws, so should actually have a fragment output type
-// that only defines a [[depth]] attribute but no color calculation.
-static const char* kNoneName = "none";
-static const char* kNoneSkSL = "outColor = half4(0.0, 0.0, 1.0, 1.0);\n";
-
-} // anonymous namespace
-
 std::tuple<SkUniquePaintParamsID, std::unique_ptr<SkUniformBlock>> ExtractPaintData(
         SkShaderCodeDictionary* dictionary,
         const PaintParams& p) {
@@ -108,42 +35,4 @@
     return { entry->uniqueID(), std::move(block) };
 }
 
-SkSpan<const SkUniform> GetUniforms(CodeSnippetID snippetID) {
-    switch (snippetID) {
-        case CodeSnippetID::kDepthStencilOnlyDraw:
-            return {nullptr, 0};
-        case CodeSnippetID::kLinearGradientShader: [[fallthrough]];
-        case CodeSnippetID::kRadialGradientShader: [[fallthrough]];
-        case CodeSnippetID::kSweepGradientShader:  [[fallthrough]];
-        case CodeSnippetID::kConicalGradientShader:
-            return SkMakeSpan(kGradientUniforms, kNumGradientUniforms);
-        case CodeSnippetID::kImageShader:
-            return SkMakeSpan(kImageUniforms, kNumImageUniforms);
-        case CodeSnippetID::kSolidColorShader:     [[fallthrough]];
-        default:
-            return SkMakeSpan(kSolidUniforms, kNumSolidUniforms);
-    }
-}
-
-// TODO: move this to the dictionary
-std::tuple<const char*, const char*> GetShaderSkSL(CodeSnippetID snippetID) {
-    switch (snippetID) {
-        case CodeSnippetID::kDepthStencilOnlyDraw:
-            return { kNoneName, kNoneSkSL};
-        case CodeSnippetID::kLinearGradientShader: [[fallthrough]];
-        case CodeSnippetID::kRadialGradientShader: [[fallthrough]];
-        case CodeSnippetID::kSweepGradientShader:  [[fallthrough]];
-        case CodeSnippetID::kConicalGradientShader:
-            return { kLinearGradient4Name, kLinearGradient4SkSL };
-        case CodeSnippetID::kImageShader:
-            return { kImageName, kImageSkSL };
-        case CodeSnippetID::kSolidColorShader:     [[fallthrough]];
-        default:
-            return { kSolidColorName, kSolidColorSkSL };
-    }
-}
-
-
-
-
 } // namespace skgpu
diff --git a/experimental/graphite/src/ContextUtils.h b/experimental/graphite/src/ContextUtils.h
index bd08d7f..443e0e1 100644
--- a/experimental/graphite/src/ContextUtils.h
+++ b/experimental/graphite/src/ContextUtils.h
@@ -26,14 +26,6 @@
 std::tuple<SkUniquePaintParamsID, std::unique_ptr<SkUniformBlock>> ExtractPaintData(
         SkShaderCodeDictionary*, const PaintParams&);
 
-SkSpan<const SkUniform> GetUniforms(CodeSnippetID);
-
-// TODO: Temporary way to get at SkSL snippet for handling the given shader type, which will be
-// embedded in the fragment function's body. It has access to the vertex output via a "interpolated"
-// variable, and must have a statement that writes to a float4 "out.color". Its uniforms (as defined
-// by GetUniforms(type)) are available as a variable named "uniforms".
-std::tuple<const char*, const char*> GetShaderSkSL(CodeSnippetID);
-
 } // namespace skgpu
 
 #endif // skgpu_ContextUtils_DEFINED
diff --git a/experimental/graphite/src/PaintParams.cpp b/experimental/graphite/src/PaintParams.cpp
index 3a1c868..1b5a8fa 100644
--- a/experimental/graphite/src/PaintParams.cpp
+++ b/experimental/graphite/src/PaintParams.cpp
@@ -48,13 +48,13 @@
     if (fShader) {
         as_SB(fShader)->addToKey(dict, backend, key, uniformBlock);
     } else {
-        SolidColorShaderBlock::AddToKey(backend, key, uniformBlock, fColor);
+        SolidColorShaderBlock::AddToKey(dict, backend, key, uniformBlock, fColor);
     }
 
     if (fBlender) {
         as_BB(fBlender)->addToKey(dict, backend, key, uniformBlock);
     } else {
-        BlendModeBlock::AddToKey(backend, key, uniformBlock, SkBlendMode::kSrcOver);
+        BlendModeBlock::AddToKey(dict, backend, key, uniformBlock, SkBlendMode::kSrcOver);
     }
 
     SkASSERT(key->sizeInBytes() > 0);
diff --git a/experimental/graphite/src/mtl/MtlGraphicsPipeline.mm b/experimental/graphite/src/mtl/MtlGraphicsPipeline.mm
index 519f5ce..9e8b3a6 100644
--- a/experimental/graphite/src/mtl/MtlGraphicsPipeline.mm
+++ b/experimental/graphite/src/mtl/MtlGraphicsPipeline.mm
@@ -169,13 +169,13 @@
     return sksl;
 }
 
-std::string get_sksl_fs(const SkShaderCodeDictionary* dictionary,
+std::string get_sksl_fs(const SkShaderCodeDictionary* dict,
                         const GraphicsPipelineDesc& desc,
                         bool* writesColor) {
     std::string sksl;
 
     SkPaintParamsKey key;
-    auto entry = dictionary->lookup(desc.paintParamsID());
+    auto entry = dict->lookup(desc.paintParamsID());
     if (entry) {
         key = entry->paintParamsKey();
     }
@@ -192,12 +192,12 @@
         }
 
         // Typedefs needed for painting
-        auto paintUniforms = GetUniforms(codeSnippetID);
+        auto paintUniforms = dict->getUniforms(codeSnippetID);
         if (!paintUniforms.empty()) {
             sksl += emit_SKSL_uniforms(2, "FS", paintUniforms);
         }
 
-        auto [name, code] = GetShaderSkSL(codeSnippetID);
+        auto [name, code] = dict->getShaderSkSL(codeSnippetID);
 
         sksl += code;
         sksl += "layout(location = 0, index = 0) out half4 sk_FragColor;\n";
diff --git a/include/private/SkPaintParamsKey.h b/include/private/SkPaintParamsKey.h
index 726da83..c827a51 100644
--- a/include/private/SkPaintParamsKey.h
+++ b/include/private/SkPaintParamsKey.h
@@ -19,7 +19,7 @@
 };
 
 // TODO: this needs to be expanded into a more flexible dictionary (esp. for user-supplied SkSL)
-// TODO: should this enum actually be in ShaderCodeDictionary.h?
+// TODO: rename to SkBuiltInCodeSnippetID and move to its own header
 enum class CodeSnippetID : uint8_t {
     // TODO: It seems like this requires some refinement. Fundamentally this doesn't seem like a
     // draw that originated from a PaintParams.
diff --git a/include/private/SkShaderCodeDictionary.h b/include/private/SkShaderCodeDictionary.h
index 1eb86fb..4890d4e 100644
--- a/include/private/SkShaderCodeDictionary.h
+++ b/include/private/SkShaderCodeDictionary.h
@@ -8,11 +8,14 @@
 #ifndef SkShaderCodeDictionary_DEFINED
 #define SkShaderCodeDictionary_DEFINED
 
+#include <array>
 #include <unordered_map>
+#include "include/core/SkSpan.h"
 #include "include/private/SkPaintParamsKey.h"
 #include "include/private/SkSpinlock.h"
 #include "include/private/SkUniquePaintParamsID.h"
 #include "src/core/SkArenaAlloc.h"
+#include "src/core/SkUniform.h"
 
 class SkShaderCodeDictionary {
 public:
@@ -44,6 +47,9 @@
 
     const Entry* lookup(SkUniquePaintParamsID) const SK_EXCLUDES(fSpinLock);
 
+    SkSpan<const SkUniform> getUniforms(CodeSnippetID) const;
+    std::tuple<const char*, const char*> getShaderSkSL(CodeSnippetID) const;
+
 private:
     Entry* makeEntry(const SkPaintParamsKey&);
 
@@ -51,6 +57,14 @@
         size_t operator()(const SkPaintParamsKey&) const;
     };
 
+    struct SnippetEntry {
+        SkSpan<const SkUniform> fUniforms;
+        const char* fName;
+        const char* fCode;
+    };
+
+    std::array<SnippetEntry, kCodeSnippetIDCount> fCodeSnippets;
+
     // TODO: can we do something better given this should have write-seldom/read-often behavior?
     mutable SkSpinlock fSpinLock;
 
diff --git a/src/core/SkBlendModeBlender.cpp b/src/core/SkBlendModeBlender.cpp
index c3503b9..8ddc62d 100644
--- a/src/core/SkBlendModeBlender.cpp
+++ b/src/core/SkBlendModeBlender.cpp
@@ -66,9 +66,9 @@
                              SkUniformBlock* uniformBlock) const {
 
     if (std::optional<SkBlendMode> bm = as_BB(this)->asBlendMode(); bm.has_value()) {
-        BlendModeBlock::AddToKey(backend, key, uniformBlock, bm.value());
+        BlendModeBlock::AddToKey(dict, backend, key, uniformBlock, bm.value());
     } else {
-        BlendModeBlock::AddToKey(backend, key, uniformBlock, SkBlendMode::kSrcOver);
+        BlendModeBlock::AddToKey(dict, backend, key, uniformBlock, SkBlendMode::kSrcOver);
     }
 }
 
diff --git a/src/core/SkKeyHelpers.cpp b/src/core/SkKeyHelpers.cpp
index 44c0979..35a1951 100644
--- a/src/core/SkKeyHelpers.cpp
+++ b/src/core/SkKeyHelpers.cpp
@@ -8,6 +8,7 @@
 #include "src/core/SkKeyHelpers.h"
 
 #include "include/private/SkPaintParamsKey.h"
+#include "include/private/SkShaderCodeDictionary.h"
 #include "src/core/SkDebugUtils.h"
 #include "src/core/SkUniform.h"
 #include "src/core/SkUniformData.h"
@@ -17,10 +18,6 @@
 #include "experimental/graphite/src/UniformManager.h"
 #endif
 
-namespace skgpu {
-SkSpan<const SkUniform> GetUniforms(CodeSnippetID snippetID);
-}
-
 namespace {
 
 #if defined(SK_DEBUG) && defined(SK_GRAPHITE_ENABLED)
@@ -68,7 +65,8 @@
 
 static const int kBlockDataSize = 0;
 
-void AddToKey(SkBackend /* backend */,
+void AddToKey(SkShaderCodeDictionary* /* dict */,
+              SkBackend /* backend */,
               SkPaintParamsKey* key,
               SkUniformBlock* /* uniformBlock */) {
     int headerOffset = key->beginBlock(CodeSnippetID::kDepthStencilOnlyDraw);
@@ -97,10 +95,10 @@
 #ifdef SK_GRAPHITE_ENABLED
 static const int kBlockDataSize = 0;
 
-sk_sp<SkUniformData> make_solid_uniform_data(SkColor4f color) {
+sk_sp<SkUniformData> make_solid_uniform_data(SkShaderCodeDictionary* dict, SkColor4f color) {
     static constexpr size_t kExpectedNumUniforms = 1;
 
-    SkSpan<const SkUniform> uniforms = skgpu::GetUniforms(CodeSnippetID::kSolidColorShader);
+    SkSpan<const SkUniform> uniforms = dict->getUniforms(CodeSnippetID::kSolidColorShader);
     SkASSERT(uniforms.size() == kExpectedNumUniforms);
 
     skgpu::UniformManager mgr(skgpu::Layout::kMetal);
@@ -118,7 +116,8 @@
 
 } // anonymous namespace
 
-void AddToKey(SkBackend backend,
+void AddToKey(SkShaderCodeDictionary* dict,
+              SkBackend backend,
               SkPaintParamsKey* key,
               SkUniformBlock* uniformBlock,
               const SkColor4f& color) {
@@ -132,7 +131,7 @@
                               CodeSnippetID::kSolidColorShader, kBlockDataSize);
 
         if (uniformBlock) {
-            uniformBlock->add(make_solid_uniform_data(color));
+            uniformBlock->add(make_solid_uniform_data(dict, color));
         }
         return;
     }
@@ -183,9 +182,10 @@
     return result;
 }
 
-sk_sp<SkUniformData> make_linear_gradient_uniform_data(const GradientData& gradData) {
+sk_sp<SkUniformData> make_linear_gradient_uniform_data(SkShaderCodeDictionary* dict,
+                                                       const GradientData& gradData) {
 
-    SkSpan<const SkUniform> uniforms = skgpu::GetUniforms(CodeSnippetID::kLinearGradientShader);
+    SkSpan<const SkUniform> uniforms = dict->getUniforms(CodeSnippetID::kLinearGradientShader);
     SkASSERT(uniforms.size() == kExpectedNumGradientUniforms);
 
     const void* srcs[kExpectedNumGradientUniforms] = {
@@ -200,9 +200,10 @@
     return make_gradient_uniform_data_common(uniforms, srcs);
 };
 
-sk_sp<SkUniformData> make_radial_gradient_uniform_data(const GradientData& gradData) {
+sk_sp<SkUniformData> make_radial_gradient_uniform_data(SkShaderCodeDictionary* dict,
+                                                       const GradientData& gradData) {
 
-    SkSpan<const SkUniform> uniforms = skgpu::GetUniforms(CodeSnippetID::kRadialGradientShader);
+    SkSpan<const SkUniform> uniforms = dict->getUniforms(CodeSnippetID::kRadialGradientShader);
     SkASSERT(uniforms.size() == kExpectedNumGradientUniforms);
 
     const void* srcs[kExpectedNumGradientUniforms] = {
@@ -217,9 +218,10 @@
     return make_gradient_uniform_data_common(uniforms, srcs);
 };
 
-sk_sp<SkUniformData> make_sweep_gradient_uniform_data(const GradientData& gradData) {
+sk_sp<SkUniformData> make_sweep_gradient_uniform_data(SkShaderCodeDictionary* dict,
+                                                      const GradientData& gradData) {
 
-    SkSpan<const SkUniform> uniforms = skgpu::GetUniforms(CodeSnippetID::kSweepGradientShader);
+    SkSpan<const SkUniform> uniforms = dict->getUniforms(CodeSnippetID::kSweepGradientShader);
     SkASSERT(uniforms.size() == kExpectedNumGradientUniforms);
 
     const void* srcs[kExpectedNumGradientUniforms] = {
@@ -234,9 +236,10 @@
     return make_gradient_uniform_data_common(uniforms, srcs);
 };
 
-sk_sp<SkUniformData> make_conical_gradient_uniform_data(const GradientData& gradData) {
+sk_sp<SkUniformData> make_conical_gradient_uniform_data(SkShaderCodeDictionary* dict,
+                                                        const GradientData& gradData) {
 
-    SkSpan<const SkUniform> uniforms = skgpu::GetUniforms(CodeSnippetID::kConicalGradientShader);
+    SkSpan<const SkUniform> uniforms = dict->getUniforms(CodeSnippetID::kConicalGradientShader);
     SkASSERT(uniforms.size() == kExpectedNumGradientUniforms);
 
     const void* srcs[kExpectedNumGradientUniforms] = {
@@ -300,7 +303,8 @@
     }
 }
 
-void AddToKey(SkBackend backend,
+void AddToKey(SkShaderCodeDictionary* dict,
+              SkBackend backend,
               SkPaintParamsKey *key,
               SkUniformBlock* uniformBlock,
               const GradientData& gradData) {
@@ -312,25 +316,25 @@
             case SkShader::kLinear_GradientType:
                 codeSnippetID = CodeSnippetID::kLinearGradientShader;
                 if (uniformBlock) {
-                    uniformBlock->add(make_linear_gradient_uniform_data(gradData));
+                    uniformBlock->add(make_linear_gradient_uniform_data(dict, gradData));
                 }
                 break;
             case SkShader::kRadial_GradientType:
                 codeSnippetID = CodeSnippetID::kRadialGradientShader;
                 if (uniformBlock) {
-                    uniformBlock->add(make_radial_gradient_uniform_data(gradData));
+                    uniformBlock->add(make_radial_gradient_uniform_data(dict, gradData));
                 }
                 break;
             case SkShader::kSweep_GradientType:
                 codeSnippetID = CodeSnippetID::kSweepGradientShader;
                 if (uniformBlock) {
-                    uniformBlock->add(make_sweep_gradient_uniform_data(gradData));
+                    uniformBlock->add(make_sweep_gradient_uniform_data(dict, gradData));
                 }
                 break;
             case SkShader::GradientType::kConical_GradientType:
                 codeSnippetID = CodeSnippetID::kConicalGradientShader;
                 if (uniformBlock) {
-                    uniformBlock->add(make_conical_gradient_uniform_data(gradData));
+                    uniformBlock->add(make_conical_gradient_uniform_data(dict, gradData));
                 }
                 break;
             case SkShader::GradientType::kColor_GradientType:
@@ -354,7 +358,7 @@
 
     if (backend == SkBackend::kSkVM || backend == SkBackend::kGanesh) {
         // TODO: add implementation of other backends
-        SolidColorShaderBlock::AddToKey(backend, key, uniformBlock, SkColors::kRed);
+        SolidColorShaderBlock::AddToKey(dict, backend, key, uniformBlock, SkColors::kRed);
     }
 }
 
@@ -440,10 +444,11 @@
 }
 #endif // SK_DEBUG
 
-sk_sp<SkUniformData> make_image_uniform_data(const ImageData& imgData) {
+sk_sp<SkUniformData> make_image_uniform_data(SkShaderCodeDictionary* dict,
+                                             const ImageData& imgData) {
     SkDEBUGCODE(static constexpr size_t kExpectedNumUniforms = 0;)
 
-    SkSpan<const SkUniform> uniforms = skgpu::GetUniforms(CodeSnippetID::kImageShader);
+    SkSpan<const SkUniform> uniforms = dict->getUniforms(CodeSnippetID::kImageShader);
     SkASSERT(uniforms.size() == kExpectedNumUniforms);
 
     skgpu::UniformManager mgr(skgpu::Layout::kMetal);
@@ -462,7 +467,8 @@
 
 } // anonymous namespace
 
-void AddToKey(SkBackend backend,
+void AddToKey(SkShaderCodeDictionary* dict,
+              SkBackend backend,
               SkPaintParamsKey* key,
               SkUniformBlock* uniformBlock,
               const ImageData& imgData) {
@@ -482,7 +488,7 @@
         SkASSERT(imgData == ExtractFromKey(*key, headerOffset));
 
         if (uniformBlock) {
-            uniformBlock->add(make_image_uniform_data(imgData));
+            uniformBlock->add(make_image_uniform_data(dict, imgData));
         }
         return;
     }
@@ -490,7 +496,7 @@
 
     if (backend == SkBackend::kSkVM || backend == SkBackend::kGanesh) {
         // TODO: add implementation for other backends
-        SolidColorShaderBlock::AddToKey(backend, key, uniformBlock, SkColors::kRed);
+        SolidColorShaderBlock::AddToKey(dict, backend, key, uniformBlock, SkColors::kRed);
     }
 }
 
@@ -513,7 +519,8 @@
 //--------------------------------------------------------------------------------------------------
 namespace BlendShaderBlock {
 
-void AddToKey(SkBackend backend,
+void AddToKey(SkShaderCodeDictionary* dict,
+              SkBackend backend,
               SkPaintParamsKey *key,
               SkUniformBlock* uniformBlock,
               const BlendData& blendData) {
@@ -524,11 +531,11 @@
 
         add_blendmode_to_key(key, blendData.fBM);
         int start = key->sizeInBytes();
-        as_SB(blendData.fDst)->addToKey(nullptr, backend, key, uniformBlock);
+        as_SB(blendData.fDst)->addToKey(dict, backend, key, uniformBlock);
         int firstShaderSize = key->sizeInBytes() - start;
 
         start = key->sizeInBytes();
-        as_SB(blendData.fSrc)->addToKey(nullptr, backend, key, uniformBlock);
+        as_SB(blendData.fSrc)->addToKey(dict, backend, key, uniformBlock);
         int secondShaderSize = key->sizeInBytes() - start;
 
         key->endBlock(headerOffset, CodeSnippetID::kBlendShader);
@@ -541,7 +548,7 @@
 
     if (backend == SkBackend::kSkVM || backend == SkBackend::kGanesh) {
         // TODO: add implementation for other backends
-        SolidColorShaderBlock::AddToKey(backend, key, uniformBlock, SkColors::kRed);
+        SolidColorShaderBlock::AddToKey(dict, backend, key, uniformBlock, SkColors::kRed);
     }
 }
 
@@ -582,7 +589,8 @@
 static const int kBlockDataSize = 1;
 #endif
 
-void AddToKey(SkBackend backend,
+void AddToKey(SkShaderCodeDictionary* dict,
+              SkBackend backend,
               SkPaintParamsKey *key,
               SkUniformBlock* uniformBlock,
               SkBlendMode bm) {
@@ -601,7 +609,7 @@
 
     if (backend == SkBackend::kSkVM || backend == SkBackend::kGanesh) {
         // TODO: add implementation for other backends
-        SolidColorShaderBlock::AddToKey(backend, key, uniformBlock, SkColors::kRed);
+        SolidColorShaderBlock::AddToKey(dict, backend, key, uniformBlock, SkColors::kRed);
     }
 }
 
@@ -631,7 +639,8 @@
 
 //--------------------------------------------------------------------------------------------------
 #ifdef SK_GRAPHITE_ENABLED
-SkPaintParamsKey CreateKey(SkBackend backend,
+SkPaintParamsKey CreateKey(SkShaderCodeDictionary* dict,
+                           SkBackend backend,
                            skgpu::ShaderCombo::ShaderType s,
                            SkTileMode tm,
                            SkBlendMode bm) {
@@ -639,30 +648,30 @@
 
     switch (s) {
         case skgpu::ShaderCombo::ShaderType::kNone:
-            DepthStencilOnlyBlock::AddToKey(backend, &key, nullptr);
+            DepthStencilOnlyBlock::AddToKey(dict, backend, &key, nullptr);
             break;
         case skgpu::ShaderCombo::ShaderType::kSolidColor:
-            SolidColorShaderBlock::AddToKey(backend, &key, nullptr, SkColors::kRed);
+            SolidColorShaderBlock::AddToKey(dict, backend, &key, nullptr, SkColors::kRed);
             break;
         case skgpu::ShaderCombo::ShaderType::kLinearGradient:
-            GradientShaderBlocks::AddToKey(backend, &key, nullptr,
+            GradientShaderBlocks::AddToKey(dict, backend, &key, nullptr,
                                            { SkShader::kLinear_GradientType, tm, 0 });
             break;
         case skgpu::ShaderCombo::ShaderType::kRadialGradient:
-            GradientShaderBlocks::AddToKey(backend, &key, nullptr,
+            GradientShaderBlocks::AddToKey(dict, backend, &key, nullptr,
                                            { SkShader::kRadial_GradientType, tm, 0 });
             break;
         case skgpu::ShaderCombo::ShaderType::kSweepGradient:
-            GradientShaderBlocks::AddToKey(backend, &key, nullptr,
+            GradientShaderBlocks::AddToKey(dict, backend, &key, nullptr,
                                            { SkShader::kSweep_GradientType, tm, 0 });
             break;
         case skgpu::ShaderCombo::ShaderType::kConicalGradient:
-            GradientShaderBlocks::AddToKey(backend, &key, nullptr,
+            GradientShaderBlocks::AddToKey(dict, backend, &key, nullptr,
                                            { SkShader::kConical_GradientType, tm, 0 });
             break;
     }
 
-    BlendModeBlock::AddToKey(backend, &key, nullptr, bm);
+    BlendModeBlock::AddToKey(dict, backend, &key, nullptr, bm);
     return key;
 }
 #endif
diff --git a/src/core/SkKeyHelpers.h b/src/core/SkKeyHelpers.h
index 1277cf3..5603ce9 100644
--- a/src/core/SkKeyHelpers.h
+++ b/src/core/SkKeyHelpers.h
@@ -18,13 +18,14 @@
 
 enum class SkBackend : uint8_t;
 class SkPaintParamsKey;
+class SkShaderCodeDictionary;
 class SkUniformBlock;
 
 // The KeyHelpers can be used to manually construct an SkPaintParamsKey
 
 namespace DepthStencilOnlyBlock {
 
-    void AddToKey(SkBackend, SkPaintParamsKey*, SkUniformBlock*);
+    void AddToKey(SkShaderCodeDictionary*, SkBackend, SkPaintParamsKey*, SkUniformBlock*);
 #ifdef SK_DEBUG
     void Dump(const SkPaintParamsKey&, int headerOffset);
 #endif
@@ -33,7 +34,8 @@
 
 namespace SolidColorShaderBlock {
 
-    void AddToKey(SkBackend,
+    void AddToKey(SkShaderCodeDictionary*,
+                  SkBackend,
                   SkPaintParamsKey*,
                   SkUniformBlock*,
                   const SkColor4f&);
@@ -89,7 +91,8 @@
         float                  fOffsets[kMaxStops];
     };
 
-    void AddToKey(SkBackend,
+    void AddToKey(SkShaderCodeDictionary*,
+                  SkBackend,
                   SkPaintParamsKey*,
                   SkUniformBlock*,
                   const GradientData&);
@@ -113,7 +116,8 @@
         SkTileMode fTileModes[2];
     };
 
-    void AddToKey(SkBackend,
+    void AddToKey(SkShaderCodeDictionary*,
+                  SkBackend,
                   SkPaintParamsKey*,
                   SkUniformBlock*,
                   const ImageData&);
@@ -132,7 +136,8 @@
         SkBlendMode fBM;
     };
 
-    void AddToKey(SkBackend,
+    void AddToKey(SkShaderCodeDictionary*,
+                  SkBackend,
                   SkPaintParamsKey*,
                   SkUniformBlock*,
                   const BlendData&);
@@ -144,7 +149,11 @@
 
 namespace BlendModeBlock {
 
-    void AddToKey(SkBackend, SkPaintParamsKey*, SkUniformBlock*, SkBlendMode);
+    void AddToKey(SkShaderCodeDictionary*,
+                  SkBackend,
+                  SkPaintParamsKey*,
+                  SkUniformBlock*,
+                  SkBlendMode);
 #ifdef SK_DEBUG
     void Dump(const SkPaintParamsKey&, int headerOffset);
 #endif
@@ -153,7 +162,11 @@
 
 #ifdef SK_GRAPHITE_ENABLED
 // Bridge between the combinations system and the SkPaintParamsKey
-SkPaintParamsKey CreateKey(SkBackend, skgpu::ShaderCombo::ShaderType, SkTileMode, SkBlendMode);
+SkPaintParamsKey CreateKey(SkShaderCodeDictionary*,
+                           SkBackend,
+                           skgpu::ShaderCombo::ShaderType,
+                           SkTileMode,
+                           SkBlendMode);
 #endif
 
 #endif // SkKeyHelpers_DEFINED
diff --git a/src/core/SkPaintPriv.cpp b/src/core/SkPaintPriv.cpp
index 14e7893..5433472 100644
--- a/src/core/SkPaintPriv.cpp
+++ b/src/core/SkPaintPriv.cpp
@@ -138,13 +138,13 @@
         if (paint.getShader()) {
             as_SB(paint.getShader())->addToKey(dict, backend, &key, nullptr);
         } else {
-            SolidColorShaderBlock::AddToKey(backend, &key, nullptr, paint.getColor4f());
+            SolidColorShaderBlock::AddToKey(dict, backend, &key, nullptr, paint.getColor4f());
         }
 
         if (paint.getBlender()) {
             as_BB(paint.getBlender())->addToKey(dict, backend, &key, nullptr);
         } else {
-            BlendModeBlock::AddToKey(backend, &key, nullptr, SkBlendMode::kSrcOver);
+            BlendModeBlock::AddToKey(dict, backend, &key, nullptr, SkBlendMode::kSrcOver);
         }
 
         SkASSERT(key.sizeInBytes() > 0);
diff --git a/src/core/SkShaderCodeDictionary.cpp b/src/core/SkShaderCodeDictionary.cpp
index d56af71..88e437d 100644
--- a/src/core/SkShaderCodeDictionary.cpp
+++ b/src/core/SkShaderCodeDictionary.cpp
@@ -8,11 +8,6 @@
 #include "include/private/SkShaderCodeDictionary.h"
 #include "src/core/SkOpts.h"
 
-SkShaderCodeDictionary::SkShaderCodeDictionary() {
-    // The 0th index is reserved as invalid
-    fEntryVector.push_back(nullptr);
-}
-
 SkShaderCodeDictionary::Entry* SkShaderCodeDictionary::makeEntry(const SkPaintParamsKey& key) {
     return fArena.make([&](void *ptr) { return new(ptr) Entry(key); });
 }
@@ -21,7 +16,8 @@
     return SkOpts::hash_fn(key.data(), key.sizeInBytes(), 0);
 }
 
-const SkShaderCodeDictionary::Entry* SkShaderCodeDictionary::findOrCreate(const SkPaintParamsKey& key) {
+const SkShaderCodeDictionary::Entry* SkShaderCodeDictionary::findOrCreate(
+        const SkPaintParamsKey& key) {
     SkAutoSpinlock lock{fSpinLock};
 
     auto iter = fHash.find(key);
@@ -51,3 +47,130 @@
 
     return fEntryVector[codeID.asUInt()];
 }
+
+SkSpan<const SkUniform> SkShaderCodeDictionary::getUniforms(CodeSnippetID id) const {
+    return fCodeSnippets[(int) id].fUniforms;
+}
+
+std::tuple<const char*, const char*> SkShaderCodeDictionary::getShaderSkSL(
+        CodeSnippetID id) const {
+    if (fCodeSnippets[(int) id].fCode) {
+        return { fCodeSnippets[(int) id].fName, fCodeSnippets[(int) id].fCode };
+    }
+
+    // If we're missing a code snippet just draw solid blue
+    return this->getShaderSkSL(CodeSnippetID::kDepthStencilOnlyDraw);
+}
+
+//--------------------------------------------------------------------------------------------------
+namespace {
+
+static constexpr int kFourStopGradient = 4;
+
+// TODO: For the sprint we unify all the gradient uniforms into a standard set of 6:
+//   kMaxStops colors
+//   kMaxStops offsets
+//   2 points
+//   2 radii
+static constexpr int kNumGradientUniforms = 6;
+static constexpr SkUniform kGradientUniforms[kNumGradientUniforms] = {
+        { "colors",  SkSLType::kHalf4, kFourStopGradient },
+        { "offsets", SkSLType::kFloat, kFourStopGradient },
+        { "point0",  SkSLType::kFloat2 },
+        { "point1",  SkSLType::kFloat2 },
+        { "radius0", SkSLType::kFloat },
+        { "radius1", SkSLType::kFloat },
+};
+
+static const char *kLinearGradient4Name = "linear_grad_4_shader";
+static const char *kLinearGradient4SkSL =
+        // TODO: This should use local coords
+        "half4 linear_grad_4_shader() {\n"
+        "    float2 pos = sk_FragCoord.xy;\n"
+        "    float2 delta = point1 - point0;\n"
+        "    float2 pt = pos - point0;\n"
+        "    float t = dot(pt, delta) / dot(delta, delta);\n"
+        "    float4 result = colors[0];\n"
+        "    result = mix(result, colors[1],\n"
+        "                 clamp((t-offsets[0])/(offsets[1]-offsets[0]),\n"
+        "                       0, 1));\n"
+        "    result = mix(result, colors[2],\n"
+        "                 clamp((t-offsets[1])/(offsets[2]-offsets[1]),\n"
+        "                       0, 1));\n"
+        "    result = mix(result, colors[3],\n"
+        "                 clamp((t-offsets[2])/(offsets[3]-offsets[2]),\n"
+        "                 0, 1));\n"
+        "    return half4(result);\n"
+        "}\n";
+
+//--------------------------------------------------------------------------------------------------
+static constexpr int kNumSolidUniforms = 1;
+static constexpr SkUniform kSolidUniforms[kNumSolidUniforms] = {
+        { "color", SkSLType::kFloat4 }
+};
+
+static const char* kSolidColorName = "solid_shader";
+static const char* kSolidColorSkSL =
+        "half4 solid_shader() {\n"
+        "    return half4(color);\n"
+        "}\n";
+
+//--------------------------------------------------------------------------------------------------
+static constexpr int kNumImageUniforms = 0;
+
+static const char* kImageName = "image_shader";
+static const char* kImageSkSL =
+        "half4 image_shader() {\n"
+        "    float r = fract(abs(sk_FragCoord.x/10.0));\n"
+        "    return half4(r, 0.0, 0.0, 1.0);\n"
+        "}\n";
+
+//--------------------------------------------------------------------------------------------------
+// TODO: kNone is for depth-only draws, so should actually have a fragment output type
+// that only defines a [[depth]] attribute but no color calculation.
+static const char* kNoneName = "none";
+static const char* kNoneSkSL =
+        "half4 none() {\n"
+        "    return half4(0.0, 0.0, 1.0, 1.0);\n"
+        "}\n";
+
+} // anonymous namespace
+
+SkShaderCodeDictionary::SkShaderCodeDictionary() {
+    // The 0th index is reserved as invalid
+    fEntryVector.push_back(nullptr);
+
+    fCodeSnippets[(int) CodeSnippetID::kDepthStencilOnlyDraw] = {
+            {}, kNoneName, kNoneSkSL
+    };
+    fCodeSnippets[(int) CodeSnippetID::kSolidColorShader] = {
+            SkMakeSpan(kSolidUniforms, kNumSolidUniforms),
+            kSolidColorName, kSolidColorSkSL
+    };
+    fCodeSnippets[(int) CodeSnippetID::kLinearGradientShader] = {
+            SkMakeSpan(kGradientUniforms, kNumGradientUniforms),
+            kLinearGradient4Name, kLinearGradient4SkSL
+    };
+    fCodeSnippets[(int) CodeSnippetID::kRadialGradientShader] = {
+            SkMakeSpan(kGradientUniforms, kNumGradientUniforms),
+            kLinearGradient4Name, kLinearGradient4SkSL
+    };
+    fCodeSnippets[(int) CodeSnippetID::kSweepGradientShader] = {
+            SkMakeSpan(kGradientUniforms, kNumGradientUniforms),
+            kLinearGradient4Name, kLinearGradient4SkSL
+    };
+    fCodeSnippets[(int) CodeSnippetID::kConicalGradientShader] = {
+            SkMakeSpan(kGradientUniforms, kNumGradientUniforms),
+            kLinearGradient4Name, kLinearGradient4SkSL
+    };
+    fCodeSnippets[(int) CodeSnippetID::kImageShader] = {
+            { nullptr, kNumImageUniforms },
+            kImageName, kImageSkSL
+    };
+    fCodeSnippets[(int) CodeSnippetID::kBlendShader] = {
+            {}, nullptr, nullptr
+    };
+    fCodeSnippets[(int) CodeSnippetID::kSimpleBlendMode] = {
+            {}, nullptr, nullptr
+    };
+}
diff --git a/src/shaders/SkColorShader.cpp b/src/shaders/SkColorShader.cpp
index 4af9ecb..28091ac 100644
--- a/src/shaders/SkColorShader.cpp
+++ b/src/shaders/SkColorShader.cpp
@@ -141,12 +141,12 @@
                              SkBackend backend,
                              SkPaintParamsKey* key,
                              SkUniformBlock* uniformBlock) const {
-    SolidColorShaderBlock::AddToKey(backend, key, uniformBlock, SkColor4f::FromColor(fColor));
+    SolidColorShaderBlock::AddToKey(dict, backend, key, uniformBlock, SkColor4f::FromColor(fColor));
 }
 
 void SkColor4Shader::addToKey(SkShaderCodeDictionary* dict,
                               SkBackend backend,
                               SkPaintParamsKey* key,
                               SkUniformBlock* uniformBlock) const {
-    SolidColorShaderBlock::AddToKey(backend, key, uniformBlock, fColor);
+    SolidColorShaderBlock::AddToKey(dict, backend, key, uniformBlock, fColor);
 }
diff --git a/src/shaders/SkComposeShader.cpp b/src/shaders/SkComposeShader.cpp
index 5ee76f6..224740d 100644
--- a/src/shaders/SkComposeShader.cpp
+++ b/src/shaders/SkComposeShader.cpp
@@ -196,5 +196,5 @@
     // TODO: add blender support
     SkASSERT(!fBlender);
 
-    BlendShaderBlock::AddToKey(backend, key, uniformBlock, { fDst.get(), fSrc.get(), fMode });
+    BlendShaderBlock::AddToKey(dict, backend, key, uniformBlock, { fDst.get(), fSrc.get(), fMode });
 }
diff --git a/src/shaders/SkImageShader.cpp b/src/shaders/SkImageShader.cpp
index d202a81..e373e9d 100755
--- a/src/shaders/SkImageShader.cpp
+++ b/src/shaders/SkImageShader.cpp
@@ -377,7 +377,7 @@
                              SkBackend backend,
                              SkPaintParamsKey* key,
                              SkUniformBlock* uniformBlock) const {
-    ImageShaderBlock::AddToKey(backend, key, uniformBlock, { fTileModeX, fTileModeY });
+    ImageShaderBlock::AddToKey(dict, backend, key, uniformBlock, { fTileModeX, fTileModeY });
 }
 
 ///////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/src/shaders/SkShader.cpp b/src/shaders/SkShader.cpp
index 820f090..74f9842 100644
--- a/src/shaders/SkShader.cpp
+++ b/src/shaders/SkShader.cpp
@@ -150,11 +150,11 @@
 }
 
 // TODO: add implementations for derived classes
-void SkShaderBase::addToKey(SkShaderCodeDictionary* dictionary,
+void SkShaderBase::addToKey(SkShaderCodeDictionary* dict,
                             SkBackend backend,
                             SkPaintParamsKey* key,
                             SkUniformBlock* uniformBlock) const {
-    SolidColorShaderBlock::AddToKey(backend, key, uniformBlock, SkColors::kRed);
+    SolidColorShaderBlock::AddToKey(dict, backend, key, uniformBlock, SkColors::kRed);
 }
 
 sk_sp<SkShader> SkShaders::Empty() { return sk_make_sp<SkEmptyShader>(); }
diff --git a/src/shaders/gradients/SkLinearGradient.cpp b/src/shaders/gradients/SkLinearGradient.cpp
index 0fc3b48..b339fc0 100644
--- a/src/shaders/gradients/SkLinearGradient.cpp
+++ b/src/shaders/gradients/SkLinearGradient.cpp
@@ -116,5 +116,5 @@
                                             fOrigColors4f,
                                             fOrigPos);
 
-    GradientShaderBlocks::AddToKey(backend, key, uniformBlock, data);
+    GradientShaderBlocks::AddToKey(dict, backend, key, uniformBlock, data);
 }
diff --git a/src/shaders/gradients/SkRadialGradient.cpp b/src/shaders/gradients/SkRadialGradient.cpp
index 24aa76a..d0cc51c 100644
--- a/src/shaders/gradients/SkRadialGradient.cpp
+++ b/src/shaders/gradients/SkRadialGradient.cpp
@@ -95,5 +95,5 @@
                                             fOrigColors4f,
                                             fOrigPos);
 
-    GradientShaderBlocks::AddToKey(backend, key, uniformBlock, data);
+    GradientShaderBlocks::AddToKey(dict, backend, key, uniformBlock, data);
 }
diff --git a/src/shaders/gradients/SkSweepGradient.cpp b/src/shaders/gradients/SkSweepGradient.cpp
index 93ef8d1..ce07c79 100644
--- a/src/shaders/gradients/SkSweepGradient.cpp
+++ b/src/shaders/gradients/SkSweepGradient.cpp
@@ -119,5 +119,5 @@
                                             fOrigColors4f,
                                             fOrigPos);
 
-    GradientShaderBlocks::AddToKey(backend, key, uniformBlock, data);
+    GradientShaderBlocks::AddToKey(dict, backend, key, uniformBlock, data);
 }
diff --git a/src/shaders/gradients/SkTwoPointConicalGradient.cpp b/src/shaders/gradients/SkTwoPointConicalGradient.cpp
index 042376e..e73ac05 100644
--- a/src/shaders/gradients/SkTwoPointConicalGradient.cpp
+++ b/src/shaders/gradients/SkTwoPointConicalGradient.cpp
@@ -289,5 +289,5 @@
                                             fOrigColors4f,
                                             fOrigPos);
 
-    GradientShaderBlocks::AddToKey(backend, key, uniformBlock, data);
+    GradientShaderBlocks::AddToKey(dict, backend, key, uniformBlock, data);
 }
diff --git a/tests/graphite/CommandBufferTest.cpp b/tests/graphite/CommandBufferTest.cpp
index 11366eb..8140fb7 100644
--- a/tests/graphite/CommandBufferTest.cpp
+++ b/tests/graphite/CommandBufferTest.cpp
@@ -237,6 +237,7 @@
 #endif
     auto recorder = context->makeRecorder();
     auto resourceProvider = recorder->priv().resourceProvider();
+    auto dict = resourceProvider->shaderCodeDictionary();
     auto commandBuffer = resourceProvider->createCommandBuffer();
 
     SkISize textureSize = { kTextureWidth, kTextureHeight };
@@ -254,12 +255,12 @@
     TextureInfo textureInfo;
 #endif
 
-    SkPaintParamsKey key = CreateKey(SkBackend::kGraphite,
+    SkPaintParamsKey key = CreateKey(dict,
+                                     SkBackend::kGraphite,
                                      ShaderCombo::ShaderType::kSolidColor,
                                      SkTileMode::kClamp,
                                      SkBlendMode::kSrc);
 
-    auto dict = resourceProvider->shaderCodeDictionary();
     auto entry = dict->findOrCreate(key);
 
     auto target = sk_sp<TextureProxy>(new TextureProxy(textureSize, textureInfo));
diff --git a/tests/graphite/UniformTest.cpp b/tests/graphite/UniformTest.cpp
index f06bede..6dfbb20 100644
--- a/tests/graphite/UniformTest.cpp
+++ b/tests/graphite/UniformTest.cpp
@@ -73,6 +73,7 @@
     using namespace skgpu;
 
     auto recorder = context->makeRecorder();
+    auto dict = recorder->priv().resourceProvider()->shaderCodeDictionary();
 
     // Intentionally does not include ShaderType::kNone, which represents no fragment shading stage
     // and is thus not relevant to uniform extraction/caching.
@@ -90,10 +91,9 @@
             }
 
             for (auto bm : { SkBlendMode::kSrc, SkBlendMode::kSrcOver }) {
-                SkPaintParamsKey expected = CreateKey(SkBackend::kGraphite, s, tm, bm);
+                SkPaintParamsKey expected = CreateKey(dict, SkBackend::kGraphite, s, tm, bm);
 
                 auto [ p, expectedNumUniforms ] = create_paint(s, tm, bm);
-                auto dict = recorder->priv().resourceProvider()->shaderCodeDictionary();
                 auto [ actualID, uniformBlock] = ExtractPaintData(dict, PaintParams(p));
                 int actualNumUniforms = uniformBlock->count();