/*========================== begin_copyright_notice ============================

Copyright (C) 2017-2024 Intel Corporation

SPDX-License-Identifier: MIT

============================= end_copyright_notice ===========================*/

#pragma once
#include "common/LLVMWarningsPush.hpp"
#include <llvm/IR/InstVisitor.h>
#include <llvm/IR/Instructions.h>
#include <llvm/Pass.h>
#include "common/LLVMWarningsPop.hpp"

#include <vector>
#include <memory>

namespace IGC
{
    enum StatusPrivArr2Reg
    {
        OK,
        IsDynamicAlloca,
        CannotUseSOALayout,
        IsNotNativeType,
        OutOfAllocSizeLimit,
        OutOfMaxGRFPressure
    };

    /// Tries to promote array in private memory to indexable vector
    /// Uses register pressure to make sure it won't cause spilling
    llvm::FunctionPass* createPromotePrivateArrayToReg();

    struct SOALayoutInfo
    {
        /// Signifies if SOA layout be applied to this variable.
        bool canUseSOALayout;
        /// The base type that all memory instructions in the use graph operate
        /// upon. Can be a compound type as long as no instruction in the chain
        /// reads from/stores into the underlying smaller strides of the
        /// primitive type.
        llvm::Type* baseType;
        /// Signifies if all memory instructions that operate on a
        /// compound-typed variable are vector-typed.
        bool allUsesAreVector;
        // Partition size for new algorithm
        uint32_t SOAPartitionBytes;

        SOALayoutInfo() : canUseSOALayout(false), baseType(nullptr),
                          allUsesAreVector(false), SOAPartitionBytes(4) {}
        SOALayoutInfo(bool canUseSOALayout, llvm::Type* baseType,
                      bool allUsesAreVector, uint32_t Size) :
            canUseSOALayout(canUseSOALayout), baseType(baseType),
            allUsesAreVector(false), SOAPartitionBytes (Size) {}
        SOALayoutInfo(SOALayoutInfo&) = default;
        ~SOALayoutInfo() = default;
    };

    /// Conceptually, this class is quite similar to 'llvm::PtrUseVisitor', but
    /// looks to ensure that actual boolean values are returned by 'visit()'
    /// calls. Through this depth-first traversal of the alloca's use-graph, we
    /// aim to determine if "structure of arrays" layout can be applied to the
    /// initialized variable.
    ///
    /// TODO: Consider generalizing the traversal logic as a parent interface
    /// (e.g. 'InstUseVisitorBase') so that it could be re-used across IGC.
    /// Long-term, we could also propose this general use-based visitor to the
    /// LLVM community as an extension/replacement for 'PtrUseVisitor'.
    ///
    /// TODO: Given that PrivateMemoryResolution also relies on SOA layout
    /// evaluation, it would be better to factor this out into a separate
    /// analysis pass.
    class SOALayoutChecker : public llvm::InstVisitor<SOALayoutChecker, bool>
    {
    public:
        friend llvm::InstVisitor<SOALayoutChecker, bool>;

        // isOCL is for testing, it will be removed once testing is done.
        SOALayoutChecker(llvm::AllocaInst& allocaToCheck, bool isOCL);
        SOALayoutChecker() = delete;
        ~SOALayoutChecker() = default;
        SOALayoutChecker(SOALayoutChecker&) = delete;

        SOALayoutInfo getOrGatherInfo();

        // for new algo
        bool useNewAlgo(llvm::Type* baseTy) const {
            return (newAlgoControl > 1 ||
                (newAlgoControl == 1 && baseTy->isStructTy()));
        }

    private:
        llvm::AllocaInst& allocaRef;
        const llvm::DataLayout* pDL;
        std::unique_ptr<SOALayoutInfo> pInfo;

        // ===== fields for new algo =====
        // todo: combine the new and old together
        //
        // SOAPartitionBytes : the size of chunk used to divide a buffer
        //   into a sequence of chunks. It should be a power-of-2 number
        //   with minimum value being 4. It is selected to be the larger
        //   of 4 and size of any scalar element type.
        uint32_t SOAPartitionBytes = 4;
        // newAlgoControl
        //   The old algo : array of DW[xn]
        //   the new algo : array of DW{xn], array of QW[xn],
        //                  array of structs.
        //    0 : disable new algorithm
        //    1 : enable new algorithm for array of simple struct
        //    2 : enable new algorithm for array of simple struct
        //        array of dw[xn], array of qw[xn]
        //        (not splitting vector, intend to replace the old algo)
        //    3 : 2 plus array of more complicated structs.
        int newAlgoControl = 0;
        uint32_t selectPartitionSize(llvm::Type* Ty);
        // Return true if struct can be transposed
        bool checkStruct(llvm::StructType* StTy);
        // ===== end of fields for new algo =====

        bool isVectorSOA = true;
        llvm::Instruction* parentLevelInst = nullptr;

        /// This method visits the direct users of the instruction, abrupting
        /// traversal if a false value is returned by any of the visits.
        ///
        /// TODO: Consider a worklist-based implementation instead. If achieved, it
        /// would make sense to rename the method into `enqueueUsers', in adherence with
        /// the convention set by llvm::PtrUseVisitor.
        bool checkUsers(llvm::Instruction& I);

        /// Given the purpose of the analysis, we can only handle memory
        /// instructions. Default to 'false'.
        bool visitInstruction(llvm::Instruction& I) { return false; }

        bool visitBitCastInst(llvm::BitCastInst&);
        bool visitGetElementPtrInst(llvm::GetElementPtrInst&);
        bool visitIntrinsicInst(llvm::IntrinsicInst&);
        bool visitLoadInst(llvm::LoadInst&);
        bool visitStoreInst(llvm::StoreInst&);
    };

    class TransposeHelper
    {
    public:
        TransposeHelper(const llvm::DataLayout& DL, bool vectorIndex)
            : m_vectorIndex(vectorIndex)
            , m_DL(DL)
        {}
        void HandleAllocaSources(
            llvm::Instruction* v,
            llvm::Value* idx);
        void handleGEPInst(
            llvm::GetElementPtrInst* pGEP,
            llvm::Value* idx);
        // Temporary, this is to replace HandleGEPInst
        void handleGEPInstNew(
            llvm::GetElementPtrInst* pGEP,
            llvm::Value* idx);
        virtual void handleLoadInst(
            llvm::LoadInst* pLoad,
            llvm::Value* pScalarizedIdx) = 0;
        virtual void handleStoreInst(
            llvm::StoreInst* pStore,
            llvm::Value* pScalarizedIdx) = 0;
        virtual void handleLifetimeMark(llvm::IntrinsicInst* inst) = 0;
        // For select handleGEPinstNew
        virtual bool useNewAlgo() { return false; }
        void EraseDeadCode();
    protected:
        const llvm::DataLayout& m_DL;
        std::vector<llvm::Instruction*> m_toBeRemovedGEP;
    private:
        bool m_vectorIndex;
        std::pair<unsigned int, llvm::Type*> getArrSizeAndEltType(llvm::Type* T);
    };
}
