fixed early return breaking things in .fp files

Bug: skia:
Change-Id: If91e8cb2f312cb7546967bd9af0d1ae180177870
Reviewed-on: https://skia-review.googlesource.com/146644
Reviewed-by: Greg Daniel <egdaniel@google.com>
Commit-Queue: Ethan Nicholas <ethannicholas@google.com>
diff --git a/src/sksl/SkSLCPPCodeGenerator.cpp b/src/sksl/SkSLCPPCodeGenerator.cpp
index 863b8c3..ef640ed 100644
--- a/src/sksl/SkSLCPPCodeGenerator.cpp
+++ b/src/sksl/SkSLCPPCodeGenerator.cpp
@@ -284,6 +284,13 @@
     INHERITED::writeIfStatement(s);
 }
 
+void CPPCodeGenerator::writeReturnStatement(const ReturnStatement& s) {
+    if (fInMain) {
+        fErrors.error(s.fOffset, "fragmentProcessor main() may not contain return statements");
+    }
+    INHERITED::writeReturnStatement(s);
+}
+
 void CPPCodeGenerator::writeSwitchStatement(const SwitchStatement& s) {
     if (s.fIsStatic) {
         this->write("@");
@@ -339,10 +346,12 @@
         OutputStream* oldOut = fOut;
         StringStream buffer;
         fOut = &buffer;
+        fInMain = true;
         for (const auto& s : ((Block&) *f.fBody).fStatements) {
             this->writeStatement(*s);
             this->writeLine();
         }
+        fInMain = false;
 
         fOut = oldOut;
         this->write(fFunctionHeader);
diff --git a/src/sksl/SkSLCPPCodeGenerator.h b/src/sksl/SkSLCPPCodeGenerator.h
index 40d9982..014f917 100644
--- a/src/sksl/SkSLCPPCodeGenerator.h
+++ b/src/sksl/SkSLCPPCodeGenerator.h
@@ -49,6 +49,8 @@
 
     void writeIfStatement(const IfStatement& s) override;
 
+    void writeReturnStatement(const ReturnStatement& s) override;
+
     void writeSwitchStatement(const SwitchStatement& s) override;
 
     void writeFunctionCall(const FunctionCall& c) override;
@@ -92,6 +94,7 @@
     std::set<int> fWrittenTransformedCoords;
     // if true, we are writing a C++ expression instead of a GLSL expression
     bool fCPPMode = false;
+    bool fInMain = false;
 
     typedef GLSLCodeGenerator INHERITED;
 };
diff --git a/src/sksl/SkSLGLSLCodeGenerator.h b/src/sksl/SkSLGLSLCodeGenerator.h
index 6346ac6..f6d78aa 100644
--- a/src/sksl/SkSLGLSLCodeGenerator.h
+++ b/src/sksl/SkSLGLSLCodeGenerator.h
@@ -190,7 +190,7 @@
 
     virtual void writeSwitchStatement(const SwitchStatement& s);
 
-    void writeReturnStatement(const ReturnStatement& r);
+    virtual void writeReturnStatement(const ReturnStatement& r);
 
     virtual void writeProgramElement(const ProgramElement& e);