Fixed some fixups not being applied to DSL code

Change-Id: I11c473a7a683282d6bd4f73d2c90f23de63f8cc7
Reviewed-on: https://skia-review.googlesource.com/c/skia/+/426557
Reviewed-by: Brian Osman <brianosman@google.com>
Reviewed-by: John Stiles <johnstiles@google.com>
Commit-Queue: Ethan Nicholas <ethannicholas@google.com>

include/sksl/DSLBlock.h

@@ -16,8 +16,8 @@
 
 namespace SkSL {
 
+class Block;
 class SymbolTable;
-class Statement;
 
 namespace dsl {
 
@@ -50,7 +50,7 @@ public:
 
     void append(DSLStatement stmt);
 
-    std::unique_ptr<SkSL::Statement> release();
+    std::unique_ptr<SkSL::Block> release();
 
 private:
     SkSL::StatementArray fStatements;

include/sksl/DSLCore.h

@@ -63,6 +63,8 @@ DSLVar sk_FragColor();
 
 DSLVar sk_FragCoord();
 
+DSLExpression sk_Position();
+
 /**
  * break;
  */

src/sksl/SkSLIRGenerator.cpp

@@ -771,7 +771,8 @@ void IRGenerator::CheckModifiers(const Context& context,
     SkASSERT(layoutFlags == 0);
 }
 
-void IRGenerator::finalizeFunction(const FunctionDeclaration& funcDecl, Statement* body) {
+std::unique_ptr<Block> IRGenerator::finalizeFunction(const FunctionDeclaration& funcDecl,
+                                                     std::unique_ptr<Block> body) {
     class Finalizer : public ProgramWriter {
     public:
         Finalizer(IRGenerator* irGenerator, const FunctionDeclaration* function)
@@ -869,6 +870,16 @@ void IRGenerator::finalizeFunction(const FunctionDeclaration& funcDecl, Statemen
         using INHERITED = ProgramWriter;
     };
 
+    bool isMain = funcDecl.isMain();
+    bool needInvocationIDWorkaround = fInvocations != -1 && isMain &&
+                                      !this->caps().gsInvocationsSupport();
+    if (needInvocationIDWorkaround) {
+        body = this->applyInvocationIDWorkaround(std::move(body));
+    }
+    if (ProgramKind::kVertex == this->programKind() && isMain && fRTAdjust) {
+        body->children().push_back(this->getNormalizeSkPositionCode());
+    }
+
     Finalizer finalizer{this, &funcDecl};
     finalizer.visitStatement(*body);
 
@@ -876,6 +887,7 @@ void IRGenerator::finalizeFunction(const FunctionDeclaration& funcDecl, Statemen
         this->errorReporter().error(funcDecl.fOffset, "function '" + funcDecl.name() +
                                                       "' can exit without returning a value");
     }
+    return body;
 }
 
 void IRGenerator::convertFunction(const ASTNode& f) {
@@ -888,8 +900,6 @@ void IRGenerator::convertFunction(const ASTNode& f) {
         return;
     }
     const ASTNode::FunctionData& funcData = f.getFunctionData();
-    bool isMain = (funcData.fName == "main");
-
     std::vector<std::unique_ptr<Variable>> parameters;
     parameters.reserve(funcData.fParameterCount);
     for (size_t i = 0; i < funcData.fParameterCount; ++i) {
@@ -952,19 +962,11 @@ void IRGenerator::convertFunction(const ASTNode& f) {
         for (const Variable* param : decl->parameters()) {
             fSymbolTable->addWithoutOwnership(param);
         }
-        bool needInvocationIDWorkaround = fInvocations != -1 && isMain &&
-                                          !this->caps().gsInvocationsSupport();
         std::unique_ptr<Block> body = this->convertBlock(*iter);
         if (!body) {
             return;
         }
-        if (needInvocationIDWorkaround) {
+        body = this->finalizeFunction(*decl, std::move(body));
-            body = this->applyInvocationIDWorkaround(std::move(body));
-        }
-        if (ProgramKind::kVertex == this->programKind() && isMain && fRTAdjust) {
-            body->children().push_back(this->getNormalizeSkPositionCode());
-        }
-        this->finalizeFunction(*decl, body.get());
         auto result = std::make_unique<FunctionDefinition>(
                 f.fOffset, decl, fIsBuiltinCode, std::move(body), std::move(fReferencedIntrinsics));
         decl->setDefinition(result.get());

src/sksl/SkSLIRGenerator.h

@@ -232,8 +232,10 @@ private:
     void copyIntrinsicIfNeeded(const FunctionDeclaration& function);
     void findAndDeclareBuiltinVariables();
     bool detectVarDeclarationWithoutScope(const Statement& stmt);
-    // Coerces returns to correct type and detects invalid break / continue placement
+    // Coerces returns to correct type, detects invalid break / continue placement, and otherwise
-    void finalizeFunction(const FunctionDeclaration& funcDecl, Statement* body);
+    // massages the function into its final form
+    std::unique_ptr<Block> finalizeFunction(const FunctionDeclaration& funcDecl,
+                                            std::unique_ptr<Block> body);
 
     // Runtime effects (and the interpreter, which uses the same CPU runtime) require adherence to
     // the strict rules from The OpenGL ES Shading Language Version 1.00. (Including Appendix A).

src/sksl/dsl/DSLBlock.cpp

@@ -31,11 +31,11 @@ DSLBlock::~DSLBlock() {
         // This will convert our Block into a DSLStatement, which is then immediately freed.
         // If an FP is being generated, this will naturally incorporate the Block's Statement into
         // our FP. If not, this will assert that unused code wasn't incorporated into the program.
-        DSLStatement(this->release());
+        DSLStatement(std::move(*this));
     }
 }
 
-std::unique_ptr<SkSL::Statement> DSLBlock::release() {
+std::unique_ptr<SkSL::Block> DSLBlock::release() {
     return std::make_unique<SkSL::Block>(/*offset=*/-1, std::move(fStatements),
                                          std::move(fSymbols));
 }

src/sksl/dsl/DSLCore.cpp

@@ -8,6 +8,7 @@
 #include "include/sksl/DSLCore.h"
 
 #include "include/private/SkSLDefines.h"
+#include "include/sksl/DSLSymbols.h"
 #include "src/sksl/SkSLCompiler.h"
 #include "src/sksl/SkSLIRGenerator.h"
 #include "src/sksl/dsl/priv/DSLWriter.h"
@@ -97,6 +98,10 @@ public:
         return DSLVar("sk_FragCoord");
     }
 
+    static DSLExpression sk_Position() {
+        return DSLExpression(Symbol("sk_Position"));
+    }
+
     template <typename... Args>
     static DSLPossibleExpression Call(const char* name, Args... args) {
         SkSL::IRGenerator& ir = DSLWriter::IRGenerator();
@@ -145,7 +150,7 @@ public:
             // sk_FragColor can end up with a null declaration despite no error occurring due to
             // specific treatment in the compiler. Ignore the null and just grab the existing
             // variable from the symbol table.
-            const Symbol* alreadyDeclared = (*DSLWriter::SymbolTable())[var.fName];
+            const SkSL::Symbol* alreadyDeclared = (*DSLWriter::SymbolTable())[var.fName];
             if (alreadyDeclared && alreadyDeclared->is<Variable>()) {
                 var.fVar = &alreadyDeclared->as<Variable>();
             }
@@ -279,6 +284,10 @@ DSLVar sk_FragCoord() {
     return DSLCore::sk_FragCoord();
 }
 
+DSLExpression sk_Position() {
+    return DSLCore::sk_Position();
+}
+
 DSLStatement Break() {
     return DSLCore::Break();
 }

src/sksl/dsl/DSLFunction.cpp

@@ -83,8 +83,8 @@ void DSLFunction::define(DSLBlock block) {
         return;
     }
     SkASSERTF(!fDecl->definition(), "function '%s' already defined", fDecl->description().c_str());
-    std::unique_ptr<Statement> body = block.release();
+    std::unique_ptr<Block> body = block.release();
-    DSLWriter::IRGenerator().finalizeFunction(*fDecl, body.get());
+    body = DSLWriter::IRGenerator().finalizeFunction(*fDecl, std::move(body));
     auto function = std::make_unique<SkSL::FunctionDefinition>(/*offset=*/-1, fDecl,
                                                                /*builtin=*/false, std::move(body));
     DSLWriter::ReportErrors();

tests/SkSLDSLTest.cpp

@@ -1961,3 +1961,20 @@ DEF_GPUTEST_FOR_MOCK_CONTEXT(DSLWrapper, r, ctxInfo) {
     vars.emplace_back(DSLVar(kInt_Type, "x"));
     REPORTER_ASSERT(r, DSLWriter::Var(*vars[0])->name() == "x");
 }
+
+DEF_GPUTEST_FOR_MOCK_CONTEXT(DSLRTAdjust, r, ctxInfo) {
+    AutoDSLContext context(ctxInfo.directContext()->priv().getGpu(), no_mark_vars_declared(),
+                           SkSL::ProgramKind::kVertex);
+    DSLVar rtAdjust(kUniform_Modifier, kFloat4_Type, "sk_RTAdjust");
+    DeclareGlobal(rtAdjust);
+    DSLFunction(kVoid_Type, "main").define(
+        sk_Position() = Half4(0)
+    );
+    REPORTER_ASSERT(r, DSLWriter::ProgramElements().size() == 2);
+    EXPECT_EQUAL(*DSLWriter::ProgramElements()[1],
+        "void main() {"
+        "(sk_PerVertex.sk_Position = float4(0.0));"
+        "(sk_PerVertex.sk_Position = float4(((sk_PerVertex.sk_Position.xy * sk_RTAdjust.xz) + "
+        "(sk_PerVertex.sk_Position.ww * sk_RTAdjust.yw)), 0.0, sk_PerVertex.sk_Position.w));"
+        "}");
+}