[sksl][wgsl] Support for trivial ternary expressions

WGSL does not have ternary expressions. The most trivial cases, which
have no side-effects (and can safely be evaluated without
short-circuiting) and evaluate to a scalar or vector can be translated
to a call to the "select" intrinsic, which is what this CL implements.

All other situations need more sophisticated handling in which any
ternary expressions need to be hoisted out as an if-else statement in
the immediately surrounding block, which will be implemented later.

Bug: skia:13092
Change-Id: I0101e4d3199a9f2bb72038d9ba153ef33f551ac8
Reviewed-on: https://skia-review.googlesource.com/c/skia/+/628852
Reviewed-by: John Stiles <johnstiles@google.com>
Commit-Queue: Arman Uguray <armansito@google.com>
diff --git a/gn/sksl_tests.gni b/gn/sksl_tests.gni
index 5d1031c..c8b9ee0 100644
--- a/gn/sksl_tests.gni
+++ b/gn/sksl_tests.gni
@@ -388,6 +388,7 @@
   "shared/HelloWorld.sksl",
   "shared/InstanceID.vert",
   "shared/Structs.sksl",
+  "shared/TernaryExpression.sksl",
   "shared/VertexID.vert",
   "wgsl/BuiltinFragmentStageIO.sksl",
   "wgsl/BuiltinVertexStageIO.vert",
diff --git a/resources/sksl/BUILD.bazel b/resources/sksl/BUILD.bazel
index 86cef7c..655e1c4 100644
--- a/resources/sksl/BUILD.bazel
+++ b/resources/sksl/BUILD.bazel
@@ -1015,6 +1015,7 @@
         "shared/HelloWorld.sksl",
         "shared/InstanceID.vert",
         "shared/Structs.sksl",
+        "shared/TernaryExpression.sksl",
         "shared/VertexID.vert",
     ],
 )
diff --git a/src/sksl/codegen/SkSLWGSLCodeGenerator.cpp b/src/sksl/codegen/SkSLWGSLCodeGenerator.cpp
index 53ec54e..dd9d414 100644
--- a/src/sksl/codegen/SkSLWGSLCodeGenerator.cpp
+++ b/src/sksl/codegen/SkSLWGSLCodeGenerator.cpp
@@ -51,6 +51,7 @@
 #include "src/sksl/ir/SkSLStructDefinition.h"
 #include "src/sksl/ir/SkSLSwizzle.h"
 #include "src/sksl/ir/SkSLSymbolTable.h"
+#include "src/sksl/ir/SkSLTernaryExpression.h"
 #include "src/sksl/ir/SkSLType.h"
 #include "src/sksl/ir/SkSLVarDeclarations.h"
 #include "src/sksl/ir/SkSLVariable.h"
@@ -790,6 +791,9 @@
         case Expression::Kind::kSwizzle:
             this->writeSwizzle(e.as<Swizzle>());
             break;
+        case Expression::Kind::kTernary:
+            this->writeTernaryExpression(e.as<TernaryExpression>(), parentPrecedence);
+            break;
         case Expression::Kind::kVariableReference:
             this->writeVariableReference(e.as<VariableReference>());
             break;
@@ -876,6 +880,51 @@
     }
 }
 
+void WGSLCodeGenerator::writeTernaryExpression(const TernaryExpression& t,
+                                               Precedence parentPrecedence) {
+    bool needParens = Precedence::kTernary >= parentPrecedence;
+    if (needParens) {
+        this->write("(");
+    }
+
+    // The trivial case is when neither branch has side effects and evaluate to a scalar or vector
+    // type. This can be represented with a call to the WGSL `select` intrinsic although it doesn't
+    // support short-circuiting.
+    if ((t.type().isScalar() || t.type().isVector()) && !Analysis::HasSideEffects(*t.ifTrue()) &&
+        !Analysis::HasSideEffects(*t.ifFalse())) {
+        this->write("select(");
+        this->writeExpression(*t.ifFalse(), Precedence::kTernary);
+        this->write(", ");
+        this->writeExpression(*t.ifTrue(), Precedence::kTernary);
+        this->write(", ");
+
+        bool isVector = t.type().isVector();
+        if (isVector) {
+            // Splat the condition expression into a vector.
+            this->write(String::printf("vec%d<bool>", t.type().columns()));
+            this->write("(");
+        }
+        this->writeExpression(*t.test(), Precedence::kTernary);
+        if (isVector) {
+            this->write(")");
+        }
+        this->write(")");
+
+        if (needParens) {
+            this->write(")");
+        }
+        return;
+    }
+
+    // TODO(skia:13092): WGSL does not support ternary expressions. To replicate the required
+    // short-circuting behavior we need to hoist the expression out into the surrounding block,
+    // convert it into an if statement that writes the result to a synthesized variable, and replace
+    // the original expression with a reference to that variable.
+    //
+    // Once hoisting is supported, we may want to use that for vector type expressions as well,
+    // since select above does a component-wise select
+}
+
 void WGSLCodeGenerator::writeVariableReference(const VariableReference& r) {
     // TODO(skia:13902): Correctly handle function parameters.
     // TODO(skia:13902): Correctly handle RTflip for built-ins.
diff --git a/src/sksl/codegen/SkSLWGSLCodeGenerator.h b/src/sksl/codegen/SkSLWGSLCodeGenerator.h
index 306bcbc..6d534ad 100644
--- a/src/sksl/codegen/SkSLWGSLCodeGenerator.h
+++ b/src/sksl/codegen/SkSLWGSLCodeGenerator.h
@@ -46,6 +46,7 @@
 class ReturnStatement;
 class Statement;
 class StructDefinition;
+class TernaryExpression;
 class VarDeclaration;
 class VariableReference;
 enum class OperatorPrecedence : uint8_t;
@@ -180,6 +181,7 @@
     void writeFieldAccess(const FieldAccess& f);
     void writeLiteral(const Literal& l);
     void writeSwizzle(const Swizzle& swizzle);
+    void writeTernaryExpression(const TernaryExpression& t, Precedence parentPrecedence);
     void writeVariableReference(const VariableReference& r);
 
     // Constructor expressions
diff --git a/tests/sksl/shared/TernaryExpression.wgsl b/tests/sksl/shared/TernaryExpression.wgsl
new file mode 100644
index 0000000..7be404f
--- /dev/null
+++ b/tests/sksl/shared/TernaryExpression.wgsl
@@ -0,0 +1,23 @@
+struct FSIn {
+    @builtin(front_facing) sk_Clockwise: bool,
+    @builtin(position) sk_FragCoord: vec4<f32>,
+};
+struct FSOut {
+    @location(0) sk_FragColor: vec4<f32>,
+};
+struct _GlobalUniforms {
+    colorGreen: vec4<f32>,
+    colorRed: vec4<f32>,
+};
+@binding(0) @group(0) var<uniform> _globalUniforms: _GlobalUniforms;
+fn main(coords: vec2<f32>) -> vec4<f32> {
+    var ok: bool = true;
+    ok = ok && (select(false, true, _globalUniforms.colorGreen.y == 1.0));
+    ok = ok && (select(true, false, _globalUniforms.colorGreen.x == 1.0));
+    return select(_globalUniforms.colorRed, _globalUniforms.colorGreen, vec4<bool>(ok));
+}
+@fragment fn fragmentMain(_stageIn: FSIn) -> FSOut {
+    var _stageOut: FSOut;
+    _stageOut.sk_FragColor = main(_stageIn.sk_FragCoord.xy);
+    return _stageOut;
+}