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

#include <unordered_set>

#include "llvm/ADT/EquivalenceClasses.h"
#include "llvm/ADT/GraphTraits.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SetVector.h"
#include "llvm/ADT/SmallPtrSet.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/IR/CFG.h"

#include "revng/ADT/GenericGraph.h"
#include "revng/ADT/RecursiveCoroutine.h"
#include "revng/DataLayoutAnalysis/DLATypeSystem.h"

#include "DLAStep.h"

using namespace llvm;

static Logger Log("dla-merge-pointees-of-ptr-union");

namespace dla {

using LTSN = LayoutTypeSystemNode;
using NeighborsConstIterator = LTSN::NeighborsSet::const_iterator;

static bool hasOutgoingPointerEdge(const LayoutTypeSystemNode *N) {
  using CPointerT = EdgeFilteredGraph<const dla::LayoutTypeSystemNode *,
                                      dla::isPointerEdge>;
  using PointerGraph = llvm::GraphTraits<CPointerT>;
  auto It = PointerGraph::child_begin(N);
  auto End = PointerGraph::child_end(N);
  return It != End;
};

static bool isWellFormedPointer(const LTSN *Pointer) {
  return Pointer->Successors.size() == 1 and hasOutgoingPointerEdge(Pointer);
}

static LTSN *getPointee(LTSN *Pointer) {
  revng_assert(isWellFormedPointer(Pointer));
  return Pointer->Successors.begin()->first;
}

bool MergePointeesOfPointerUnion::runOnTypeSystem(LayoutTypeSystem &TS) {
  bool Changed = false;

  revng_log(Log, "MergePointeesOfPointerUnion");
  LoggerIndent StepIndent{ Log };

  if (VerifyLog.isEnabled())
    revng_assert(TS.verifyDAG() and TS.verifyLeafs());

  // Initialize a vector of nodes before iterating.
  // The algorithm iterates over all nodes in the graph, but it can end merging
  // the current node (and a bunch of others) with another one, and that would
  // invalidate the iterators if we iterate on llvm::nodes(&TS) directly.
  std::vector<LTSN *> Nodes{ llvm::nodes(&TS).begin(), llvm::nodes(&TS).end() };
  std::unordered_set<LTSN *> Erased;

  while (not Nodes.empty()) {
    LTSN *Node = Nodes.back();
    Nodes.pop_back();

    revng_log(Log, "Analyzing Node: " << Node->ID);
    LoggerIndent Indent{ Log };

    if (Erased.contains(Node)) {
      revng_log(Log, "merged by a previous iteration");
      continue;
    }

    if (isInstanceLeaf(Node)) {
      revng_log(Log, "no instance children");
      continue;
    }

    llvm::EquivalenceClasses<LTSN *> PointersToMerge;

    auto ChildEnd = Node->Successors.end();
    for (auto AChildIt = Node->Successors.begin(); AChildIt != ChildEnd;
         ++AChildIt) {

      const auto &AEdge = *AChildIt;
      if (not isInstanceEdge(AEdge))
        continue;

      const auto &[APointer, ATag] = AEdge;
      if (not hasOutgoingPointerEdge(APointer))
        continue;
      revng_assert(isWellFormedPointer(APointer));

      for (auto BChildIt = std::next(AChildIt); BChildIt != ChildEnd;
           ++BChildIt) {

        const auto &BEdge = *BChildIt;
        const auto &[BPointer, BTag] = BEdge;
        if (ATag != BTag)
          continue;

        if (not hasOutgoingPointerEdge(BPointer))
          continue;
        revng_assert(isWellFormedPointer(BPointer));

        // Here we're sure that A and B are connected to Node with the same kind
        // of instance edge. And that they are both pointer nodes.

        revng_log(Log,
                  "has a pair of instance children at the same offset that are "
                  "pointer nodes:");
        revng_log(Log, "A: " << APointer->ID << ", B:" << BPointer->ID);
        revng_log(Log, "ATag: " << *ATag << ", BTag: " << *BTag);

        revng_assert(APointer->Successors.size() == 1,
                     std::to_string(APointer->ID).c_str());
        revng_assert(BPointer->Successors.size() == 1,
                     std::to_string(BPointer->ID).c_str());
        PointersToMerge.unionSets(APointer, BPointer);
      }
    }

    if (not PointersToMerge.empty()) {
      revng_log(Log, "Merging children");
      LoggerIndent MoreIndent{ Log };

      // Iterate over all of the equivalence sets.
      for (auto I = PointersToMerge.begin(), E = PointersToMerge.end(); I != E;
           ++I) {
        // Ignore non-leader sets.
        if (not I->isLeader())
          continue;

        // Loop over members in this set to select the node that we want to
        // merge the others into.
        const auto NotErasedPointer =
          [&Erased = std::as_const(Erased)](LayoutTypeSystemNode *N) {
            revng_log(Log, "NotErased ID: " << N->ID);
            LoggerIndent Indent{ Log };

            // Skip erased nodes.
            if (Erased.contains(N)) {
              revng_log(Log, "Erased");
              return false;
            }

            // Skip non-pointer nodes.
            // Nodes in PointersToMerge, being scalar, can be pointed to by
            // other pointers whose pointees are also scalars to merge.
            // This means that in some cases a pointer node can be merged with
            // its scalar pointee, turning it first into a scalar with an
            // outgoing pointer edge to self, and then into a non-pointer scalar
            // (because self edges are cancelled).
            // In those cases, we might end up with a node that was in
            // PointersToMerge being turned into a non-pointer.
            // If that happens we just ignore it, because it's not a pointer
            // anymore, and we just look at others.
            if (not hasOutgoingPointerEdge(N)) {
              revng_log(Log, "NOT hasOutgoingPointerEdge");
              return false;
            }

            // If the pointee was erased, the outgoing pointer edge from N would
            // have been erased too, and we would have fallen in the check
            // above.
            revng_assert(not Erased.contains(getPointee(N)));
            return true;
          };
        using llvm::make_filter_range;
        auto Pointers = make_filter_range(llvm::make_range(PointersToMerge
                                                             .member_begin(I),
                                                           PointersToMerge
                                                             .member_end()),
                                          NotErasedPointer);

        if (Log.isEnabled()) {
          revng_log(Log, "Preparing to merge pointees:");
          LoggerIndent EvenMoreIndent{ Log };
          for (LayoutTypeSystemNode *N : Pointers) {
            const LTSN *Pointee = getPointee(N);
            std::string ScalarString = Pointee->NonScalar ? "Aggregate" :
                                                            "Scalar";
            revng_log(Log,
                      N->ID << " with pointee " << Pointee->ID << " ("
                            << ScalarString << ", size: " << getPointee(N)->Size
                            << ")");
          }
          revng_log(Log, "----");
        }

        llvm::SmallSetVector<LTSN *, 8> UniquedScalars;
        llvm::SmallPtrSet<LTSN *, 8> PointersToScalars;
        llvm::SmallSetVector<LTSN *, 8> UniquedAggregates;
        llvm::SmallPtrSet<LTSN *, 8> PointersToAggregates;
        llvm::SmallSetVector<LTSN *, 8> UniquedFromModel;
        llvm::SmallPtrSet<LTSN *, 8> PointersToFromModel;

        for (LayoutTypeSystemNode *Pointer : Pointers) {
          revng_assert(not Erased.contains(Pointer));
          LTSN *Pointee = getPointee(Pointer);
          revng_assert(not Erased.contains(Pointee));

          if (Pointee->NonScalar) {
            PointersToFromModel.insert(Pointer);
            UniquedFromModel.insert(Pointee);
          } else if (Pointee->Successors.empty()
                     or hasOutgoingPointerEdge(Pointee)) {
            revng_assert(not hasOutgoingPointerEdge(Pointee)
                         or isWellFormedPointer(Pointee));
            PointersToScalars.insert(Pointer);
            UniquedScalars.insert(Pointee);
          } else {
            PointersToAggregates.insert(Pointer);
            UniquedAggregates.insert(Pointee);
          }
        }

        // Sort scalars and aggregates so that the first is the node with the
        // lowerst ID among the nodes with largest size.

        llvm::SmallVector<LTSN *> Scalars = UniquedScalars.takeVector();
        llvm::SmallVector<LTSN *> Aggregates = UniquedAggregates.takeVector();
        llvm::SmallVector<LTSN *> FromModel = UniquedAggregates.takeVector();

        const auto Ordering = [](const LTSN *LHS, const LTSN *RHS) {
          auto LSize = LHS->Size;
          auto RSize = RHS->Size;
          if (LSize > RSize)
            return true;

          if (LSize == RSize)
            return LHS->ID < RHS->ID;

          return false;
        };

        llvm::sort(Scalars, Ordering);
        llvm::sort(Aggregates, Ordering);
        llvm::sort(FromModel, Ordering);

        // Merge all the scalars together.
        LTSN *MergedScalar = nullptr;
        if (not Scalars.empty()) {
          if (Log.isEnabled()) {
            revng_log(Log, "merging Scalars:");
            LoggerIndent MoreMoreIndent{ Log };
            for (const LTSN *N : Scalars)
              revng_log(Log, N->ID);
          }
          TS.mergeNodes(Scalars);
          Erased.insert(std::next(Scalars.begin()), Scalars.end());
          MergedScalar = Scalars.front();

          revng_log(Log, "MergedScalar: " << MergedScalar->ID);

          // Check if we merged more than one scalar that also was a pointer.
          // In that case we have to create a new union of their pointees,
          // enqueue it for further analysis
          llvm::SmallVector<LTSN::NeighborIterator> PointerEdges;
          {
            LTSN::NeighborIterator ChildIt = MergedScalar->Successors.begin();
            LTSN::NeighborIterator ChildEnd = MergedScalar->Successors.end();
            for (; ChildIt != ChildEnd; ++ChildIt)
              if (isPointerEdge(*ChildIt))
                PointerEdges.push_back(ChildIt);

            revng_assert(PointerEdges.empty()
                         or MergedScalar->Size == PointerSize);
          }

          if (PointerEdges.size() > 1) {
            revng_log(Log,
                      "Merged scalar is a union of pointers: "
                        << MergedScalar->ID);
            for (LTSN::NeighborIterator &PointerEdgeIt : PointerEdges) {
              LTSN *NewPointer = TS.createArtificialLayoutType();
              NewPointer->Size = PointerSize;
              TS.moveEdgeSource(MergedScalar, NewPointer, PointerEdgeIt, 0);
              TS.addInstanceLink(MergedScalar,
                                 NewPointer,
                                 OffsetExpression{ 0 });
            }
            Nodes.push_back(MergedScalar);
          }
        }

        LTSN *MergedAggregate = nullptr;
        if (not Aggregates.empty()) {
          if (Log.isEnabled()) {
            revng_log(Log, "merging Aggregates:");
            LoggerIndent MoreMoreIndent{ Log };
            for (const LTSN *N : Aggregates)
              revng_log(Log, N->ID);
          }
          MergedAggregate = Aggregates.front();
          revng_log(Log, "MergedAggregate: " << MergedAggregate->ID);
          TS.mergeNodes(Aggregates);
          Erased.insert(std::next(Aggregates.begin()), Aggregates.end());
        }

        if (MergedAggregate and not FromModel.empty()) {
          revng_log(Log, "merged aggregate from model");
          size_t AggregateSize = MergedAggregate->Size;
          auto It = llvm::find_if(FromModel, [AggregateSize](LTSN *N) {
            return N->Size < AggregateSize;
          });
          auto SmallFromModel = llvm::make_range(It, FromModel.end());

          for (LTSN *AggregateFromModel : SmallFromModel) {
            revng_assert(AggregateFromModel->Size < AggregateSize);
            revng_log(Log,
                      "AggregateFromModel: " << AggregateFromModel->ID
                                             << " Size: "
                                             << AggregateFromModel->Size);
            // We want to move the pointers that point to AggregateFromModel,
            // and make it point to MergedAggregate.
            const auto PredecessorPointsToAggregate =
              [&PointersToFromModel](const auto &Pair) {
                return isPointerEdge(Pair)
                       and PointersToFromModel.contains(Pair.first);
              };
            auto InverseEdgeIt = llvm::find_if(AggregateFromModel->Predecessors,
                                               PredecessorPointsToAggregate);
            TS.moveEdgeTarget(AggregateFromModel,
                              MergedAggregate,
                              InverseEdgeIt,
                              0);

            // We want to put an instance link at offset 0 from MergedAggregate
            // to each small AggregateFromModel.
            TS.addInstanceLink(MergedAggregate,
                               AggregateFromModel,
                               OffsetExpression{ 0 });
          }

          if (1 == std::distance(FromModel.begin(), It)) {
            LTSN *UniqueFromModelToPreserve = *FromModel.begin();

            revng_log(Log,
                      "UniqueFromModelToPreserve: "
                        << (UniqueFromModelToPreserve ?
                              std::to_string(UniqueFromModelToPreserve->ID) :
                              "none"));
            TS.mergeNodes({ UniqueFromModelToPreserve, MergedAggregate });
            Erased.insert(MergedAggregate);
            MergedAggregate = UniqueFromModelToPreserve;
          }
        }

        if (MergedAggregate and MergedScalar) {
          revng_log(Log, "MergedAggregate and MergedScalar");

          // First, we want all pointers that point to MergedScalar to actually
          // start pointing to MergedAggregate.
          for (LTSN *Pointer : PointersToScalars) {
            revng_log(Log, "PointerToScalar: " << Pointer->ID);
            if (Erased.contains(Pointer)) {
              // This can happen if one of the PointersToScalars was also itself
              // a pointee of another pointer to scalar.
              revng_log(Log, "was erased");
              continue;
            }

            if (Pointer == MergedScalar
                and not hasOutgoingPointerEdge(Pointer)) {
              revng_log(Log,
                        "Pointer was merged with its pointee scalar, and is no "
                        "longer a pointer");
              continue;
            }

            const auto &[Pointee, PointerTag] = *Pointer->Successors.begin();
            revng_log(Log, "Pointee: " << Pointee->ID);
            auto InverseEdgeIt = Pointee->Predecessors.find({ Pointer,
                                                              PointerTag });
            TS.moveEdgeTarget(Pointee, MergedAggregate, InverseEdgeIt, 0);
          }
          revng_log(Log, "Fixed up PointersToScalars");

          // Second, we want to inject an instance of MergedScalar at offset 0
          // inside MergedAggregate.
          // If MergedAggregate is larger than MergedScalar we're fine.
          if (MergedAggregate->Size >= MergedScalar->Size) {
            TS.addInstanceLink(MergedAggregate,
                               MergedScalar,
                               OffsetExpression{ 0 });
          } else if (not MergedAggregate->NonScalar) {
            MergedAggregate->Size = MergedScalar->Size;
            TS.addInstanceLink(MergedAggregate,
                               MergedScalar,
                               OffsetExpression{ 0 });
          } else {
            revng_abort();
          }
        }
      }
    }
  }

  return Changed;
}

} // end namespace dla