ES31: support ssbo as an argument of aggregate type

src/compiler/translator/tree_ops/RewriteExpressionsWithShaderStorageBlock.cpp

@@ -13,6 +13,7 @@
 
 #include "compiler/translator/tree_ops/RewriteExpressionsWithShaderStorageBlock.h"
 
+#include "compiler/translator/Symbol.h"
 #include "compiler/translator/tree_util/IntermNode_util.h"
 #include "compiler/translator/tree_util/IntermTraverse.h"
 #include "compiler/translator/util.h"
@@ -82,6 +83,19 @@ bool IsReadonlyBinaryOperatorNotInSSBOAccessChain(TOperator op)
     }
 }
 
+bool HasSSBOAsFunctionArgument(TIntermSequence *arguments)
+{
+    for (TIntermNode *arg : *arguments)
+    {
+        TIntermTyped *typedArg = arg->getAsTyped();
+        if (IsInShaderStorageBlock(typedArg))
+        {
+            return true;
+        }
+    }
+    return false;
+}
+
 class RewriteExpressionsWithShaderStorageBlockTraverser : public TIntermTraverser
 {
   public:
@@ -91,14 +105,16 @@ class RewriteExpressionsWithShaderStorageBlockTraverser : public TIntermTraverse
 
   private:
     bool visitBinary(Visit, TIntermBinary *node) override;
+    bool visitAggregate(Visit visit, TIntermAggregate *node) override;
     TIntermSymbol *insertInitStatementAndReturnTempSymbol(TIntermTyped *node,
                                                           TIntermSequence *insertions);
+
     bool mFoundSSBO;
 };
 
 RewriteExpressionsWithShaderStorageBlockTraverser::
     RewriteExpressionsWithShaderStorageBlockTraverser(TSymbolTable *symbolTable)
-    : TIntermTraverser(true, false, false, symbolTable), mFoundSSBO(false)
+    : TIntermTraverser(true, true, false, symbolTable), mFoundSSBO(false)
 {}
 
 TIntermSymbol *
@@ -114,8 +130,15 @@ RewriteExpressionsWithShaderStorageBlockTraverser::insertInitStatementAndReturnT
     return CreateTempSymbolNode(tempVariable);
 }
 
-bool RewriteExpressionsWithShaderStorageBlockTraverser::visitBinary(Visit, TIntermBinary *node)
+bool RewriteExpressionsWithShaderStorageBlockTraverser::visitBinary(Visit visit,
+                                                                    TIntermBinary *node)
 {
+    // Make sure that the expression is caculated from left to right.
+    if (visit != InVisit)
+    {
+        return true;
+    }
+
     if (mFoundSSBO)
     {
         return false;
@@ -130,16 +153,9 @@ bool RewriteExpressionsWithShaderStorageBlockTraverser::visitBinary(Visit, TInte
 
     // case 1: Compound assigment operator
     //  original:
-    //      lssbo += expr_no_ssbo;
+    //      lssbo += expr;
     //  new:
-    //      var temp = lssbo;
-    //      temp += expr_no_ssbo;
-    //      lssbo = temp;
-    //
-    //  original:
-    //      lssbo += rssbo;
-    //  new:
-    //      var rvalue = rssbo;
+    //      var rvalue = expr;
     //      var temp = lssbo;
     //      temp += rvalue;
     //      lssbo = temp;
@@ -153,12 +169,8 @@ bool RewriteExpressionsWithShaderStorageBlockTraverser::visitBinary(Visit, TInte
     {
         mFoundSSBO = true;
         TIntermSequence insertions;
-        TIntermTyped *rightNode = node->getRight();
-        if (rightSSBO)
-        {
-            rightNode = insertInitStatementAndReturnTempSymbol(node->getRight(), &insertions);
-        }
-
+        TIntermTyped *rightNode =
+            insertInitStatementAndReturnTempSymbol(node->getRight(), &insertions);
         if (leftSSBO)
         {
             TIntermSymbol *tempSymbol =
@@ -175,9 +187,9 @@ bool RewriteExpressionsWithShaderStorageBlockTraverser::visitBinary(Visit, TInte
         else
         {
             insertStatementsInParentBlock(insertions);
-            TIntermBinary *assignTempValueToLValue =
+            TIntermBinary *compoundAssignRValueToLValue =
                 new TIntermBinary(node->getOp(), node->getLeft(), rightNode);
-            queueReplacement(assignTempValueToLValue, OriginalNode::IS_DROPPED);
+            queueReplacement(compoundAssignRValueToLValue, OriginalNode::IS_DROPPED);
         }
     }
     // case 2: Readonly binary operator
@@ -210,6 +222,96 @@ bool RewriteExpressionsWithShaderStorageBlockTraverser::visitBinary(Visit, TInte
     return !mFoundSSBO;
 }
 
+// case 3: ssbo as the argument of aggregate type
+//  original:
+//      foo(ssbo);
+//  new:
+//      var tempArg = ssbo;
+//      foo(tempArg);
+//      ssbo = tempArg;  (Optional based on whether ssbo is an out|input argument)
+//
+//  original:
+//      foo(ssbo) * expr;
+//  new:
+//      var tempArg = ssbo;
+//      var tempReturn = foo(tempArg);
+//      ssbo = tempArg;  (Optional based on whether ssbo is an out|input argument)
+//      tempReturn * expr;
+bool RewriteExpressionsWithShaderStorageBlockTraverser::visitAggregate(Visit visit,
+                                                                       TIntermAggregate *node)
+{
+    // Make sure that visitAggregate is only executed once for same node.
+    if (visit != PreVisit)
+    {
+        return true;
+    }
+
+    if (mFoundSSBO)
+    {
+        return false;
+    }
+
+    if (IsAtomicFunction(node->getOp()))
+    {
+        return true;
+    }
+
+    if (!HasSSBOAsFunctionArgument(node->getSequence()))
+    {
+        return true;
+    }
+
+    mFoundSSBO = true;
+    TIntermSequence insertions;
+    TIntermSequence readBackToSSBOs;
+    TIntermSequence *originalArguments = node->getSequence();
+    for (size_t i = 0; i < node->getChildCount(); ++i)
+    {
+        TIntermTyped *ssboArgument = (*originalArguments)[i]->getAsTyped();
+        if (IsInShaderStorageBlock(ssboArgument))
+        {
+            TIntermSymbol *argumentCopy =
+                insertInitStatementAndReturnTempSymbol(ssboArgument, &insertions);
+            if (node->getFunction() != nullptr)
+            {
+                TQualifier qual = node->getFunction()->getParam(i)->getType().getQualifier();
+                if (qual == EvqInOut || qual == EvqOut)
+                {
+                    TIntermBinary *readBackToSSBO = new TIntermBinary(
+                        EOpAssign, ssboArgument->deepCopy(), argumentCopy->deepCopy());
+                    readBackToSSBOs.push_back(readBackToSSBO);
+                }
+            }
+            node->replaceChildNode(ssboArgument, argumentCopy);
+        }
+    }
+
+    TIntermBlock *parentBlock = getParentNode()->getAsBlock();
+    if (parentBlock)
+    {
+        // Aggregate node is as a single sentence.
+        insertions.push_back(node);
+        if (!readBackToSSBOs.empty())
+        {
+            insertions.insert(insertions.end(), readBackToSSBOs.begin(), readBackToSSBOs.end());
+        }
+        mMultiReplacements.push_back(NodeReplaceWithMultipleEntry(parentBlock, node, insertions));
+    }
+    else
+    {
+        // Aggregate node is inside an expression.
+        TIntermSymbol *tempSymbol = insertInitStatementAndReturnTempSymbol(node, &insertions);
+        if (!readBackToSSBOs.empty())
+        {
+            insertions.insert(insertions.end(), readBackToSSBOs.begin(), readBackToSSBOs.end());
+        }
+        insertStatementsInParentBlock(insertions);
+        queueReplacement(tempSymbol, OriginalNode::IS_DROPPED);
+    }
+
+    return false;
+}
+
 void RewriteExpressionsWithShaderStorageBlockTraverser::nextIteration()
 {
     mFoundSSBO = false;

src/tests/gl_tests/ShaderStorageBufferTest.cpp

@@ -1660,6 +1660,67 @@ TEST_P(ShaderStorageBufferTest31, ReadonlyBinaryOperator)
     EXPECT_GL_NO_ERROR();
 }
 
+// Test that ssbo as an argument of a function can be translated.
+TEST_P(ShaderStorageBufferTest31, SSBOAsFunctionArgument)
+{
+    // http://anglebug.com/2990
+    ANGLE_SKIP_TEST_IF(IsD3D11());
+
+    constexpr char kComputeShaderSource[] =
+        R"(#version 310 es
+layout(local_size_x = 1) in;
+
+layout(std430, binding = 0) buffer Block
+{
+    uint var1;
+    uint var2;
+};
+
+bool compare(uint a, uint b)
+{
+    return a == b;
+}
+
+uint increase(inout uint a)
+{
+    a++;
+    return a;
+}
+
+void main(void)
+{
+    bool isEqual = compare(var1, 2u);
+    if (isEqual)
+    {
+        var2 += increase(var1);
+    }
+}
+)";
+
+    ANGLE_GL_COMPUTE_PROGRAM(program, kComputeShaderSource);
+    glUseProgram(program);
+
+    constexpr unsigned int kBytesPerComponent = sizeof(unsigned int);
+    constexpr unsigned int kInputValues[2]    = {2u, 2u};
+    constexpr unsigned int kExpectedValues[2] = {3u, 5u};
+    GLBuffer shaderStorageBuffer;
+    glBindBuffer(GL_SHADER_STORAGE_BUFFER, shaderStorageBuffer);
+    glBufferData(GL_SHADER_STORAGE_BUFFER, 2 * kBytesPerComponent, &kInputValues, GL_STATIC_DRAW);
+
+    glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 0, shaderStorageBuffer);
+
+    glDispatchCompute(1, 1, 1);
+    glFinish();
+
+    glBindBuffer(GL_SHADER_STORAGE_BUFFER, shaderStorageBuffer);
+    const GLuint *ptr = reinterpret_cast<const GLuint *>(
+        glMapBufferRange(GL_SHADER_STORAGE_BUFFER, 0, 2 * kBytesPerComponent, GL_MAP_READ_BIT));
+    EXPECT_EQ(kExpectedValues[0], *ptr);
+    EXPECT_EQ(kExpectedValues[1], *(ptr + 1));
+
+    EXPECT_GL_NO_ERROR();
+}
+
 ANGLE_INSTANTIATE_TEST(ShaderStorageBufferTest31, ES31_OPENGL(), ES31_OPENGLES(), ES31_D3D11());
 
 }  // namespace