//===- BPSectionOrdererBase.inc ---------------------------------*- 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
//
//===----------------------------------------------------------------------===//
//
// This file defines the common BPSectionOrderer interface using the Curiously
// Recurring Template Pattern and dispatches to the BalancedPartitioning
// algorithm implemented in LLVMSupport. The optimized section layout attempts
// to group similar sections together (resulting in a smaller compressed app
// size) and utilize a temporal profile file to reduce page faults during
// program startup.
//
// Clients should derive from BPOrderer, providing concrete implementations for
// section and symbol representations. Include this file in a .cpp file to
// specialize the template for the derived class.
//
//===----------------------------------------------------------------------===//

#include "lld/Common/BPSectionOrdererBase.h"
#include "lld/Common/ErrorHandler.h"
#include "lld/Common/Utils.h"
#include "llvm/ADT/CachedHashString.h"
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/MapVector.h"
#include "llvm/ADT/SetVector.h"
#include "llvm/ADT/SmallSet.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/StringMap.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/ADT/Twine.h"
#include "llvm/ProfileData/InstrProfReader.h"
#include "llvm/Support/BalancedPartitioning.h"
#include "llvm/Support/TimeProfiler.h"
#include "llvm/Support/VirtualFileSystem.h"
#include <memory>
#include <optional>
#include <set>

#define DEBUG_TYPE "bp-section-orderer"

using namespace llvm;
using namespace lld;

namespace lld {
template <class D> struct BPOrdererTraits;

template <class D> struct BPOrderer {
  using Section = typename BPOrdererTraits<D>::Section;
  using Defined = typename BPOrdererTraits<D>::Defined;

  // Compute a section order using the Balanced Partitioning algorithm.
  //
  // * compressionSortSpecs: Improve Lempel-Ziv compression by grouping
  //   similar sections together. The specs select which sections participate,
  //   divide them into distinct groups, and run BP independently for each
  //   group, then place the groups according to their layout priority.
  // * profilePath: Utilize a temporal profile file to reduce page faults during
  //   program startup.
  // * compressionSortStartupFunctions: if profilePath is specified, allocate
  //   extra utility vertices to prioritize nearby function similarity.
  auto computeOrder(llvm::StringRef profilePath,
                    llvm::ArrayRef<BPCompressionSortSpec> compressionSortSpecs,
                    bool forFunctionCompression, bool forDataCompression,
                    bool compressionSortStartupFunctions, bool verbose,
                    llvm::ArrayRef<Section *> sections,
                    const DenseMap<CachedHashStringRef, std::set<unsigned>>
                        &rootSymbolToSectionIdxs)
      -> llvm::DenseMap<const Section *, int>;

  std::optional<StringRef> static getResolvedLinkageName(StringRef name) {
    return {};
  }

  std::string static getCompressionSubgroupKey(const Section &sec) {
    return "";
  }
};
} // namespace lld

using UtilityNodes = SmallVector<BPFunctionNode::UtilityNodeT>;

template <class D>
static SmallVector<std::pair<unsigned, UtilityNodes>> getUnsForCompression(
    ArrayRef<const typename D::Section *> sections,
    const DenseMap<const void *, uint64_t> &sectionToIdx,
    ArrayRef<unsigned> sectionIdxs,
    DenseMap<unsigned, SmallVector<unsigned, 0>> *duplicateSectionIdxs,
    BPFunctionNode::UtilityNodeT &maxUN) {
  TimeTraceScope timeScope("Build nodes for compression");

  SmallVector<std::pair<unsigned, SmallVector<uint64_t>>> sectionHashes;
  sectionHashes.reserve(sectionIdxs.size());
  SmallVector<uint64_t> hashes;

  for (unsigned sectionIdx : sectionIdxs) {
    const auto *isec = sections[sectionIdx];
    D::getSectionHashes(*isec, hashes, sectionToIdx);
    sectionHashes.emplace_back(sectionIdx, std::move(hashes));
    hashes.clear();
  }

  MapVector<uint64_t, unsigned> hashFrequency;
  for (auto &[sectionIdx, hashes] : sectionHashes)
    for (auto hash : hashes)
      ++hashFrequency[hash];

  if (duplicateSectionIdxs) {
    // Merge sections that are nearly identical
    SmallVector<std::pair<unsigned, SmallVector<uint64_t>>> newSectionHashes;
    DenseMap<uint64_t, unsigned> wholeHashToSectionIdx;
    unsigned threshold = sectionHashes.size() > 10000 ? 5 : 0;
    for (auto &[sectionIdx, hashes] : sectionHashes) {
      uint64_t wholeHash = 0;
      for (auto hash : hashes)
        if (hashFrequency[hash] > threshold)
          wholeHash ^= hash;
      auto [it, wasInserted] =
          wholeHashToSectionIdx.insert(std::make_pair(wholeHash, sectionIdx));
      if (wasInserted) {
        newSectionHashes.emplace_back(sectionIdx, hashes);
      } else {
        (*duplicateSectionIdxs)[it->getSecond()].push_back(sectionIdx);
      }
    }
    sectionHashes = newSectionHashes;

    // Recompute hash frequencies
    hashFrequency.clear();
    for (auto &[sectionIdx, hashes] : sectionHashes)
      for (auto hash : hashes)
        ++hashFrequency[hash];
  }

  // Filter rare and common hashes and assign each a unique utility node that
  // doesn't conflict with the trace utility nodes
  DenseMap<uint64_t, BPFunctionNode::UtilityNodeT> hashToUN;
  for (auto &[hash, frequency] : hashFrequency) {
    if (frequency <= 1 || frequency * 2 > sectionHashes.size())
      continue;
    hashToUN[hash] = ++maxUN;
  }

  SmallVector<std::pair<unsigned, UtilityNodes>> sectionUns;
  for (auto &[sectionIdx, hashes] : sectionHashes) {
    UtilityNodes uns;
    for (auto &hash : hashes) {
      auto it = hashToUN.find(hash);
      if (it != hashToUN.end())
        uns.push_back(it->second);
    }
    sectionUns.emplace_back(sectionIdx, uns);
  }
  return sectionUns;
}

template <class D>
auto BPOrderer<D>::computeOrder(
    StringRef profilePath, ArrayRef<BPCompressionSortSpec> compressionSortSpecs,
    bool forFunctionCompression, bool forDataCompression,
    bool compressionSortStartupFunctions, bool verbose,
    ArrayRef<Section *> sections,
    const DenseMap<CachedHashStringRef, std::set<unsigned>>
        &rootSymbolToSectionIdxs) -> DenseMap<const Section *, int> {
  TimeTraceScope timeScope("Setup Balanced Partitioning");
  DenseMap<const void *, uint64_t> sectionToIdx;
  for (auto [i, isec] : llvm::enumerate(sections))
    sectionToIdx.try_emplace(isec, i);

  BPFunctionNode::UtilityNodeT maxUN = 0;
  MapVector<unsigned, UtilityNodes> startupSectionIdxUNs;
  // Used to define the initial order for startup functions.
  DenseMap<unsigned, size_t> sectionIdxToTimestamp;
  std::unique_ptr<InstrProfReader> reader;
  if (!profilePath.empty()) {
    auto fs = vfs::getRealFileSystem();
    auto readerOrErr = InstrProfReader::create(profilePath, *fs);
    if (!readerOrErr) {
      lld::error(toString(readerOrErr.takeError()));
      return DenseMap<const Section *, int>{};
    }

    reader = std::move(readerOrErr.get());
    for (auto &entry : *reader) {
      // Read all entries
      (void)entry;
    }
    auto &traces = reader->getTemporalProfTraces();

    MapVector<unsigned, BPFunctionNode::UtilityNodeT> sectionIdxToFirstUN;
    for (size_t traceIdx = 0; traceIdx < traces.size(); traceIdx++) {
      uint64_t currentSize = 0, cutoffSize = 1;
      size_t cutoffTimestamp = 1;
      auto &trace = traces[traceIdx].FunctionNameRefs;
      for (size_t timestamp = 0; timestamp < trace.size(); timestamp++) {
        auto [_, parsedFuncName] = getParsedIRPGOName(
            reader->getSymtab().getFuncOrVarName(trace[timestamp]));
        parsedFuncName = lld::utils::getRootSymbol(parsedFuncName);

        auto sectionIdxsIt =
            rootSymbolToSectionIdxs.find(CachedHashStringRef(parsedFuncName));
        if (sectionIdxsIt == rootSymbolToSectionIdxs.end())
          continue;
        auto &sectionIdxs = sectionIdxsIt->second;
        // If the same symbol is found in multiple sections, they might be
        // identical, so we arbitrarily use the size from the first section.
        currentSize += D::getSize(*sections[*sectionIdxs.begin()]);

        // Since BalancedPartitioning is sensitive to the initial order, we need
        // to explicitly define it to be ordered by earliest timestamp.
        for (unsigned sectionIdx : sectionIdxs) {
          auto [it, wasInserted] =
              sectionIdxToTimestamp.try_emplace(sectionIdx, timestamp);
          if (!wasInserted)
            it->getSecond() = std::min<size_t>(it->getSecond(), timestamp);
        }

        if (timestamp >= cutoffTimestamp || currentSize >= cutoffSize) {
          ++maxUN;
          cutoffSize = 2 * currentSize;
          cutoffTimestamp = 2 * cutoffTimestamp;
        }
        for (unsigned sectionIdx : sectionIdxs)
          sectionIdxToFirstUN.try_emplace(sectionIdx, maxUN);
      }
      for (auto &[sectionIdx, firstUN] : sectionIdxToFirstUN)
        for (auto un = firstUN; un <= maxUN; ++un)
          startupSectionIdxUNs[sectionIdx].push_back(un);
      ++maxUN;
      sectionIdxToFirstUN.clear();
    }
  }

  // Compression grouping: for each section, find the winning spec among
  // matching --bp-compression-sort-section globs. Match resolution:
  // - Explicit matchPriority always beats positional.
  // - Among explicit: lowest matchPriority wins.
  // - Among positional: last match wins.
  // Sections matched by the same winning globString are grouped together and
  // may be further subdivided by getCompressionSubgroupKey(). Layout order is
  // (layoutPriority, globString, subgroupKey).
  struct GlobGroup {
    unsigned layoutPriority;
    std::map<std::string, SmallVector<unsigned>> subgroups;
  };
  StringMap<GlobGroup> globGroups;
  SmallVector<unsigned> sectionIdxsForFunctionCompression,
      sectionIdxsForDataCompression;
  for (unsigned sectionIdx = 0; sectionIdx < sections.size(); sectionIdx++) {
    if (startupSectionIdxUNs.contains(sectionIdx))
      continue;
    const auto *isec = sections[sectionIdx];
    auto sectionName = D::getSectionName(*isec);

    const BPCompressionSortSpec *winner = nullptr;
    for (const auto &spec : compressionSortSpecs) {
      if (!spec.glob.match(sectionName))
        continue;
      if (!winner || spec.matchPriority.value_or(UINT_MAX) <=
                         winner->matchPriority.value_or(UINT_MAX))
        winner = &spec;
    }
    if (winner) {
      auto &group = globGroups[winner->globString];
      group.layoutPriority = winner->layoutPriority;
      auto subgroupKey = D::getCompressionSubgroupKey(*isec);
      group.subgroups[subgroupKey].push_back(sectionIdx);
      continue;
    }
    bool isCode = D::isCodeSection(*isec);
    if (forFunctionCompression && isCode)
      sectionIdxsForFunctionCompression.push_back(sectionIdx);
    if (forDataCompression && !isCode)
      sectionIdxsForDataCompression.push_back(sectionIdx);
  }

  // Flatten groups sorted by (layoutPriority, globString, subgroupKey).
  struct CompressionGroup {
    std::string key;
    SmallVector<unsigned> sectionIdxs;
  };
  SmallVector<CompressionGroup> groupsForCompression;
  {
    struct SortedCompressionGroup {
      std::tuple<unsigned, StringRef, StringRef> sortKey;
      SmallVector<unsigned> *sectionIdxs;
    };
    SmallVector<SortedCompressionGroup> sortedGroups;
    for (auto &[globString, group] : globGroups)
      for (auto &[subKey, idxs] : group.subgroups)
        sortedGroups.push_back(
            {{group.layoutPriority, globString, subKey}, &idxs});
    llvm::sort(sortedGroups, [](const auto &a, const auto &b) {
      return a.sortKey < b.sortKey;
    });
    for (auto &group : sortedGroups) {
      auto &[prio, globString, subKey] = group.sortKey;
      std::string fullKey =
          subKey.empty() ? globString.str() : (globString + ":" + subKey).str();
      groupsForCompression.push_back(
          {std::move(fullKey), std::move(*group.sectionIdxs)});
    }
  }
  if (forFunctionCompression)
    groupsForCompression.push_back(
        {"legacy:function", std::move(sectionIdxsForFunctionCompression)});
  if (forDataCompression)
    groupsForCompression.push_back(
        {"legacy:data", std::move(sectionIdxsForDataCompression)});

  if (compressionSortStartupFunctions) {
    SmallVector<unsigned> startupIdxs;
    for (auto &[sectionIdx, uns] : startupSectionIdxUNs)
      startupIdxs.push_back(sectionIdx);
    auto unsForStartupFunctionCompression =
        getUnsForCompression<D>(sections, sectionToIdx, startupIdxs,
                                /*duplicateSectionIdxs=*/nullptr, maxUN);
    for (auto &[sectionIdx, compressionUns] :
         unsForStartupFunctionCompression) {
      auto &uns = startupSectionIdxUNs[sectionIdx];
      uns.append(compressionUns);
      llvm::sort(uns);
      uns.erase(llvm::unique(uns), uns.end());
    }
  }

  // Map a section index (ordered directly) to a list of duplicate section
  // indices (not ordered directly).
  DenseMap<unsigned, SmallVector<unsigned, 0>> duplicateSectionIdxs;
  SmallVector<std::vector<BPFunctionNode>> nodesForCompression;
  for (auto &group : groupsForCompression) {
    auto &nodes = nodesForCompression.emplace_back();
    auto uns =
        getUnsForCompression<D>(sections, sectionToIdx, group.sectionIdxs,
                                &duplicateSectionIdxs, maxUN);
    for (auto &[sectionIdx, uns] : uns)
      nodes.emplace_back(sectionIdx, uns);
    // Sort compression nodes by their Id (which is the section index) because
    // the input linker order tends to be not bad.
    llvm::sort(nodes, [](auto &L, auto &R) { return L.Id < R.Id; });
  }

  std::vector<BPFunctionNode> nodesForStartup;
  for (auto &[sectionIdx, uns] : startupSectionIdxUNs)
    nodesForStartup.emplace_back(sectionIdx, uns);

  // Use the first timestamp to define the initial order for startup nodes.
  llvm::sort(nodesForStartup, [&sectionIdxToTimestamp](auto &L, auto &R) {
    return std::make_pair(sectionIdxToTimestamp[L.Id], L.Id) <
           std::make_pair(sectionIdxToTimestamp[R.Id], R.Id);
  });

  {
    TimeTraceScope timeScope("Balanced Partitioning");
    BalancedPartitioningConfig config;
    BalancedPartitioning bp(config);
    bp.run(nodesForStartup);
    for (auto &nodes : nodesForCompression)
      bp.run(nodes);
  }

  unsigned numStartupSections = 0, startupSize = 0;
  unsigned numCompressionSections = 0, compressionSize = 0;
  unsigned numDuplicateCompressionSections = 0, duplicateCompressionSize = 0;
  SetVector<const Section *> orderedSections;
  // Order startup functions,
  for (auto &node : nodesForStartup) {
    const auto *isec = sections[node.Id];
    if (orderedSections.insert(isec)) {
      startupSize += D::getSize(*isec);
      ++numStartupSections;
    }
  }
  // then sections for compression.
  for (const auto &nodes : nodesForCompression) {
    for (auto &node : nodes) {
      const auto *isec = sections[node.Id];
      if (orderedSections.insert(isec)) {
        compressionSize += D::getSize(*isec);
        ++numCompressionSections;
      }
      auto It = duplicateSectionIdxs.find(node.Id);
      if (It == duplicateSectionIdxs.end())
        continue;
      for (auto dupSecIdx : It->getSecond()) {
        const auto *dupIsec = sections[dupSecIdx];
        if (orderedSections.insert(dupIsec)) {
          duplicateCompressionSize += D::getSize(*dupIsec);
          ++numDuplicateCompressionSections;
        }
      }
    }
  }

  if (verbose) {
    unsigned numTotalOrderedSections = numStartupSections +
                                       numCompressionSections +
                                       numDuplicateCompressionSections;
    unsigned totalOrderedSize =
        startupSize + compressionSize + duplicateCompressionSize;
    dbgs() << "Ordered " << numTotalOrderedSections << " sections ("
           << totalOrderedSize << " bytes) using balanced partitioning:\n";
    dbgs() << "  Functions for startup: " << numStartupSections << " ("
           << startupSize << " bytes)\n";
    dbgs() << "  Sections for compression: " << numCompressionSections << " ("
           << compressionSize << " bytes)\n";
    dbgs() << "  Compression groups: " << groupsForCompression.size() << "\n";
    for (const auto &group : groupsForCompression)
      dbgs() << "    " << group.key << ": " << group.sectionIdxs.size()
             << " sections\n";
    dbgs() << "  Duplicate compression sections: "
           << numDuplicateCompressionSections << " ("
           << duplicateCompressionSize << " bytes)\n";

    if (!profilePath.empty()) {
      // Evaluate this function order for startup
      StringMap<std::pair<uint64_t, uint64_t>> symbolToPageNumbers;
      const uint64_t pageSize = (1 << 14);
      uint64_t currentAddress = 0;
      for (const auto *isec : orderedSections) {
        for (auto *sym : static_cast<D *>(this)->getSymbols(*isec)) {
          uint64_t startAddress = currentAddress + D::getSymValue(*sym);
          uint64_t endAddress = startAddress + D::getSymSize(*sym);
          uint64_t firstPage = startAddress / pageSize;
          // I think the kernel might pull in a few pages when one it touched,
          // so it might be more accurate to force lastPage to be aligned by
          // 4?
          uint64_t lastPage = endAddress / pageSize;
          StringRef rootSymbol = D::getSymName(*sym);
          rootSymbol = lld::utils::getRootSymbol(rootSymbol);
          symbolToPageNumbers.try_emplace(rootSymbol, firstPage, lastPage);
          if (auto resolvedLinkageName = D::getResolvedLinkageName(rootSymbol))
            symbolToPageNumbers.try_emplace(resolvedLinkageName.value(),
                                            firstPage, lastPage);
        }
        currentAddress += D::getSize(*isec);
      }

      // The area under the curve F where F(t) is the total number of page
      // faults at step t.
      unsigned area = 0;
      for (auto &trace : reader->getTemporalProfTraces()) {
        SmallSet<uint64_t, 0> touchedPages;
        for (unsigned step = 0; step < trace.FunctionNameRefs.size(); step++) {
          auto traceId = trace.FunctionNameRefs[step];
          auto [Filename, ParsedFuncName] =
              getParsedIRPGOName(reader->getSymtab().getFuncOrVarName(traceId));
          ParsedFuncName = lld::utils::getRootSymbol(ParsedFuncName);
          auto it = symbolToPageNumbers.find(ParsedFuncName);
          if (it != symbolToPageNumbers.end()) {
            auto &[firstPage, lastPage] = it->getValue();
            for (uint64_t i = firstPage; i <= lastPage; i++)
              touchedPages.insert(i);
          }
          area += touchedPages.size();
        }
      }
      dbgs() << "Total area under the page fault curve: " << (float)area
             << "\n";
    }
  }

  DenseMap<const Section *, int> sectionPriorities;
  int prio = -orderedSections.size();
  for (const auto *isec : orderedSections)
    sectionPriorities[isec] = prio++;
  return sectionPriorities;
}
