//===- lib/ReaderWriter/PECOFF/EdataPass.cpp ------------------------------===//
//
//                             The LLVM Linker
//
// This file is distributed under the University of Illinois Open Source
// License. See LICENSE.TXT for details.
//
//===----------------------------------------------------------------------===//

#include "Pass.h"
#include "EdataPass.h"
#include "lld/Core/File.h"
#include "lld/Core/Pass.h"
#include "lld/Core/Simple.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/Path.h"
#include <climits>
#include <ctime>
#include <utility>

using lld::pecoff::edata::EdataAtom;
using lld::pecoff::edata::TableEntry;
using llvm::object::export_address_table_entry;
using llvm::object::export_directory_table_entry;

namespace lld {
namespace pecoff {

typedef PECOFFLinkingContext::ExportDesc ExportDesc;

// dedupExports removes duplicate export entries. If two exports are
// referring the same symbol, they are considered duplicates.
// This could happen if the same symbol name is specified as an argument
// to /export more than once, or an unmangled and mangled name of the
// same symbol are given to /export. In the latter case, we choose
// unmangled (shorter) name.
static void dedupExports(PECOFFLinkingContext &ctx) {
  std::vector<ExportDesc> &exports = ctx.getDllExports();
  // Pass 1: find duplicate entries
  std::set<const ExportDesc *> dup;
  std::map<StringRef, ExportDesc *> map;
  for (ExportDesc &exp : exports) {
    if (!exp.externalName.empty())
      continue;
    StringRef symbol = exp.getRealName();
    auto it = map.find(symbol);
    if (it == map.end()) {
      map[symbol] = &exp;
    } else if (symbol.size() < it->second->getRealName().size()) {
      map[symbol] = &exp;
      dup.insert(it->second);
    } else {
      dup.insert(&exp);
    }
  }
  // Pass 2: remove duplicate entries
  auto pred = [&](const ExportDesc &exp) {
    return dup.count(&exp) == 1;
  };
  exports.erase(std::remove_if(exports.begin(), exports.end(), pred),
                exports.end());
}

static void assignOrdinals(PECOFFLinkingContext &ctx) {
  std::vector<ExportDesc> &exports = ctx.getDllExports();
  int maxOrdinal = -1;
  for (ExportDesc &desc : exports)
    maxOrdinal = std::max(maxOrdinal, desc.ordinal);

  std::sort(exports.begin(), exports.end(),
            [](const ExportDesc &a, const ExportDesc &b) {
    return a.getExternalName().compare(b.getExternalName()) < 0;
  });

  int nextOrdinal = (maxOrdinal == -1) ? 1 : (maxOrdinal + 1);
  for (ExportDesc &desc : exports)
    if (desc.ordinal == -1)
      desc.ordinal = nextOrdinal++;
}

static bool getExportedAtoms(PECOFFLinkingContext &ctx, const SimpleFile &file,
                             std::vector<TableEntry> &ret) {
  std::map<StringRef, const DefinedAtom *> definedAtoms;
  for (const DefinedAtom *atom : file.defined())
    definedAtoms[atom->name()] = atom;

  for (PECOFFLinkingContext::ExportDesc &desc : ctx.getDllExports()) {
    auto it = definedAtoms.find(desc.getRealName());
    if (it == definedAtoms.end()) {
      llvm::errs() << "Symbol <" << desc.name
                   << "> is exported but not defined.\n";
      return false;
    }
    const DefinedAtom *atom = it->second;

    // One can export a symbol with a different name than the symbol
    // name used in DLL. If such name is specified, use it in the
    // .edata section.
    ret.push_back(TableEntry(ctx.undecorateSymbol(desc.getExternalName()),
                             desc.ordinal, atom, desc.noname));
  }
  std::sort(ret.begin(), ret.end(),
            [](const TableEntry &a, const TableEntry &b) {
    return a.exportName.compare(b.exportName) < 0;
  });

  return true;
}

static std::pair<int, int> getOrdinalBase(std::vector<TableEntry> &entries) {
  int ordinalBase = INT_MAX;
  int maxOrdinal = -1;
  for (TableEntry &e : entries) {
    ordinalBase = std::min(ordinalBase, e.ordinal);
    maxOrdinal = std::max(maxOrdinal, e.ordinal);
  }
  return std::pair<int, int>(ordinalBase, maxOrdinal);
}

edata::EdataAtom *
EdataPass::createAddressTable(const std::vector<TableEntry> &entries,
                              int ordinalBase, int maxOrdinal) {
  EdataAtom *addressTable =
      new (_alloc) EdataAtom(_file, sizeof(export_address_table_entry) *
                                        (maxOrdinal - ordinalBase + 1));

  for (const TableEntry &e : entries) {
    int index = e.ordinal - ordinalBase;
    size_t offset = index * sizeof(export_address_table_entry);
    addDir32NBReloc(addressTable, e.atom, _ctx.getMachineType(), offset);
  }
  return addressTable;
}

edata::EdataAtom *
EdataPass::createNamePointerTable(const PECOFFLinkingContext &ctx,
                                  const std::vector<TableEntry> &entries,
                                  SimpleFile *file) {
  EdataAtom *table =
      new (_alloc) EdataAtom(_file, sizeof(uint32_t) * entries.size());

  size_t offset = 0;
  for (const TableEntry &e : entries) {
    auto *stringAtom = new (_alloc) COFFStringAtom(
        _file, _stringOrdinal++, ".edata", e.exportName);
    file->addAtom(*stringAtom);
    addDir32NBReloc(table, stringAtom, _ctx.getMachineType(), offset);
    offset += sizeof(uint32_t);
  }
  return table;
}

edata::EdataAtom *EdataPass::createExportDirectoryTable(
    const std::vector<edata::TableEntry> &namedEntries, int ordinalBase,
    int maxOrdinal) {
  EdataAtom *ret =
      new (_alloc) EdataAtom(_file, sizeof(export_directory_table_entry));
  auto *data = ret->getContents<export_directory_table_entry>();
  data->TimeDateStamp = time(nullptr);
  data->OrdinalBase = ordinalBase;
  data->AddressTableEntries = maxOrdinal - ordinalBase + 1;
  data->NumberOfNamePointers = namedEntries.size();
  return ret;
}

edata::EdataAtom *
EdataPass::createOrdinalTable(const std::vector<TableEntry> &entries,
                              int ordinalBase) {
  EdataAtom *ret =
      new (_alloc) EdataAtom(_file, sizeof(uint16_t) * entries.size());
  uint16_t *data = ret->getContents<uint16_t>();
  int i = 0;
  for (const TableEntry &e : entries)
    data[i++] = e.ordinal - ordinalBase;
  return ret;
}

std::error_code EdataPass::perform(SimpleFile &file) {
  dedupExports(_ctx);
  assignOrdinals(_ctx);

  std::vector<TableEntry> entries;
  if (!getExportedAtoms(_ctx, file, entries))
    return std::error_code();
  if (entries.empty())
    return std::error_code();

  int ordinalBase, maxOrdinal;
  std::tie(ordinalBase, maxOrdinal) = getOrdinalBase(entries);

  std::vector<TableEntry> namedEntries;
  for (TableEntry &e : entries)
    if (!e.noname)
      namedEntries.push_back(e);

  EdataAtom *table =
      createExportDirectoryTable(namedEntries, ordinalBase, maxOrdinal);
  file.addAtom(*table);

  COFFStringAtom *dllName =
      new (_alloc) COFFStringAtom(_file, _stringOrdinal++, ".edata",
                                  llvm::sys::path::filename(_ctx.outputPath()));
  file.addAtom(*dllName);
  addDir32NBReloc(table, dllName, _ctx.getMachineType(),
                  offsetof(export_directory_table_entry, NameRVA));

  EdataAtom *addressTable =
      createAddressTable(entries, ordinalBase, maxOrdinal);
  file.addAtom(*addressTable);
  addDir32NBReloc(
      table, addressTable, _ctx.getMachineType(),
      offsetof(export_directory_table_entry, ExportAddressTableRVA));

  EdataAtom *namePointerTable =
      createNamePointerTable(_ctx, namedEntries, &file);
  file.addAtom(*namePointerTable);
  addDir32NBReloc(table, namePointerTable, _ctx.getMachineType(),
                  offsetof(export_directory_table_entry, NamePointerRVA));

  EdataAtom *ordinalTable = createOrdinalTable(namedEntries, ordinalBase);
  file.addAtom(*ordinalTable);
  addDir32NBReloc(table, ordinalTable, _ctx.getMachineType(),
                  offsetof(export_directory_table_entry, OrdinalTableRVA));

  return std::error_code();
}

} // namespace pecoff
} // namespace lld