//=- IndirectToDirectCallAnalysis.h - Analyze indirect to direct calls --==//
//
//===----------------------------------------------------------------------===//
//                     The LLVM Compiler Infrastructure
//
// This file is distributed under the University of Illinois Open Source
// License. See LICENSE.TXT for details.
//
//===----------------------------------------------------------------------===//
//
// 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.
//
//===----------------------------------------------------------------------===//
#include "llvm/Analysis/AliasAnalysis.h"
#include "llvm/Analysis/AliasSetTracker.h"
#include "llvm/Analysis/TargetLibraryInfo.h"
#include "llvm/Analysis/VectorUtils.h"
#include "llvm/Support/Debug.h"
#include "llvm/Transforms/IPO.h"
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/SmallPtrSet.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/IR/Constants.h"
#include "llvm/IR/DebugInfo.h"
#include "llvm/IR/DerivedTypes.h"
#include "llvm/IR/Instructions.h"
#include "llvm/IR/InstIterator.h"
#include "llvm/IR/Module.h"
#include "llvm/IR/TypeFinder.h"
#include "llvm/IR/ValueSymbolTable.h"
#include "llvm/InitializePasses.h"
#include "llvm/Pass.h"
#include "llvm/Transforms/Utils/BasicBlockUtils.h"
#include "llvm/Transforms/Utils/Cloning.h"
#include "llvm/Transforms/Utils/Local.h"
#include "llvm/Transforms/Utils/ValueMapper.h"
#include <iostream>
#include <list>
#include <map>
#include <utility>
#include <memory>

using namespace llvm;

namespace llvm {

// Both Analysis & Transform
// Return formal argument at specified position.
Value *functionArgAtPosition(Function *F, unsigned argPos);

class AnalysisGraph {
public:
  int argumentPosition;
  bool hasCall;
  bool subTreeHasIndirectCall;
  bool isRecursive;
  Function *F;
  SmallVector<AnalysisGraph *, 8> Child;
  SmallVector<bool,8> ChildLeased;

  AnalysisGraph(int ArgPos, Function *Func)
      : argumentPosition(ArgPos), hasCall(false),
        subTreeHasIndirectCall(false), isRecursive(false), F(Func) {}

  ~AnalysisGraph() {
    for (unsigned i = 0, e = Child.size(); i != e; ++i) {
      if (!ChildLeased[i]) delete Child[i];
    }
  }

  void setNoIndirectCallInSubTree() { subTreeHasIndirectCall = false; }

  void setIndirectCallInSubTree() { subTreeHasIndirectCall = true; }

  void hasIndirectCall() { hasCall = true; }

  void print();
};

class AnalysisGraphInfo {
public:
  Function *TargetFunction;
  AnalysisGraph *Tree;
  const CallInst *Call;
  AnalysisGraphInfo(Function *TF, AnalysisGraph *AG, const CallInst *C)
      : TargetFunction(TF), Tree(AG), Call(C) {}
  ~AnalysisGraphInfo() { delete Tree; }
};

struct AGNodeKey {
  Function *FrameFunc;
  Function *TargetFunc;
  unsigned argPos;
  struct Hash {
    long operator()(AGNodeKey const &k) const {
      return cantor(cantor((long)k.FrameFunc, (long)k.TargetFunc), k.argPos);
    }
    long cantor(long a, long b) const { return (a + b + 1) * (a + b) / 2 + b; }
  };
  bool operator==(AGNodeKey const &k) const { return FrameFunc == k.FrameFunc && TargetFunc == k.TargetFunc && argPos == k.argPos; }
};
class IndirectToDirectCallAnalysis {
public:
  ModulePass *CurPass;
  AnalysisManager<Function> *FAM;
  AliasAnalysis *AA;
  AliasSetTracker *CurAST;

  SmallVector<std::shared_ptr<AnalysisGraphInfo> , 16> IdentifiedAnalysisGraph;
  std::map<Function *, bool> IndirectRecursionTracer;
  std::map<CallInst *, SmallVector<Function *, 16>> IndirectCallPossiblities;
  std::unordered_map<AGNodeKey, AnalysisGraph*, AGNodeKey::Hash> AnalysisGraphsVisited;
  IndirectToDirectCallAnalysis(ModulePass *CurPass,
                               AnalysisManager<Function> *FAM) :
                               CurPass(CurPass), FAM(FAM), AA(nullptr),
                               CurAST(nullptr) {
    IdentifiedAnalysisGraph.clear();
    IndirectRecursionTracer.clear();
    IndirectCallPossiblities.clear();
  }
  
  void analyze(Module &);
  void analyze(Function &);
  void analyze(Instruction &);
  bool processModule(Module &M);

  AnalysisGraph *buildAnalysisGraph(Function *FrameFunc, Function *TargetFunc,
                                    unsigned argPos, unsigned depth);

  void storeAnalysisGraph(Function *TargetFunction, AnalysisGraph *Tree,
                          const CallInst *Call);

  bool isAccessSafeInFunction(Function *F, Value *V);
  bool isIndirectRecursion(Function *F) {
    std::map<Function *, bool>::iterator MItr = IndirectRecursionTracer.find(F);
    if (MItr == IndirectRecursionTracer.end())
      return false;
    return true;
  }

  void updateIndirectCallPossiblities(CallInst *Call, Function *F) {
    std::map<CallInst *, SmallVector<Function *, 16>>::iterator MItr =
        IndirectCallPossiblities.find(Call);
    if (MItr != IndirectCallPossiblities.end()) {
      for (auto &VItr : MItr->second) {
        if (VItr == F)
          return;
      }
    }
    IndirectCallPossiblities[Call].push_back(F);
  }

  const SmallVector<Function *, 16> *
  findIndirectCallPossiblities(CallInst *Call) {
    std::map<CallInst *, SmallVector<Function *, 16>>::iterator MItr =
        IndirectCallPossiblities.find(Call);
    if (MItr == IndirectCallPossiblities.end())
      return nullptr;

    return &MItr->second;
  }

  void printCallPossiblities() {
    for (auto &Itr : IndirectCallPossiblities) {
      dbgs() << "Indirect Call Instruction: \n";
      Itr.first->dump();
      dbgs() << "Indirect Call Instruction TYPE: \n";
      Itr.first->getType()->dump();
      dbgs() << "Has Following Target Possiblities : \n";
      for (auto &VItr : Itr.second)
        dbgs() << "    " << VItr->getName() << "\n";
    }
  }
};

class IndirectToDirectCallAnalysisLegacy : public ModulePass {
public:
  static char ID; // Pass identification, replacement for typeid
  IndirectToDirectCallAnalysis *IDCA;
  explicit IndirectToDirectCallAnalysisLegacy() : ModulePass(ID),
                                                  IDCA(nullptr) {
    initializeIndirectToDirectCallAnalysisLegacyPass(
        *PassRegistry::getPassRegistry());
  }
 
  ~IndirectToDirectCallAnalysisLegacy() {
    if (IDCA)
      delete IDCA;
  }

  bool runOnModule(Module &M) override;
  IndirectToDirectCallAnalysis *getAnalysis() { return IDCA;}

  void getAnalysisUsage(AnalysisUsage &AU) const override {
    AU.addRequired<TargetLibraryInfoWrapperPass>();
    AU.addRequired<AssumptionCacheTracker>();
    AU.setPreservesAll();
  }
};

class IndirectToDirectCallAnalysisInfo : public AnalysisInfoMixin<IndirectToDirectCallAnalysisInfo> {
public:
    friend AnalysisInfoMixin<IndirectToDirectCallAnalysisInfo>;
    static AnalysisKey Key;
    typedef IndirectToDirectCallAnalysis Result;
    Result run(Module& M, ModuleAnalysisManager& AM);
};

}
