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

Copyright (C) 2024 Intel Corporation

SPDX-License-Identifier: MIT

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

#pragma once

#include "AdaptorCommon/ImplicitArgs.hpp"
#include "Compiler/Optimizer/OpenCLPasses/KernelArgs/KernelArgs.hpp"
#include "Compiler/MetaDataUtilsWrapper.h"

#include "common/LLVMWarningsPush.hpp"
#include "llvmWrapper/IR/Module.h"
#include <llvm/Pass.h>
#include <llvm/IR/InstVisitor.h>
#include <llvm/IR/IRBuilder.h>
#include "common/LLVMWarningsPop.hpp"

namespace IGC
{
    class BufferBoundsChecking : public llvm::ModulePass, public llvm::InstVisitor<BufferBoundsChecking>
    {
    public:
        static char ID;

        BufferBoundsChecking();
        ~BufferBoundsChecking() = default;

        virtual void getAnalysisUsage(llvm::AnalysisUsage& analysisUsage) const override
        {
            analysisUsage.addRequired<MetaDataUtilsWrapper>();
            analysisUsage.addRequired<CodeGenContextWrapper>();
        }

        virtual llvm::StringRef getPassName() const override
        {
            return "BufferBoundsChecking";
        }

        virtual bool runOnModule(llvm::Module& M) override;

        void visitLoadInst(llvm::LoadInst& load);
        void visitStoreInst(llvm::StoreInst& store);

    private:
        struct AccessInfo
        {
            llvm::StringRef filename;
            int line;
            int column;
            llvm::StringRef bufferName;
            llvm::Value* bufferAddress;
            llvm::Value* bufferOffsetInBytes;
            uint32_t implicitArgBufferSizeIndex;
            llvm::Value* elementSizeInBytes;
        };

        bool modified = false;

        IGCMD::MetaDataUtils* metadataUtils = nullptr;
        ModuleMetaData* moduleMetadata = nullptr;
        ImplicitArgs* implicitArgs = nullptr;
        KernelArgs* kernelArgs = nullptr;
        llvm::SmallVector<llvm::Instruction*, 4> loadsAndStoresToCheck;
        llvm::DenseMap<llvm::StringRef, llvm::GlobalVariable*> stringsCache;
        llvm::DICompileUnit* compileUnit = nullptr;
        llvm::Function* bufferSizePlaceholderFunction = nullptr;
        llvm::Value* localId0;
        llvm::Value* localId1;
        llvm::Value* localId2;
        llvm::Value* globalId0;
        llvm::Value* globalId1;
        llvm::Value* globalId2;

        std::tuple<llvm::Value*, llvm::Value*, llvm::Value*> getLocalIds(llvm::Function& function);
        std::tuple<llvm::Value*, llvm::Value*, llvm::Value*> getGlobalIds(llvm::Function& function);

        void handleLoadStore(llvm::Instruction* instruction);

        bool argumentQualifiesForChecking(const llvm::Argument* argument);

        llvm::Value* createBoundsCheckingCondition(const AccessInfo& accessInfo, llvm::Instruction* insertBefore);
        llvm::Value* createLoadStoreReplacement(llvm::Instruction* instruction, llvm::Instruction* insertBefore);
        void createAssertCall(const AccessInfo& accessInfo, llvm::Instruction* insertBefore);
        llvm::SmallVector<llvm::Value*, 4> createAssertArgs(const AccessInfo& accessInfo, llvm::Instruction* insertBefore);
        llvm::GlobalVariable* getOrCreateGlobalConstantString(llvm::Module* M, llvm::StringRef format);
        void createBoundsCheckingCode(llvm::Instruction* instruction, const AccessInfo& accessInfo);

        const IGC::KernelArg* getKernelArg(llvm::Value* value);
        const IGC::KernelArg* getKernelArgFromPtr(const llvm::PointerType& pointerType, llvm::Value* value);
        llvm::Argument* getBufferSizeArg(llvm::Function* function, uint32_t n);
        AccessInfo getAccessInfo(llvm::Instruction* instruction, llvm::Value* value);

        llvm::Value* createBufferSizePlaceholder(uint32_t implicitArgBufferSizeIndex, llvm::Instruction* insertBefore);
    };
}
