//
// Copyright 2018 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.
//
// RewriteStructSamplers: Extract samplers from structs.
//

#include "compiler/translator/tree_ops/RewriteStructSamplers.h"

#include "compiler/translator/ImmutableStringBuilder.h"
#include "compiler/translator/SymbolTable.h"
#include "compiler/translator/tree_util/IntermNode_util.h"
#include "compiler/translator/tree_util/IntermTraverse.h"

namespace sh
{
namespace
{

// Used to map one structure type to another (one where the samplers are removed).
struct StructureData
{
    // The structure this was replaced with.  If nullptr, it means the structure is removed (because
    // it had all samplers).
    const TStructure *modified;
    // Indexed by the field index of original structure, to get the field index of the modified
    // structure.  For example:
    //
    //     struct Original
    //     {
    //         sampler2D s1;
    //         vec4 f1;
    //         sampler2D s2;
    //         sampler2D s3;
    //         vec4 f2;
    //     };
    //
    //     struct Modified
    //     {
    //         vec4 f1;
    //         vec4 f2;
    //     };
    //
    //     fieldMap:
    //         0 -> Invalid
    //         1 -> 0
    //         2 -> Invalid
    //         3 -> Invalid
    //         4 -> 1
    //
    TVector<int> fieldMap;
};

using StructureMap        = angle::HashMap<const TStructure *, StructureData>;
using StructureUniformMap = angle::HashMap<const TVariable *, const TVariable *>;
using ExtractedSamplerMap = angle::HashMap<std::string, const TVariable *>;

TIntermTyped *RewriteModifiedStructFieldSelectionExpression(
    TCompiler *compiler,
    TIntermBinary *node,
    const StructureMap &structureMap,
    const StructureUniformMap &structureUniformMap,
    const ExtractedSamplerMap &extractedSamplers);

TIntermTyped *RewriteExpressionVisitBinaryHelper(TCompiler *compiler,
                                                 TIntermBinary *node,
                                                 const StructureMap &structureMap,
                                                 const StructureUniformMap &structureUniformMap,
                                                 const ExtractedSamplerMap &extractedSamplers)
{
    // Only interested in EOpIndexDirectStruct binary nodes.
    if (node->getOp() != EOpIndexDirectStruct)
    {
        return nullptr;
    }

    const TStructure *structure = node->getLeft()->getType().getStruct();
    ASSERT(structure);

    // If the result of the index is not a sampler and the struct is not replaced, there's nothing
    // to do.
    if (!node->getType().isSampler() && structureMap.find(structure) == structureMap.end())
    {
        return nullptr;
    }

    // Otherwise, replace the whole expression such that:
    //
    // - if sampler, it's indexed with whatever indices the parent structs were indexed with,
    // - otherwise, the chain of field selections is rewritten by modifying the base uniform so all
    //   the intermediate nodes would have the correct type (and therefore fields).
    ASSERT(structureMap.find(structure) != structureMap.end());

    return RewriteModifiedStructFieldSelectionExpression(compiler, node, structureMap,
                                                         structureUniformMap, extractedSamplers);
}

// Given an expression, this traverser calculates a new expression where sampler-in-structs are
// replaced with their extracted ones, and field indices are adjusted for the rest of the fields.
// In particular, this is run on the right node of EOpIndexIndirect binary nodes, so that the
// expression in the index gets a chance to go through this transformation.
class RewriteExpressionTraverser final : public TIntermTraverser
{
  public:
    explicit RewriteExpressionTraverser(TCompiler *compiler,
                                        const StructureMap &structureMap,
                                        const StructureUniformMap &structureUniformMap,
                                        const ExtractedSamplerMap &extractedSamplers)
        : TIntermTraverser(true, false, false),
          mCompiler(compiler),
          mStructureMap(structureMap),
          mStructureUniformMap(structureUniformMap),
          mExtractedSamplers(extractedSamplers)
    {}

    bool visitBinary(Visit visit, TIntermBinary *node) override
    {
        TIntermTyped *rewritten = RewriteExpressionVisitBinaryHelper(
            mCompiler, node, mStructureMap, mStructureUniformMap, mExtractedSamplers);

        if (rewritten == nullptr)
        {
            return true;
        }

        queueReplacement(rewritten, OriginalNode::IS_DROPPED);

        // Don't iterate as the expression is rewritten.
        return false;
    }

    void visitSymbol(TIntermSymbol *node) override
    {
        // It's impossible to reach here with a symbol that needs replacement.
        // MonomorphizeUnsupportedFunctions makes sure that whole structs containing
        // samplers are not passed to functions, so any instance of the struct uniform is
        // necessarily indexed right away.  visitBinary should have already taken care of it.
        ASSERT(mStructureUniformMap.find(&node->variable()) == mStructureUniformMap.end());
    }

  private:
    TCompiler *mCompiler;

    // See RewriteStructSamplersTraverser.
    const StructureMap &mStructureMap;
    const StructureUniformMap &mStructureUniformMap;
    const ExtractedSamplerMap &mExtractedSamplers;
};

// Rewrite the index of an EOpIndexIndirect expression.  The root can never need replacing, because
// it cannot be a sampler itself or of a struct type.
void RewriteIndexExpression(TCompiler *compiler,
                            TIntermTyped *expression,
                            const StructureMap &structureMap,
                            const StructureUniformMap &structureUniformMap,
                            const ExtractedSamplerMap &extractedSamplers)
{
    RewriteExpressionTraverser traverser(compiler, structureMap, structureUniformMap,
                                         extractedSamplers);
    expression->traverse(&traverser);
    bool valid = traverser.updateTree(compiler, expression);
    ASSERT(valid);
}

// Given an expression such as the following:
//
//                                                    EOpIndexDirectStruct (sampler)
//                                                    /                  \
//                                               EOpIndex*           field index
//                                              /        \
//                                EOpIndexDirectStruct   index 2
//                                /                  \
//                           EOpIndex*           field index
//                          /        \
//            EOpIndexDirectStruct   index 1
//            /                  \
//     Uniform Struct           field index
//
// produces:
//
//                                EOpIndex*
//                                /      \
//                           EOpIndex*  index 2
//                          /        \
//                      sampler    index 1
//
// Alternatively, if the expression is as such:
//
//                                                    EOpIndexDirectStruct
//                                                    /                  \
//                        (modified struct type) EOpIndex*           field index
//                                              /        \
//                                EOpIndexDirectStruct   index 2
//                                /                  \
//                           EOpIndex*           field index
//                          /        \
//            EOpIndexDirectStruct   index 1
//            /                  \
//     Uniform Struct           field index
//
// produces:
//
//                                                    EOpIndexDirectStruct
//                                                    /                  \
//                                               EOpIndex*     mapped field index
//                                              /        \
//                                EOpIndexDirectStruct   index 2
//                                /                  \
//                           EOpIndex*      mapped field index
//                          /        \
//            EOpIndexDirectStruct   index 1
//            /                  \
//     Uniform Struct     mapped field index
//
TIntermTyped *RewriteModifiedStructFieldSelectionExpression(
    TCompiler *compiler,
    TIntermBinary *node,
    const StructureMap &structureMap,
    const StructureUniformMap &structureUniformMap,
    const ExtractedSamplerMap &extractedSamplers)
{
    ASSERT(node->getOp() == EOpIndexDirectStruct);

    const bool isSampler = node->getType().isSampler();

    TIntermSymbol *baseUniform = nullptr;
    std::string samplerName;

    TVector<TIntermBinary *> indexNodeStack;

    // Iterate once and build the name of the sampler.
    TIntermBinary *iter = node;
    while (baseUniform == nullptr)
    {
        indexNodeStack.push_back(iter);
        baseUniform = iter->getLeft()->getAsSymbolNode();

        if (isSampler)
        {
            if (iter->getOp() == EOpIndexDirectStruct)
            {
                // When indexed into a struct, get the field name instead and construct the sampler
                // name.
                samplerName.insert(0, iter->getIndexStructFieldName().data());
                samplerName.insert(0, "_");
            }

            if (baseUniform)
            {
                // If left is a symbol, we have reached the end of the chain.  Use the struct name
                // to finish building the name of the sampler.
                samplerName.insert(0, baseUniform->variable().name().data());
            }
        }

        iter = iter->getLeft()->getAsBinaryNode();
    }

    TIntermTyped *rewritten = nullptr;

    if (isSampler)
    {
        ASSERT(extractedSamplers.find(samplerName) != extractedSamplers.end());
        rewritten = new TIntermSymbol(extractedSamplers.at(samplerName));
    }
    else
    {
        const TVariable *baseUniformVar = &baseUniform->variable();
        ASSERT(structureUniformMap.find(baseUniformVar) != structureUniformMap.end());
        rewritten = new TIntermSymbol(structureUniformMap.at(baseUniformVar));
    }

    // Iterate again and build the expression from bottom up.
    for (auto it = indexNodeStack.rbegin(); it != indexNodeStack.rend(); ++it)
    {
        TIntermBinary *indexNode = *it;

        switch (indexNode->getOp())
        {
            case EOpIndexDirectStruct:
                if (!isSampler)
                {
                    // Remap the field.
                    const TStructure *structure = indexNode->getLeft()->getType().getStruct();
                    ASSERT(structureMap.find(structure) != structureMap.end());

                    TIntermConstantUnion *asConstantUnion =
                        indexNode->getRight()->getAsConstantUnion();
                    ASSERT(asConstantUnion);

                    const int fieldIndex = asConstantUnion->getIConst(0);
                    ASSERT(fieldIndex <
                           static_cast<int>(structureMap.at(structure).fieldMap.size()));

                    const int mappedFieldIndex = structureMap.at(structure).fieldMap[fieldIndex];

                    rewritten = new TIntermBinary(EOpIndexDirectStruct, rewritten,
                                                  CreateIndexNode(mappedFieldIndex));
                }
                break;

            case EOpIndexDirect:
                rewritten = new TIntermBinary(EOpIndexDirect, rewritten, indexNode->getRight());
                break;

            case EOpIndexIndirect:
            {
                // Run RewriteExpressionTraverser on the right node.  It may itself be an expression
                // with a sampler inside that needs to be rewritten, or simply use a field of a
                // struct that's remapped.
                TIntermTyped *indexExpression = indexNode->getRight();
                RewriteIndexExpression(compiler, indexExpression, structureMap, structureUniformMap,
                                       extractedSamplers);
                rewritten = new TIntermBinary(EOpIndexIndirect, rewritten, indexExpression);
                break;
            }

            default:
                UNREACHABLE();
                break;
        }
    }

    return rewritten;
}

class RewriteStructSamplersTraverser final : public TIntermTraverser
{
  public:
    explicit RewriteStructSamplersTraverser(TCompiler *compiler, TSymbolTable *symbolTable)
        : TIntermTraverser(true, false, false, symbolTable),
          mCompiler(compiler),
          mRemovedUniformsCount(0)
    {}

    int removedUniformsCount() const { return mRemovedUniformsCount; }

    // Each struct sampler declaration is stripped of its samplers. New uniforms are added for each
    // stripped struct sampler.
    bool visitDeclaration(Visit visit, TIntermDeclaration *decl) override
    {
        if (!mInGlobalScope)
        {
            return true;
        }

        const TIntermSequence &sequence = *(decl->getSequence());
        TIntermTyped *declarator        = sequence.front()->getAsTyped();
        const TType &type               = declarator->getType();

        if (!type.isStructureContainingSamplers())
        {
            return false;
        }

        TIntermSequence newSequence;

        if (type.isStructSpecifier())
        {
            // If this is just a struct definition (not a uniform variable declaration of a
            // struct type), just remove the samplers.  They are not instantiated yet.
            const TStructure *structure = type.getStruct();
            ASSERT(structure && mStructureMap.find(structure) == mStructureMap.end());

            stripStructSpecifierSamplers(structure, &newSequence);
        }
        else
        {
            const TStructure *structure = type.getStruct();

            // If the structure is defined at the same time, create the mapping to the stripped
            // version first.
            if (mStructureMap.find(structure) == mStructureMap.end())
            {
                stripStructSpecifierSamplers(structure, &newSequence);
            }

            // Then, extract the samplers from the struct and create global-scope variables instead.
            TIntermSymbol *asSymbol = declarator->getAsSymbolNode();
            ASSERT(asSymbol);
            const TVariable &variable = asSymbol->variable();
            ASSERT(variable.symbolType() != SymbolType::Empty);

            extractStructSamplerUniforms(variable, structure, &newSequence);
        }

        mMultiReplacements.emplace_back(getParentNode()->getAsBlock(), decl,
                                        std::move(newSequence));

        return false;
    }

    // Same implementation as in RewriteExpressionTraverser.  That traverser cannot replace root.
    bool visitBinary(Visit visit, TIntermBinary *node) override
    {
        TIntermTyped *rewritten = RewriteExpressionVisitBinaryHelper(
            mCompiler, node, mStructureMap, mStructureUniformMap, mExtractedSamplers);

        if (rewritten == nullptr)
        {
            return true;
        }

        queueReplacement(rewritten, OriginalNode::IS_DROPPED);

        // Don't iterate as the expression is rewritten.
        return false;
    }

    // Same implementation as in RewriteExpressionTraverser.  That traverser cannot replace root.
    void visitSymbol(TIntermSymbol *node) override
    {
        ASSERT(mStructureUniformMap.find(&node->variable()) == mStructureUniformMap.end());
    }

  private:
    // Removes all samplers from a struct specifier.
    void stripStructSpecifierSamplers(const TStructure *structure, TIntermSequence *newSequence)
    {
        TFieldList *newFieldList = new TFieldList;
        ASSERT(structure->containsSamplers());

        // Add this struct to the struct map
        ASSERT(mStructureMap.find(structure) == mStructureMap.end());
        StructureData *modifiedData = &mStructureMap[structure];

        modifiedData->modified = nullptr;
        modifiedData->fieldMap.resize(structure->fields().size(), std::numeric_limits<int>::max());

        for (size_t fieldIndex = 0; fieldIndex < structure->fields().size(); ++fieldIndex)
        {
            const TField *field    = structure->fields()[fieldIndex];
            const TType &fieldType = *field->type();

            // If the field is a sampler, or a struct that's entirely removed, skip it.
            if (!fieldType.isSampler() && !isRemovedStructType(fieldType))
            {
                TType *newType = nullptr;

                // Otherwise, if it's a struct that's replaced, create a new field of the replaced
                // type.
                if (fieldType.isStructureContainingSamplers())
                {
                    const TStructure *fieldStruct = fieldType.getStruct();
                    ASSERT(mStructureMap.find(fieldStruct) != mStructureMap.end());

                    const TStructure *modifiedStruct = mStructureMap[fieldStruct].modified;
                    ASSERT(modifiedStruct);

                    newType = new TType(modifiedStruct, true);
                    if (fieldType.isArray())
                    {
                        newType->makeArrays(fieldType.getArraySizes());
                    }
                }
                else
                {
                    // If not, duplicate the field as is.
                    newType = new TType(fieldType);
                }

                // Record the mapping of the field indices, so future EOpIndexDirectStruct's into
                // this struct can be fixed up.
                modifiedData->fieldMap[fieldIndex] = static_cast<int>(newFieldList->size());

                TField *newField =
                    new TField(newType, field->name(), field->line(), field->symbolType());
                newFieldList->push_back(newField);
            }
        }

        // Prune empty structs.
        if (newFieldList->empty())
        {
            return;
        }

        // Declare a new struct with the same name and the new fields.
        modifiedData->modified =
            new TStructure(mSymbolTable, structure->name(), newFieldList, structure->symbolType());
        TType *newStructType = new TType(modifiedData->modified, true);
        TVariable *newStructVar =
            new TVariable(mSymbolTable, kEmptyImmutableString, newStructType, SymbolType::Empty);
        TIntermSymbol *newStructRef = new TIntermSymbol(newStructVar);

        TIntermDeclaration *structDecl = new TIntermDeclaration;
        structDecl->appendDeclarator(newStructRef);

        newSequence->push_back(structDecl);
    }

    // Returns true if the type is a struct that was removed because we extracted all the members.
    bool isRemovedStructType(const TType &type) const
    {
        const TStructure *structure = type.getStruct();
        if (structure == nullptr)
        {
            // Not a struct
            return false;
        }

        // A struct is removed if it is in the map, but doesn't have a replacement struct.
        auto iter = mStructureMap.find(structure);
        return iter != mStructureMap.end() && iter->second.modified == nullptr;
    }

    // Removes samplers from struct uniforms. For each sampler removed also adds a new globally
    // defined sampler uniform.
    void extractStructSamplerUniforms(const TVariable &variable,
                                      const TStructure *structure,
                                      TIntermSequence *newSequence)
    {
        ASSERT(structure->containsSamplers());
        ASSERT(mStructureMap.find(structure) != mStructureMap.end());

        const TType &type = variable.getType();
        enterArray(type);

        for (const TField *field : structure->fields())
        {
            extractFieldSamplers(variable.name().data(), field, newSequence);
        }

        // If there's a replacement structure (because there are non-sampler fields in the struct),
        // add a declaration with that type.
        const TStructure *modified = mStructureMap[structure].modified;
        if (modified != nullptr)
        {
            TType *newType = new TType(modified, false);
            if (type.isArray())
            {
                newType->makeArrays(type.getArraySizes());
            }
            newType->setQualifier(EvqUniform);
            const TVariable *newVariable =
                new TVariable(mSymbolTable, variable.name(), newType, variable.symbolType());

            TIntermDeclaration *newDecl = new TIntermDeclaration();
            newDecl->appendDeclarator(new TIntermSymbol(newVariable));

            newSequence->push_back(newDecl);

            ASSERT(mStructureUniformMap.find(&variable) == mStructureUniformMap.end());
            mStructureUniformMap[&variable] = newVariable;
        }
        else
        {
            mRemovedUniformsCount++;
        }

        exitArray(type);
    }

    // Extracts samplers from a field of a struct. Works with nested structs and arrays.
    void extractFieldSamplers(const std::string &prefix,
                              const TField *field,
                              TIntermSequence *newSequence)
    {
        const TType &fieldType = *field->type();
        if (fieldType.isSampler() || fieldType.isStructureContainingSamplers())
        {
            std::string newPrefix = prefix + "_" + field->name().data();

            if (fieldType.isSampler())
            {
                extractSampler(newPrefix, fieldType, newSequence);
            }
            else
            {
                enterArray(fieldType);
                const TStructure *structure = fieldType.getStruct();
                for (const TField *nestedField : structure->fields())
                {
                    extractFieldSamplers(newPrefix, nestedField, newSequence);
                }
                exitArray(fieldType);
            }
        }
    }

    void GenerateArraySizesFromStack(TVector<unsigned int> *sizesOut)
    {
        sizesOut->reserve(mArraySizeStack.size());

        for (auto it = mArraySizeStack.rbegin(); it != mArraySizeStack.rend(); ++it)
        {
            sizesOut->push_back(*it);
        }
    }

    // Extracts a sampler from a struct. Declares the new extracted sampler.
    void extractSampler(const std::string &newName,
                        const TType &fieldType,
                        TIntermSequence *newSequence)
    {
        ASSERT(fieldType.isSampler());

        TType *newType = new TType(fieldType);

        // Add array dimensions accumulated so far due to struct arrays.  Note that to support
        // nested arrays, mArraySizeStack has the outermost size in the front.  |makeArrays| thus
        // expects this in reverse order.
        TVector<unsigned int> parentArraySizes;
        GenerateArraySizesFromStack(&parentArraySizes);
        newType->makeArrays(parentArraySizes);

        ImmutableStringBuilder nameBuilder(newName.size() + 1);
        nameBuilder << newName;

        newType->setQualifier(EvqUniform);
        TVariable *newVariable =
            new TVariable(mSymbolTable, nameBuilder, newType, SymbolType::AngleInternal);
        TIntermSymbol *newSymbol = new TIntermSymbol(newVariable);

        TIntermDeclaration *samplerDecl = new TIntermDeclaration;
        samplerDecl->appendDeclarator(newSymbol);

        newSequence->push_back(samplerDecl);

        // TODO: Use a temp name instead of generating a name as currently done.  There is no
        // guarantee that these generated names cannot clash.  Create a mapping from the previous
        // name to the name assigned to the temp variable so ShaderVariable::mappedName can be
        // updated post-transformation.  http://anglebug.com/4301
        ASSERT(mExtractedSamplers.find(newName) == mExtractedSamplers.end());
        mExtractedSamplers[newName] = newVariable;
    }

    void enterArray(const TType &arrayType)
    {
        const TSpan<const unsigned int> &arraySizes = arrayType.getArraySizes();
        for (auto it = arraySizes.rbegin(); it != arraySizes.rend(); ++it)
        {
            unsigned int arraySize = *it;
            mArraySizeStack.push_back(arraySize);
        }
    }

    void exitArray(const TType &arrayType)
    {
        mArraySizeStack.resize(mArraySizeStack.size() - arrayType.getNumArraySizes());
    }

    TCompiler *mCompiler;
    int mRemovedUniformsCount;

    // Map structures with samplers to ones that have their samplers removed.
    StructureMap mStructureMap;

    // Map uniform variables of structure type that are replaced with another variable.
    StructureUniformMap mStructureUniformMap;

    // Map a constructed sampler name to its variable.  Used to replace an expression that uses this
    // sampler with the extracted one.
    ExtractedSamplerMap mExtractedSamplers;

    // A stack of array sizes.  Used to figure out the array dimensions of the extracted sampler,
    // for example when it's nested in an array of structs in an array of structs.
    TVector<unsigned int> mArraySizeStack;
};
}  // anonymous namespace

bool RewriteStructSamplers(TCompiler *compiler,
                           TIntermBlock *root,
                           TSymbolTable *symbolTable,
                           int *removedUniformsCountOut)
{
    RewriteStructSamplersTraverser traverser(compiler, symbolTable);
    root->traverse(&traverser);
    *removedUniformsCountOut = traverser.removedUniformsCount();
    return traverser.updateTree(compiler, root);
}
}  // namespace sh