1 files changed, 198 insertions, 0 deletions
diff --git a/gfx/angle/checkout/src/compiler/translator/tree_util/RewriteSampleMaskVariable.cpp b/gfx/angle/checkout/src/compiler/translator/tree_util/RewriteSampleMaskVariable.cpp
new file mode 100644
index 0000000000..a81d230492
--- /dev/null
+++ b/gfx/angle/checkout/src/compiler/translator/tree_util/RewriteSampleMaskVariable.cpp
@@ -0,0 +1,198 @@
+//
+// Copyright 2020 The ANGLE Project Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+//
+// RewriteSampleMaskVariable.cpp: Find any references to gl_SampleMask and gl_SampleMaskIn, and
+// rewrite it with ANGLESampleMask or ANGLESampleMaskIn.
+//
+
+#include "compiler/translator/tree_util/RewriteSampleMaskVariable.h"
+
+#include "common/bitset_utils.h"
+#include "common/debug.h"
+#include "common/utilities.h"
+#include "compiler/translator/Compiler.h"
+#include "compiler/translator/SymbolTable.h"
+#include "compiler/translator/tree_util/BuiltIn.h"
+#include "compiler/translator/tree_util/IntermNode_util.h"
+#include "compiler/translator/tree_util/IntermTraverse.h"
+#include "compiler/translator/tree_util/RunAtTheEndOfShader.h"
+
+namespace sh
+{
+namespace
+{
+constexpr int kMaxIndexForSampleMaskVar = 0;
+constexpr int kFullSampleMask           = 0xFFFFFFFF;
+
+// Traverse the tree and collect the redeclaration and replace all non constant index references of
+// gl_SampleMask or gl_SampleMaskIn with constant index references
+class GLSampleMaskRelatedReferenceTraverser : public TIntermTraverser
+{
+  public:
+    GLSampleMaskRelatedReferenceTraverser(const TIntermSymbol **redeclaredSymOut,
+                                          const ImmutableString &targetStr)
+        : TIntermTraverser(true, false, false),
+          mRedeclaredSym(redeclaredSymOut),
+          mTargetStr(targetStr)
+    {
+        *mRedeclaredSym = nullptr;
+    }
+
+    bool visitDeclaration(Visit visit, TIntermDeclaration *node) override
+    {
+        // If gl_SampleMask is redeclared, we need to collect its information
+        const TIntermSequence &sequence = *(node->getSequence());
+
+        if (sequence.size() != 1)
+        {
+            return true;
+        }
+
+        TIntermTyped *variable = sequence.front()->getAsTyped();
+        TIntermSymbol *symbol  = variable->getAsSymbolNode();
+        if (symbol == nullptr || symbol->getName() != mTargetStr)
+        {
+            return true;
+        }
+
+        *mRedeclaredSym = symbol;
+
+        return true;
+    }
+
+    bool visitBinary(Visit visit, TIntermBinary *node) override
+    {
+        TOperator op = node->getOp();
+        if (op != EOpIndexDirect && op != EOpIndexIndirect)
+        {
+            return true;
+        }
+        TIntermSymbol *left = node->getLeft()->getAsSymbolNode();
+        if (!left)
+        {
+            return true;
+        }
+        if (left->getName() != mTargetStr)
+        {
+            return true;
+        }
+        const TConstantUnion *constIdx = node->getRight()->getConstantValue();
+        if (!constIdx)
+        {
+            if (node->getRight()->hasSideEffects())
+            {
+                insertStatementInParentBlock(node->getRight());
+            }
+
+            queueReplacementWithParent(node, node->getRight(),
+                                       CreateIndexNode(kMaxIndexForSampleMaskVar),
+                                       OriginalNode::IS_DROPPED);
+        }
+
+        return true;
+    }
+
+  private:
+    const TIntermSymbol **mRedeclaredSym;
+    const ImmutableString mTargetStr;
+};
+
+}  // anonymous namespace
+
+[[nodiscard]] bool RewriteSampleMask(TCompiler *compiler,
+                                     TIntermBlock *root,
+                                     TSymbolTable *symbolTable,
+                                     const TIntermTyped *numSamplesUniform)
+{
+    const TIntermSymbol *redeclaredGLSampleMask = nullptr;
+    GLSampleMaskRelatedReferenceTraverser indexTraverser(&redeclaredGLSampleMask,
+                                                         ImmutableString("gl_SampleMask"));
+
+    root->traverse(&indexTraverser);
+    if (!indexTraverser.updateTree(compiler, root))
+    {
+        return false;
+    }
+
+    // Retrieve gl_SampleMask variable reference
+    // Search user redeclared it first
+    const TVariable *glSampleMaskVar = nullptr;
+    if (redeclaredGLSampleMask)
+    {
+        glSampleMaskVar = &redeclaredGLSampleMask->variable();
+    }
+    else
+    {
+        // User defined not found, find in built-in table
+        glSampleMaskVar = static_cast<const TVariable *>(symbolTable->findBuiltIn(
+            ImmutableString("gl_SampleMask"), compiler->getShaderVersion()));
+    }
+    if (!glSampleMaskVar)
+    {
+        return false;
+    }
+
+    // Current ANGLE assumes that the maximum number of samples is less than or equal to
+    // VK_SAMPLE_COUNT_32_BIT. So, the size of gl_SampleMask array is always one.
+    const unsigned int arraySizeOfSampleMask = glSampleMaskVar->getType().getOutermostArraySize();
+    ASSERT(arraySizeOfSampleMask == 1);
+
+    TIntermSymbol *glSampleMaskSymbol = new TIntermSymbol(glSampleMaskVar);
+
+    // if (ANGLEUniforms.numSamples == 1)
+    // {
+    //     gl_SampleMask[0] = int(0xFFFFFFFF);
+    // }
+    TIntermConstantUnion *singleSampleCount = CreateUIntNode(1);
+    TIntermBinary *equalTo =
+        new TIntermBinary(EOpEqual, numSamplesUniform->deepCopy(), singleSampleCount);
+
+    TIntermBlock *trueBlock = new TIntermBlock();
+
+    TIntermBinary *sampleMaskVar = new TIntermBinary(EOpIndexDirect, glSampleMaskSymbol->deepCopy(),
+                                                     CreateIndexNode(kMaxIndexForSampleMaskVar));
+    TIntermConstantUnion *fullSampleMask = CreateIndexNode(kFullSampleMask);
+    TIntermBinary *assignment = new TIntermBinary(EOpAssign, sampleMaskVar, fullSampleMask);
+
+    trueBlock->appendStatement(assignment);
+
+    TIntermIfElse *multiSampleOrNot = new TIntermIfElse(equalTo, trueBlock, nullptr);
+
+    return RunAtTheEndOfShader(compiler, root, multiSampleOrNot, symbolTable);
+}
+
+[[nodiscard]] bool RewriteSampleMaskIn(TCompiler *compiler,
+                                       TIntermBlock *root,
+                                       TSymbolTable *symbolTable)
+{
+    const TIntermSymbol *redeclaredGLSampleMaskIn = nullptr;
+    GLSampleMaskRelatedReferenceTraverser indexTraverser(&redeclaredGLSampleMaskIn,
+                                                         ImmutableString("gl_SampleMaskIn"));
+
+    root->traverse(&indexTraverser);
+    if (!indexTraverser.updateTree(compiler, root))
+    {
+        return false;
+    }
+
+    // Retrieve gl_SampleMaskIn variable reference
+    const TVariable *glSampleMaskInVar = nullptr;
+    glSampleMaskInVar                  = static_cast<const TVariable *>(
+        symbolTable->findBuiltIn(ImmutableString("gl_SampleMaskIn"), compiler->getShaderVersion()));
+    if (!glSampleMaskInVar)
+    {
+        return false;
+    }
+
+    // Current ANGLE assumes that the maximum number of samples is less than or equal to
+    // VK_SAMPLE_COUNT_32_BIT. So, the size of gl_SampleMask array is always one.
+    const unsigned int arraySizeOfSampleMaskIn =
+        glSampleMaskInVar->getType().getOutermostArraySize();
+    ASSERT(arraySizeOfSampleMaskIn == 1);
+
+    return true;
+}
+
+}  // namespace sh