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

Copyright (C) 2022-2023 Intel Corporation

SPDX-License-Identifier: MIT

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

//===----------------------------------------------------------------------===//
//
// The purpose of this pass is to add a limit for each loop,
// so no infinite loops occur.
//===----------------------------------------------------------------------===//

#include <vector>
#include <set>

#include "CapLoopIterationsPass.h"
#include "IGCIRBuilder.h"
#include "Compiler/IGCPassSupport.h"

#include "common/LLVMWarningsPush.hpp"
#include "llvm/Analysis/LoopInfo.h"
#include "llvm/IR/Dominators.h"
#include "common/LLVMWarningsPop.hpp"

using namespace llvm;
using namespace IGC;

// Register pass to igc-opt
#define PASS_FLAG "cap-loop-iterations-pass"
#define PASS_DESCRIPTION "Limit the number of iterations in each loop to UINT_MAX. Prevents inifinite loops"
#define PASS_CFG_ONLY false
#define PASS_ANALYSIS false
IGC_INITIALIZE_PASS_BEGIN(CapLoopIterations, PASS_FLAG, PASS_DESCRIPTION, PASS_CFG_ONLY, PASS_ANALYSIS)
IGC_INITIALIZE_PASS_DEPENDENCY(LoopInfoWrapperPass)
IGC_INITIALIZE_PASS_END(CapLoopIterations, PASS_FLAG, PASS_DESCRIPTION, PASS_CFG_ONLY, PASS_ANALYSIS)

char CapLoopIterations::ID = 0;

CapLoopIterations::CapLoopIterations() : FunctionPass(ID), m_iterationLimit(UINT_MAX)
{
    initializeCapLoopIterationsPass(*PassRegistry::getPassRegistry());
}

CapLoopIterations::CapLoopIterations(uint32_t iterationLimit) : FunctionPass(ID), m_iterationLimit(iterationLimit)
{
    initializeCapLoopIterationsPass(*PassRegistry::getPassRegistry());
}

void CapLoopIterations::getAnalysisUsage(llvm::AnalysisUsage& AU) const
{
    AU.addRequired<LoopInfoWrapperPass>();
}

bool CapLoopIterations::runOnFunction(Function& F)
{
    auto changed = false;
    auto& loopInfo = getAnalysis<LoopInfoWrapperPass>().getLoopInfo();
    auto IRB = IRBuilder(F.getContext());

    for (auto loop : loopInfo.getLoopsInPreorder())
    {
        if (!loop->isLoopSimplifyForm()) // preheader is required
            continue;

        IRB.SetInsertPoint(loop->getHeader()->getFirstNonPHI());

        auto counterphi = IRB.CreatePHI(IRB.getInt32Ty(), 2, "counterphi");
        counterphi->addIncoming(IRB.getInt32(0), loop->getLoopPreheader());

        auto counter = IRB.CreateAdd(counterphi, IRB.getInt32(1), "counter");
        counterphi->addIncoming(counter, loop->getLoopLatch());

        auto forceexit = IRB.CreateICmpEQ(counter, IRB.getInt32(m_iterationLimit), "forceloopexit");

        llvm::SmallVector<std::pair<BasicBlock*, BasicBlock*>, 4> exitedges;

        loop->getExitEdges(exitedges);

        for (const auto& [exitingbb, exitbb] : exitedges)
        {
            if (auto br = dyn_cast<BranchInst>(exitingbb->getTerminator()))
            {
                if (!br->isConditional())
                    continue;

                IRB.SetInsertPoint(br);

                if (br->getSuccessor(0) == exitbb) // br i1 %cond, %exit, %notexit
                {
                    br->setCondition(
                        IRB.CreateOr(forceexit, br->getCondition())
                    );
                }
                else // br i1 %cond, %notexit, %exit
                {
                    br->setCondition(
                        IRB.CreateAnd(IRB.CreateNot(forceexit), br->getCondition())
                    );
                }
            }
        }

        changed = true;
    }
    return changed;
}
