//
// This file is distributed under the MIT License. See LICENSE.md for details.
//

#include <functional>
#include <optional>

#include "llvm/ADT/StringExtras.h"
#include "llvm/IR/Dominators.h"
#include "llvm/IR/InstIterator.h"
#include "llvm/IR/Instruction.h"
#include "llvm/IR/Instructions.h"
#include "llvm/IR/LegacyPassManager.h"
#include "llvm/IR/Verifier.h"
#include "llvm/Pass.h"
#include "llvm/Support/Casting.h"
#include "llvm/Support/FormatVariadic.h"

#include "revng/ABI/ModelHelpers.h"
#include "revng/Lift/LoadBinaryPass.h"
#include "revng/Model/Architecture.h"
#include "revng/Model/Binary.h"
#include "revng/Model/FunctionTags.h"
#include "revng/Model/LoadModelPass.h"
#include "revng/Model/NameBuilder.h"
#include "revng/Model/RawBinaryView.h"
#include "revng/Pipeline/Context.h"
#include "revng/Pipeline/Contract.h"
#include "revng/Pipeline/Kind.h"
#include "revng/Pipeline/LLVMContainer.h"
#include "revng/Pipeline/RegisterPipe.h"
#include "revng/Pipes/FileContainer.h"
#include "revng/Pipes/FunctionPass.h"
#include "revng/Pipes/Kinds.h"
#include "revng/RemoveLiftingArtifacts/MakeSegmentRef.h"
#include "revng/Support/Debug.h"
#include "revng/Support/IRHelpers.h"
#include "revng/Support/MetaAddress.h"
#include "revng/Support/OpaqueFunctionsPool.h"

using namespace llvm;

// Generating function pointers so far relied on the fact that all the
// isolated functions were in the same LLVM module. Since in the new pipeline
// this is no longer the case (there's one isolated function per LLVM module)
// the generating of function pointers needs to be reworked. As a temporary fix
// the new pipeline will emit pointers to segments. To allow comparison between
// the old and new pipeline, this option allows the old pipeline to also have
// the same behaviour.
static cl::opt<bool> DisableFunctionPointersOpt("disable-function-pointers",
                                                cl::desc("Disable emitting "
                                                         "pointers to "
                                                         "functions through "
                                                         "segments"),
                                                cl::init(false));

struct MakeSegmentRefPassImpl : public pipeline::FunctionPassImpl {
private:
  const model::Binary &Binary;
  llvm::Module &M;
  llvm::LLVMContext &Context;
  OpaqueFunctionsPool<FunctionTags::SegmentRefPoolKey> SegmentRefPool;
  OpaqueFunctionsPool<FunctionTags::TypePair> AddressOfPool;
  OpaqueFunctionsPool<FunctionTags::StringLiteralPoolKey> StringLiteralPool;
  RawBinaryView &BinaryView;
  bool DisableFunctionPointers = false;

public:
  MakeSegmentRefPassImpl(llvm::ModulePass &Pass,
                         const model::Binary &Binary,
                         llvm::Module &M) :
    pipeline::FunctionPassImpl(Pass),
    Binary(Binary),
    M(M),
    Context(M.getContext()),
    SegmentRefPool(FunctionTags::SegmentRef.getPool(M)),
    AddressOfPool(FunctionTags::AddressOf.getPool(M)),
    StringLiteralPool(FunctionTags::StringLiteral.getPool(M)),
    BinaryView(getAnalysis<LoadBinaryWrapperPass>().get()),
    DisableFunctionPointers(DisableFunctionPointersOpt) {}

  MakeSegmentRefPassImpl(const model::Binary &Binary,
                         llvm::Module &M,
                         RawBinaryView &BinaryView) :
    pipeline::FunctionPassImpl(),
    Binary(Binary),
    M(M),
    Context(M.getContext()),
    SegmentRefPool(FunctionTags::SegmentRef.getPool(M)),
    AddressOfPool(FunctionTags::AddressOf.getPool(M)),
    StringLiteralPool(FunctionTags::StringLiteral.getPool(M)),
    BinaryView(BinaryView),
    // TODO: once functions pointers are derived post-lift this can be removed
    DisableFunctionPointers(true) {}

  bool runOnFunction(const model::Function &ModelFunction,
                     llvm::Function &Function) override;

public:
  static void getAnalysisUsage(llvm::AnalysisUsage &AU);
};

void MakeSegmentRefPassImpl::getAnalysisUsage(llvm::AnalysisUsage &AU) {
  AU.addRequired<LoadBinaryWrapperPass>();
  AU.addRequired<LoadModelWrapperPass>();
}

static std::optional<std::pair<MetaAddress, uint64_t>>
findSegmentContainingLiteral(const model::Binary &Binary, uint64_t Literal) {
  std::optional<std::pair<MetaAddress, uint64_t>> Result = std::nullopt;
  auto Architecture = Binary.Architecture();

  for (const auto &Segment : Binary.Segments()) {
    if (Segment.contains(MetaAddress::fromGeneric(Architecture, Literal))) {
      revng_assert(not Result.has_value());
      Result = { { Segment.StartAddress(), Segment.VirtualSize() } };
    }
  }

  return Result;
}

static std::optional<llvm::StringRef>
getStringLiteral(RawBinaryView &BinaryView,
                 MetaAddress SegmentAddress,
                 uint64_t SegmentVirtualSize,
                 uint64_t StringOffsetInSegment) {
  // If the segment is not read only it's not a string literal
  const bool IsReadOnly = BinaryView.isReadOnly(SegmentAddress,
                                                SegmentVirtualSize);
  if (not IsReadOnly)
    return std::nullopt;

  MetaAddress StringAddress = SegmentAddress + StringOffsetInSegment;
  uint64_t SegmentSizeFromString = SegmentVirtualSize - StringOffsetInSegment;

  // If it's not there, it's not a string literal
  auto DataOrNone = BinaryView.getStringByAddress(StringAddress,
                                                  SegmentSizeFromString);
  if (not DataOrNone.has_value())
    return std::nullopt;

  llvm::StringRef StringView = DataOrNone.value();

  // If it doesn't end with \0 it's not a string literal
  auto NullTerminatorPos = StringView.find('\0');
  if (NullTerminatorPos == llvm::StringLiteral::npos)
    return std::nullopt;

  StringView = StringView.take_front(NullTerminatorPos);

  // If some of the characters are not printable, is not a string literal
  const auto IsPrintableInCString = [](const char C) {
    return not llvm::isPrint(C) and not llvm::isSpace(C);
  };
  if (llvm::any_of(StringView, IsPrintableInCString))
    return std::nullopt;

  return StringView;
}

static bool isBitcast(const ConstantExpr *E) {
  auto OpCode = E->getOpcode();
  return OpCode == Instruction::BitCast or OpCode == Instruction::IntToPtr
         or OpCode == Instruction::PtrToInt;
}

static Value *skipConstexprBitcasts(Value *V) {
  while (isa<ConstantExpr>(V) and isBitcast(cast<ConstantExpr>(V)))
    V = cast<User>(V)->getOperand(0);
  return V;
}

bool MakeSegmentRefPassImpl::runOnFunction(const model::Function &ModelFunction,
                                           llvm::Function &F) {
  model::CNameBuilder NameBuilder(Binary);

  using llvm::DominatorTree;
  DominatorTree DT;
  DT.recalculate(F);

  std::map<MetaAddress, MetaAddress> FunctionEntries;
  for (const model::Function &Function : Binary.Functions())
    FunctionEntries[Function.Entry().toGeneric()] = Function.Entry();

  // TODO: the checks should be enabled conditionally based on the user.
  revng::NonDebugInfoCheckingIRBuilder IRB(Context);

  bool Changed = false;
  llvm::Type *PtrSizedInteger = getPointerSizedInteger(Context,
                                                       Binary.Architecture());

  for (Instruction &I : instructions(F)) {

    if (isCallToTagged(&I, FunctionTags::AllocatesLocalVariable))
      continue;

    // In the following loop we're going to look at each constant operand of I,
    // and see if it can be replaced either with a string literal or with a
    // reference to a segment, or with a reference to a function.
    //
    // If we do the replacement blindly, and the same constant is used many
    // times in the same instruction, we will end up replacing that constant
    // many times in the same instruction, replacing it with redundant
    // instructions that could be CSE's.
    //
    // In addition, if I is a PHINode, we could end up generating invalid IR, if
    // the PHI has the same constant incoming value many times coming from the
    // same predecessor block. In that case, the IR would be valid, but if we
    // replace the same incoming constant with many different expressions (even
    // if equivalent), the IR doesn't verify anymore, because now we have many
    // different incoming values (the replacement instructions) from the same
    // predecessor block.
    //
    // To avoid this problem, we need map to track how we replace each constant
    // operand in a given instruction. In this way, if an instruction uses the
    // same constant many times, we we'll emit a single expression as a
    // replacement, that will have many uses in the user instruction.
    // Given that we have to do this for PHINodes for verification reasons, we
    // might as well do it for all instructions, to keep the generated IR
    // somewhat smaller.
    //
    // This is correct only if all the Values we generate as a replacement
    // don't have any side effects.
    // The key is Constant, but we expect them to be only ConstantInt and
    // cast-like ConstantExpr.
    DenseMap<Constant *, Value *> ConstantOpReplacements;

    for (Use &Op : I.operands()) {

      // Operands of calls to OpaqueExtractValue should never be promoted to
      // references to segments
      if (isCallToTagged(Op.getUser(), FunctionTags::OpaqueExtractValue))
        continue;

      if (isa<SwitchInst>(&I) && (cast<SwitchInst>(&I)->getCondition() != Op))
        continue;

      if (auto *BI = dyn_cast<BranchInst>(&I))
        if (BI->isConditional())
          if (BI->getCondition() != Op)
            continue;

      // If Op is not a constant it cannot be a reference to a segment, nor
      // a string literal, nor a function.
      // TODO: this will change when we properly support PIE code. But when that
      // happens this part of the code should be gone already.
      if (not isa<ConstantInt>(Op)
          and not isa<ConstantInt>(skipConstexprBitcasts(Op)))
        continue;

      ConstantInt *LeafConstOp = cast<ConstantInt>(skipConstexprBitcasts(Op));

      // Skip stuff that is not pointer-sized
      {
        using namespace model::Architecture;
        auto PointerSize = getPointerSize(Binary.Architecture());
        if (LeafConstOp->getBitWidth() != (8 * PointerSize))
          continue;
      }

      auto *ConstOp = cast<Constant>(Op);
      // We've already this very same constant used as another operand of this
      // very same instruction. Just reuse the replacement that we've computed
      // earlier.
      if (auto It = ConstantOpReplacements.find(ConstOp);
          It != ConstantOpReplacements.end()) {
        I.setOperand(Op.getOperandNo(), It->second);
        Changed = true;
        revng::verify(&M);
        continue;
      }

      uint64_t ConstantAddress = LeafConstOp->getZExtValue();

      if (auto Segment = findSegmentContainingLiteral(Binary, ConstantAddress);
          Segment) {

        const auto &[StartAddress, VirtualSize] = *Segment;
        auto OffsetInSegment = ConstantAddress - StartAddress.address();
        MetaAddress Address = StartAddress + OffsetInSegment;

        auto *PHI = dyn_cast<PHINode>(&I);
        if (PHI) {
          BasicBlock
            *InsertionBlock = DT[PHI->getParent()]->getIDom()->getBlock();
          IRB.SetInsertPoint(InsertionBlock->getTerminator());
        } else {
          IRB.SetInsertPoint(&I);
        }

        Value *ConstantOpReplacement = nullptr;

        // If Address matches a function entry, the constant operand must be
        // replaced with a function reference.
        auto It = FunctionEntries.find(Address);
        if (It != FunctionEntries.end() and not DisableFunctionPointers) {
          auto Name = llvmName(Binary.Functions().at(It->second));

          auto *ReferencedFunction = M.getFunction(Name);
          revng_assert(ReferencedFunction != nullptr);
          ConstantOpReplacement = IRB.CreatePtrToInt(ReferencedFunction,
                                                     Op->getType());
        } else {

          // Then we we check to see if the address matches a string literal.

          // Check if the use of this constant is a icmp. If it is we cannot
          // replace it with a string literal, because comparisons between
          // string literals are undefined behavior in C.
          bool UseIsComparison = llvm::isa<llvm::CmpInst>(Op.getUser());

          // Check if the Op is large as a pointer. If it isn't it can't be a
          // string literal.
          // See if we can find a string literal there.
          auto OptString = getStringLiteral(BinaryView,
                                            StartAddress,
                                            VirtualSize,
                                            OffsetInSegment);

          if (not UseIsComparison and OptString.has_value()) {
            auto Str = OptString.value();
            Constant *ConstExpr = getUniqueString(&M, Str);

            auto RealOpType = Op->getType();
            FunctionTags::StringLiteralPoolKey Key = {
              StartAddress, VirtualSize, OffsetInSegment, RealOpType
            };

            auto LiteralFunction = StringLiteralPool
                                     .get(Key,
                                          RealOpType,
                                          { ConstExpr->getType() },
                                          "cstringLiteral");

            if (not hasStringLiteralMetadata(*LiteralFunction)) {
              setStringLiteralMetadata(*LiteralFunction,
                                       StartAddress,
                                       VirtualSize,
                                       OffsetInSegment,
                                       Str.size(),
                                       RealOpType);
            }

            CallInst *Call = IRB.CreateCall(LiteralFunction, { ConstExpr });
            ConstantOpReplacement = Call;

          } else {
            // If it cannot be emitted as a string literal we emit it as a
            // reference to a segment.

            IntegerType *OperandType = LeafConstOp->getType();
            FunctionTags::SegmentRefPoolKey Key = { StartAddress,
                                                    VirtualSize,
                                                    OperandType };
            auto *SegmentRefFunction = SegmentRefPool.get(Key,
                                                          OperandType,
                                                          {},
                                                          "segmentRef");

            if (not hasSegmentKeyMetadata(*SegmentRefFunction)) {
              setSegmentKeyMetadata(*SegmentRefFunction,
                                    StartAddress,
                                    VirtualSize);
            }

            // Inject call to SegmentRef
            auto *SegmentRefCall = IRB.CreateCall(SegmentRefFunction);

            // Inject call to AddressOf
            auto SegmentRefType = SegmentRefCall->getType();
            auto *AddressOfFunctionType = getAddressOfType(PtrSizedInteger,
                                                           SegmentRefType);
            auto *AddressOfFunction = AddressOfPool.get({ PtrSizedInteger,
                                                          SegmentRefType },
                                                        AddressOfFunctionType,
                                                        "AddressOf");
            Constant *ModelTypeString = nullptr;
            if (const auto &SegmentType = Binary.Segments()
                                            .at({ StartAddress, VirtualSize })
                                            .Type()) {
              ModelTypeString = toLLVMString(SegmentType, M);
            } else {
              auto Byte = model::PrimitiveType::makeGeneric(1);
              auto Arr = model::ArrayType::make(std::move(Byte), VirtualSize);
              ModelTypeString = toLLVMString(Arr, M);
            }
            Value *AddressOfCall = IRB.CreateCall(AddressOfFunction,
                                                  { ModelTypeString,
                                                    SegmentRefCall });
            AddressOfCall = IRB.CreateZExtOrTrunc(AddressOfCall, OperandType);

            // Inject LLVM add
            if (OffsetInSegment > 0) {
              auto *Offset = ConstantInt::get(OperandType, OffsetInSegment);
              AddressOfCall = IRB.CreateAdd(AddressOfCall, Offset);
            }

            if (auto *CE = dyn_cast<ConstantExpr>(Op))
              if (CE->getType()->isPointerTy())
                AddressOfCall = IRB.CreateIntToPtr(AddressOfCall,
                                                   CE->getType());

            ConstantOpReplacement = AddressOfCall;
          }
        }

        revng_assert(nullptr != ConstantOpReplacement);
        I.setOperand(Op.getOperandNo(), ConstantOpReplacement);
        Changed = true;
        revng::verify(&M);

        ConstantOpReplacements[ConstOp] = ConstantOpReplacement;
      }
    }
  }

  return Changed;
}

template<>
char pipeline::FunctionPass<MakeSegmentRefPassImpl>::ID = 0;

namespace revng::pipes {

class MakeSegmentRef {
public:
  static constexpr auto Name = "make-segment-ref";

  std::array<pipeline::ContractGroup, 1> getContract() const {
    pipeline::Contract BinaryPart(kinds::Binary, 0, kinds::Binary, 0);
    pipeline::Contract FunctionsPart(kinds::StackAccessesSegregated,
                                     1,
                                     kinds::StackAccessesSegregated,
                                     1);
    return { pipeline::ContractGroup({ BinaryPart, FunctionsPart }) };
  }

  void run(pipeline::ExecutionContext &EC,
           const BinaryFileContainer &SourceBinary,
           pipeline::LLVMContainer &Output);
};

void MakeSegmentRef::run(pipeline::ExecutionContext &EC,
                         const BinaryFileContainer &SourceBinary,
                         pipeline::LLVMContainer &ModuleContainer) {
  if (not SourceBinary.exists())
    return;

  const TupleTree<model::Binary> &Model = getModelFromContext(EC);
  auto BufferOrError = llvm::MemoryBuffer::getFileOrSTDIN(*SourceBinary.path());
  auto Buffer = cantFail(errorOrToExpected(std::move(BufferOrError)));
  RawBinaryView RawBinary(*Model, Buffer->getBuffer());

  llvm::legacy::PassManager PM;
  PM.add(new pipeline::LoadExecutionContextPass(&EC, ModuleContainer.name()));
  PM.add(new LoadModelWrapperPass(Model));
  PM.add(new LoadBinaryWrapperPass(Buffer->getBuffer()));
  PM.add(new pipeline::FunctionPass<MakeSegmentRefPassImpl>);
  PM.run(ModuleContainer.getModule());
}

} // namespace revng::pipes

static pipeline::RegisterPipe<revng::pipes::MakeSegmentRef> RegMSRPipe;

namespace revng::pypeline::piperuns {

// TODO: merge MakeSegmentRefPassImpl into MakeSegmentRef once we dismiss the
//       old pipeline
void MakeSegmentRef::runOnLLVMFunction(const model::Function &Function,
                                       llvm::Function &LLVMFunction) {
  MakeSegmentRefPassImpl Impl(Binary, *LLVMFunction.getParent(), BinaryView);
  Impl.runOnFunction(Function, LLVMFunction);
}

} // namespace revng::pypeline::piperuns