sksl can now fold constant vector or matrix equality expressions

Bug: skia:
Change-Id: Icaddae68e53ed3629bcdc04b5f0b541d9e4398e2
Reviewed-on: https://skia-review.googlesource.com/14260
Commit-Queue: Ethan Nicholas <ethannicholas@google.com>
Reviewed-by: Ben Wagner <benjaminwagner@google.com>

6 files changed, 131 insertions, 0 deletions

diff --git a/src/sksl/SkSLIRGenerator.cpp b/src/sksl/SkSLIRGenerator.cpp
index 56858a92ca..523b7a009d 100644
--- a/src/sksl/SkSLIRGenerator.cpp
+++ b/src/sksl/SkSLIRGenerator.cpp
@@ -1043,6 +1043,12 @@ std::unique_ptr<Expression> IRGenerator::constantFold(const Expression& left,
             return std::unique_ptr<Expression>(new Constructor(Position(), left.fType, \
                                                                std::move(args)));
         switch (op) {
+            case Token::EQEQ:
+                return std::unique_ptr<Expression>(new BoolLiteral(fContext, Position(),
+                                                            left.compareConstant(fContext, right)));
+            case Token::NEQ:
+                return std::unique_ptr<Expression>(new BoolLiteral(fContext, Position(),
+                                                           !left.compareConstant(fContext, right)));
             case Token::PLUS:  RETURN_VEC_COMPONENTWISE_RESULT(+);
             case Token::MINUS: RETURN_VEC_COMPONENTWISE_RESULT(-);
             case Token::STAR:  RETURN_VEC_COMPONENTWISE_RESULT(*);
@@ -1050,6 +1056,20 @@ std::unique_ptr<Expression> IRGenerator::constantFold(const Expression& left,
             default:           return nullptr;
         }
     }
+    if (left.fType.kind() == Type::kMatrix_Kind &&
+        right.fType.kind() == Type::kMatrix_Kind &&
+        left.fKind == right.fKind) {
+        switch (op) {
+            case Token::EQEQ:
+                return std::unique_ptr<Expression>(new BoolLiteral(fContext, Position(),
+                                                            left.compareConstant(fContext, right)));
+            case Token::NEQ:
+                return std::unique_ptr<Expression>(new BoolLiteral(fContext, Position(),
+                                                           !left.compareConstant(fContext, right)));
+            default:
+                return nullptr;
+        }
+    }
     #undef RESULT
     return nullptr;
 }
diff --git a/src/sksl/ir/SkSLBoolLiteral.h b/src/sksl/ir/SkSLBoolLiteral.h
index 13203a4e55..a4151b8b35 100644
--- a/src/sksl/ir/SkSLBoolLiteral.h
+++ b/src/sksl/ir/SkSLBoolLiteral.h
@@ -33,6 +33,11 @@ struct BoolLiteral : public Expression {
         return true;
     }
 
+    bool compareConstant(const Context& context, const Expression& other) const override {
+        BoolLiteral& b = (BoolLiteral&) other;
+        return fValue == b.fValue;
+    }
+
     const bool fValue;
 
     typedef Expression INHERITED;
diff --git a/src/sksl/ir/SkSLConstructor.h b/src/sksl/ir/SkSLConstructor.h
index 208031abba..05f409649a 100644
--- a/src/sksl/ir/SkSLConstructor.h
+++ b/src/sksl/ir/SkSLConstructor.h
@@ -81,6 +81,44 @@ struct Constructor : public Expression {
         return true;
     }
 
+    bool compareConstant(const Context& context, const Expression& other) const override {
+        ASSERT(other.fKind == Expression::kConstructor_Kind && other.fType == fType);
+        Constructor& c = (Constructor&) other;
+        if (c.fType.kind() == Type::kVector_Kind) {
+            for (int i = 0; i < fType.columns(); i++) {
+                if (!this->getVecComponent(i).compareConstant(context, c.getVecComponent(i))) {
+                    return false;
+                }
+            }
+            return true;
+        }
+        // shouldn't be possible to have a constant constructor that isn't a vector or matrix;
+        // a constant scalar constructor should have been collapsed down to the appropriate
+        // literal
+        ASSERT(fType.kind() == Type::kMatrix_Kind);
+        const FloatLiteral fzero(context, Position(), 0);
+        const IntLiteral izero(context, Position(), 0);
+        const Expression* zero;
+        if (fType.componentType() == *context.fFloat_Type) {
+            zero = &fzero;
+        } else {
+            ASSERT(fType.componentType() == *context.fInt_Type);
+            zero = &izero;
+        }
+        for (int col = 0; col < fType.columns(); col++) {
+            for (int row = 0; row < fType.rows(); row++) {
+                const Expression* component1 = getMatComponent(col, row);
+                const Expression* component2 = c.getMatComponent(col, row);
+                if (!(component1 ? component1 : zero)->compareConstant(
+                                                                context,
+                                                                component2 ? *component2 : *zero)) {
+                    return false;
+                }
+            }
+        }
+        return true;
+    }
+
     const Expression& getVecComponent(int index) const {
         ASSERT(fType.kind() == Type::kVector_Kind);
         if (fArguments.size() == 1 && fArguments[0]->fType.kind() == Type::kScalar_Kind) {
@@ -118,6 +156,51 @@ struct Constructor : public Expression {
         return ((IntLiteral&) c).fValue;
     }
 
+    // null return should be interpreted as zero
+    const Expression* getMatComponent(int col, int row) const {
+        ASSERT(this->isConstant());
+        ASSERT(fType.kind() == Type::kMatrix_Kind);
+        ASSERT(col < fType.columns() && row < fType.rows());
+        if (fArguments.size() == 1) {
+            if (fArguments[0]->fType.kind() == Type::kScalar_Kind) {
+                // single scalar argument, so matrix is of the form:
+                // x 0 0
+                // 0 x 0
+                // 0 0 x
+                // return x if col == row
+                return col == row ? fArguments[0].get() : nullptr;
+            }
+            if (fArguments[0]->fType.kind() == Type::kMatrix_Kind) {
+                ASSERT(fArguments[0]->fKind == Expression::kConstructor_Kind);
+                // single matrix argument. make sure we're within the argument's bounds.
+                const Type& argType = ((Constructor&) *fArguments[0]).fType;
+                if (col < argType.columns() && row < argType.rows()) {
+                    // within bounds, defer to argument
+                    return ((Constructor&) *fArguments[0]).getMatComponent(col, row);
+                }
+                // out of bounds, return 0
+                return nullptr;
+            }
+        }
+        int currentIndex = 0;
+        int targetIndex = col * fType.rows() + row;
+        for (const auto& arg : fArguments) {
+            ASSERT(targetIndex >= currentIndex);
+            ASSERT(arg->fType.rows() == 1);
+            if (currentIndex + arg->fType.columns() > targetIndex) {
+                if (arg->fType.columns() == 1) {
+                    return arg.get();
+                } else {
+                    ASSERT(arg->fType.kind() == Type::kVector_Kind);
+                    ASSERT(arg->fKind == Expression::kConstructor_Kind);
+                    return &((Constructor&) *arg).getVecComponent(targetIndex - currentIndex);
+                }
+            }
+            currentIndex += arg->fType.columns();
+        }
+        ABORT("can't happen, matrix component out of bounds");
+    }
+
     std::vector<std::unique_ptr<Expression>> fArguments;
 
     typedef Expression INHERITED;
diff --git a/src/sksl/ir/SkSLExpression.h b/src/sksl/ir/SkSLExpression.h
index 5db9ddf96f..07dad1d7df 100644
--- a/src/sksl/ir/SkSLExpression.h
+++ b/src/sksl/ir/SkSLExpression.h
@@ -48,11 +48,24 @@ struct Expression : public IRNode {
     , fKind(kind)
     , fType(std::move(type)) {}
 
+    /**
+     * Returns true if this expression is constant. compareConstant must be implemented for all
+     * constants!
+     */
     virtual bool isConstant() const {
         return false;
     }
 
     /**
+     * Compares this constant expression against another constant expression of the same type. It is
+     * an error to call this on non-constant expressions, or if the types of the expressions do not
+     * match.
+     */
+    virtual bool compareConstant(const Context& context, const Expression& other) const {
+        ABORT("cannot call compareConstant on this type");
+    }
+
+    /**
      * Returns true if evaluating the expression potentially has side effects. Expressions may never
      * return false if they actually have side effects, but it is legal (though suboptimal) to
      * return true if there are not actually any side effects.
diff --git a/src/sksl/ir/SkSLFloatLiteral.h b/src/sksl/ir/SkSLFloatLiteral.h
index 8f83e2866c..21a485fb0a 100644
--- a/src/sksl/ir/SkSLFloatLiteral.h
+++ b/src/sksl/ir/SkSLFloatLiteral.h
@@ -34,6 +34,11 @@ struct FloatLiteral : public Expression {
         return true;
     }
 
+    bool compareConstant(const Context& context, const Expression& other) const override {
+        FloatLiteral& f = (FloatLiteral&) other;
+        return fValue == f.fValue;
+    }
+
     const double fValue;
 
     typedef Expression INHERITED;
diff --git a/src/sksl/ir/SkSLIntLiteral.h b/src/sksl/ir/SkSLIntLiteral.h
index 3a95ed65ba..d8eba5573a 100644
--- a/src/sksl/ir/SkSLIntLiteral.h
+++ b/src/sksl/ir/SkSLIntLiteral.h
@@ -35,6 +35,11 @@ struct IntLiteral : public Expression {
         return true;
     }
 
+    bool compareConstant(const Context& context, const Expression& other) const override {
+        IntLiteral& i = (IntLiteral&) other;
+        return fValue == i.fValue;
+    }
+
     const int64_t fValue;
 
     typedef Expression INHERITED;