Add SkRuntimeShaderBuilder, clean up SkRuntimeEffect API a bit

Utility class for getting named access to uniforms and children of an
SkRuntimeEffect (also functions as an example of using the
SkRuntimeEffect public API).

Moved several internal SkRuntimeEffect functions to private, and added
findInput/findChild helpers.

Change-Id: I8c2e7745ea81670a49b7ab2f51ce44a8d8169278
Reviewed-on: https://skia-review.googlesource.com/c/skia/+/286516
Commit-Queue: Brian Osman <brianosman@google.com>
Reviewed-by: Brian Salomon <bsalomon@google.com>
diff --git a/include/effects/SkRuntimeEffect.h b/include/effects/SkRuntimeEffect.h
index 5b6a7c9..943bbb7 100644
--- a/include/effects/SkRuntimeEffect.h
+++ b/include/effects/SkRuntimeEffect.h
@@ -8,6 +8,7 @@
 #ifndef SkRuntimeEffect_DEFINED
 #define SkRuntimeEffect_DEFINED
 
+#include "include/core/SkData.h"
 #include "include/core/SkString.h"
 
 #include <vector>
@@ -32,8 +33,8 @@
 /*
  * SkRuntimeEffect supports creating custom SkShader and SkColorFilter objects using Skia's SkSL
  * shading language.
- * *
- * This API is experimental and subject to change.
+ *
+ * NOTE: This API is experimental and subject to change.
  */
 class SK_API SkRuntimeEffect : public SkRefCnt {
 public:
@@ -87,13 +88,16 @@
     // [Effect, ErrorText]
     // If successful, Effect != nullptr, otherwise, ErrorText contains the reason for failure.
     using EffectResult = std::tuple<sk_sp<SkRuntimeEffect>, SkString>;
-
     static EffectResult Make(SkString sksl);
 
-    sk_sp<SkShader> makeShader(sk_sp<SkData> inputs, sk_sp<SkShader> children[], size_t childCount,
-                               const SkMatrix* localMatrix, bool isOpaque);
+    sk_sp<SkShader> makeShader(sk_sp<SkData> inputs,
+                               sk_sp<SkShader> children[],
+                               size_t childCount,
+                               const SkMatrix* localMatrix,
+                               bool isOpaque);
 
-    sk_sp<SkColorFilter> makeColorFilter(sk_sp<SkData> inputs, sk_sp<SkColorFilter> children[],
+    sk_sp<SkColorFilter> makeColorFilter(sk_sp<SkData> inputs,
+                                         sk_sp<SkColorFilter> children[],
                                          size_t childCount);
     sk_sp<SkColorFilter> makeColorFilter(sk_sp<SkData> inputs);
 
@@ -119,29 +123,17 @@
     // makeShader, provide an SkData of this size, containing values for all of those variables.
     size_t inputSize() const;
 
-    // Combined size of just the 'uniform' variables.
-    size_t uniformSize() const { return fUniformSize; }
-
     ConstIterable<Variable> inputs() const { return ConstIterable<Variable>(fInAndUniformVars); }
     ConstIterable<SkString> children() const { return ConstIterable<SkString>(fChildren); }
     ConstIterable<Varying> varyings() const { return ConstIterable<Varying>(fVaryings); }
 
-#if SK_SUPPORT_GPU
-    // This re-compiles the program from scratch, using the supplied shader caps.
-    // This is necessary to get the correct values of settings.
-    bool toPipelineStage(const void* inputs, const GrShaderCaps* shaderCaps,
-                         GrContextOptions::ShaderErrorHandler* errorHandler,
-                         SkSL::PipelineStageArgs* outArgs);
-#endif
+    // Returns pointer to the named in/uniform variable's description, or nullptr if not found
+    const Variable* findInput(const char* name) const;
 
-    // [ByteCode, ErrorText]
-    // If successful, ByteCode != nullptr, otherwise, ErrorText contains the reason for failure.
-    using ByteCodeResult = std::tuple<std::unique_ptr<SkSL::ByteCode>, SkString>;
-
-    ByteCodeResult toByteCode(const void* inputs) const;
+    // Returns index of the named child, or -1 if not found
+    int findChild(const char* name) const;
 
     static void RegisterFlattenables();
-
     ~SkRuntimeEffect();
 
 private:
@@ -153,6 +145,28 @@
     SpecializeResult specialize(SkSL::Program& baseProgram, const void* inputs,
                                 const SkSL::SharedCompiler&) const;
 
+#if SK_SUPPORT_GPU
+    friend class GrSkSLFP;  // toPipelineStage
+
+    // This re-compiles the program from scratch, using the supplied shader caps.
+    // This is necessary to get the correct values of settings.
+    bool toPipelineStage(const void* inputs, const GrShaderCaps* shaderCaps,
+                         GrContextOptions::ShaderErrorHandler* errorHandler,
+                         SkSL::PipelineStageArgs* outArgs);
+#endif
+
+    friend class SkRTShader;            // toByteCode & uniformSize
+    friend class SkRuntimeColorFilter;  //
+
+    // [ByteCode, ErrorText]
+    // If successful, ByteCode != nullptr, otherwise, ErrorText contains the reason for failure.
+    using ByteCodeResult = std::tuple<std::unique_ptr<SkSL::ByteCode>, SkString>;
+    ByteCodeResult toByteCode(const void* inputs) const;
+
+    // Combined size of just the 'uniform' variables.
+    size_t uniformSize() const { return fUniformSize; }
+
+
     uint32_t fHash;
     SkString fSkSL;
 
@@ -164,4 +178,69 @@
     size_t fUniformSize;
 };
 
+/**
+ * SkRuntimeShaderBuilder is a utility to simplify creating SkShader objects from SkRuntimeEffects.
+ *
+ * NOTE: Like SkRuntimeEffect, this API is experimental and subject to change!
+ *
+ * Given an SkRuntimeEffect, the SkRuntimeShaderBuilder manages creating an input data block and
+ * provides named access to the 'in' and 'uniform' variables in that block, as well as named access
+ * to a list of child shader slots. Usage:
+ *
+ *   sk_sp<SkRuntimeEffect> effect = ...;
+ *   SkRuntimeShaderBuilder builder(effect);
+ *   builder.input("some_uniform_float")  = 3.14f;
+ *   builder.input("some_uniform_matrix") = SkM44::Rotate(...);
+ *   builder.child("some_child_effect")   = mySkImage->makeShader(...);
+ *   ...
+ *   sk_sp<SkShader> shader = builder.makeShader(nullptr, false);
+ *
+ * Note that SkRuntimeShaderBuilder is built entirely on the public API of SkRuntimeEffect,
+ * so can be used as-is or serve as inspiration for other interfaces or binding techniques.
+ */
+struct SkRuntimeShaderBuilder {
+    SkRuntimeShaderBuilder(sk_sp<SkRuntimeEffect>);
+    ~SkRuntimeShaderBuilder();
+
+    struct BuilderInput {
+        // Copy 'val' to this variable. No type conversion is performed - 'val' must be same
+        // size as expected by the effect. Information about the variable can be queried by
+        // looking at fVar. If the size is incorrect, no copy will be performed, and debug
+        // builds will abort. If this is the result of querying a missing variable, fVar will
+        // be nullptr, and assigning will also do nothing (and abort in debug builds).
+        template <typename T>
+        std::enable_if_t<std::is_trivially_copyable<T>::value, BuilderInput&> operator=(
+                const T& val) {
+            if (!fVar) {
+                SkDEBUGFAIL("Assigning to missing variable");
+            } else if (sizeof(val) != fVar->sizeInBytes()) {
+                SkDEBUGFAIL("Incorrect value size");
+            } else {
+                memcpy(SkTAddOffset<void>(fOwner->fInputs->writable_data(), fVar->fOffset),
+                        &val, sizeof(val));
+            }
+            return *this;
+        }
+
+        SkRuntimeShaderBuilder*          fOwner;
+        const SkRuntimeEffect::Variable* fVar;    // nullptr if the variable was not found
+    };
+
+    struct BuilderChild {
+        BuilderChild& operator=(const sk_sp<SkShader>& val);
+
+        SkRuntimeShaderBuilder* fOwner;
+        int                     fIndex;  // -1 if the child was not found
+    };
+
+    BuilderInput input(const char* name) { return { this, fEffect->findInput(name) }; }
+    BuilderChild child(const char* name) { return { this, fEffect->findChild(name) }; }
+
+    sk_sp<SkShader> makeShader(const SkMatrix* localMatrix, bool isOpaque);
+
+    sk_sp<SkRuntimeEffect>       fEffect;
+    sk_sp<SkData>                fInputs;
+    std::vector<sk_sp<SkShader>> fChildren;
+};
+
 #endif
diff --git a/samplecode/Sample3D.cpp b/samplecode/Sample3D.cpp
index 7525c48..51e3bae 100644
--- a/samplecode/Sample3D.cpp
+++ b/samplecode/Sample3D.cpp
@@ -93,12 +93,6 @@
 
         canvas->concat(viewport * perspective * camera * inv(viewport));
     }
-
-    SkM44 localToWorld(SkCanvas* canvas) {
-        SkM44 worldToDevice;
-        SkAssertResult(canvas->findMarkedCTM(kLocalToWorld, &worldToDevice));
-        return inv(worldToDevice) * canvas->getLocalToDevice();
-    }
 };
 
 struct Face {
@@ -413,19 +407,16 @@
             return;
         }
 
-        struct Uniforms {
-            SkM44  fLocalToWorld;        // Automatically populated, via layout(marker)
-            SkM44  fLocalToWorldAdjInv;  // Ditto
-            SkV3   fLightPos;
-        } uni;
-        uni.fLightPos = fLight.computeWorldPos(fSphere);
+        SkRuntimeShaderBuilder builder(fEffect);
+        builder.input("lightPos") = fLight.computeWorldPos(fSphere);
+        // localToWorld matrices are automatically populated, via layout(marker)
 
-        sk_sp<SkData> data = SkData::MakeWithCopy(&uni, sizeof(uni));
-        sk_sp<SkShader> children[] = { fImgShader, fBmpShader };
+        builder.child("color_map")  = fImgShader;
+        builder.child("normal_map") = fBmpShader;
 
         SkPaint paint;
         paint.setColor(color);
-        paint.setShader(fEffect->makeShader(data, children, 2, nullptr, true));
+        paint.setShader(builder.makeShader(nullptr, true));
 
         canvas->drawRRect(fRR, paint);
     }
@@ -470,12 +461,14 @@
 
         const char code[] = R"(
             varying float3 vtx_normal;
-            uniform float4x4 localToWorld;
+
+            layout (marker=local_to_world)          uniform float4x4 localToWorld;
+            layout (marker=normals(local_to_world)) uniform float4x4 localToWorldAdjInv;
             uniform float3   lightPos;
 
             void main(float2 p, inout half4 color) {
                 float3 norm = normalize(vtx_normal);
-                float3 plane_norm = normalize(localToWorld * float4(norm, 0)).xyz;
+                float3 plane_norm = normalize(localToWorldAdjInv * float4(norm, 0)).xyz;
 
                 float3 plane_pos = (localToWorld * float4(p, 0, 1)).xyz;
                 float3 light_dir = normalize(lightPos - plane_pos);
@@ -499,18 +492,12 @@
             return;
         }
 
-        struct Uniforms {
-            SkM44  fLocalToWorld;
-            SkV3   fLightPos;
-        } uni;
-        uni.fLocalToWorld = this->localToWorld(canvas);
-        uni.fLightPos = fLight.computeWorldPos(fSphere);
-
-        sk_sp<SkData> data = SkData::MakeWithCopy(&uni, sizeof(uni));
+        SkRuntimeShaderBuilder builder(fEffect);
+        builder.input("lightPos") = fLight.computeWorldPos(fSphere);
 
         SkPaint paint;
         paint.setColor(color);
-        paint.setShader(fEffect->makeShader(data, nullptr, 0, nullptr, true));
+        paint.setShader(builder.makeShader(nullptr, true));
 
         canvas->drawVertices(fVertices, paint);
     }
diff --git a/src/core/SkRuntimeEffect.cpp b/src/core/SkRuntimeEffect.cpp
index 387b31c..b8456b3 100644
--- a/src/core/SkRuntimeEffect.cpp
+++ b/src/core/SkRuntimeEffect.cpp
@@ -30,6 +30,8 @@
 #include "src/gpu/effects/generated/GrMatrixEffect.h"
 #endif
 
+#include <algorithm>
+
 namespace SkSL {
 class SharedCompiler {
 public:
@@ -298,6 +300,18 @@
                                                 fInAndUniformVars.back().sizeInBytes());
 }
 
+const SkRuntimeEffect::Variable* SkRuntimeEffect::findInput(const char* name) const {
+    auto iter = std::find_if(fInAndUniformVars.begin(), fInAndUniformVars.end(),
+                             [name](const Variable& v) { return v.fName.equals(name); });
+    return iter == fInAndUniformVars.end() ? nullptr : &(*iter);
+}
+
+int SkRuntimeEffect::findChild(const char* name) const {
+    auto iter = std::find_if(fChildren.begin(), fChildren.end(),
+                             [name](const SkString& s) { return s.equals(name); });
+    return iter == fChildren.end() ? -1 : static_cast<int>(iter - fChildren.begin());
+}
+
 SkRuntimeEffect::SpecializeResult
 SkRuntimeEffect::specialize(SkSL::Program& baseProgram,
                             const void* inputs,
@@ -1082,3 +1096,24 @@
     SK_REGISTER_FLATTENABLE(SkRuntimeColorFilter);
     SK_REGISTER_FLATTENABLE(SkRTShader);
 }
+
+SkRuntimeShaderBuilder::SkRuntimeShaderBuilder(sk_sp<SkRuntimeEffect> effect)
+    : fEffect(std::move(effect))
+    , fInputs(SkData::MakeUninitialized(fEffect->inputSize()))
+    , fChildren(fEffect->children().count()) {}
+
+SkRuntimeShaderBuilder::~SkRuntimeShaderBuilder() = default;
+
+sk_sp<SkShader> SkRuntimeShaderBuilder::makeShader(const SkMatrix* localMatrix, bool isOpaque) {
+    return fEffect->makeShader(fInputs, fChildren.data(), fChildren.size(), localMatrix, isOpaque);
+}
+
+SkRuntimeShaderBuilder::BuilderChild&
+SkRuntimeShaderBuilder::BuilderChild::operator=(const sk_sp<SkShader>& val) {
+    if (fIndex < 0) {
+        SkDEBUGFAIL("Assigning to missing child");
+    } else {
+        fOwner->fChildren[fIndex] = val;
+    }
+    return *this;
+}
diff --git a/tests/SkRuntimeEffectTest.cpp b/tests/SkRuntimeEffectTest.cpp
index ef4f7ad..d28b3e4 100644
--- a/tests/SkRuntimeEffectTest.cpp
+++ b/tests/SkRuntimeEffectTest.cpp
@@ -10,6 +10,7 @@
 #include "include/core/SkSurface.h"
 #include "include/effects/SkRuntimeEffect.h"
 #include "include/gpu/GrContext.h"
+#include "src/core/SkTLazy.h"
 #include "tests/Test.h"
 
 #include <algorithm>
@@ -67,33 +68,16 @@
             return;
         }
 
-        fEffect = std::move(effect);
-        fInputs = SkData::MakeUninitialized(fEffect->inputSize());
+        fBuilder.init(std::move(effect));
     }
 
-    struct InputVar {
-        template <typename T> InputVar& operator=(const T& val) {
-            SkASSERT(sizeof(T) == fVar.sizeInBytes());
-            memcpy(SkTAddOffset<void>(fOwner->fInputs->writable_data(), fVar.fOffset), &val,
-                   sizeof(T));
-            return *this;
-        }
-        TestEffect* fOwner;
-        const SkRuntimeEffect::Variable& fVar;
-    };
-
-    InputVar operator[](const char* name) {
-        auto input = std::find_if(fEffect->inputs().begin(), fEffect->inputs().end(),
-                                  [name](const auto& v) { return v.fName.equals(name); });
-        SkASSERT(input != fEffect->inputs().end());
-        return {this, *input};
+    SkRuntimeShaderBuilder::BuilderInput operator[](const char* name) {
+        return fBuilder->input(name);
     }
 
     void test(skiatest::Reporter* r, sk_sp<SkSurface> surface,
               uint32_t TL, uint32_t TR, uint32_t BL, uint32_t BR) {
-        if (!fEffect) { return; }
-
-        auto shader = fEffect->makeShader(fInputs, nullptr, 0, nullptr, false);
+        auto shader = fBuilder->makeShader(nullptr, false);
         if (!shader) {
             REPORT_FAILURE(r, "shader", SkString("Effect didn't produce a shader"));
             return;
@@ -119,7 +103,7 @@
                                           "Got     : [ %08x %08x %08x %08x ]\n"
                                           "SkSL:\n%s\n",
                                           TL, TR, BL, BR, actual[0], actual[1], actual[2],
-                                          actual[3], fEffect->source().c_str()));
+                                          actual[3], fBuilder->fEffect->source().c_str()));
         }
     }
 
@@ -128,8 +112,7 @@
     }
 
 private:
-    sk_sp<SkRuntimeEffect> fEffect;
-    sk_sp<SkData> fInputs;
+    SkTLazy<SkRuntimeShaderBuilder> fBuilder;
 };
 
 static void test_RuntimeEffect_Shaders(skiatest::Reporter* r, GrContext* context) {