//
// This file is distributed under the MIT License. See LICENSE.md for details.
//

#include "llvm/ADT/DepthFirstIterator.h"
#include "llvm/ADT/GenericCycleImpl.h"
#include "llvm/ADT/GenericCycleInfo.h"
#include "llvm/ADT/GraphTraits.h"
#include "llvm/IR/CFG.h"
#include "llvm/IR/Function.h"
#include "llvm/IR/SSAContext.h"

#include "revng/RestructureCFG/GenericRegionInfo.h"
#include "revng/RestructureCFG/ScopeGraphGraphTraits.h"
#include "revng/Support/Debug.h"
#include "revng/Support/GraphAlgorithms.h"

using namespace llvm;

// Debug logger
Logger GenericRegionInfoLogger("generic-region-info");

/// Helper function which mimics the `at` behavior for a `llvm::SmallDenseMap`
template<class KeyT, class ValueT>
static ValueT mapAt(llvm::SmallDenseMap<KeyT, ValueT> &Map, KeyT Key) {
  auto MapIt = Map.find(Key);
  revng_assert(MapIt != Map.end());
  return MapIt->second;
}

/// Helper function to obtain a `GenericCycleInfo` analysis
template<class GraphT>
static GenericCycleInfo<SSAContext, GraphT> getGenericCycleInfo(GraphT &F) {
  // We instantiate the `GenericCycle` analysis and wrap the results in
  // the region objects
  GenericCycleInfo<SSAContext, GraphT> GCI;
  GCI.compute(*F);

  return GCI;
}

/// Template function specialization to obtain the `GenericCycleInfo` analysis
/// starting from a `Scope<llvm::Function *>` parameter, since we need to unwrap
/// the `Graph` object from the `Scope` wrapper class
template<>
GenericCycleInfo<SSAContext, Scope<llvm::Function *>>
getGenericCycleInfo(Scope<llvm::Function *> &SG) {
  // We instantiate the `GenericCycle` analysis and wrap the results in
  // the region objects
  GenericCycleInfo<SSAContext, Scope<llvm::Function *>> GCI;
  GCI.compute(*SG.Graph);

  return GCI;
}

template<class GraphT, class GT>
void GenericRegionInfo<GraphT, GT>::initializeRegions(GraphT F) {

  // Obtain the `GenericCycleInfo` analysis
  auto GCI = getGenericCycleInfo(F);

  using CycleT = GenericCycleInfo<SSAContext, GraphT>::CycleT;
  using Region = GenericRegion<NodeT>;
  llvm::SmallDenseMap<const CycleT *, Region *> CycleToRegionMap;

  // Populate the `Regions` with the identified regions
  for (const auto *TLC : GCI.toplevel_cycles()) {
    for (const auto *Cycle : depth_first(TLC)) {

      // Create a new `Region`
      Regions.push_back(std::make_unique<Region>());
      Region *CurrentRegion = Regions.back().get();

      // Populate the mapping between the `CycleT` object and our custom
      // `Region`
      CycleToRegionMap[Cycle] = CurrentRegion;

      // Iterate over all the blocks and insert them in the `CurrentRegion`
      for (auto *Block : Cycle->blocks()) {
        CurrentRegion->insertBlock(Block);
      }
    }
  }

  // Populate the children regions. We need to perform this operation in a
  // separate step in order to have already all the created regions in the step
  // above
  for (const auto *TLC : GCI.toplevel_cycles()) {
    for (const auto *Cycle : depth_first(TLC)) {
      auto *Region = mapAt(CycleToRegionMap, Cycle);
      for (const auto *Child : Cycle->children()) {
        auto *ChildRegion = mapAt(CycleToRegionMap, Child);
        Region->addChild(ChildRegion);
      }
    }
  }
}

/// Helper static function to compute the shortest distance from the entry block
template<class GraphT>
static llvm::SmallDenseMap<typename llvm::GraphTraits<GraphT>::NodeRef, size_t>
computeShortesPath(GraphT F) {
  using NodeT = llvm::GraphTraits<GraphT>::NodeRef;
  llvm::SmallDenseMap<NodeT, size_t> ShortestPathFromEntry;

  for (auto BFSIt = bf_begin(F); BFSIt != bf_end(F); BFSIt++) {
    NodeT Block = *BFSIt;
    size_t Depth = BFSIt.getLevel();

    // Obtain the insertion iterator for the `Depth` block element
    auto ShortestIt = ShortestPathFromEntry.insert({ Block, Depth });

    // If we already had in the map an entry for the current block, we need to
    // assert that the previously found value for the `Depth` is less or equal
    // of the `Depth` we are inserting
    if (ShortestIt.second == false) {
      revng_assert(ShortestIt.first->second <= Depth);
    }
  }

  return ShortestPathFromEntry;
}

template<class GraphT, class GT>
void GenericRegionInfo<GraphT, GT>::electHead(GraphT F) {

  // For each `Region`, we perform the election of the head node
  // 1) Compute the reverse post order
  llvm::SmallVector<NodeT> RPOT;
  llvm::copy(llvm::post_order(F), std::back_inserter(RPOT));

  // 2) We use the shortest distance from the entry block of the function just
  //    as a tie breaker. Therefore, we delay its computation until it is
  //    necessary.
  std::optional<llvm::SmallDenseMap<NodeT, size_t>>
    ShortestPathFromEntry = std::nullopt;

  // 3) Perform the head election for each `Region`
  for (auto &TopLevelRegion : top_level_regions()) {
    for (auto &CurrentRegion : depth_first(&TopLevelRegion)) {

      // During the `Head` election phase, we now introduce the following
      // additional criterion:
      // When processing a `GenericRegion` nested into an
      // outer one(its `ParentRegion`), if the inner `Region` contains the block
      // that has been elected as `Head` of the `ParentRegion`, we also force
      // that block to be the `Head` of the inner `GenericRegion`.
      // This criterion is justified by the following observation:
      // Suppose that we elect for the outer `Region` A as `Head`. If A is
      // also contained in the inner child region, and we elect another
      // block, say B, as its `Head`, it would mean that A becomes a late
      // entry for the inner region, causing it to be disconnected (late
      // entry edges are transformed into `goto` edges). This would clearly
      // break the assumption that all the blocks remain connected to the
      // entry in the `ScopeGraph`. If this criterion does not apply, we
      // continue with the standard criterion election.
      auto *ParentRegion = CurrentRegion->getParent();
      if (ParentRegion) {
        NodeT ParentHead = ParentRegion->getHead();
        revng_assert(ParentHead);

        if (CurrentRegion->containsBlock(ParentHead)) {
          CurrentRegion->setHead(ParentHead);
          continue;
        }
      }

      // All the blocks which have an incoming edge from a block not part of the
      // region itself, are considered as head candidates
      llvm::SmallMapVector<NodeT, size_t, 4> HeadCandidates;
      for (NodeT Block : CurrentRegion->blocks()) {
        for (NodeT Predecessor : graph_predecessors(Block)) {
          if (not CurrentRegion->containsBlock(Predecessor)) {
            HeadCandidates[Block]++;
          }
        }
      }

      // Elect the `Head` as the candidate head with the largest number of
      // incoming edges from outside the region.
      // If there is a tie, i.e., there are 2 or more candidate heads with the
      // same number of incoming edges from outside the region itself, we select
      // the entry with the minimal shortest path from entry. If it is still a
      // tie, i.e., there are 2 or more candidate heads with, also, the same
      // minimal shortest path from entry, then we disambiguate by picking the
      // head that comes first in RPOT.
      NodeT Head = HeadCandidates.begin()->first;
      {
        size_t MaxNHead = HeadCandidates.begin()->second;
        auto HeadEnd = HeadCandidates.end();
        for (NodeT Block : RPOT) {
          auto HeadIt = HeadCandidates.find(Block);
          if (HeadIt != HeadEnd) {
            const auto &[HeadCandidate, NumIncoming] = *HeadIt;
            if (NumIncoming > MaxNHead) {
              Head = HeadCandidate;
            } else if (NumIncoming == MaxNHead) {

              // Compute the `ShortestPathFromEntry` map since we need to break
              // a tie here
              if (not ShortestPathFromEntry.has_value()) {
                ShortestPathFromEntry = computeShortesPath(F);
              }
              size_t CurrentShortest = mapAt(*ShortestPathFromEntry, Head);
              size_t CandidateShortest = mapAt(*ShortestPathFromEntry,
                                               HeadCandidate);
              if (CandidateShortest < CurrentShortest) {
                Head = HeadCandidate;
              }
            }
          }
        }
      }

      // Verify that we found a `Head`
      revng_assert(Head != nullptr);

      // Set the `Head` for the `Region`
      CurrentRegion->setHead(Head);
    }
  }
}

template<class GraphT, class GT>
void GenericRegionInfo<GraphT, GT>::compute(GraphT F) {

  initializeRegions(F);
  electHead(F);

  // Print the `GenericRegionInfo` results, when the respective Logger is
  // activated. This is used both for debugging purposes and for testing with
  // `FileCheck`.
  revng_log(GenericRegionInfoLogger, print());
}

template<class GraphT, class GT>
std::string GenericRegionInfo<GraphT, GT>::print() const {
  std::string Output;

  // Print each top level `GenericRegion`, and for each one explore it using a
  // DFS strategy
  size_t RegionIndex = 0;
  Output += "\nGeneric Region Info Results:\n";
  for (auto &TopLevelRegion : top_level_regions()) {
    for (auto *Region : llvm::depth_first(&TopLevelRegion)) {
      Output += "\nRegion " + std::to_string(RegionIndex) + ":\n";
      Output += "Elected head: " + Region->getHead()->getName().str() + "\n";
      for (auto &Block : Region->blocks()) {
        Output += Block->getName().str() + "\n";
      }
      RegionIndex++;
    }
  }

  return Output;
}

template class GenericRegionInfo<llvm::Function *>;
template class GenericRegionInfo<Scope<llvm::Function *>>;