Separate SkVM's SampleChild into two callbacks.

Previously, one callback was used and the caller needed to infer from
the passed-in index what type of effect to invoke, and which arguments
were valid.

Now, each type of effect has a separate callback with a unique
signature. (In particular, SkColorFilters lack coordinates entirely.)
In a later CL, SkBlenders will also be supported, and those will have
a different signature again, as they take two colors and no coords.

Change-Id: Ibcb39c91c35256c5339e4d2790ed0c36e434f191
Bug: skia:12257
Reviewed-on: https://skia-review.googlesource.com/c/skia/+/431656
Reviewed-by: Brian Osman <brianosman@google.com>
Auto-Submit: John Stiles <johnstiles@google.com>
Commit-Queue: Brian Osman <brianosman@google.com>
diff --git a/src/core/SkRuntimeEffect.cpp b/src/core/SkRuntimeEffect.cpp
index a69e4b4..58e7415 100644
--- a/src/core/SkRuntimeEffect.cpp
+++ b/src/core/SkRuntimeEffect.cpp
@@ -586,7 +586,7 @@
         return x.r.id == y.r.id && x.g.id == y.g.id && x.b.id == y.b.id && x.a.id == y.a.id;
     };
     bool allSampleCallsSupported = true;
-    auto sampleChild = [&](int ix, skvm::Coord, skvm::Color c) {
+    auto sampleColorFilter = [&](int ix, skvm::Color c) {
         skvm::Color result = p.uniformColor(/*placeholder*/ SkColors::kWhite, &childColorUniforms);
         SkFilterColorProgram::SampleCall call;
         call.fChild = ix;
@@ -623,7 +623,8 @@
                                              /*local=*/zeroCoord,
                                              inputColor,
                                              inputColor,
-                                             sampleChild);
+                                             /*sampleShader=*/nullptr,
+                                             sampleColorFilter);
 
     // Then store the result to the *third* arg ptr
     p.store({skvm::PixelFormat::FLOAT, 32, 32, 32, 32, 0, 32, 64, 96}, p.arg(16), result);
@@ -808,27 +809,30 @@
         sk_sp<SkData> inputs = get_xformed_uniforms(fEffect.get(), fUniforms, dst.colorSpace());
         SkASSERT(inputs);
 
-        // There should be no way for the color filter to use device coords, but we need to supply
-        // something. (Uninitialized values can trigger asserts in skvm::Builder).
-        skvm::Coord zeroCoord = { p->splat(0.0f), p->splat(0.0f) };
-
-        auto sampleChild = [&](int ix, skvm::Coord coord, skvm::Color color) {
-            if (fChildren[ix].shader) {
+        auto sampleShader = [&](int ix, skvm::Coord coord, skvm::Color color) {
+            if (SkShader* shader = fChildren[ix].shader.get()) {
                 SkSimpleMatrixProvider mats{SkMatrix::I()};
-                return as_SB(fChildren[ix].shader)
-                        ->program(p, coord, coord, color, mats, nullptr, dst, uniforms, alloc);
-            } else if (fChildren[ix].colorFilter) {
-                return as_CFB(fChildren[ix].colorFilter)->program(p, color, dst, uniforms, alloc);
-            } else {
-                return color;
+                return as_SB(shader)->program(p, coord, coord, color, mats, /*localM=*/nullptr,
+                                              dst, uniforms, alloc);
             }
+            return color;
+        };
+        auto sampleColorFilter = [&](int ix, skvm::Color color) {
+            if (SkColorFilter* colorFilter = fChildren[ix].colorFilter.get()) {
+                return as_CFB(colorFilter)->program(p, color, dst, uniforms, alloc);
+            }
+            return color;
         };
 
         std::vector<skvm::Val> uniform = make_skvm_uniforms(p, uniforms, fEffect->uniformSize(),
                                                             *inputs);
 
+        // There should be no way for the color filter to use device coords, but we need to supply
+        // something. (Uninitialized values can trigger asserts in skvm::Builder).
+        skvm::Coord zeroCoord = { p->splat(0.0f), p->splat(0.0f) };
         return SkSL::ProgramToSkVM(*fEffect->fBaseProgram, fEffect->fMain, p, SkMakeSpan(uniform),
-                                   /*device=*/zeroCoord, /*local=*/zeroCoord, c, c, sampleChild);
+                                   /*device=*/zeroCoord, /*local=*/zeroCoord, c, c, sampleShader,
+                                   sampleColorFilter);
     }
 
     SkPMColor4f onFilterColor4f(const SkPMColor4f& color, SkColorSpace* dstCS) const override {
@@ -993,23 +997,26 @@
         }
         local = SkShaderBase::ApplyMatrix(p,inv,local,uniforms);
 
-        auto sampleChild = [&](int ix, skvm::Coord coord, skvm::Color color) {
-            if (fChildren[ix].shader) {
+        auto sampleShader = [&](int ix, skvm::Coord coord, skvm::Color color) {
+            if (SkShader* shader = fChildren[ix].shader.get()) {
                 SkOverrideDeviceMatrixProvider mats{matrices, SkMatrix::I()};
-                return as_SB(fChildren[ix].shader)
-                        ->program(p, device, coord, color, mats, nullptr, dst, uniforms, alloc);
-            } else if (fChildren[ix].colorFilter) {
-                return as_CFB(fChildren[ix].colorFilter)->program(p, color, dst, uniforms, alloc);
-            } else {
-                return color;
+                return as_SB(shader)->program(p, device, coord, color, mats, /*localM=*/nullptr,
+                                              dst, uniforms, alloc);
             }
+            return color;
+        };
+        auto sampleColorFilter = [&](int ix, skvm::Color color) {
+            if (SkColorFilter* colorFilter = fChildren[ix].colorFilter.get()) {
+                return as_CFB(colorFilter)->program(p, color, dst, uniforms, alloc);
+            }
+            return color;
         };
 
         std::vector<skvm::Val> uniform = make_skvm_uniforms(p, uniforms, fEffect->uniformSize(),
                                                             *inputs);
 
         return SkSL::ProgramToSkVM(*fEffect->fBaseProgram, fEffect->fMain, p, SkMakeSpan(uniform),
-                                   device, local, paint, paint, sampleChild);
+                                   device, local, paint, paint, sampleShader, sampleColorFilter);
     }
 
     void flatten(SkWriteBuffer& buffer) const override {
@@ -1110,15 +1117,17 @@
                                                     colorInfo.colorSpace());
         SkASSERT(inputs);
 
-        auto sampleChild = [&](int index, skvm::Coord coord, skvm::Color color) {
-            const SkRuntimeEffect::ChildPtr& effect = fChildren[index];
-            if (effect.shader) {
+        auto sampleShader = [&](int ix, skvm::Coord coord, skvm::Color color) {
+            if (SkShader* shader = fChildren[ix].shader.get()) {
                 SkSimpleMatrixProvider mats{SkMatrix::I()};
-                return as_SB(effect.shader)->program(p, coord, coord, color, mats,
-                                                    /*localM=*/nullptr, colorInfo, uniforms, alloc);
+                return as_SB(shader)->program(p, coord, coord, color, mats, /*localM=*/nullptr,
+                                              colorInfo, uniforms, alloc);
             }
-            if (effect.colorFilter) {
-                return as_CFB(effect.colorFilter)->program(p, color, colorInfo, uniforms, alloc);
+            return color;
+        };
+        auto sampleColorFilter = [&](int ix, skvm::Color color) {
+            if (SkColorFilter* colorFilter = fChildren[ix].colorFilter.get()) {
+                return as_CFB(colorFilter)->program(p, color, colorInfo, uniforms, alloc);
             }
             return color;
         };
@@ -1129,7 +1138,8 @@
         // Emit the blend function as an SkVM program.
         skvm::Coord zeroCoord = {p->splat(0.0f), p->splat(0.0f)};
         return SkSL::ProgramToSkVM(*fEffect->fBaseProgram, fEffect->fMain, p, SkMakeSpan(uniform),
-                                  /*device=*/zeroCoord, /*local=*/zeroCoord, src, dst, sampleChild);
+                                  /*device=*/zeroCoord, /*local=*/zeroCoord, src, dst,
+                                  sampleShader, sampleColorFilter);
     }
 
 #if SK_SUPPORT_GPU
diff --git a/src/sksl/codegen/SkSLVMCodeGenerator.cpp b/src/sksl/codegen/SkSLVMCodeGenerator.cpp
index 82cd65b..436c290 100644
--- a/src/sksl/codegen/SkSLVMCodeGenerator.cpp
+++ b/src/sksl/codegen/SkSLVMCodeGenerator.cpp
@@ -117,7 +117,8 @@
                   skvm::Coord device,
                   skvm::Coord local,
                   skvm::Color inputColor,
-                  SampleChildFn sampleChild);
+                  SampleShaderFn sampleShader,
+                  SampleColorFilterFn sampleColorFilter);
 
     void writeFunction(const FunctionDefinition& function,
                        SkSpan<skvm::Val> arguments,
@@ -219,7 +220,8 @@
 
     const skvm::Coord fLocalCoord;
     const skvm::Color fInputColor;
-    const SampleChildFn fSampleChild;
+    const SampleShaderFn fSampleShader;
+    const SampleColorFilterFn fSampleColorFilter;
 
     // [Variable, first slot in fSlots]
     std::unordered_map<const Variable*, size_t> fVariableMap;
@@ -277,12 +279,14 @@
                              skvm::Coord device,
                              skvm::Coord local,
                              skvm::Color inputColor,
-                             SampleChildFn sampleChild)
+                             SampleShaderFn sampleShader,
+                             SampleColorFilterFn sampleColorFilter)
         : fProgram(program)
         , fBuilder(builder)
         , fLocalCoord(local)
         , fInputColor(inputColor)
-        , fSampleChild(std::move(sampleChild)) {
+        , fSampleShader(std::move(sampleShader))
+        , fSampleColorFilter(std::move(sampleColorFilter)) {
     fConditionMask = fLoopMask = fBuilder->splat(0xffff'ffff);
 
     // Now, add storage for each global variable (including uniforms) to fSlots, and entries in
@@ -297,7 +301,7 @@
             SkASSERT(fVariableMap.find(&var) == fVariableMap.end());
 
             // For most variables, fVariableMap stores an index into fSlots, but for children,
-            // fVariableMap stores the index to pass to fSampleChild().
+            // fVariableMap stores the index to pass to fSample(Shader|ColorFilter)
             if (var.type().isEffectChild()) {
                 fVariableMap[&var] = fpCount++;
                 continue;
@@ -879,22 +883,22 @@
 
         // Shaders require a coordinate argument. Color filters require a color argument.
         // When we call sampleChild, the other value remains the incoming default.
-        skvm::Color inColor = fInputColor;
-        skvm::Coord coord = fLocalCoord;
         const Expression* arg = c.arguments()[1].get();
         Value argVal = this->writeExpression(*arg);
+        skvm::Color color;
 
         if (child->type().typeKind() == Type::TypeKind::kShader) {
             SkASSERT(arg->type() == *fProgram.fContext->fTypes.fFloat2);
-            coord = {f32(argVal[0]), f32(argVal[1])};
+            skvm::Coord coord = {f32(argVal[0]), f32(argVal[1])};
+            color = fSampleShader(fp_it->second, coord, fInputColor);
         } else {
             SkASSERT(child->type().typeKind() == Type::TypeKind::kColorFilter);
             SkASSERT(arg->type() == *fProgram.fContext->fTypes.fHalf4 ||
                      arg->type() == *fProgram.fContext->fTypes.fFloat4);
-            inColor = {f32(argVal[0]), f32(argVal[1]), f32(argVal[2]), f32(argVal[3])};
+            skvm::Color inColor = {f32(argVal[0]), f32(argVal[1]), f32(argVal[2]), f32(argVal[3])};
+            color = fSampleColorFilter(fp_it->second, inColor);
         }
 
-        skvm::Color color = fSampleChild(fp_it->second, coord, inColor);
         Value result(4);
         result[0] = color.r;
         result[1] = color.g;
@@ -1547,7 +1551,8 @@
                           skvm::Coord local,
                           skvm::Color inputColor,
                           skvm::Color destColor,
-                          SampleChildFn sampleChild) {
+                          SampleShaderFn sampleShader,
+                          SampleColorFilterFn sampleColorFilter) {
     skvm::Val zero = builder->splat(0.0f).id;
     skvm::Val result[4] = {zero,zero,zero,zero};
 
@@ -1584,8 +1589,8 @@
     }
     SkASSERT(argSlots <= SK_ARRAY_COUNT(args));
 
-    SkVMGenerator generator(
-            program, builder, uniforms, device, local, inputColor, std::move(sampleChild));
+    SkVMGenerator generator(program, builder, uniforms, device, local, inputColor,
+                            std::move(sampleShader), std::move(sampleColorFilter));
     generator.writeFunction(function, {args, argSlots}, SkMakeSpan(result));
 
     return skvm::Color{{builder, result[0]},
@@ -1627,7 +1632,8 @@
     skvm::Coord zeroCoord = {zero, zero};
     skvm::Color zeroColor = {zero, zero, zero, zero};
     SkVMGenerator generator(program, b, uniforms, /*device=*/zeroCoord, /*local=*/zeroCoord,
-                            /*inputColor=*/zeroColor, /*sampleChild=*/{});
+                            /*inputColor=*/zeroColor, /*sampleShader=*/nullptr,
+                            /*sampleColorFilter=*/nullptr);
     generator.writeFunction(function, SkMakeSpan(argVals), SkMakeSpan(returnVals));
 
     // generateCode has updated the contents of 'argVals' for any 'out' or 'inout' parameters.
@@ -1744,7 +1750,7 @@
         children.push_back({uniforms.pushPtr(nullptr), builder->uniform32(uniforms.push(0))});
     }
 
-    auto sampleChild = [&](int i, skvm::Coord coord, skvm::Color) {
+    auto sampleShader = [&](int i, skvm::Coord coord, skvm::Color) {
         skvm::PixelFormat pixelFormat = skvm::SkColorType_to_PixelFormat(kRGBA_F32_SkColorType);
         skvm::I32 index  = trunc(coord.x);
                   index += trunc(coord.y) * children[i].rowBytesAsPixels;
@@ -1760,7 +1766,8 @@
     skvm::Color destColor = builder->uniformColor(SkColors::kBlack, &uniforms);
 
     skvm::Color result = SkSL::ProgramToSkVM(program, *main, builder, SkMakeSpan(uniformVals),
-                                             device, local, inColor, destColor, sampleChild);
+                                             device, local, inColor, destColor, sampleShader,
+                                             /*sampleColorFilter=*/nullptr);
 
     storeF(builder->varying<float>(), result.r);
     storeF(builder->varying<float>(), result.g);
diff --git a/src/sksl/codegen/SkSLVMCodeGenerator.h b/src/sksl/codegen/SkSLVMCodeGenerator.h
index 76a63df..9017e6b 100644
--- a/src/sksl/codegen/SkSLVMCodeGenerator.h
+++ b/src/sksl/codegen/SkSLVMCodeGenerator.h
@@ -20,7 +20,8 @@
 class FunctionDefinition;
 struct Program;
 
-using SampleChildFn = std::function<skvm::Color(int, skvm::Coord, skvm::Color)>;
+using SampleShaderFn = std::function<skvm::Color(int, skvm::Coord, skvm::Color)>;
+using SampleColorFilterFn = std::function<skvm::Color(int, skvm::Color)>;
 
 // Convert 'function' to skvm instructions in 'builder', for use by blends, shaders, & color filters
 skvm::Color ProgramToSkVM(const Program& program,
@@ -31,7 +32,8 @@
                           skvm::Coord local,
                           skvm::Color inputColor,
                           skvm::Color destColor,
-                          SampleChildFn sampleChild);
+                          SampleShaderFn sampleShader,
+                          SampleColorFilterFn sampleColorFilter);
 
 struct SkVMSignature {
     size_t fParameterSlots = 0;