//===-- llvm-dwp.cpp - Split DWARF merging tool for llvm ------------------===//
//
// 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
//
//===----------------------------------------------------------------------===//
//
// A utility for merging DWARF 5 Split DWARF .dwo files into .dwp (DWARF
// package files).
//
//===----------------------------------------------------------------------===//
#include "llvm/DWP/DWP.h"
#include "llvm/DWP/DWPError.h"
#include "llvm/Object/ObjectFile.h"
#include "llvm/Option/ArgList.h"
#include "llvm/Option/Option.h"
#include "llvm/Support/CommandLine.h"
#include "llvm/Support/FileSystem.h"
#include "llvm/Support/LLVMDriver.h"
#include "llvm/Support/MemoryBuffer.h"
#include "llvm/Support/ToolOutputFile.h"
#include <optional>

using namespace llvm;
using namespace llvm::object;

// Command-line option boilerplate.
namespace {
enum ID {
  OPT_INVALID = 0, // This is not an option ID.
#define OPTION(...) LLVM_MAKE_OPT_ID(__VA_ARGS__),
#include "Opts.inc"
#undef OPTION
};

#define OPTTABLE_STR_TABLE_CODE
#include "Opts.inc"
#undef OPTTABLE_STR_TABLE_CODE

#define OPTTABLE_PREFIXES_TABLE_CODE
#include "Opts.inc"
#undef OPTTABLE_PREFIXES_TABLE_CODE

using namespace llvm::opt;
static constexpr opt::OptTable::Info InfoTable[] = {
#define OPTION(...) LLVM_CONSTRUCT_OPT_INFO(__VA_ARGS__),
#include "Opts.inc"
#undef OPTION
};

class DwpOptTable : public opt::GenericOptTable {
public:
  DwpOptTable()
      : GenericOptTable(OptionStrTable, OptionPrefixesTable, InfoTable) {}
};
} // end anonymous namespace

// Options
static std::vector<std::string> ExecFilenames;
static std::string OutputFilename;
static std::string ContinueOption;

static Expected<SmallVector<std::string, 16>>
getDWOFilenames(StringRef ExecFilename) {
  auto ErrOrObj = object::ObjectFile::createObjectFile(ExecFilename);
  if (!ErrOrObj)
    return ErrOrObj.takeError();

  const ObjectFile &Obj = *ErrOrObj.get().getBinary();
  std::unique_ptr<DWARFContext> DWARFCtx = DWARFContext::create(Obj);

  SmallVector<std::string, 16> DWOPaths;
  for (const auto &CU : DWARFCtx->compile_units()) {
    const DWARFDie &Die = CU->getUnitDIE();
    std::string DWOName = dwarf::toString(
        Die.find({dwarf::DW_AT_dwo_name, dwarf::DW_AT_GNU_dwo_name}), "");
    if (DWOName.empty())
      continue;
    std::string DWOCompDir =
        dwarf::toString(Die.find(dwarf::DW_AT_comp_dir), "");
    if (!DWOCompDir.empty()) {
      SmallString<16> DWOPath(DWOName);
      sys::path::make_absolute(DWOCompDir, DWOPath);
      if (!sys::fs::exists(DWOPath) && sys::fs::exists(DWOName))
        DWOPaths.push_back(std::move(DWOName));
      else
        DWOPaths.emplace_back(DWOPath.data(), DWOPath.size());
    } else {
      DWOPaths.push_back(std::move(DWOName));
    }
  }
  return std::move(DWOPaths);
}

static int error(const Twine &Error, const Twine &Context) {
  errs() << Twine("while processing ") + Context + ":\n";
  errs() << Twine("error: ") + Error + "\n";
  return 1;
}

int llvm_dwp_main(int argc, char **argv, const llvm::ToolContext &) {
  DwpOptTable Tbl;
  llvm::BumpPtrAllocator A;
  llvm::StringSaver Saver{A};
  OnCuIndexOverflow OverflowOptValue = OnCuIndexOverflow::HardStop;
  Dwarf64StrOffsetsPromotion Dwarf64StrOffsetsValue =
      Dwarf64StrOffsetsPromotion::Disabled;

  opt::InputArgList Args =
      Tbl.parseArgs(argc, argv, OPT_UNKNOWN, Saver, [&](StringRef Msg) {
        llvm::errs() << Msg << '\n';
        std::exit(1);
      });

  if (Args.hasArg(OPT_help)) {
    Tbl.printHelp(llvm::outs(), "llvm-dwp [options] <input files>",
                  "merge split dwarf (.dwo) files");
    std::exit(0);
  }

  if (Args.hasArg(OPT_version)) {
    llvm::cl::PrintVersionMessage();
    std::exit(0);
  }

  OutputFilename = Args.getLastArgValue(OPT_outputFileName, "");
  if (Arg *Arg = Args.getLastArg(OPT_continueOnCuIndexOverflow,
                                 OPT_continueOnCuIndexOverflow_EQ)) {
    if (Arg->getOption().matches(OPT_continueOnCuIndexOverflow)) {
      OverflowOptValue = OnCuIndexOverflow::Continue;
    } else {
      ContinueOption = Arg->getValue();
      if (ContinueOption == "soft-stop") {
        OverflowOptValue = OnCuIndexOverflow::SoftStop;
      } else if (ContinueOption == "continue") {
        OverflowOptValue = OnCuIndexOverflow::Continue;
      } else {
        llvm::errs() << "invalid value for --continue-on-cu-index-overflow"
                     << ContinueOption << '\n';
        exit(1);
      }
    }
  }

  if (Arg *Arg = Args.getLastArg(OPT_dwarf64StringOffsets,
                                 OPT_dwarf64StringOffsets_EQ)) {
    if (Arg->getOption().matches(OPT_dwarf64StringOffsets)) {
      Dwarf64StrOffsetsValue = Dwarf64StrOffsetsPromotion::Enabled;
    } else {
      std::string OptValue = Arg->getValue();
      if (OptValue == "disabled") {
        Dwarf64StrOffsetsValue = Dwarf64StrOffsetsPromotion::Disabled;
      } else if (OptValue == "enabled") {
        Dwarf64StrOffsetsValue = Dwarf64StrOffsetsPromotion::Enabled;
      } else if (OptValue == "always") {
        Dwarf64StrOffsetsValue = Dwarf64StrOffsetsPromotion::Always;
      } else {
        llvm::errs()
            << "invalid value for --dwarf64-str-offsets-promotion. Valid "
               "values are one of: \"enabled\", \"disabled\" or \"always\".\n";
        exit(1);
      }
    }
  }

  for (const llvm::opt::Arg *A : Args.filtered(OPT_execFileNames))
    ExecFilenames.emplace_back(A->getValue());

  std::vector<std::string> DWOFilenames;
  for (const llvm::opt::Arg *A : Args.filtered(OPT_INPUT))
    DWOFilenames.emplace_back(A->getValue());

  for (const auto &ExecFilename : ExecFilenames) {
    auto DWOs = getDWOFilenames(ExecFilename);
    if (!DWOs) {
      logAllUnhandledErrors(
          handleErrors(DWOs.takeError(),
                       [&](std::unique_ptr<ECError> EC) -> Error {
                         return createFileError(ExecFilename,
                                                Error(std::move(EC)));
                       }),
          WithColor::error());
      return 1;
    }
    DWOFilenames.insert(DWOFilenames.end(),
                        std::make_move_iterator(DWOs->begin()),
                        std::make_move_iterator(DWOs->end()));
  }

  if (DWOFilenames.empty()) {
    WithColor::defaultWarningHandler(make_error<DWPError>(
        "executable file does not contain any references to dwo files"));
    return 0;
  }

  StringRef DiscardPrefix = Args.getLastArgValue(OPT_prioritizeDiscardPath, "");
  if (OverflowOptValue == OnCuIndexOverflow::SoftStop &&
      !DiscardPrefix.empty()) {
    SmallString<256> CanonicalDiscardPrefix(DiscardPrefix);
    if (std::error_code EC =
            sys::fs::real_path(DiscardPrefix, CanonicalDiscardPrefix)) {
      WithColor::warning() << "invalid --prioritize-discard-path '"
                           << DiscardPrefix << "': " << EC.message()
                           << "; ignoring option.\n";
    } else {
      StringRef PrefixRef(CanonicalDiscardPrefix);
      auto IsNonDiscarded = [&](const std::string &Name) {
        SmallString<256> CanonicalDWO;
        if (sys::fs::real_path(Name, CanonicalDWO))
          return true;
        StringRef DWORef(CanonicalDWO);
        if (!DWORef.starts_with(PrefixRef))
          return true;
        if (DWORef.size() == PrefixRef.size())
          return false;
        if (sys::path::is_separator(DWORef[PrefixRef.size()]))
          return false;
        return true;
      };
      std::stable_partition(DWOFilenames.begin(), DWOFilenames.end(),
                            IsNonDiscarded);
    }
  }

  // Create the output file.
  std::error_code EC;
  ToolOutputFile OutFile(OutputFilename, EC, sys::fs::OF_None);
  std::optional<buffer_ostream> BOS;
  raw_pwrite_stream *OS;
  if (EC)
    return error(Twine(OutputFilename) + ": " + EC.message(),
                 "dwp output init");
  if (OutFile.os().supportsSeeking()) {
    OS = &OutFile.os();
  } else {
    BOS.emplace(OutFile.os());
    OS = &*BOS;
  }

  // Use DWPWriter for direct ELF output
  DWPWriter Writer;

  if (auto Err = write(Writer, DWOFilenames, OverflowOptValue,
                       Dwarf64StrOffsetsValue, OS)) {
    logAllUnhandledErrors(std::move(Err), WithColor::error());
    return 1;
  }

  OutFile.keep();
  return 0;
}