Support input color argument to process() function in sksl .fp files

--

This expands sksl's capabilities with .fp files. Previously, it was possible to declare "in fragmentProcessor foo" and emit it automatically when "process(foo);" was called. This adds a variant of process that takes a second argument, which must be a half4 expression. This argument specifies the value, or dynamic expression calculated earlier in the parent shader, to use as sk_InColor by the child.

The CL is longer than anticipated because of properly handling dependencies between previous sksl statements and the input to the child. The original writeEmitCode() collected all extra emission code (the calls to build->emitChild) and put them before any call to codeAppendf. This makes it impossible to use a parent's variable, or the output of another child, as the input for process.

To solve this, there is now a flushEmittedCode() function that takes over the logic of outputting the extra emission code and the necessary codeAppendf calls. When invoked, it (by default) only appends completed sksl statements, and places any current expression back at the beginning of the output stream. It now updates fFormatArgs and fExtraEmitCodeCode as it consumes their contents. This allows writeFunctionCall() for a call to "process" to flush all previous statements before it adds its emit child code.

Bug: skia:
Change-Id: I63c41af6f3e0620aa890d10d14436ee6244f0051
Reviewed-on: https://skia-review.googlesource.com/148395
Commit-Queue: Michael Ludwig <michaelludwig@google.com>
Reviewed-by: Ethan Nicholas <ethannicholas@google.com>
diff --git a/src/sksl/README b/src/sksl/README
index 5062ed9..efe9be6 100644
--- a/src/sksl/README
+++ b/src/sksl/README
@@ -136,6 +136,11 @@
   the program's key. Matrix variables additionally support 'key=identity',
   which causes the key to consider only whether or not the matrix is an
   identity matrix.
+* child processors can be declared with 'in fragmentProcessor <name>;', and can
+  be invoked by calling 'process(<name>)' or 'process(<name>, <inputColor>)'.
+  The first variant emits the child with a solid white input color. The second
+  variant emits the child with the result of the 2nd argument's expression,
+  which must evaluate to a half4. The process function returns a half4.
 
 
 Creating a new .fp file
diff --git a/src/sksl/SkSLCPPCodeGenerator.cpp b/src/sksl/SkSLCPPCodeGenerator.cpp
index 7e8f492..548089f 100644
--- a/src/sksl/SkSLCPPCodeGenerator.cpp
+++ b/src/sksl/SkSLCPPCodeGenerator.cpp
@@ -10,6 +10,8 @@
 #include "SkSLCompiler.h"
 #include "SkSLHCodeGenerator.h"
 
+#include <algorithm>
+
 namespace SkSL {
 
 static bool needs_uniform_var(const Variable& var) {
@@ -309,8 +311,22 @@
 
 void CPPCodeGenerator::writeFunctionCall(const FunctionCall& c) {
     if (c.fFunction.fBuiltin && c.fFunction.fName == "process") {
-        SkASSERT(c.fArguments.size() == 1);
-        SkASSERT(Expression::kVariableReference_Kind == c.fArguments[0]->fKind);
+        // Sanity checks that are detected by function definition in sksl_fp.inc
+        SkASSERT(c.fArguments.size() == 1 || c.fArguments.size() == 2);
+        SkASSERT("fragmentProcessor" == c.fArguments[0]->fType.name());
+
+        // Actually fail during compilation if arguments with valid types are
+        // provided that are not variable references, since process() is a
+        // special function that impacts code emission.
+        if (c.fArguments[0]->fKind != Expression::kVariableReference_Kind) {
+            fErrors.error(c.fArguments[0]->fOffset,
+                    "process()'s fragmentProcessor argument must be a variable reference\n");
+            return;
+        }
+        if (c.fArguments.size() > 1) {
+            // Second argument must also be a half4 expression
+            SkASSERT("half4" == c.fArguments[1]->fType.name());
+        }
         int index = 0;
         bool found = false;
         for (const auto& p : fProgram) {
@@ -330,10 +346,35 @@
             }
         }
         SkASSERT(found);
+
+        // Flush all previous statements so that this emitted child can
+        // depend upon any declared variables within that section
+        this->flushEmittedCode();
+
+        // Set to the empty string when no input color parameter should be emitted,
+        // which means this must be properly formatted with a prefixed comma
+        // when the parameter should be inserted into the emitChild() parameter list.
+        String inputArg;
+        if (c.fArguments.size() > 1) {
+            SkASSERT(c.fArguments.size() == 2);
+            // Use the emitChild() variant that accepts an input color, so convert
+            // the 2nd argument's expression into C++ code that produces sksl
+            // stored in an SkString.
+
+            String inputName = "_input" + to_string(index);
+            fExtraEmitCodeCode += "        " + convertSKSLExpressionToCPP(
+                    *c.fArguments[1], inputName);
+
+            // emitChild() needs a char*
+            inputArg = ", " + inputName + ".c_str()";
+        }
+
+        // Write the output handling after the possible input handling
         String childName = "_child" + to_string(index);
-        fExtraEmitCodeCode += "        SkString " + childName + "(\"" + childName + "\");\n" +
-                              "        this->emitChild(" + to_string(index) + ", &" + childName +
-                              ", args);\n";
+        fExtraEmitCodeCode += "        SkString " + childName + "(\"" + childName + "\");\n";
+        fExtraEmitCodeCode += "        this->emitChild(" + to_string(index) + inputArg +
+                              ", &" + childName + ", args);\n";
+
         this->write("%s");
         fFormatArgs.push_back(childName + ".c_str()");
         return;
@@ -503,6 +544,52 @@
            Type::kOther_Kind != var.fType.kind();
 }
 
+void CPPCodeGenerator::flushEmittedCode(bool flushAll) {
+    if (fCPPBuffer == nullptr) {
+        // Not actually within writeEmitCode() so nothing to flush
+        return;
+    }
+
+    StringStream* skslBuffer = static_cast<StringStream*>(fOut);
+
+    String sksl = skslBuffer->str();
+    // Empty the accumulation buffer; if there are any partial statements in
+    // the extracted sksl string they will be re-added later
+    skslBuffer->reset();
+
+    if (!flushAll) {
+        // Find the last ';', '{', or '}' and leave everything after that in the buffer
+        int lastStatementEnd = sksl.findLastOf(';');
+        int lastBlockOpen = sksl.findLastOf('{');
+        int lastBlockEnd = sksl.findLastOf('}');
+
+        int flushPoint = std::max(lastStatementEnd, std::max(lastBlockEnd, lastBlockOpen));
+
+        // NOTE: this does the right thing when flushPoint = -1
+        if (flushPoint < (int) sksl.size() - 1) {
+            // There is partial sksl content that can't be flushed yet so put
+            // that back into the sksl buffer and remove it from the string
+            skslBuffer->writeText(sksl.c_str() + flushPoint + 1);
+            sksl = String(sksl.c_str(), flushPoint + 1);
+        }
+    }
+
+    // Switch back to the CPP buffer for the actual code appending statements
+    fOut = fCPPBuffer;
+    if (fExtraEmitCodeCode.size() > 0) {
+        this->writef("%s", fExtraEmitCodeCode.c_str());
+        fExtraEmitCodeCode.reset();
+    }
+    // writeCodeAppend automatically removes the format args that it consumed,
+    // so fFormatArgs will be in a valid state for any partial statements left
+    // in the sksl buffer.
+    this->writeCodeAppend(sksl);
+
+    // After appending, switch back to an sksl buffer that contains any
+    // remaining partial statements that couldn't be appended
+    fOut = skslBuffer;
+}
+
 void CPPCodeGenerator::writeCodeAppend(const String& code) {
     // codeAppendf can only handle appending 1024 bytes at a time, so we need to break the string
     // into chunks. Unfortunately we can't tell exactly how long the string is going to end up,
@@ -539,6 +626,51 @@
         argStart += argCount;
         start = index;
     }
+
+    // argStart is equal to the number of fFormatArgs that were consumed
+    // so they should be removed from the list
+    if (argStart > 0) {
+        fFormatArgs.erase(fFormatArgs.begin(), fFormatArgs.begin() + argStart);
+    }
+}
+
+String CPPCodeGenerator::convertSKSLExpressionToCPP(const Expression& e,
+                                                    const String& cppVar) {
+    // To do this conversion, we temporarily switch the sksl output stream
+    // to an empty stringstream and reset the format args to empty.
+    OutputStream* oldSKSL = fOut;
+    StringStream exprBuffer;
+    fOut = &exprBuffer;
+
+    std::vector<String> oldArgs(fFormatArgs);
+    fFormatArgs.clear();
+
+    // Convert the argument expression into a format string and args
+    this->writeExpression(e, Precedence::kTopLevel_Precedence);
+    std::vector<String> newArgs(fFormatArgs);
+    String exprFormat = exprBuffer.str();
+
+    // After generating, restore the original output stream and format args
+    fFormatArgs = oldArgs;
+    fOut = oldSKSL;
+
+    // Now build the final C++ code snippet from the format string and args
+    String cppExpr;
+    if (newArgs.size() == 0) {
+        // This was a static expression, so we can simplify the input
+        // color declaration in the emitted code to just a static string
+        cppExpr = "SkString " + cppVar + "(\"" + exprFormat + "\");\n";
+    } else {
+        // String formatting must occur dynamically, so have the C++ declaration
+        // use SkStringPrintf with the format args that were accumulated
+        // when the expression was written.
+        cppExpr = "SkString " + cppVar + " = SkStringPrintf(\"" + exprFormat + "\"";
+        for (size_t i = 0; i < newArgs.size(); i++) {
+            cppExpr += ", " + newArgs[i];
+        }
+        cppExpr += ");\n";
+    }
+    return cppExpr;
 }
 
 bool CPPCodeGenerator::writeEmitCode(std::vector<const Variable*>& uniforms) {
@@ -568,13 +700,19 @@
         this->addUniform(*u);
     }
     this->writeSection(EMIT_CODE_SECTION);
-    OutputStream* old = fOut;
-    StringStream mainBuffer;
-    fOut = &mainBuffer;
+
+    // Save original buffer as the CPP buffer for flushEmittedCode()
+    fCPPBuffer = fOut;
+    StringStream skslBuffer;
+    fOut = &skslBuffer;
+
     bool result = INHERITED::generateCode();
-    fOut = old;
-    this->writef("%s", fExtraEmitCodeCode.c_str());
-    this->writeCodeAppend(mainBuffer.str());
+    // Final flush in case there is anything extra, forcing it to emit everything
+    this->flushEmittedCode(true);
+
+    // Then restore the original CPP buffer and close the function
+    fOut = fCPPBuffer;
+    fCPPBuffer = nullptr;
     this->write("    }\n");
     return result;
 }
diff --git a/src/sksl/SkSLCPPCodeGenerator.h b/src/sksl/SkSLCPPCodeGenerator.h
index 39255d6..6c67bad 100644
--- a/src/sksl/SkSLCPPCodeGenerator.h
+++ b/src/sksl/SkSLCPPCodeGenerator.h
@@ -23,6 +23,21 @@
     bool generateCode() override;
 
 private:
+    // When inside writeEmitCode(), certain SkSL elements need to control
+    // when fragBuilder->codeAppendf is added to the function block. This
+    // takes all completed statements in the SkSL buffer, and their corresponding
+    // format args, and writes them into the emitCode()'s statement block
+    // using writeCodeAppend().
+    //
+    // This control is necessary for handling special functions in SkSL, like
+    // process(), which need to intermix the current FP's SkSL with that of
+    // an emitted child.
+    //
+    //  :forceAll - If false, only the completed statements (terminated by ;),
+    //     will be flushed and the sksl buffer will be set to any partial
+    //     statements that remain. If true, everything is flushed, regardless.
+    void flushEmittedCode(bool forceAll = false);
+
     void writef(const char* s, va_list va) SKSL_PRINTF_LIKE(2, 0);
 
     void writef(const char* s, ...) SKSL_PRINTF_LIKE(2, 3);
@@ -88,6 +103,15 @@
 
     void writeTest();
 
+    // If the returned C++ is included in the generated code, then the variable
+    // name stored in cppVar will refer to a valid SkString that matches the
+    // Expression. Successful returns leave the output buffer (and related state)
+    // unmodified.
+    //
+    // In the simplest cases, this will return "SkString {cppVar}(\"{e}\");",
+    // while more advanced cases will properly insert format arguments.
+    String convertSKSLExpressionToCPP(const Expression& e, const String& cppVar);
+
     String fName;
     String fFullName;
     SectionAndParameterHelper fSectionAndParameterHelper;
@@ -98,6 +122,11 @@
     bool fCPPMode = false;
     bool fInMain = false;
 
+    // if not null, we are accumulating SkSL for emitCode into fOut, which
+    // replaced the original buffer with a StringStream. The original buffer is
+    // stored here for restoration.
+    OutputStream* fCPPBuffer = nullptr;
+
     typedef GLSLCodeGenerator INHERITED;
 };
 
diff --git a/src/sksl/sksl_fp.inc b/src/sksl/sksl_fp.inc
index c5edc00..7a3963d 100644
--- a/src/sksl/sksl_fp.inc
+++ b/src/sksl/sksl_fp.inc
@@ -24,5 +24,5 @@
 layout(builtin=10012) half sk_Height;
 
 half4 process(fragmentProcessor fp);
-
+half4 process(fragmentProcessor fp, half4 input);
 )
diff --git a/tests/SkSLFPTest.cpp b/tests/SkSLFPTest.cpp
index 4208246..20a17bb 100644
--- a/tests/SkSLFPTest.cpp
+++ b/tests/SkSLFPTest.cpp
@@ -381,3 +381,72 @@
             "this->registerChildProcessor(src.childProcessor(1).clone());"
          });
 }
+
+DEF_TEST(SkSLFPChildProcessorsWithInput, r) {
+    test(r,
+         "in fragmentProcessor child1;"
+         "in fragmentProcessor child2;"
+         "void main() {"
+         "    half4 childIn = sk_InColor;"
+         "    half4 childOut1 = process(child1, childIn);"
+         "    half4 childOut2 = process(child2, childOut1);"
+         "    sk_OutColor = childOut2;"
+         "}",
+         *SkSL::ShaderCapsFactory::Default(),
+         {
+            "this->registerChildProcessor(std::move(child1));",
+            "this->registerChildProcessor(std::move(child2));"
+         },
+         {
+            "SkString _input0(\"childIn\");",
+            "SkString _child0(\"_child0\");",
+            "this->emitChild(0, _input0.c_str(), &_child0, args);",
+            "SkString _input1(\"childOut1\");",
+            "SkString _child1(\"_child1\");",
+            "this->emitChild(1, _input1.c_str(), &_child1, args);",
+            "this->registerChildProcessor(src.childProcessor(0).clone());",
+            "this->registerChildProcessor(src.childProcessor(1).clone());"
+         });
+}
+
+DEF_TEST(SkSLFPChildProcessorWithInputExpression, r) {
+    test(r,
+         "in fragmentProcessor child;"
+         "void main() {"
+         "    sk_OutColor = process(child, sk_InColor * half4(0.5));"
+         "}",
+         *SkSL::ShaderCapsFactory::Default(),
+         {
+            "this->registerChildProcessor(std::move(child));",
+         },
+         {
+            "SkString _input0 = SkStringPrintf(\"%s * half4(0.5)\", args.fInputColor);",
+            "SkString _child0(\"_child0\");",
+            "this->emitChild(0, _input0.c_str(), &_child0, args);",
+            "this->registerChildProcessor(src.childProcessor(0).clone());",
+         });
+}
+
+DEF_TEST(SkSLFPNestedChildProcessors, r) {
+    test(r,
+         "in fragmentProcessor child1;"
+         "in fragmentProcessor child2;"
+         "void main() {"
+         "    sk_OutColor = process(child2, sk_InColor * process(child1, sk_InColor * half4(0.5)));"
+         "}",
+         *SkSL::ShaderCapsFactory::Default(),
+         {
+            "this->registerChildProcessor(std::move(child1));",
+            "this->registerChildProcessor(std::move(child2));"
+         },
+         {
+            "SkString _input0 = SkStringPrintf(\"%s * half4(0.5)\", args.fInputColor);",
+            "SkString _child0(\"_child0\");",
+            "this->emitChild(0, _input0.c_str(), &_child0, args);",
+            "SkString _input1 = SkStringPrintf(\"%s * %s\", args.fInputColor, _child0.c_str());",
+            "SkString _child1(\"_child1\");",
+            "this->emitChild(1, _input1.c_str(), &_child1, args);",
+            "this->registerChildProcessor(src.childProcessor(0).clone());",
+            "this->registerChildProcessor(src.childProcessor(1).clone());"
+         });
+}