Fix codegen errors with Metal return statements.

The Metal return type from main() diverges from the SkSL source, so we
patch it in the Metal code generator. This CL improves the patching
process in multiple ways:

- A `return` statement from a fragment processor main() is rewritten to:
    return *_out;

- A `return` statement from a vertex processor main() is rewritten to:
    return (_out->sk_Position.y = -_out->sk_Position.y, *_out);

- We avoid emitting a duplicate `return *_out;` statement if we can
  determine that main() already ends in a return statement. This is
  harmless either way so it doesn't necessarily catch everything. (e.g.
  it doesn't detect an if/else which returns at the end of both blocks.)

Also added a unit test which returns from the middle of a vertex shader,
since we didn't test this anywhere and we need to verify that
sk_Position.y will be negated. (This didn't work properly before.)

Change-Id: I14cf18375894fc712fa6c6466df3888ebaeba7c8
Bug: skia:10903
Reviewed-on: https://skia-review.googlesource.com/c/skia/+/339636
Reviewed-by: Ethan Nicholas <ethannicholas@google.com>
Commit-Queue: John Stiles <johnstiles@google.com>
Auto-Submit: John Stiles <johnstiles@google.com>

gn/sksl_tests.gni

@@ -284,6 +284,7 @@ sksl_shared_tests = [
   "$_tests/sksl/shared/UnusedVariables.sksl",
   "$_tests/sksl/shared/VectorConstructors.sksl",
   "$_tests/sksl/shared/VectorFolding.sksl",
+  "$_tests/sksl/shared/VertexEarlyReturn.vert",
   "$_tests/sksl/shared/VertexID.vert",
   "$_tests/sksl/shared/Width.sksl",
 ]

src/sksl/SkSLMetalCodeGenerator.cpp

@@ -7,6 +7,7 @@
 
 #include "src/sksl/SkSLMetalCodeGenerator.h"
 
+#include "src/core/SkScopeExit.h"
 #include "src/sksl/SkSLCompiler.h"
 #include "src/sksl/SkSLMemoryLayout.h"
 #include "src/sksl/ir/SkSLExpressionStatement.h"
@@ -1142,6 +1143,27 @@ void MetalCodeGenerator::writeFunctionPrototype(const FunctionPrototype& f) {
     this->writeLine(";");
 }
 
+static bool is_block_ending_with_return(const Statement* stmt) {
+    // This function detects (potentially nested) blocks that end in a return statement.
+    if (!stmt->is<Block>()) {
+        return false;
+    }
+    const StatementArray& block = stmt->as<Block>().children();
+    for (int index = block.count(); index--; ) {
+        const Statement& stmt = *block[index];
+        if (stmt.is<ReturnStatement>()) {
+            return true;
+        }
+        if (stmt.is<Block>()) {
+            return is_block_ending_with_return(&stmt);
+        }
+        if (!stmt.is<Nop>()) {
+            break;
+        }
+    }
+    return false;
+}
+
 void MetalCodeGenerator::writeFunction(const FunctionDefinition& f) {
     SkASSERT(!fProgram.fSettings.fFragColorIsInOut);
 
@@ -1149,6 +1171,9 @@ void MetalCodeGenerator::writeFunction(const FunctionDefinition& f) {
         return;
     }
 
+    fCurrentFunction = &f.declaration();
+    SkScopeExit clearCurrentFunction([&] { fCurrentFunction = nullptr; });
+
     this->writeLine(" {");
 
     if (f.declaration().name() == "main") {
@@ -1169,16 +1194,10 @@ void MetalCodeGenerator::writeFunction(const FunctionDefinition& f) {
         }
     }
     if (f.declaration().name() == "main") {
-        switch (fProgram.fKind) {
-            case Program::kFragment_Kind:
-                this->writeLine("return *_out;");
-                break;
-            case Program::kVertex_Kind:
-                this->writeLine("_out->sk_Position.y = -_out->sk_Position.y;");
-                this->writeLine("return *_out;"); // FIXME - detect if function already has return
-                break;
-            default:
-                SkDEBUGFAIL("unsupported kind of program");
+        // If the main function doesn't end with a return, we need to synthesize one here.
+        if (!is_block_ending_with_return(f.body().get())) {
+            this->writeReturnStatementFromMain();
+            this->writeLine("");
         }
     }
     fIndentation--;
@@ -1467,7 +1486,29 @@ void MetalCodeGenerator::writeSwitchStatement(const SwitchStatement& s) {
     this->write("}");
 }
 
+void MetalCodeGenerator::writeReturnStatementFromMain() {
+    // main functions in Metal return a magic _out parameter that doesn't exist in SkSL.
+    switch (fProgram.fKind) {
+        case Program::kFragment_Kind:
+            this->write("return *_out;");
+            break;
+        case Program::kVertex_Kind:
+            this->write("return (_out->sk_Position.y = -_out->sk_Position.y, *_out);");
+            break;
+        default:
+            SkDEBUGFAIL("unsupported kind of program");
+    }
+}
+
 void MetalCodeGenerator::writeReturnStatement(const ReturnStatement& r) {
+    if (fCurrentFunction && fCurrentFunction->name() == "main") {
+        if (r.expression()) {
+            fErrors.error(r.fOffset, "Metal does not support returning values from main()");
+        }
+        this->writeReturnStatementFromMain();
+        return;
+    }
+
     this->write("return");
     if (r.expression()) {
         this->write(" ");

src/sksl/SkSLMetalCodeGenerator.h

@@ -257,6 +257,8 @@ protected:
 
     void writeSwitchStatement(const SwitchStatement& s);
 
+    void writeReturnStatementFromMain();
+
     void writeReturnStatement(const ReturnStatement& r);
 
     void writeProgramElement(const ProgramElement& e);
@@ -296,6 +298,7 @@ protected:
     std::unordered_set<String> fHelpers;
     int fUniformBuffer = -1;
     String fRTHeightName;
+    const FunctionDeclaration* fCurrentFunction = nullptr;
 
     using INHERITED = CodeGenerator;
 };

tests/sksl/shared/Control.sksl

@@ -4,7 +4,9 @@ void main() {
     while (i < 10) { sk_FragColor *= 0.5; i++; }
     do { sk_FragColor += 0.25; } while (sk_FragColor.x < 0.75);
     for (int i = 0; i < 10; i++) {
-        if (i % 2 == 1) break; else continue;
+        if (i % 2 == 1) break;
+        else if (i > 100) return;
+        else continue;
     }
     return;
 }

tests/sksl/shared/VertexEarlyReturn.vert

@@ -0,0 +1,7 @@
+layout(set=0) uniform half zoom;
+
+void main() {
+    sk_Position = half4(1);
+    if (zoom == 1) return;
+    sk_Position *= zoom;
+}

tests/sksl/shared/golden/Control.asm.frag

@@ -37,6 +37,7 @@ OpDecorate %49 RelaxedPrecision
 %int_1 = OpConstant %int 1
 %float_0_25 = OpConstant %float 0.25
 %int_2 = OpConstant %int 2
+%int_100 = OpConstant %int 100
 %main = OpFunction %void None %11
 %12 = OpLabel
 %i = OpVariable %_ptr_Function_int Function
@@ -107,13 +108,22 @@ OpBranchConditional %63 %64 %65
 %64 = OpLabel
 OpBranch %57
 %65 = OpLabel
+%67 = OpLoad %int %i_0
+%69 = OpSGreaterThan %bool %67 %int_100
+OpSelectionMerge %72 None
+OpBranchConditional %69 %70 %71
+%70 = OpLabel
+OpReturn
+%71 = OpLabel
 OpBranch %56
+%72 = OpLabel
+OpBranch %66
 %66 = OpLabel
 OpBranch %56
 %56 = OpLabel
-%67 = OpLoad %int %i_0
-%68 = OpIAdd %int %67 %int_1
-OpStore %i_0 %68
+%73 = OpLoad %int %i_0
+%74 = OpIAdd %int %73 %int_1
+OpStore %i_0 %74
 OpBranch %53
 %57 = OpLabel
 OpReturn

tests/sksl/shared/golden/Control.glsl

@@ -15,7 +15,7 @@ void main() {
         sk_FragColor += 0.25;
     } while (sk_FragColor.x < 0.75);
     for (int i = 0;i < 10; i++) {
-        if (i % 2 == 1) break; else continue;
+        if (i % 2 == 1) break; else if (i > 100) return; else continue;
     }
     return;
 }

tests/sksl/shared/golden/Control.metal

@@ -23,8 +23,7 @@ fragment Outputs fragmentMain(Inputs _in [[stage_in]], bool _frontFacing [[front
         _out->sk_FragColor += 0.25;
     } while (_out->sk_FragColor.x < 0.75);
     for (int i = 0;i < 10; i++) {
-        if (i % 2 == 1) break; else continue;
+        if (i % 2 == 1) break; else if (i > 100) return *_out; else continue;
     }
-    return;
     return *_out;
 }

tests/sksl/shared/golden/DeadIfStatement.metal

@@ -9,6 +9,5 @@ struct Outputs {
 fragment Outputs fragmentMain(Inputs _in [[stage_in]], bool _frontFacing [[front_facing]], float4 _fragCoord [[position]]) {
     Outputs _outputStruct;
     thread Outputs* _out = &_outputStruct;
-    return;
     return *_out;
 }

tests/sksl/shared/golden/InstanceID.metal

@@ -13,6 +13,5 @@ vertex Outputs vertexMain(Inputs _in [[stage_in]], uint sk_VertexID [[vertex_id]
     Outputs _outputStruct;
     thread Outputs* _out = &_outputStruct;
     _out->id = sk_InstanceID;
-    _out->sk_Position.y = -_out->sk_Position.y;
-    return *_out;
+    return (_out->sk_Position.y = -_out->sk_Position.y, *_out);
 }

tests/sksl/shared/golden/NoFragCoordsPos.metal

@@ -13,6 +13,5 @@ vertex Outputs vertexMain(Inputs _in [[stage_in]], uint sk_VertexID [[vertex_id]
     Outputs _outputStruct;
     thread Outputs* _out = &_outputStruct;
     _out->sk_Position = _in.pos;
-    _out->sk_Position.y = -_out->sk_Position.y;
-    return *_out;
+    return (_out->sk_Position.y = -_out->sk_Position.y, *_out);
 }

tests/sksl/shared/golden/NoFragCoordsPosRT.metal

@@ -18,6 +18,5 @@ vertex Outputs vertexMain(Inputs _in [[stage_in]], constant Uniforms& _uniforms
     thread Outputs* _out = &_outputStruct;
     _out->sk_Position = _in.pos;
     _out->sk_Position = float4(_out->sk_Position.xy * _uniforms.sk_RTAdjust.xz + _out->sk_Position.ww * _uniforms.sk_RTAdjust.yw, 0.0, _out->sk_Position.w);
-    _out->sk_Position.y = -_out->sk_Position.y;
-    return *_out;
+    return (_out->sk_Position.y = -_out->sk_Position.y, *_out);
 }

tests/sksl/shared/golden/NormalizationVert.metal

@@ -16,6 +16,5 @@ vertex Outputs vertexMain(Inputs _in [[stage_in]], constant Uniforms& _uniforms
     thread Outputs* _out = &_outputStruct;
     _out->sk_Position = float4(1.0);
     _out->sk_Position = float4(_out->sk_Position.xy * _uniforms.sk_RTAdjust.xz + _out->sk_Position.ww * _uniforms.sk_RTAdjust.yw, 0.0, _out->sk_Position.w);
-    _out->sk_Position.y = -_out->sk_Position.y;
-    return *_out;
+    return (_out->sk_Position.y = -_out->sk_Position.y, *_out);
 }

tests/sksl/shared/golden/VertexEarlyReturn.asm.vert

@@ -0,0 +1,58 @@
+### Compilation failed:
+
+error: 1: SPIR-V validation error: Uniform OpVariable <id> '8[%zoom]' has illegal type.
+From Vulkan spec, section 14.5.2:
+Variables identified with the Uniform storage class are used to access transparent buffer backed resources. Such variables must be typed as OpTypeStruct, or an array of this type
+  %zoom = OpVariable %_ptr_Uniform_float Uniform
+
+OpCapability Shader
+%1 = OpExtInstImport "GLSL.std.450"
+OpMemoryModel Logical GLSL450
+OpEntryPoint Vertex %main "main" %3
+OpName %sk_PerVertex "sk_PerVertex"
+OpMemberName %sk_PerVertex 0 "sk_Position"
+OpMemberName %sk_PerVertex 1 "sk_PointSize"
+OpName %zoom "zoom"
+OpName %main "main"
+OpMemberDecorate %sk_PerVertex 0 BuiltIn Position
+OpMemberDecorate %sk_PerVertex 1 BuiltIn PointSize
+OpDecorate %sk_PerVertex Block
+OpDecorate %zoom RelaxedPrecision
+OpDecorate %zoom DescriptorSet 0
+OpDecorate %19 RelaxedPrecision
+OpDecorate %26 RelaxedPrecision
+%float = OpTypeFloat 32
+%v4float = OpTypeVector %float 4
+%sk_PerVertex = OpTypeStruct %v4float %float
+%_ptr_Output_sk_PerVertex = OpTypePointer Output %sk_PerVertex
+%3 = OpVariable %_ptr_Output_sk_PerVertex Output
+%_ptr_Uniform_float = OpTypePointer Uniform %float
+%zoom = OpVariable %_ptr_Uniform_float Uniform
+%void = OpTypeVoid
+%11 = OpTypeFunction %void
+%float_1 = OpConstant %float 1
+%13 = OpConstantComposite %v4float %float_1 %float_1 %float_1 %float_1
+%int = OpTypeInt 32 1
+%int_0 = OpConstant %int 0
+%_ptr_Output_v4float = OpTypePointer Output %v4float
+%bool = OpTypeBool
+%main = OpFunction %void None %11
+%12 = OpLabel
+%17 = OpAccessChain %_ptr_Output_v4float %3 %int_0
+OpStore %17 %13
+%19 = OpLoad %float %zoom
+%20 = OpFOrdEqual %bool %19 %float_1
+OpSelectionMerge %23 None
+OpBranchConditional %20 %22 %23
+%22 = OpLabel
+OpReturn
+%23 = OpLabel
+%24 = OpAccessChain %_ptr_Output_v4float %3 %int_0
+%25 = OpLoad %v4float %24
+%26 = OpLoad %float %zoom
+%27 = OpVectorTimesScalar %v4float %25 %26
+OpStore %24 %27
+OpReturn
+OpFunctionEnd
+
+1 error

tests/sksl/shared/golden/VertexEarlyReturn.glsl

@@ -0,0 +1,7 @@
+
+layout (set = 0) uniform float zoom;
+void main() {
+    gl_Position = vec4(1.0);
+    if (zoom == 1.0) return;
+    gl_Position *= zoom;
+}

tests/sksl/shared/golden/VertexEarlyReturn.metal

@@ -0,0 +1,21 @@
+#include <metal_stdlib>
+#include <simd/simd.h>
+using namespace metal;
+struct Uniforms {
+    float zoom;
+};
+struct Inputs {
+};
+struct Outputs {
+    float4 sk_Position [[position]];
+    float sk_PointSize [[point_size]];
+};
+
+vertex Outputs vertexMain(Inputs _in [[stage_in]], constant Uniforms& _uniforms [[buffer(0)]], uint sk_VertexID [[vertex_id]], uint sk_InstanceID [[instance_id]]) {
+    Outputs _outputStruct;
+    thread Outputs* _out = &_outputStruct;
+    _out->sk_Position = float4(1.0);
+    if (_uniforms.zoom == 1.0) return (_out->sk_Position.y = -_out->sk_Position.y, *_out);
+    _out->sk_Position *= _uniforms.zoom;
+    return (_out->sk_Position.y = -_out->sk_Position.y, *_out);
+}

tests/sksl/shared/golden/VertexID.metal

@@ -13,6 +13,5 @@ vertex Outputs vertexMain(Inputs _in [[stage_in]], uint sk_VertexID [[vertex_id]
     Outputs _outputStruct;
     thread Outputs* _out = &_outputStruct;
     _out->id = sk_VertexID;
-    _out->sk_Position.y = -_out->sk_Position.y;
-    return *_out;
+    return (_out->sk_Position.y = -_out->sk_Position.y, *_out);
 }