Added SkSL DSL swizzles

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

src/sksl/dsl/DSLCore.cpp

@@ -31,6 +31,26 @@ void End() {
 void SetErrorHandler(ErrorHandler* errorHandler) {
     DSLWriter::SetErrorHandler(errorHandler);
 }
+
+static char swizzle_component(SwizzleComponent c) {
+    switch (c) {
+        case X:
+            return 'x';
+        case Y:
+            return 'y';
+        case Z:
+            return 'z';
+        case W:
+            return 'w';
+        case ZERO:
+            return '0';
+        case ONE:
+            return '1';
+        default:
+            SkUNREACHABLE;
+    }
+}
+
 class DSLCore {
 public:
     template <typename... Args>
@@ -75,6 +95,29 @@ public:
                                                   ifTrue.release(), ifFalse.release());
     }
 
+    static DSLExpression Swizzle(DSLExpression base, SwizzleComponent a) {
+        char mask[] = { swizzle_component(a), 0 };
+        return DSLWriter::IRGenerator().convertSwizzle(base.release(), mask);
+    }
+
+    static DSLExpression Swizzle(DSLExpression base, SwizzleComponent a, SwizzleComponent b) {
+        char mask[] = { swizzle_component(a), swizzle_component(b), 0 };
+        return DSLWriter::IRGenerator().convertSwizzle(base.release(), mask);
+    }
+
+    static DSLExpression Swizzle(DSLExpression base, SwizzleComponent a, SwizzleComponent b,
+                                 SwizzleComponent c) {
+        char mask[] = { swizzle_component(a), swizzle_component(b), swizzle_component(c), 0 };
+        return DSLWriter::IRGenerator().convertSwizzle(base.release(), mask);
+    }
+
+    static DSLExpression Swizzle(DSLExpression base, SwizzleComponent a, SwizzleComponent b,
+                                 SwizzleComponent c, SwizzleComponent d) {
+        char mask[] = { swizzle_component(a), swizzle_component(b), swizzle_component(c),
+                        swizzle_component(d), 0 };
+        return DSLWriter::IRGenerator().convertSwizzle(base.release(), mask);
+    }
+
     static DSLExpression Ternary(DSLExpression test, DSLExpression ifTrue, DSLExpression ifFalse) {
         return DSLWriter::IRGenerator().convertTernaryExpression(test.release(), ifTrue.release(),
                                                                  ifFalse.release());
@@ -277,6 +320,24 @@ DSLExpression Step(DSLExpression edge, DSLExpression x) {
     return DSLCore::Call("step", std::move(edge), std::move(x));
 }
 
+DSLExpression Swizzle(DSLExpression base, SwizzleComponent a) {
+    return DSLCore::Swizzle(std::move(base), a);
+}
+
+DSLExpression Swizzle(DSLExpression base, SwizzleComponent a, SwizzleComponent b) {
+    return DSLCore::Swizzle(std::move(base), a, b);
+}
+
+DSLExpression Swizzle(DSLExpression base, SwizzleComponent a, SwizzleComponent b,
+                      SwizzleComponent c) {
+    return DSLCore::Swizzle(std::move(base), a, b, c);
+}
+
+DSLExpression Swizzle(DSLExpression base, SwizzleComponent a, SwizzleComponent b,
+                      SwizzleComponent c, SwizzleComponent d) {
+    return DSLCore::Swizzle(std::move(base), a, b, c, d);
+}
+
 DSLExpression Tan(DSLExpression x) {
     return DSLCore::Call("tan", std::move(x));
 }

src/sksl/dsl/DSLCore.h

@@ -79,6 +79,24 @@ DSLExpression Ternary(DSLExpression test, DSLExpression ifTrue, DSLExpression if
  */
 DSLStatement While(DSLExpression test, DSLStatement stmt);
 
+
+enum SwizzleComponent : int8_t {
+    X = 0, Y = 1, Z = 2, W = 3,
+    R = 0, G = 1, B = 2, A = 3,
+    ZERO,
+    ONE
+};
+
+DSLExpression Swizzle(DSLExpression base, SwizzleComponent a);
+
+DSLExpression Swizzle(DSLExpression base, SwizzleComponent a, SwizzleComponent b);
+
+DSLExpression Swizzle(DSLExpression base, SwizzleComponent a, SwizzleComponent b,
+                      SwizzleComponent c);
+
+DSLExpression Swizzle(DSLExpression base, SwizzleComponent a, SwizzleComponent b,
+                      SwizzleComponent c, SwizzleComponent d);
+
 /**
  * Returns the absolute value of x. If x is a vector, operates componentwise.
  */

src/sksl/dsl/DSLExpression.cpp

@@ -9,6 +9,7 @@
 
 #include "src/sksl/SkSLCompiler.h"
 #include "src/sksl/SkSLIRGenerator.h"
+#include "src/sksl/dsl/DSLCore.h"
 #include "src/sksl/dsl/DSLVar.h"
 #include "src/sksl/dsl/priv/DSLWriter.h"
 #include "src/sksl/ir/SkSLBinaryExpression.h"
@@ -65,6 +66,38 @@ std::unique_ptr<SkSL::Expression> DSLExpression::release() {
     return std::move(fExpression);
 }
 
+DSLExpression DSLExpression::x() {
+    return Swizzle(this->release(), X);
+}
+
+DSLExpression DSLExpression::y() {
+    return Swizzle(this->release(), Y);
+}
+
+DSLExpression DSLExpression::z() {
+    return Swizzle(this->release(), Z);
+}
+
+DSLExpression DSLExpression::w() {
+    return Swizzle(this->release(), W);
+}
+
+DSLExpression DSLExpression::r() {
+    return Swizzle(this->release(), R);
+}
+
+DSLExpression DSLExpression::g() {
+    return Swizzle(this->release(), G);
+}
+
+DSLExpression DSLExpression::b() {
+    return Swizzle(this->release(), B);
+}
+
+DSLExpression DSLExpression::a() {
+    return Swizzle(this->release(), A);
+}
+
 DSLExpression DSLExpression::operator=(DSLExpression right) {
     return DSLWriter::ConvertBinary(this->release(), SkSL::Token::Kind::TK_EQ, right.release());
 }

src/sksl/dsl/DSLExpression.h

@@ -67,6 +67,22 @@ public:
      */
     DSLExpression operator=(DSLExpression other);
 
+    DSLExpression x();
+
+    DSLExpression y();
+
+    DSLExpression z();
+
+    DSLExpression w();
+
+    DSLExpression r();
+
+    DSLExpression g();
+
+    DSLExpression b();
+
+    DSLExpression a();
+
     /**
      * Creates an SkSL array index expression.
      */

src/sksl/dsl/DSLVar.h

@@ -30,6 +30,38 @@ public:
 
     DSLVar(DSLVar&&) = delete;
 
+    DSLExpression x() {
+        return DSLExpression(*this).x();
+    }
+
+    DSLExpression y() {
+        return DSLExpression(*this).y();
+    }
+
+    DSLExpression z() {
+        return DSLExpression(*this).z();
+    }
+
+    DSLExpression w() {
+        return DSLExpression(*this).w();
+    }
+
+    DSLExpression r() {
+        return DSLExpression(*this).r();
+    }
+
+    DSLExpression g() {
+        return DSLExpression(*this).g();
+    }
+
+    DSLExpression b() {
+        return DSLExpression(*this).b();
+    }
+
+    DSLExpression a() {
+        return DSLExpression(*this).a();
+    }
+
     DSLExpression operator=(const DSLVar& var) {
         return this->operator=(DSLExpression(var));
     }

tests/SkSLDSLTest.cpp

@@ -954,6 +954,50 @@ DEF_GPUTEST_FOR_MOCK_CONTEXT(DSLIf, r, ctxInfo) {
     }
 }
 
+DEF_GPUTEST_FOR_MOCK_CONTEXT(DSLSwizzle, r, ctxInfo) {
+    AutoDSLContext context(ctxInfo.directContext()->priv().getGpu());
+    Var a(kFloat4, "a");
+
+    Expression e1 = a.x();
+    REPORTER_ASSERT(r, e1.release()->description() == "a.x");
+
+    Expression e2 = a.y();
+    REPORTER_ASSERT(r, e2.release()->description() == "a.y");
+
+    Expression e3 = a.z();
+    REPORTER_ASSERT(r, e3.release()->description() == "a.z");
+
+    Expression e4 = a.w();
+    REPORTER_ASSERT(r, e4.release()->description() == "a.w");
+
+    Expression e5 = a.r();
+    REPORTER_ASSERT(r, e5.release()->description() == "a.x");
+
+    Expression e6 = a.g();
+    REPORTER_ASSERT(r, e6.release()->description() == "a.y");
+
+    Expression e7 = a.b();
+    REPORTER_ASSERT(r, e7.release()->description() == "a.z");
+
+    Expression e8 = a.a();
+    REPORTER_ASSERT(r, e8.release()->description() == "a.w");
+
+    Expression e9 = Swizzle(a, R);
+    REPORTER_ASSERT(r, e9.release()->description() == "a.x");
+
+    Expression e10 = Swizzle(a, ZERO, G);
+    REPORTER_ASSERT(r, e10.release()->description() == "float2(a.y, float(0)).yx");
+
+    Expression e11 = Swizzle(a, B, G, G);
+    REPORTER_ASSERT(r, e11.release()->description() == "a.zyy");
+
+    Expression e12 = Swizzle(a, R, G, B, ONE);
+    REPORTER_ASSERT(r, e12.release()->description() == "float4(a.xyz, float(1))");
+
+    Expression e13 = Swizzle(a, R, G, B, ONE).r();
+    REPORTER_ASSERT(r, e13.release()->description() == "float4(a.xyz, float(1)).x");
+}
+
 DEF_GPUTEST_FOR_MOCK_CONTEXT(DSLTernary, r, ctxInfo) {
     AutoDSLContext context(ctxInfo.directContext()->priv().getGpu());
     Var a(kInt, "a");