//

// Copyright (C) 2016 LunarG, Inc.

//

// All rights reserved.

//

// Redistribution and use in source and binary forms, with or without

// modification, are permitted provided that the following conditions

// are met:

//

//    Redistributions of source code must retain the above copyright

//    notice, this list of conditions and the following disclaimer.

//

//    Redistributions in binary form must reproduce the above

//    copyright notice, this list of conditions and the following

//    disclaimer in the documentation and/or other materials provided

//    with the distribution.

//

//    Neither the name of 3Dlabs Inc. Ltd. nor the names of its

//    contributors may be used to endorse or promote products derived

//    from this software without specific prior written permission.

//

// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS

// "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT

// LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS

// FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE

// COPYRIGHT HOLDERS OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,

// INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,

// BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;

// LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER

// CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT

// LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN

// ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE

// POSSIBILITY OF SUCH DAMAGE.

//

#ifndef _IOMAPPER_INCLUDED

#define _IOMAPPER_INCLUDED

#include <cstdint>

#include <array>

#include <unordered_map>

#include <unordered_set>

//

// A reflection database and its interface, consistent with the OpenGL API reflection queries.

//

class TInfoSink;

namespace glslang {

class TIntermediate;

struct TVarEntryInfo;

// Base class for shared TIoMapResolver services, used by several derivations.

struct TDefaultIoResolverBase : public glslang::TIoMapResolver {

public:

    TDefaultIoResolverBase(const TIntermediate& intermediate);

    typedef std::vector<int> TSlotSet;

    typedef std::unordered_map<int, TSlotSet> TSlotSetMap;

    typedef std::array<TSlotSetMap, EResCount> TSlotSetMapResourceArray;

    // grow the reflection stage by stage

    void notifyBinding(EShLanguage, TVarEntryInfo& /*ent*/) override {}

    void notifyInOut(EShLanguage, TVarEntryInfo& /*ent*/) override {}

    void beginNotifications(EShLanguage) override {}

    void endNotifications(EShLanguage) override {}

    void beginResolve(EShLanguage) override {}

    void endResolve(EShLanguage) override {}

    void beginCollect(EShLanguage) override {}

    void endCollect(EShLanguage) override {}

    void reserverResourceSlot(TVarEntryInfo& /*ent*/, TInfoSink& /*infoSink*/) override {}

    void reserverStorageSlot(TVarEntryInfo& /*ent*/, TInfoSink& /*infoSink*/) override {}

    int getBaseBinding(EShLanguage stage, TResourceType res, unsigned int set) const;

    const std::vector<std::string>& getResourceSetBinding(EShLanguage stage) const;

    virtual TResourceType getResourceType(const glslang::TType& type) = 0;

    bool doAutoBindingMapping() const;

    bool doAutoLocationMapping() const;

    TSlotSet::iterator findSlot(int resource, int set, int slot);

    bool checkEmpty(int resource, int set, int slot);

    bool validateInOut(EShLanguage /*stage*/, TVarEntryInfo& /*ent*/) override { return true; }

    int reserveSlot(int resource, int set, int slot, int size = 1);

    int getFreeSlot(int resource, int set, int base, int size = 1);

    int resolveSet(EShLanguage /*stage*/, TVarEntryInfo& ent) override;

    int resolveUniformLocation(EShLanguage /*stage*/, TVarEntryInfo& ent) override;

    int resolveInOutLocation(EShLanguage stage, TVarEntryInfo& ent) override;

    int resolveInOutComponent(EShLanguage /*stage*/, TVarEntryInfo& ent) override;

    int resolveInOutIndex(EShLanguage /*stage*/, TVarEntryInfo& ent) override;

    void addStage(EShLanguage stage, TIntermediate& stageIntermediate) override {

        if ((unsigned)stage < EShLangCount) {

            stageMask[(unsigned)stage] = true;

            stageIntermediates[(unsigned)stage] = &stageIntermediate;

        }

    }

    uint32_t computeTypeLocationSize(const TType& type, EShLanguage stage);

    TSlotSetMapResourceArray slots;

    bool hasError = false;

protected:

    TDefaultIoResolverBase(TDefaultIoResolverBase&);

    TDefaultIoResolverBase& operator=(TDefaultIoResolverBase&);

    const TIntermediate& referenceIntermediate;

    int nextUniformLocation;

    int nextInputLocation;

    int nextOutputLocation;

    bool stageMask[EShLangCount + 1];

    const TIntermediate* stageIntermediates[EShLangCount];

    // Return descriptor set specific base if there is one, and the generic base otherwise.

    int selectBaseBinding(int base, int descriptorSetBase) const {

        return descriptorSetBase != -1 ? descriptorSetBase : base;

    }

    static int getLayoutSet(const glslang::TType& type) {

        if (type.getQualifier().hasSet())

            return type.getQualifier().layoutSet;

        else

            return 0;

    }

    static bool isSamplerType(const glslang::TType& type) {

        return type.getBasicType() == glslang::EbtSampler && type.getSampler().isPureSampler();

    }

    static bool isTextureType(const glslang::TType& type) {

        return (type.getBasicType() == glslang::EbtSampler && !type.getSampler().isCombined() &&

                (type.getSampler().isTexture() || type.getSampler().isSubpass()));

    }

    static bool isUboType(const glslang::TType& type) {

        return type.getQualifier().storage == EvqUniform && type.isStruct();

    }

    static bool isImageType(const glslang::TType& type) {

        return type.getBasicType() == glslang::EbtSampler && type.getSampler().isImage();

    }

    static bool isSsboType(const glslang::TType& type) {

        return type.getQualifier().storage == EvqBuffer;

    }

    static bool isCombinedSamplerType(const glslang::TType& type) {

        return type.getBasicType() == glslang::EbtSampler && type.getSampler().isCombined();

    }

    static bool isAsType(const glslang::TType& type) {

        return type.getBasicType() == glslang::EbtAccStruct;

    }

    static bool isTensorType(const glslang::TType& type) {

        return type.isTensorARM();

    }

    static bool isValidGlslType(const glslang::TType& type) {

        // at most one must be true

        return (isSamplerType(type) + isTextureType(type) + isUboType(type) + isImageType(type) +

            isSsboType(type) + isCombinedSamplerType(type) + isAsType(type) + isTensorType(type)) <= 1;

    }

    // Return true if this is a SRV (shader resource view) type:

    static bool isSrvType(const glslang::TType& type) {

        return isTextureType(type) || type.getQualifier().storage == EvqBuffer;

    }

    // Return true if this is a UAV (unordered access view) type:

    static bool isUavType(const glslang::TType& type) {

        if (type.getQualifier().isReadOnly())

            return false;

        return (type.getBasicType() == glslang::EbtSampler && type.getSampler().isImage()) ||

                (type.getQualifier().storage == EvqBuffer);

    }

};

// Default I/O resolver for OpenGL

struct TDefaultGlslIoResolver : public TDefaultIoResolverBase {

public:

    typedef std::map<TString, int> TVarSlotMap;  // <resourceName, location/binding>

    typedef std::map<int, TVarSlotMap> TSlotMap; // <resourceKey, TVarSlotMap>

    typedef std::array<TSlotMap, EResCount> TResourceSlotMap;

    TDefaultGlslIoResolver(const TIntermediate& intermediate);

    bool validateBinding(EShLanguage /*stage*/, TVarEntryInfo& /*ent*/) override { return true; }

    TResourceType getResourceType(const glslang::TType& type) override;

    int resolveInOutLocation(EShLanguage stage, TVarEntryInfo& ent) override;

    int resolveUniformLocation(EShLanguage /*stage*/, TVarEntryInfo& ent) override;

    int resolveBinding(EShLanguage /*stage*/, TVarEntryInfo& ent) override;

    void beginResolve(EShLanguage /*stage*/) override;

    void endResolve(EShLanguage stage) override;

    void beginCollect(EShLanguage) override;

    void endCollect(EShLanguage) override;

    void reserverStorageSlot(TVarEntryInfo& ent, TInfoSink& infoSink) override;

    void reserverResourceSlot(TVarEntryInfo& ent, TInfoSink& infoSink) override;

    // in/out symbol and uniform symbol are stored in the same resourceSlotMap, the storage key is used to identify each type of symbol.

    // We use stage and storage qualifier to construct a storage key. it can help us identify the same storage resource used in different stage.

    // if a resource is a program resource and we don't need know it usage stage, we can use same stage to build storage key.

    // Note: both stage and type must less then 0xffff.

    int buildStorageKey(EShLanguage stage, TStorageQualifier type) {

        assert(static_cast<uint32_t>(stage) <= 0x0000ffff && static_cast<uint32_t>(type) <= 0x0000ffff);

        return (stage << 16) | type;

    }

protected:

    // Use for mark pre stage, to get more interface symbol information.

    EShLanguage preStage;

    // Use for mark current shader stage for resolver

    EShLanguage currentStage;

    // Slot map for storage resource(location of uniform and interface symbol) It's a program share slot

    TResourceSlotMap resourceSlotMap;

    // Slot map for other resource(image, ubo, ssbo), It's a program share slot.

    TSlotMap storageSlotMap;

};

typedef std::map<TString, TVarEntryInfo> TVarLiveMap;

// I/O mapper for GLSL

class TGlslIoMapper : public TIoMapper {

public:

    TGlslIoMapper();

    virtual ~TGlslIoMapper();

    // If set, the uniform block with the given name will be changed to be backed by

    // push_constant if it's size is <= maxSize

    bool setAutoPushConstantBlock(const char* name, unsigned int maxSize, TLayoutPacking packing) override {

        autoPushConstantBlockName = name;

        autoPushConstantMaxSize = maxSize;

        autoPushConstantBlockPacking = packing;

        return true;

    }

    // grow the reflection stage by stage

    bool addStage(EShLanguage, TIntermediate&, TInfoSink&, TIoMapResolver*) override;

    bool doMap(TIoMapResolver*, TInfoSink&) override;

    TIntermediate* intermediates[EShLangCount];

    bool hadError = false;

    bool relaxSetBindingLimits = false;

    EProfile profile;

    int version;

private:

    TString autoPushConstantBlockName;

    unsigned int autoPushConstantMaxSize;

    TLayoutPacking autoPushConstantBlockPacking;

    TVarLiveMap *inVarMaps[EShLangCount], *outVarMaps[EShLangCount],

                *uniformVarMap[EShLangCount];

};

} // end namespace glslang

#endif // _IOMAPPER_INCLUDED