#include <thread>
#include <unordered_set>
#include "llvm/DebugInfo/DIContext.h"
#include "llvm/DebugInfo/DWARF/DWARFCompileUnit.h"
#include "llvm/DebugInfo/DWARF/DWARFContext.h"
#include "llvm/Support/Error.h"
#include "llvm/Support/ThreadPool.h"
#include "llvm/Support/raw_ostream.h"
#include "llvm/DebugInfo/GSYM/DwarfTransformer.h"
#include "llvm/DebugInfo/GSYM/FunctionInfo.h"
#include "llvm/DebugInfo/GSYM/GsymCreator.h"
#include "llvm/DebugInfo/GSYM/GsymReader.h"
#include "llvm/DebugInfo/GSYM/InlineInfo.h"
using namespace llvm;
using namespace gsym;
struct llvm::gsym::CUInfo {
  const DWARFDebugLine::LineTable *LineTable;
  const char *CompDir;
  std::vector<uint32_t> FileCache;
  uint64_t Language = 0;
  uint8_t AddrSize = 0;
  CUInfo(DWARFContext &DICtx, DWARFCompileUnit *CU) {
    LineTable = DICtx.getLineTableForUnit(CU);
    CompDir = CU->getCompilationDir();
    FileCache.clear();
    if (LineTable)
      FileCache.assign(LineTable->Prologue.FileNames.size() + 1, UINT32_MAX);
    DWARFDie Die = CU->getUnitDIE();
    Language = dwarf::toUnsigned(Die.find(dwarf::DW_AT_language), 0);
    AddrSize = CU->getAddressByteSize();
  }
          bool isHighestAddress(uint64_t Addr) const {
    if (AddrSize == 4)
      return Addr == UINT32_MAX;
    else if (AddrSize == 8)
      return Addr == UINT64_MAX;
    return false;
  }
                    uint32_t DWARFToGSYMFileIndex(GsymCreator &Gsym, uint32_t DwarfFileIdx) {
    if (!LineTable)
      return 0;
    assert(DwarfFileIdx < FileCache.size());
    uint32_t &GsymFileIdx = FileCache[DwarfFileIdx];
    if (GsymFileIdx != UINT32_MAX)
      return GsymFileIdx;
    std::string File;
    if (LineTable->getFileNameByIndex(
            DwarfFileIdx, CompDir,
            DILineInfoSpecifier::FileLineInfoKind::AbsoluteFilePath, File))
      GsymFileIdx = Gsym.insertFile(File);
    else
      GsymFileIdx = 0;
    return GsymFileIdx;
  }
};
static DWARFDie GetParentDeclContextDIE(DWARFDie &Die) {
  if (DWARFDie SpecDie =
          Die.getAttributeValueAsReferencedDie(dwarf::DW_AT_specification)) {
    if (DWARFDie SpecParent = GetParentDeclContextDIE(SpecDie))
      return SpecParent;
  }
  if (DWARFDie AbstDie =
          Die.getAttributeValueAsReferencedDie(dwarf::DW_AT_abstract_origin)) {
    if (DWARFDie AbstParent = GetParentDeclContextDIE(AbstDie))
      return AbstParent;
  }
        if (Die.getTag() == dwarf::DW_TAG_inlined_subroutine)
    return DWARFDie();
  DWARFDie ParentDie = Die.getParent();
  if (!ParentDie)
    return DWARFDie();
  switch (ParentDie.getTag()) {
  case dwarf::DW_TAG_namespace:
  case dwarf::DW_TAG_structure_type:
  case dwarf::DW_TAG_union_type:
  case dwarf::DW_TAG_class_type:
  case dwarf::DW_TAG_subprogram:
    return ParentDie;   case dwarf::DW_TAG_lexical_block:
    return GetParentDeclContextDIE(ParentDie);
  default:
    break;
  }
  return DWARFDie();
}
static Optional<uint32_t> getQualifiedNameIndex(DWARFDie &Die,
                                                uint64_t Language,
                                                GsymCreator &Gsym) {
    if (auto LinkageName =
          dwarf::toString(Die.findRecursively({dwarf::DW_AT_MIPS_linkage_name,
                                               dwarf::DW_AT_linkage_name}),
                          nullptr))
    return Gsym.insertString(LinkageName,  false);
  StringRef ShortName(Die.getName(DINameKind::ShortName));
  if (ShortName.empty())
    return llvm::None;
    if (!(Language == dwarf::DW_LANG_C_plus_plus ||
        Language == dwarf::DW_LANG_C_plus_plus_03 ||
        Language == dwarf::DW_LANG_C_plus_plus_11 ||
        Language == dwarf::DW_LANG_C_plus_plus_14 ||
        Language == dwarf::DW_LANG_ObjC_plus_plus ||
                        Language == dwarf::DW_LANG_C))
    return Gsym.insertString(ShortName,  false);
        if (ShortName.startswith("_Z") &&
      (ShortName.contains(".isra.") || ShortName.contains(".part.")))
    return Gsym.insertString(ShortName,  false);
  DWARFDie ParentDeclCtxDie = GetParentDeclContextDIE(Die);
  if (ParentDeclCtxDie) {
    std::string Name = ShortName.str();
    while (ParentDeclCtxDie) {
      StringRef ParentName(ParentDeclCtxDie.getName(DINameKind::ShortName));
      if (!ParentName.empty()) {
                                if (ParentName.front() == '<' && ParentName.back() == '>')
          Name = "{" + ParentName.substr(1, ParentName.size() - 2).str() + "}" +
                "::" + Name;
        else
          Name = ParentName.str() + "::" + Name;
      }
      ParentDeclCtxDie = GetParentDeclContextDIE(ParentDeclCtxDie);
    }
        return Gsym.insertString(Name,  true);
  }
    return Gsym.insertString(ShortName,  false);
}
static bool hasInlineInfo(DWARFDie Die, uint32_t Depth) {
  bool CheckChildren = true;
  switch (Die.getTag()) {
  case dwarf::DW_TAG_subprogram:
        CheckChildren = Depth == 0;
    break;
  case dwarf::DW_TAG_inlined_subroutine:
    return true;
  default:
    break;
  }
  if (!CheckChildren)
    return false;
  for (DWARFDie ChildDie : Die.children()) {
    if (hasInlineInfo(ChildDie, Depth + 1))
      return true;
  }
  return false;
}
static void parseInlineInfo(GsymCreator &Gsym, CUInfo &CUI, DWARFDie Die,
                            uint32_t Depth, FunctionInfo &FI,
                            InlineInfo &parent) {
  if (!hasInlineInfo(Die, Depth))
    return;
  dwarf::Tag Tag = Die.getTag();
  if (Tag == dwarf::DW_TAG_inlined_subroutine) {
        InlineInfo II;
    DWARFAddressRange FuncRange =
        DWARFAddressRange(FI.startAddress(), FI.endAddress());
    Expected<DWARFAddressRangesVector> RangesOrError = Die.getAddressRanges();
    if (RangesOrError) {
      for (const DWARFAddressRange &Range : RangesOrError.get()) {
                        if (FuncRange.LowPC <= Range.LowPC && Range.HighPC <= FuncRange.HighPC)
          II.Ranges.insert(AddressRange(Range.LowPC, Range.HighPC));
      }
    }
    if (II.Ranges.empty())
      return;
    if (auto NameIndex = getQualifiedNameIndex(Die, CUI.Language, Gsym))
      II.Name = *NameIndex;
    II.CallFile = CUI.DWARFToGSYMFileIndex(
        Gsym, dwarf::toUnsigned(Die.find(dwarf::DW_AT_call_file), 0));
    II.CallLine = dwarf::toUnsigned(Die.find(dwarf::DW_AT_call_line), 0);
        for (DWARFDie ChildDie : Die.children())
      parseInlineInfo(Gsym, CUI, ChildDie, Depth + 1, FI, II);
    parent.Children.emplace_back(std::move(II));
    return;
  }
  if (Tag == dwarf::DW_TAG_subprogram || Tag == dwarf::DW_TAG_lexical_block) {
        for (DWARFDie ChildDie : Die.children())
      parseInlineInfo(Gsym, CUI, ChildDie, Depth + 1, FI, parent);
  }
}
static void convertFunctionLineTable(raw_ostream &Log, CUInfo &CUI,
                                     DWARFDie Die, GsymCreator &Gsym,
                                     FunctionInfo &FI) {
  std::vector<uint32_t> RowVector;
  const uint64_t StartAddress = FI.startAddress();
  const uint64_t EndAddress = FI.endAddress();
  const uint64_t RangeSize = EndAddress - StartAddress;
  const object::SectionedAddress SecAddress{
      StartAddress, object::SectionedAddress::UndefSection};
  if (!CUI.LineTable->lookupAddressRange(SecAddress, RangeSize, RowVector)) {
            std::string FilePath = Die.getDeclFile(
        DILineInfoSpecifier::FileLineInfoKind::AbsoluteFilePath);
    if (FilePath.empty())
      return;
    if (auto Line =
            dwarf::toUnsigned(Die.findRecursively({dwarf::DW_AT_decl_line}))) {
      LineEntry LE(StartAddress, Gsym.insertFile(FilePath), *Line);
      FI.OptLineTable = LineTable();
      FI.OptLineTable->push(LE);
    }
    return;
  }
  FI.OptLineTable = LineTable();
  DWARFDebugLine::Row PrevRow;
  for (uint32_t RowIndex : RowVector) {
        const DWARFDebugLine::Row &Row = CUI.LineTable->Rows[RowIndex];
    const uint32_t FileIdx = CUI.DWARFToGSYMFileIndex(Gsym, Row.File);
    uint64_t RowAddress = Row.Address.Address;
                            if (!FI.Range.contains(RowAddress)) {
      if (RowAddress < FI.Range.start()) {
        Log << "error: DIE has a start address whose LowPC is between the "
          "line table Row[" << RowIndex << "] with address "
          << HEX64(RowAddress) << " and the next one.\n";
        Die.dump(Log, 0, DIDumpOptions::getForSingleDIE());
        RowAddress = FI.Range.start();
      } else {
        continue;
      }
    }
    LineEntry LE(RowAddress, FileIdx, Row.Line);
    if (RowIndex != RowVector[0] && Row.Address < PrevRow.Address) {
                                          auto FirstLE = FI.OptLineTable->first();
      if (FirstLE && *FirstLE == LE) {
        if (!Gsym.isQuiet()) {
          Log << "warning: duplicate line table detected for DIE:\n";
          Die.dump(Log, 0, DIDumpOptions::getForSingleDIE());
        }
      } else {
                Log << "error: line table has addresses that do not "
             << "monotonically increase:\n";
        for (uint32_t RowIndex2 : RowVector) {
          CUI.LineTable->Rows[RowIndex2].dump(Log);
        }
        Die.dump(Log, 0, DIDumpOptions::getForSingleDIE());
      }
      break;
    }
        auto LastLE = FI.OptLineTable->last();
    if (LastLE && LastLE->File == FileIdx && LastLE->Line == Row.Line)
        continue;
                if (Row.EndSequence) {
                              PrevRow = DWARFDebugLine::Row();
    } else {
      FI.OptLineTable->push(LE);
      PrevRow = Row;
    }
  }
      if (FI.OptLineTable->empty())
    FI.OptLineTable = llvm::None;
}
void DwarfTransformer::handleDie(raw_ostream &OS, CUInfo &CUI, DWARFDie Die) {
  switch (Die.getTag()) {
  case dwarf::DW_TAG_subprogram: {
    Expected<DWARFAddressRangesVector> RangesOrError = Die.getAddressRanges();
    if (!RangesOrError) {
      consumeError(RangesOrError.takeError());
      break;
    }
    const DWARFAddressRangesVector &Ranges = RangesOrError.get();
    if (Ranges.empty())
      break;
    auto NameIndex = getQualifiedNameIndex(Die, CUI.Language, Gsym);
    if (!NameIndex) {
      OS << "error: function at " << HEX64(Die.getOffset())
         << " has no name\n ";
      Die.dump(OS, 0, DIDumpOptions::getForSingleDIE());
      break;
    }
        for (const DWARFAddressRange &Range : Ranges) {
                                                      if (Range.LowPC >= Range.HighPC || CUI.isHighestAddress(Range.LowPC))
        break;
                              if (!Gsym.IsValidTextAddress(Range.LowPC)) {
                                        if (Range.LowPC != 0) {
          if (!Gsym.isQuiet()) {
                        OS << "warning: DIE has an address range whose start address is "
                  "not in any executable sections ("
               << *Gsym.GetValidTextRanges()
               << ") and will not be processed:\n";
            Die.dump(OS, 0, DIDumpOptions::getForSingleDIE());
          }
        }
        break;
      }
      FunctionInfo FI;
      FI.Range = {Range.LowPC, Range.HighPC};
      FI.Name = *NameIndex;
      if (CUI.LineTable) {
        convertFunctionLineTable(OS, CUI, Die, Gsym, FI);
      }
      if (hasInlineInfo(Die, 0)) {
        FI.Inline = InlineInfo();
        FI.Inline->Name = *NameIndex;
        FI.Inline->Ranges.insert(FI.Range);
        parseInlineInfo(Gsym, CUI, Die, 0, FI, *FI.Inline);
      }
      Gsym.addFunctionInfo(std::move(FI));
    }
  } break;
  default:
    break;
  }
  for (DWARFDie ChildDie : Die.children())
    handleDie(OS, CUI, ChildDie);
}
Error DwarfTransformer::convert(uint32_t NumThreads) {
  size_t NumBefore = Gsym.getNumFunctionInfos();
  auto getDie = [&](DWARFUnit &DwarfUnit) -> DWARFDie {
    DWARFDie ReturnDie = DwarfUnit.getUnitDIE(false);
    if (llvm::Optional<uint64_t> DWOId = DwarfUnit.getDWOId()) {
      DWARFUnit *DWOCU = DwarfUnit.getNonSkeletonUnitDIE(false).getDwarfUnit();
      if (!DWOCU->isDWOUnit()) {
        std::string DWOName = dwarf::toString(
            DwarfUnit.getUnitDIE().find(
                {dwarf::DW_AT_dwo_name, dwarf::DW_AT_GNU_dwo_name}),
            "");
        Log << "warning: Unable to retrieve DWO .debug_info section for "
            << DWOName << "\n";
      } else {
        ReturnDie = DWOCU->getUnitDIE(false);
      }
    }
    return ReturnDie;
  };
  if (NumThreads == 1) {
            for (const auto &CU : DICtx.compile_units()) {
      DWARFDie Die = getDie(*CU);
      CUInfo CUI(DICtx, dyn_cast<DWARFCompileUnit>(CU.get()));
      handleDie(Log, CUI, Die);
    }
  } else {
                
            for (const auto &CU : DICtx.compile_units())
      CU->getAbbreviations();
            ThreadPool pool(hardware_concurrency(NumThreads));
    for (const auto &CU : DICtx.compile_units())
      pool.async([&CU]() { CU->getUnitDIE(false ); });
    pool.wait();
        std::mutex LogMutex;
    for (const auto &CU : DICtx.compile_units()) {
      DWARFDie Die = getDie(*CU);
      if (Die) {
        CUInfo CUI(DICtx, dyn_cast<DWARFCompileUnit>(CU.get()));
        pool.async([this, CUI, &LogMutex, Die]() mutable {
          std::string ThreadLogStorage;
          raw_string_ostream ThreadOS(ThreadLogStorage);
          handleDie(ThreadOS, CUI, Die);
          ThreadOS.flush();
          if (!ThreadLogStorage.empty()) {
                        std::lock_guard<std::mutex> guard(LogMutex);
            Log << ThreadLogStorage;
          }
        });
      }
    }
    pool.wait();
  }
  size_t FunctionsAddedCount = Gsym.getNumFunctionInfos() - NumBefore;
  Log << "Loaded " << FunctionsAddedCount << " functions from DWARF.\n";
  return Error::success();
}
llvm::Error DwarfTransformer::verify(StringRef GsymPath) {
  Log << "Verifying GSYM file \"" << GsymPath << "\":\n";
  auto Gsym = GsymReader::openFile(GsymPath);
  if (!Gsym)
    return Gsym.takeError();
  auto NumAddrs = Gsym->getNumAddresses();
  DILineInfoSpecifier DLIS(
      DILineInfoSpecifier::FileLineInfoKind::AbsoluteFilePath,
      DILineInfoSpecifier::FunctionNameKind::LinkageName);
  std::string gsymFilename;
  for (uint32_t I = 0; I < NumAddrs; ++I) {
    auto FuncAddr = Gsym->getAddress(I);
    if (!FuncAddr)
        return createStringError(std::errc::invalid_argument,
                                  "failed to extract address[%i]", I);
    auto FI = Gsym->getFunctionInfo(*FuncAddr);
    if (!FI)
      return createStringError(std::errc::invalid_argument,
                            "failed to extract function info for address 0x%"
                            PRIu64, *FuncAddr);
    for (auto Addr = *FuncAddr; Addr < *FuncAddr + FI->size(); ++Addr) {
      const object::SectionedAddress SectAddr{
          Addr, object::SectionedAddress::UndefSection};
      auto LR = Gsym->lookup(Addr);
      if (!LR)
        return LR.takeError();
      auto DwarfInlineInfos =
          DICtx.getInliningInfoForAddress(SectAddr, DLIS);
      uint32_t NumDwarfInlineInfos = DwarfInlineInfos.getNumberOfFrames();
      if (NumDwarfInlineInfos == 0) {
        DwarfInlineInfos.addFrame(
            DICtx.getLineInfoForAddress(SectAddr, DLIS));
      }
            if (NumDwarfInlineInfos == 1 &&
          DwarfInlineInfos.getFrame(0).FileName == "<invalid>") {
        DwarfInlineInfos = DIInliningInfo();
        NumDwarfInlineInfos = 0;
      }
      if (NumDwarfInlineInfos > 0 &&
          NumDwarfInlineInfos != LR->Locations.size()) {
        Log << "error: address " << HEX64(Addr) << " has "
            << NumDwarfInlineInfos << " DWARF inline frames and GSYM has "
            << LR->Locations.size() << "\n";
        Log << "    " << NumDwarfInlineInfos << " DWARF frames:\n";
        for (size_t Idx = 0; Idx < NumDwarfInlineInfos; ++Idx) {
          const auto &dii = DwarfInlineInfos.getFrame(Idx);
          Log << "    [" << Idx << "]: " << dii.FunctionName << " @ "
              << dii.FileName << ':' << dii.Line << '\n';
        }
        Log << "    " << LR->Locations.size() << " GSYM frames:\n";
        for (size_t Idx = 0, count = LR->Locations.size();
              Idx < count; ++Idx) {
          const auto &gii = LR->Locations[Idx];
          Log << "    [" << Idx << "]: " << gii.Name << " @ " << gii.Dir
              << '/' << gii.Base << ':' << gii.Line << '\n';
        }
        DwarfInlineInfos = DICtx.getInliningInfoForAddress(SectAddr, DLIS);
        Gsym->dump(Log, *FI);
        continue;
      }
      for (size_t Idx = 0, count = LR->Locations.size(); Idx < count;
            ++Idx) {
        const auto &gii = LR->Locations[Idx];
        if (Idx < NumDwarfInlineInfos) {
          const auto &dii = DwarfInlineInfos.getFrame(Idx);
          gsymFilename = LR->getSourceFile(Idx);
                    if (dii.FunctionName.find(gii.Name.str()) != 0)
            Log << "error: address " << HEX64(Addr) << " DWARF function \""
                << dii.FunctionName.c_str()
                << "\" doesn't match GSYM function \"" << gii.Name << "\"\n";
                    if (dii.FileName != gsymFilename)
            Log << "error: address " << HEX64(Addr) << " DWARF path \""
                << dii.FileName.c_str() << "\" doesn't match GSYM path \""
                << gsymFilename.c_str() << "\"\n";
                    if (dii.Line != gii.Line)
            Log << "error: address " << HEX64(Addr) << " DWARF line "
                << dii.Line << " != GSYM line " << gii.Line << "\n";
        }
      }
    }
  }
  return Error::success();
}