//===- LoopRegionAnalysis.h. - Finds a Loop region starting from a call site
//---===//
//
// 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
//
//===----------------------------------------------------------------------===//
//
// Copyright (c) 2015 Advanced Micro Devices, Inc.  All rights reserved.
// Developed by: Advanced Micro Devices, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// with the Software without restriction, including without limitation the
// rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
// sell copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
// Redistributions of source code must retain the above copyright notice, this
// list of conditions and the following disclaimers.
// Redistributions in binary form must reproduce the above copyright notice,
// this list of conditions and the following disclaimers in the documentation
// and/or other materials provided with the distribution.
// Neither the names of Advanced Micro Devices, Inc., nor the names of its
// contributors may be used to endorse or promote products derived from this
// Software without specific prior written permission.
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// CONTRIBUTORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS WITH
// THE SOFTWARE.
//
//===----------------------------------------------------------------------===//
//
// This pass finds tries to find loop regions of program. It starts with
// collection of  call sites and tries to go up its call chain to its caller and
// caller's caller etc up to a certain depth.  Then it tries to find if there
// are enclosing loop regions to all these call sites.
//===----------------------------------------------------------------------===//

#ifndef LLVM_TRANSFORMS_UTILS_LOOPREGIONANALYSIS_H
#define LLVM_TRANSFORMS_UTILS_LOOPREGIONANALYSIS_H

#include "llvm/ADT/DepthFirstIterator.h"
#include "llvm/Analysis/CallGraph.h"
#include "llvm/Analysis/LoopInfo.h"
#include "llvm/IR/AbstractCallSite.h"
#include "llvm/IR/Function.h"
#include "llvm/IR/Instructions.h"
#include "llvm/IR/Module.h"
#include "llvm/Pass.h"
#include "llvm/Support/raw_ostream.h"

namespace llvm {

#define DEBUG_TYPE "loop-region"

class Node {
public:
  Loop *L;
  int incomming = 0;
  int bottom = 0;
  int top = 0;
  bool pruned = false;
  std::vector<Node *> neighbors;
  std::vector<Node *> rootlist;
  std::vector<Node *> startlist;
  Node(Loop *L) { this->L = L; }
};

class DirectedGraph {
public:
  std::vector<Node *> nodes;
  std::unordered_map<Loop *, Node *> loopNodeMap;
  std::unordered_map<CallBase *, Node *> callNodeMap;
  std::vector<Node *> roots;
  std::vector<Node *> visited;
  Node *Start;
  DirectedGraph() {}

  void addNode(Node *node) { nodes.push_back(node); }

  Node *GetNode(Loop *L) { return loopNodeMap[L]; }

  void addEdge(Node *source, Node *destination) {
    for (Node *neighbor : source->neighbors) {
      if (neighbor == destination)
        return;
    }
    source->neighbors.push_back(destination);
    destination->incomming++;
  }

  void collectRoots() {
    roots.clear();
    for (Node *node : nodes) {
      if (node->neighbors.size() == 0) {
        roots.push_back(node);
        LLVM_DEBUG(dbgs() << "Collected Root:"
                          << node->L->getHeader()->getParent()->getName()
                          << "\n");
      }
    }
  }

  void printLoopRegions() {
    for (auto rootnode : roots) {
      if (!rootnode->pruned) {
        LLVM_DEBUG(dbgs() << "Root Function:"
                          << rootnode->L->getHeader()->getParent()->getName()
                          << "\n");
        LLVM_DEBUG(dbgs() << "Root Loop:" << *(rootnode->L) << "\n");
      }
    }
  }

  bool checkReachablity() {
    collectRoots();
    for (auto entry : callNodeMap) {
      visited.clear();
      Start = entry.second;
      dfsSearch(entry.second);
      for (auto rootnode : roots) {
        if (!rootnode->pruned && std::find(visited.begin(), visited.end(),
                                           rootnode) == visited.end()) {
          return false;
        }
      }
    }
    return true;
  }

  void dfsSearch(Node *node) {
    visited.push_back(node);
    node->bottom++;
    node->startlist.push_back(Start);
    for (Node *neighbor : node->neighbors) {
      if (std::find(visited.begin(), visited.end(), neighbor) ==
          visited.end()) {
        dfsSearch(neighbor);
      }
    }

    int totalpruned = 0;
    for (Node *neighbor : node->neighbors) {
      if (neighbor->neighbors.size() == 0 && (neighbor->incomming == 1)) {
        neighbor->pruned = true;
        LLVM_DEBUG(dbgs() << "Pruning neighbour\n");
        totalpruned++;
      }
    }

    if (node->neighbors.size() == totalpruned) {
      if (std::find(roots.begin(), roots.end(), node) == roots.end())
        roots.push_back(node);
    }
  }

  void print() {
    for (Node *node : nodes) {
      if (node->L) {
        Loop *L = node->L;
        LLVM_DEBUG(dbgs() << "Node:" << *L << " ");
        LLVM_DEBUG(dbgs() << L->getHeader()->getParent()->getName());
        LLVM_DEBUG(dbgs() << "\t Pruned" << node->pruned << "\n");
        LLVM_DEBUG(dbgs() << "\t Incomming" << node->incomming << "\n");
        LLVM_DEBUG(dbgs() << "\t bottom" << node->bottom << "\n");
        LLVM_DEBUG(dbgs() << "\t start list\n");
        for (Node *n : node->startlist) {
          LLVM_DEBUG(dbgs() << "\t" << n << "\n");
        }
      } else {
        LLVM_DEBUG(dbgs() << "Call site node");
      }

      for (Node *neighbor : node->neighbors) {
        Loop *L = neighbor->L;
        LLVM_DEBUG(dbgs() << "\t Neighbors:" << *L << "\n");
        LLVM_DEBUG(dbgs() << L->getHeader()->getParent()->getName());
      }
      LLVM_DEBUG(dbgs() << "\n");
    }
  }
};

class LoopRegionAnalysis {
private:
  FunctionAnalysisManager *FM;

public:
  // The main entry point for the pass
  LoopRegionAnalysis(FunctionAnalysisManager *FM) : FM(FM) {}
  bool run(Module &M, std::vector<CallBase *> CallSites);
  void analyzeLoopRegions(std::vector<CallBase *> CallSites);
  void dfsLoopRegions(CallBase *CB, std::vector<CallBase *> &Path,
                      int loopdepth);
  void storeLoopRegion(std::vector<CallBase *> Path);
  Loop *getOuterLoopInCallSite(CallBase *CB);
  void getNewCallSites(Function *F, std::vector<CallBase *> &NewCallSites);
  std::vector<CallBase *> CallSites;
  std::vector<CallBase *> Visited;
  std::map<Function *, LoopInfo *> LoopMap;
  std::vector<CallBase *> Path;
  std::set<std::vector<Loop *>> collectedPaths;
  std::vector<Loop *> LoopRegions;
  std::vector<Loop *> getLoopRegions() {
    LoopRegions.clear();
    for (auto rootnode : G->roots)
      if (!rootnode->pruned && rootnode->L)
        LoopRegions.push_back(rootnode->L);
    return LoopRegions;
  }
  DirectedGraph *G;
};

} // namespace llvm
#endif // LLVM_TRANSFORMS_UTILS__LOOPREGIONANALYSIS_H
