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

Copyright (C) 2017-2024 Intel Corporation

SPDX-License-Identifier: MIT

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

#include "GASRetValuePropagator.h"

#include "Compiler/CISACodeGen/ShaderCodeGen.hpp"
#include "Compiler/InitializePasses.h"
#include "LLVM3DBuilder/MetadataBuilder.h"

#include "common/LLVMWarningsPush.hpp"
#include "llvm/Analysis/CallGraph.h"
#include "llvm/ADT/PostOrderIterator.h"
#include "llvm/BinaryFormat/Dwarf.h"
#include "llvmWrapper/IR/Function.h"
#include "common/LLVMWarningsPop.hpp"
#include "llvmWrapper/IR/DerivedTypes.h"
#include "DebugInfo/DwarfDebug.hpp"
#include <optional>

#define PASS_FLAG "igc-gas-ret-value-propagator"
#define PASS_DESC "Resolve generic pointer return value"
#define PASS_CFG_ONLY false
#define PASS_ANALYSIS false
IGC_INITIALIZE_PASS_BEGIN(GASRetValuePropagator, PASS_FLAG, PASS_DESC, PASS_CFG_ONLY, PASS_ANALYSIS)
IGC_INITIALIZE_PASS_DEPENDENCY(MetaDataUtilsWrapper)
IGC_INITIALIZE_PASS_DEPENDENCY(CodeGenContextWrapper)
IGC_INITIALIZE_PASS_DEPENDENCY(CallGraphWrapperPass)
IGC_INITIALIZE_PASS_DEPENDENCY(LoopInfoWrapperPass)
IGC_INITIALIZE_PASS_END(GASRetValuePropagator, PASS_FLAG, PASS_DESC, PASS_CFG_ONLY, PASS_ANALYSIS)

ModulePass* IGC::createGASRetValuePropagatorPass() { return new GASRetValuePropagator(); }

char GASRetValuePropagator::ID = 0;

bool GASRetValuePropagator::runOnModule(Module& M) {
    bool changed = false;
    m_module = &M;
    m_mdUtils = getAnalysis<MetaDataUtilsWrapper>().getMetaDataUtils();
    m_ctx = getAnalysis<CodeGenContextWrapper>().getCodeGenContext();

    CallGraph& CG = getAnalysis<CallGraphWrapperPass>().getCallGraph();
    std::vector<Function*> candidates = findCandidates(CG);

    for (auto* F : candidates)
    {
        LoopInfo& LI = getAnalysis<LoopInfoWrapperPass>(*F).getLoopInfo();
        GASPropagator ThePropagator(F->getContext(), &LI);
        m_Propagator = &ThePropagator;

        if (propagateReturnValue(F))
        {
            changed = true;
        }
    }

    return changed;
}

bool GASRetValuePropagator::propagateReturnValue(Function*& F) {
    PointerType* nonGASPtr = getRetValueNonGASType(F);

    if (!nonGASPtr) return false;

    Function* newFunc = cloneFunctionWithModifiedRetType(F, nonGASPtr);

    updateAllUsesWithNewFunction(F, newFunc);

    IGC_ASSERT(nullptr != F);
    IGC_ASSERT_MESSAGE(F->use_empty(), "All function uses should have been transfered to new function");
    F->eraseFromParent();
    F = newFunc;
    return true;
}

std::vector<Function*> GASRetValuePropagator::findCandidates(CallGraph& CG) {
    std::vector<Function*> candidates;

    auto skip = [](Function* F)
    {
        // Skip functions with variable number of arguments, e.g. printf.
        if (F->isVarArg())
            return true;

        // Only non-extern functions within the module are optimized
        if (F->hasFnAttribute("referenced-indirectly") || F->isDeclaration()
            || F->isIntrinsic() || F->user_empty())
            return true;

        return false;
    };

    auto isGenericPtrTy = [](Type* T)
    {
        return T->isPointerTy() && T->getPointerAddressSpace() == ADDRESS_SPACE_GENERIC;
    };

    // Find the candidates, which are functions returning generic pointer args.
    // Functions will be updated later in down-top ordering (starting from most nested function).
    for (auto I : post_order(&CG))
    {
        auto F = I->getFunction();
        if (F == nullptr)
            continue;
        if (skip(F))
            continue;
        if (!isGenericPtrTy(F->getReturnType()))
            continue;

        candidates.push_back(F);
    }

    return candidates;
}

std::vector<ReturnInst*> GASRetValuePropagator::getAllRetInstructions(Function& F)
{
    std::vector<ReturnInst*> retInstructions;
    for (auto& BB : F)
    {
        if (auto retInst = dyn_cast<ReturnInst>(BB.getTerminator()))
        {
            retInstructions.push_back(retInst);
        }
    }
    return retInstructions;
}

PointerType* GASRetValuePropagator::getRetValueNonGASType(Function* F)
{
    std::vector<ReturnInst*> retInstructions = getAllRetInstructions(*F);

    std::optional<unsigned> originAddrSpace = std::nullopt;
    for (auto retInst : retInstructions)
    {
        Value* retValue = retInst->getReturnValue();

        if (isa<ConstantPointerNull>(retValue))
            continue;

        if (!isa<AddrSpaceCastInst>(retValue))
            return nullptr;

        auto I = cast<AddrSpaceCastInst>(retValue);
        IGC_ASSERT(I->getDestAddressSpace() == ADDRESS_SPACE_GENERIC);

        unsigned AS = I->getSrcAddressSpace();
        if (originAddrSpace && originAddrSpace.value() != AS)
            return nullptr;

        originAddrSpace.emplace(AS);
    }

    PointerType* retTy = cast<PointerType>(F->getReturnType());

    return originAddrSpace ?
        IGCLLVM::getWithSamePointeeType(retTy, originAddrSpace.value()) :
        nullptr;
}

Function* GASRetValuePropagator::createNewFunctionDecl(Function* oldFunc, Type* newRetTy)
{
    Module* M = oldFunc->getParent();
    ArrayRef<Type*> params = oldFunc->getFunctionType()->params();
    FunctionType* newFTy = FunctionType::get(newRetTy, params, oldFunc->isVarArg());

    Function* newFunc = Function::Create(newFTy, oldFunc->getLinkage());
    newFunc->copyAttributesFrom(oldFunc);
    newFunc->setSubprogram(oldFunc->getSubprogram());
    M->getFunctionList().insert(oldFunc->getIterator(), newFunc);
    newFunc->takeName(oldFunc);
    return newFunc;
}

void GASRetValuePropagator::transferFunctionBody(Function* oldFunc, Function* newFunc)
{
    newFunc->stealArgumentListFrom(*oldFunc);
    IGCLLVM::splice(newFunc, newFunc->begin(), oldFunc);
}

void GASRetValuePropagator::updateFunctionRetInstruction(Function* F)
{
    std::vector<ReturnInst*> retInstructions = getAllRetInstructions(*F);

    for (auto retInst : retInstructions)
    {
        Value* retValue = retInst->getReturnValue();

        if (isa<ConstantPointerNull>(retValue))
        {
            retInst->setOperand(0, ConstantPointerNull::get(cast<PointerType>(F->getReturnType())));
            continue;
        }

        IGC_ASSERT(isa<AddrSpaceCastInst>(retValue));

        auto ASC = cast<AddrSpaceCastInst>(retValue);
        IGC_ASSERT(ASC->getDestAddressSpace() == ADDRESS_SPACE_GENERIC);

        retInst->setOperand(0, ASC->getPointerOperand());

        if (ASC->use_empty()) ASC->eraseFromParent();
    }
}

void GASRetValuePropagator::updateAllUsesWithNewFunction(Function* oldFunc, Function* newFunc)
{
    IGC_ASSERT(!oldFunc->use_empty());

    // Keep track of old calls and addrspacecast to be deleted later
    std::vector<CallInst*> callsToDelete;

    for (auto U : oldFunc->users())
    {
        CallInst* cInst = dyn_cast<CallInst>(U);
        if (!cInst)
        {
            IGC_ASSERT_MESSAGE(0, "Unknown function usage");
            return;
        }

        // Prepare args for new call
        std::vector<Value*> callArgs;
        for (unsigned I = 0, E = IGCLLVM::getNumArgOperands(cInst); I != E; ++I) {
            callArgs.push_back(cInst->getArgOperand(I));
        }

        // Create new call and insert it before old one
        CallInst* newCall = CallInst::Create(newFunc, callArgs, "", cInst);

        newCall->setCallingConv(newFunc->getCallingConv());
        newCall->setAttributes(cInst->getAttributes());
        newCall->setDebugLoc(cInst->getDebugLoc());

        IGC_ASSERT(oldFunc->getType()->isPointerTy() &&
            newFunc->getReturnType()->isPointerTy());

        auto* oldRetTy = dyn_cast<PointerType>(oldFunc->getReturnType());
        auto* newRetTy = dyn_cast<PointerType>(newFunc->getReturnType());

        IGC_ASSERT(
            oldRetTy->getAddressSpace() == ADDRESS_SPACE_GENERIC &&
            newRetTy->getAddressSpace() != ADDRESS_SPACE_GENERIC);

        auto ASC = CastInst::Create(Instruction::AddrSpaceCast, newCall, oldFunc->getReturnType(), "", cInst);

        cInst->replaceAllUsesWith(ASC);
        callsToDelete.push_back(cInst);

        m_Propagator->propagate(newCall);
    }

    // Delete old calls
    for (auto call : callsToDelete)
    {
        call->eraseFromParent();
    }
}

Function* GASRetValuePropagator::cloneFunctionWithModifiedRetType(Function* F, PointerType* newRetTy)
{
    Function* newFunc = createNewFunctionDecl(F, newRetTy);
    transferFunctionBody(F, newFunc);
    updateFunctionRetInstruction(newFunc);
    updateMetadata(F, newFunc);
    return newFunc;
}

void GASRetValuePropagator::updateMetadata(Function* oldFunc, Function* newFunc) {
    MetadataBuilder mbuilder(m_module);
    auto& FuncMD = m_ctx->getModuleMetaData()->FuncMD;

    auto oldFuncIter = m_mdUtils->findFunctionsInfoItem(oldFunc);
    m_mdUtils->setFunctionsInfoItem(newFunc, oldFuncIter->second);
    m_mdUtils->eraseFunctionsInfoItem(oldFuncIter);
    mbuilder.UpdateShadingRate(oldFunc, newFunc);
    updateDwarfAddressSpace(newFunc);

    auto loc = FuncMD.find(oldFunc);
    if (loc != FuncMD.end())
    {
        auto funcInfo = loc->second;
        FuncMD.erase(oldFunc);
        FuncMD[newFunc] = funcInfo;
    }

    m_mdUtils->save(m_module->getContext());
}

void GASRetValuePropagator::updateDwarfAddressSpace(Function *F) {
    DISubprogram *subprogram = F->getSubprogram();
    if (!subprogram)
        return;

    // Currently only SLM tag needed
    if (F->getReturnType()->getPointerAddressSpace() != ADDRESS_SPACE_LOCAL)
        return;

    DISubroutineType *subtype = subprogram->getType();
    IGC_ASSERT_MESSAGE(subtype, "Type field must point at DISubroutineType");

    DITypeRefArray functionTypes = subtype->getTypeArray();
    IGC_ASSERT_MESSAGE(functionTypes.size() > 0, "DITypeRefArray can't be empty");
    IGC_ASSERT_MESSAGE(functionTypes[0], "Null return value not expected");

    DIDerivedType *returnType = cast<DIDerivedType>(functionTypes[0]);

    auto isPtrOrReferenceTag = [](unsigned tag) {
      return tag == llvm::dwarf::DW_TAG_pointer_type ||
             tag == llvm::dwarf::DW_TAG_reference_type;
    };

    DIDerivedType *prevType = nullptr;
    while (!isPtrOrReferenceTag(returnType->getTag())) {
        prevType = returnType;
        returnType = cast<DIDerivedType>(returnType->getBaseType());
    }

    DIDerivedType *newType = getDIDerivedTypeWithDwarfAddrspace(
        returnType, DwarfLocalAddressSpaceTag);

    if (prevType) {
        IGC_ASSERT(prevType->getOperand(3).get() &&
                   isa<DIDerivedType>(prevType->getOperand(3).get()));
        prevType->replaceOperandWith(3, newType);
    }
    else
    {
        IGC_ASSERT(functionTypes.get()->getOperand(0).get() &&
                   isa<DIDerivedType>(functionTypes.get()->getOperand(0).get()));
        functionTypes.get()->replaceOperandWith(0, newType);
    }
}

DIDerivedType *
GASRetValuePropagator::getDIDerivedTypeWithDwarfAddrspace(DIDerivedType* type,
    unsigned dwarfTag) {
    return DIDerivedType::get(
        type->getContext(), type->getTag(), type->getName(), type->getFile(),
        type->getLine(), type->getScope(), type->getBaseType(),
        type->getSizeInBits(), type->getAlignInBits(), type->getOffsetInBits(),
        dwarfTag, type->getFlags(), type->getExtraData());
}
