//===--- DXILDebugInfo.cpp - analysis&lowering for Debug info -*- C++ -*- ---=//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//

#include "DXILDebugInfo.h"
#include "DXILAttributes.h"
#include "llvm/BinaryFormat/Dwarf.h"
#include "llvm/IR/AttributeMask.h"
#include "llvm/IR/Attributes.h"
#include "llvm/IR/DebugInfo.h"
#include "llvm/IR/Instructions.h"
#include "llvm/IR/IntrinsicsDirectX.h"
#include "llvm/IR/Module.h"
#include "llvm/Transforms/Utils/BasicBlockUtils.h"

#define DEBUG_TYPE "dx-debug-info"

using namespace llvm;
using namespace llvm::dxil;

// llvm.dbg.value has an additional "offset" operand in DXIL.
static void replaceDbgVariableIntr(DbgVariableIntrinsic *DVI, Function *NewF,
                                   DXILDebugInfoMap &Res) {
  if (DVI->getIntrinsicID() != Intrinsic::dbg_value) {
    return;
  }

  Type *Int64Ty = Type::getInt64Ty(DVI->getContext());
  Constant *ZeroOffset = ConstantInt::get(Int64Ty, 0);

  Value *NewOps[] = {
      DVI->getOperand(0),
      ZeroOffset,
      DVI->getOperand(1),
      DVI->getOperand(2),
  };

  CallInst *NewI = CallInst::Create(NewF->getFunctionType(), NewF, NewOps);
  NewI->setTailCall(DVI->isTailCall());
  NewI->setDebugLoc(DVI->getDebugLoc());
  Res.InstReplace.insert({DVI, decltype(Res.InstReplace)::mapped_type(NewI)});
}

static void replaceDbgValue(Module &M, DXILDebugInfoMap &Res) {
  Function *F = getDeclarationIfExists(&M, Intrinsic::dbg_value);
  if (!F)
    return;

  FunctionType *FT = F->getFunctionType();
  Type *Int64Ty = Type::getInt64Ty(F->getContext());
  FunctionType *NewFT = FunctionType::get(
      FT->getReturnType(),
      {FT->getParamType(0), Int64Ty, FT->getParamType(1), FT->getParamType(2)},
      /*isVarArg=*/false);
  Function *NewF = Function::Create(NewFT, F->getLinkage(), F->getName());
  NewF->copyAttributesFrom(F);
  Res.FuncReplace.insert({F, decltype(Res.FuncReplace)::mapped_type(NewF)});

  for (User *U : F->users()) {
    auto *DVI = cast<DbgVariableIntrinsic>(U);
    replaceDbgVariableIntr(DVI, NewF, Res);
  }
}

DXILDebugInfoMap DXILDebugInfoPass::run(Module &M) {
  M.convertFromNewDbgValues();

  DXILDebugInfoMap Res;
  DebugInfoFinder DIF;
  DIF.processModule(M);

  {
    Function *DVDecl = nullptr;

    // Logically these should be variables in the
    //  for (BasicBlock &BB : F) loop.
    // They are defined here and cleared at the start of the loop body to avoid
    // the cost of deconstruction and reconstruction.
    DenseMap<DILocalVariable *, std::pair<Instruction *, DbgValueInst *>>
        DbgValues;
    DenseMap<std::pair<DILocalVariable *, DIExpression *>,
             std::pair<Instruction *, DbgValueInst *>>
        DbgValueFragments;
    // Likewise, logically, this should be a variable in the
    // for (Function &F : M) loop.
    DenseSet<DILocalVariable *> DbgVariablesSeen;

    const AttributeMask &AttrMask = getNonDXILAttributeMask();

    for (Function &F : M) {
      F.removeFnAttrs(AttrMask);
      F.removeRetAttrs(AttrMask);
      for (unsigned ArgNo = 0; ArgNo != F.arg_size(); ++ArgNo)
        F.removeParamAttrs(ArgNo, AttrMask);

      bool IsEntryBlock = true;
      DbgVariablesSeen.clear();
      for (BasicBlock &BB : F) {
        Instruction *NextNonDebugInst = nullptr;
        DbgValues.clear();
        DbgValueFragments.clear();
        for (Instruction &I : make_early_inc_range(reverse(BB))) {
          I.eraseMetadataIf([](unsigned KindID, MDNode *) {
            return KindID == LLVMContext::MD_DIAssignID;
          });
          if (!isa<DbgInfoIntrinsic>(I)) {
            NextNonDebugInst = &I;
            continue;
          }
          if (auto *DL = dyn_cast<DbgLabelInst>(&I)) {
            DL->eraseFromParent();
            continue;
          }
          // Process both llvm.dbg.value and llvm.dbg.assign here. We convert
          // llvm.dbg.assign to llvm.dbg.value by dropping the last arguments,
          // and remove redundant llvm.dbg.values.
          if (auto *DV = dyn_cast<DbgValueInst>(&I)) {
            // Keep track of the last location where we saw any debug value for
            // a variable.
            DILocalVariable *V = DV->getVariable();
            DIExpression *E = DV->getExpression();
            // If this is already an llvm.dbg.value instruction that we can
            // keep, just do that, otherwise convert it.
            auto *Val = cast<MetadataAsValue>(DV->getArgOperand(0));
            auto *Var = cast<MetadataAsValue>(DV->getArgOperand(1));
            auto *Expr = cast<MetadataAsValue>(DV->getArgOperand(2));
            bool Replace = DV->getIntrinsicID() != Intrinsic::dbg_value;
            if (!isa<ValueAsMetadata>(Val->getMetadata())) {
              // This may be a DIArgList which is not supported in LLVM 3.7. If
              // it is, we cannot record the new value, but we still need to
              // kill any old value. Do this by poison. We do not know the
              // correct type to use here and arbitrarily use i1.
              // This should never be anything other than ValueAsMetadata or
              // DIArgList, but in manually constructed LLVM IR, it can be.
              // Handle this gracefully by also replacing it with poison.
              Val = MetadataAsValue::get(
                  M.getContext(), ConstantAsMetadata::get(PoisonValue::get(
                                      Type::getInt1Ty(M.getContext()))));
              E = DIExpression::get(M.getContext(), {});
              Expr = MetadataAsValue::get(M.getContext(), E);
              Replace = true;
            }
            std::pair<Instruction *, DbgValueInst *> &DbgValue = DbgValues[V];
            std::pair<Instruction *, DbgValueInst *> &DbgValueFragment =
                DbgValueFragments[{V, E}];
            if (DbgValue.second) {
              // If there is a later value of the same fragment at the same
              // location, this value is redundant.
              if (DbgValueFragment.first == NextNonDebugInst) {
                DV->eraseFromParent();
                continue;
              }
              // If there is a later identical value of the same fragment at a
              // later point, and there have been no intervening values of
              // different possibly overlapping fragments, that later value is
              // redundant.
              if (DbgValueFragment.second &&
                  DbgValueFragment.second == DbgValue.second &&
                  DbgValueFragment.second->getValue() == DV->getValue()) {
                DbgValue.second->eraseFromParent();
              }
            }
            DbgValueInst *NewDV;
            if (Replace) {
              if (!DVDecl) {
                DVDecl =
                    Intrinsic::getOrInsertDeclaration(&M, Intrinsic::dbg_value);
                AttributeMask AM;
                for (Attribute A : DVDecl->getAttributes().getFnAttrs())
                  if (A.isStringAttribute() ||
                      (A.getKindAsEnum() != Attribute::NoUnwind &&
                       A.getKindAsEnum() != Attribute::Memory))
                    AM.addAttribute(A);
                DVDecl->removeFnAttrs(AM);
              }
              NewDV = cast<DbgValueInst>(
                  CallInst::Create(DVDecl, {Val, Var, Expr}, {}, "",
                                   std::next(DV->getIterator())));
              NewDV->setTailCall();
              NewDV->setDebugLoc(DV->getDebugLoc());
              DV->eraseFromParent();
            } else {
              NewDV = DV;
            }
            DbgValue = DbgValueFragment = {NextNonDebugInst, NewDV};
            continue;
          }
        }
        // If this is the entry block, if the first value we see for each debug
        // value is undef, it is redundant.
        if (IsEntryBlock) {
          for (Instruction &I : make_early_inc_range(BB)) {
            auto *DV = dyn_cast<DbgValueInst>(&I);
            if (!DV || DbgVariablesSeen.contains(DV->getVariable()))
              continue;
            if (isa<UndefValue>(DV->getValue())) {
              DV->eraseFromParent();
              continue;
            }
            DbgVariablesSeen.insert(DV->getVariable());
          }
        }
        IsEntryBlock = false;
      }
    }
  }

  for (DISubprogram *SP : DIF.subprograms()) {
    if (MDTuple *RN = cast_or_null<MDTuple>(SP->getRawRetainedNodes())) {
      SmallVector<Metadata *> MDs(RN->operands());
      MDs.erase(std::remove_if(MDs.begin(), MDs.end(),
                               [](Metadata *M) { return isa<DILabel>(M); }),
                MDs.end());
      SP->replaceRetainedNodes(MDTuple::get(M.getContext(), MDs));
    }
  }

  // Re-scan the module to account for removed metadata.
  DIF.reset();
  DIF.processModule(M);

  // Replace llvm.dbg.value with equivalent DXIL intrinsics.
  replaceDbgValue(M, Res);

  for (DICompileUnit *CU : DIF.compile_units()) {
    DISourceLanguageName Lang = CU->getSourceLanguage();
    if (Lang.hasVersionedName()) {
      auto LangName = static_cast<dwarf::SourceLanguageName>(Lang.getName());
      Lang = dwarf::toDW_LANG(LangName, Lang.getVersion())
                 .value_or(dwarf::SourceLanguage{});
      auto *NewCU = DICompileUnit::getDistinct(
          M.getContext(), Lang, CU->getFile(), CU->getProducer(),
          CU->isOptimized(), CU->getFlags(), CU->getRuntimeVersion(),
          CU->getSplitDebugFilename(), CU->getEmissionKind(),
          CU->getEnumTypes(), CU->getRetainedTypes(), CU->getGlobalVariables(),
          CU->getImportedEntities(), CU->getMacros(), CU->getDWOId(),
          CU->getSplitDebugInlining(), CU->getDebugInfoForProfiling(),
          CU->getNameTableKind(), CU->getRangesBaseAddress(), CU->getSysRoot(),
          CU->getSDK());
      Res.MDReplace.insert({CU, NewCU});
    }
  }

  std::vector<std::pair<const DICompileUnit *, const Metadata *>> CUSubprograms;

  for (const Function &F : M) {
    if (const DISubprogram *SP = F.getSubprogram()) {
      auto *FunctionMD = ConstantAsMetadata::get(const_cast<Function *>(&F));
      Res.MDExtra.insert({SP, FunctionMD});
    }
  }

  for (const DISubprogram *SP : DIF.subprograms()) {
    const DISubprogram *NewSP = SP;

    static constexpr auto SupportedDIFlags =
        static_cast<DISubprogram::DIFlags>(DISubprogram::FlagExportSymbols - 1);
    static constexpr auto SupportedDISPFlags =
        static_cast<DISubprogram::DISPFlags>(DISubprogram::SPFlagPure - 1);
    if (SP->isDistinct() || SP->getFlags() & ~SupportedDIFlags ||
        SP->getSPFlags() & ~SupportedDISPFlags) {
      NewSP = DISubprogram::get(
          M.getContext(), SP->getScope(), SP->getName(), SP->getLinkageName(),
          SP->getFile(), SP->getLine(), SP->getType(), SP->getScopeLine(),
          SP->getContainingType(), SP->getVirtualIndex(),
          SP->getThisAdjustment(), SP->getFlags() & SupportedDIFlags,
          SP->getSPFlags() & SupportedDISPFlags, SP->getUnit(),
          SP->getTemplateParams(), SP->getDeclaration(), SP->getRetainedNodes(),
          SP->getThrownTypes(), SP->getAnnotations(), SP->getTargetFuncName(),
          SP->getKeyInstructionsEnabled());

      Res.MDReplace.insert({SP, NewSP});

      if (auto It = Res.MDExtra.find(SP); It != Res.MDExtra.end()) {
        const Metadata *FunctionMD = It->second;
        Res.MDExtra.erase(It);
        Res.MDExtra.insert({NewSP, FunctionMD});
      }
    }

    if (SP->getUnit())
      CUSubprograms.push_back(
          {SP->getUnit(), static_cast<const Metadata *>(SP)});
  }

  std::stable_sort(
      CUSubprograms.begin(), CUSubprograms.end(), [](auto &&A, auto &&B) {
        return std::less<const DICompileUnit *>()(A.first, B.first);
      });
  for (auto It = CUSubprograms.begin(), End = CUSubprograms.end(); It != End;) {
    const DICompileUnit *CU = It->first;
    const DICompileUnit *NewCU =
        cast<DICompileUnit>(Res.MDReplace.lookup_or(CU, CU));
    SmallVector<Metadata *, 16> Subprograms;
    do {
      Subprograms.push_back(const_cast<Metadata *>(It->second));
    } while (++It != End && It->first == CU);
    const auto *SubprogramsMD = MDTuple::get(M.getContext(), Subprograms);
    Res.MDExtra.insert({NewCU, SubprogramsMD});
  }

  for (const GlobalVariable &GV : M.globals()) {
    SmallVector<DIGlobalVariableExpression *, 4> GVEs;
    GV.getDebugInfo(GVEs);
    for (DIGlobalVariableExpression *GVE : GVEs) {
      if (GVE->getExpression()->getNumElements())
        continue;
      auto [It, Inserted] = Res.MDExtra.insert(
          {GVE->getVariable(),
           ValueAsMetadata::get(const_cast<GlobalVariable *>(&GV))});
      if (!Inserted)
        It->second = nullptr;
    }
  }

  for (DIGlobalVariableExpression *GVE : DIF.global_variables())
    Res.MDReplace.insert({GVE, GVE->getVariable()});

  for (DIScope *S : DIF.scopes()) {
    if (auto *CB = dyn_cast<DICommonBlock>(S)) {
      const Metadata *Scope = CB->getScope();
      Scope = Res.MDReplace.lookup_or(Scope, Scope);
      Res.MDReplace.insert({CB, Scope});
    }
  }

  for (DIType *T : DIF.types()) {
    if (auto *SR = dyn_cast<DISubrangeType>(T)) {
      DIType *BT = SR->getBaseType();
      if (!BT)
        BT = DIBasicType::get(SR->getContext(), dwarf::DW_TAG_base_type,
                              SR->getName(), SR->getSizeInBits(),
                              SR->getAlignInBits(), dwarf::DW_ATE_unsigned,
                              SR->getNumExtraInhabitants(),
                              /*DataSizeInBits=*/0, SR->getFlags());
      Res.MDReplace.insert({T, BT});
    }
  }

  return Res;
}