Improve index-folding of arrays and matrices.

Yesterday's implementation was close but I realized later that it wasn't
quite ideal.
- Array index-folding was gated on `isCompileTimeConstant`, which is too
  strict. The real limitation is `hasSideEffects`. If an array contains
  a side-effecting expression, we should leave it alone. Otherwise it
  is safe to pluck out an element from the array and toss the rest.
- Matrix index-folding was gated on `getConstantSubexpression` for the
  extracted elements, but did not check the other elements at all. This
  was too lenient; we now only proceed to the folding step if
  `hasSideEffects` returns false.

I added some tests to verify the final behavior and also discovered a
small related issue. Diagonal matrices were not substituting literals
in for constant-values, which inhibited folding as well and would break
constant-expression evaluation. This is now fixed.

Change-Id: Idda32fd8643c1f32ba21475251cd4d4dd7cea94c
Reviewed-on: https://skia-review.googlesource.com/c/skia/+/470396
Auto-Submit: John Stiles <johnstiles@google.com>
Reviewed-by: Brian Osman <brianosman@google.com>
Commit-Queue: John Stiles <johnstiles@google.com>
diff --git a/resources/sksl/folding/ArrayFolding.sksl b/resources/sksl/folding/ArrayFolding.sksl
index b253a5b..7feb107 100644
--- a/resources/sksl/folding/ArrayFolding.sksl
+++ b/resources/sksl/folding/ArrayFolding.sksl
@@ -17,15 +17,16 @@
     const int b [x[x[z]]]    = int[2](1, 2);
     const int c [x[x[x[z]]]] = int[3](1, 2, 3);
 
-    // Constant-expression arrays can be optimized.
-    int flatten0 = (int[3](1, 2, 3))[0];
-    int flatten1 = (int[3](1, 2, 3))[1];
-    int flatten2 = (int[3](1, 2, 3))[2];
+    // Arrays with elements lacking side-effects can be optimized.
+    int two = 2;
+    int flatten0 = (int[3](x[0], two, 3))[0];
+    int flatten1 = (int[3](x[0], two, 3))[1];
+    int flatten2 = (int[3](x[0], two, 3))[2];
 
-    // Non-constant-expression arrays are not eligible for optimization.
-    int noFlatten0 = (int[3](1, side_effecting(2), 3))[0];
+    // Arrays with elements that have side-effects are not eligible for optimization.
+    int noFlatten0 = (int[3](--two, side_effecting(2), 3))[0];
     int noFlatten1 = (int[3](side_effecting(1), 2, 3))[1];
-    int noFlatten2 = (int[3](1, 2, side_effecting(3)))[2];
+    int noFlatten2 = (int[3](1, ++two, 3))[2];
 
     return (x == xx) && !(x != xx) && (x != y) && !(x == y) &&
            (x[0] == y[0]) && (c == x) &&
diff --git a/resources/sksl/folding/MatrixFoldingES2.sksl b/resources/sksl/folding/MatrixFoldingES2.sksl
index e5c2223..ab880ef 100644
--- a/resources/sksl/folding/MatrixFoldingES2.sksl
+++ b/resources/sksl/folding/MatrixFoldingES2.sksl
@@ -61,6 +61,36 @@
     ok = ok &&  (m[2][1] == 8);
     ok = ok &&  (m[2][2] == 9);
 
+    {
+        // This `five` is constant and should always fold.
+        const float five = 5.0;
+        ok = ok &&  (float2x2(five)[0] == float2(five, 0.0));
+        ok = ok &&  (float2x2(five)[1] == float2(0.0, five));
+
+        ok = ok &&  (float2x2(five)[0][0] == five);
+        ok = ok &&  (float2x2(five)[0][1] == 0.0);
+        ok = ok &&  (float2x2(five)[1][0] == 0.0);
+        ok = ok &&  (float2x2(five)[1][1] == five);
+
+        ok = ok &&  (float3x3(1, 2, 3, 4, five, 6, 7, 8, 9)[0] == float3(1, 2, 3));
+        ok = ok &&  (float3x3(1, 2, 3, 4, five, 6, 7, 8, 9)[1] == float3(4, five, 6));
+        ok = ok &&  (float3x3(1, 2, 3, 4, five, 6, 7, 8, 9)[2] == float3(7, 8, 9));
+    }
+    {
+        // This `five` cannot be folded, but the first and third columns should still be foldable.
+        float five = 5.0;
+        ok = ok &&  (float3x3(1, 2, 3, 4, five, 6, 7, 8, 9)[0] == float3(1, 2, 3));
+        ok = ok &&  (float3x3(1, 2, 3, 4, five, 6, 7, 8, 9)[1] == float3(4, five, 6));
+        ok = ok &&  (float3x3(1, 2, 3, 4, five, 6, 7, 8, 9)[2] == float3(7, 8, 9));
+    }
+    {
+        // Side-effecting expressions should never be folded away.
+        float num = 6.0;
+        ok = ok &&  (float3x3(1, 2, 3, 4, 5, num++, 7, 8, 9)[0] == float3(1, 2, 3));
+        ok = ok &&  (float3x3(1, 2, 3, 4, 5, 6, num++, 8, 9)[1] == float3(4, 5, 6));
+        ok = ok &&  (float3x3(1, 2, 3, 4, 5, 6, 7, num++, 9)[2] == float3(7, 8, 9));
+    }
+
     return ok;
 }
 
diff --git a/src/sksl/ir/SkSLConstructorDiagonalMatrix.cpp b/src/sksl/ir/SkSLConstructorDiagonalMatrix.cpp
index fd2b479..09f0513 100644
--- a/src/sksl/ir/SkSLConstructorDiagonalMatrix.cpp
+++ b/src/sksl/ir/SkSLConstructorDiagonalMatrix.cpp
@@ -7,6 +7,7 @@
 
 #include "src/sksl/ir/SkSLConstructorDiagonalMatrix.h"
 
+#include "src/sksl/SkSLConstantFolder.h"
 #include "src/sksl/ir/SkSLConstructor.h"
 #include "src/sksl/ir/SkSLType.h"
 
@@ -20,6 +21,11 @@
     SkASSERT(type.isAllowedInES2(context));
     SkASSERT(arg->type().isScalar());
     SkASSERT(arg->type() == type.componentType());
+
+    // Look up the value of constant variables. This allows constant-expressions like `mat4(five)`
+    // to be replaced with `mat4(5.0)`.
+    arg = ConstantFolder::MakeConstantValueForVariable(std::move(arg));
+
     return std::make_unique<ConstructorDiagonalMatrix>(line, type, std::move(arg));
 }
 
diff --git a/src/sksl/ir/SkSLIndexExpression.cpp b/src/sksl/ir/SkSLIndexExpression.cpp
index 64676a4..7308c9a 100644
--- a/src/sksl/ir/SkSLIndexExpression.cpp
+++ b/src/sksl/ir/SkSLIndexExpression.cpp
@@ -104,11 +104,11 @@
                 return Swizzle::Make(context, std::move(base), ComponentArray{(int8_t)indexValue});
             }
 
-            if (baseType.isArray()) {
+            if (baseType.isArray() && !base->hasSideEffects()) {
                 // Indexing an constant array constructor with a constant index can just pluck out
                 // the requested value from the array.
                 const Expression* baseExpr = ConstantFolder::GetConstantValueForVariable(*base);
-                if (baseExpr->is<ConstructorArray>() && baseExpr->isCompileTimeConstant()) {
+                if (baseExpr->is<ConstructorArray>()) {
                     const ConstructorArray& arrayCtor = baseExpr->as<ConstructorArray>();
                     const ExpressionArray& arguments = arrayCtor.arguments();
                     SkASSERT(arguments.count() == baseType.columns());
@@ -117,7 +117,7 @@
                 }
             }
 
-            if (baseType.isMatrix()) {
+            if (baseType.isMatrix() && !base->hasSideEffects()) {
                 // Matrices can be constructed with vectors that don't line up on column boundaries,
                 // so extracting out the values from the constructor can be tricky. Fortunately, we
                 // can reconstruct an equivalent vector using `getConstantSubexpression`. If we
diff --git a/tests/sksl/folding/ArrayFolding.glsl b/tests/sksl/folding/ArrayFolding.glsl
index 1dfa5f3..bdf5d70 100644
--- a/tests/sksl/folding/ArrayFolding.glsl
+++ b/tests/sksl/folding/ArrayFolding.glsl
@@ -8,11 +8,12 @@
     return value;
 }
 vec4 main() {
-    int _7_flatten0 = 1;
-    int _8_flatten1 = 2;
-    int _9_flatten2 = 3;
-    int _10_noFlatten0 = int[3](1, side_effecting_ii(2), 3)[0];
-    int _11_noFlatten1 = int[3](side_effecting_ii(1), 2, 3)[1];
-    int _12_noFlatten2 = int[3](1, 2, side_effecting_ii(3))[2];
-    return (_7_flatten0 == _10_noFlatten0 && _8_flatten1 == _11_noFlatten1) && _9_flatten2 == _12_noFlatten2 ? colorGreen : colorRed;
+    int _7_two = 2;
+    int _8_flatten0 = 1;
+    int _9_flatten1 = _7_two;
+    int _10_flatten2 = 3;
+    int _11_noFlatten0 = int[3](--_7_two, side_effecting_ii(2), 3)[0];
+    int _12_noFlatten1 = int[3](side_effecting_ii(1), 2, 3)[1];
+    int _13_noFlatten2 = int[3](1, ++_7_two, 3)[2];
+    return (_8_flatten0 == _11_noFlatten0 && _9_flatten1 == _12_noFlatten1) && _10_flatten2 == _13_noFlatten2 ? colorGreen : colorRed;
 }
diff --git a/tests/sksl/folding/MatrixFoldingES2.glsl b/tests/sksl/folding/MatrixFoldingES2.glsl
index d428cfe..580bdb9 100644
--- a/tests/sksl/folding/MatrixFoldingES2.glsl
+++ b/tests/sksl/folding/MatrixFoldingES2.glsl
@@ -9,5 +9,15 @@
     _0_ok = _0_ok && mat3(unknownInput) == mat3(mat2(1.0));
     _0_ok = _0_ok && mat3(9.0, 0.0, 0.0, 0.0, 9.0, 0.0, 0.0, 0.0, unknownInput) == mat3(mat2(9.0));
     _0_ok = _0_ok && vec4(testMatrix2x2) == vec4(1.0, 2.0, 3.0, 4.0);
+    {
+        float _3_five = 5.0;
+        _0_ok = _0_ok && mat3(1.0, 2.0, 3.0, 4.0, _3_five, 6.0, 7.0, 8.0, 9.0)[1] == vec3(4.0, _3_five, 6.0);
+    }
+    {
+        float _4_num = 6.0;
+        _0_ok = _0_ok && mat3(1.0, 2.0, 3.0, 4.0, 5.0, _4_num++, 7.0, 8.0, 9.0)[0] == vec3(1.0, 2.0, 3.0);
+        _0_ok = _0_ok && mat3(1.0, 2.0, 3.0, 4.0, 5.0, 6.0, _4_num++, 8.0, 9.0)[1] == vec3(4.0, 5.0, 6.0);
+        _0_ok = _0_ok && mat3(1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, _4_num++, 9.0)[2] == vec3(7.0, 8.0, 9.0);
+    }
     return _0_ok ? colorGreen : colorRed;
 }