/*
 * Call graph construction
 *
 * Copyright (C) 2012 Xi Wang, Haogang Chen, Nickolai Zeldovich
 * Copyright (C) 2015 - 2016 Chengyu Song 
 * Copyright (C) 2016 Kangjie Lu
 *
 * For licensing details see LICENSE
 */


#include <llvm/IR/DebugInfo.h>
#include <llvm/Pass.h>
#include <llvm/IR/Instructions.h>
#include <llvm/Support/Debug.h>
#include <llvm/IR/InstIterator.h>
#include <llvm/IR/Module.h>
#include <llvm/IR/Constants.h>
#include <llvm/ADT/StringExtras.h>
#include <llvm/Analysis/CallGraph.h>

#include "CallGraph.h"

using namespace llvm;

static std::string getScopeName(const Function *F) {
  if (llvm::GlobalValue::isExternalLinkage(F->getLinkage())) {
    return F->getName().str();
  } else {
    llvm::StringRef ModuleName = llvm::sys::path::stem(
        F->getParent()->getModuleIdentifier());
    return "_" + ModuleName.str() + "." + F->getName().str();
  }
}

Function* CallGraphPass::getFuncDef(Function *F) {
  FuncMap::iterator it = Ctx->Funcs.find(getScopeName(F));
  if (it != Ctx->Funcs.end())
    return it->second;
  else
    return F;
} 

bool CallGraphPass::isCompatibleType(Type *T1, Type *T2) {
  return false;
}

bool CallGraphPass::findCalleesByType(CallInst *CI, FuncSet &FS) {

  for (Function *F : Ctx->AddressTakenFuncs) {

    // just compare known args
    if (F->getFunctionType()->isVarArg()) {
      continue;
      //report_fatal_error("VarArg address taken function\n");
    } else if (F->arg_size() != CI->arg_size()) {
      continue;
    } else if (!isCompatibleType(F->getReturnType(), CI->getType())) {
      continue;
    }

    if (F->isIntrinsic()) {
      continue;
    }

    // type matching on args
    bool Matched = true;
    auto AI = CI->arg_begin();
    for (auto FI = F->arg_begin(), FE = F->arg_end(); FI != FE; ++FI, ++AI) {
      // check type mis-match
      Type *FormalTy = FI->getType();
      Type *ActualTy = (*AI)->getType();

      if (isCompatibleType(FormalTy, ActualTy)) {
        continue;
      } else {
        Matched = false;
        break;
      }
    }

    if (Matched)
      FS.insert(F);
  }

  return false;
}

bool CallGraphPass::runOnFunction(Function *F) {

  for (inst_iterator i = inst_begin(F), e = inst_end(F); i != e; ++i) {
       Instruction *I = &*i;
    // map callsite to possible callees
    if (CallInst *CI = dyn_cast<CallInst>(I)) {
      // ignore inline asm
      if (CI->isInlineAsm())
        continue;
      
      Function *CF = CI->getCalledFunction();
      if (CF) // ignore direct call
        continue;

      // save called values for point-to analysis
      Ctx->IndirectCallInsts.push_back(CI);

      // must be an indirect call, find all possible callees
      FuncSet &FS = Ctx->Callees[CI];
      findCalleesByType(CI, FS);
    }
  }

  return false;
}

bool CallGraphPass::doInitialization(Module *M) {

  for (Function &F : *M) { 
    // collect address-taken functions
    if (F.hasAddressTaken())
      Ctx->AddressTakenFuncs.insert(&F);
  }

  return false;
}

bool CallGraphPass::doFinalization(Module *M) {

  // update callee mapping
  for (Function &F : *M) {
    for (inst_iterator i = inst_begin(F), e = inst_end(F); i != e; ++i) {
      // map callsite to possible callees
      if (CallInst *CI = dyn_cast<CallInst>(&*i)) {
        FuncSet &FS = Ctx->Callees[CI];
        // calculate the caller info here
        for (Function *CF : FS) {
          CallInstSet &CIS = Ctx->Callers[CF];
          CIS.insert(CI);
        }
      }
    }
  }

  dumpCallees();

  return false;
}

bool CallGraphPass::doModulePass(Module *M) {
  bool Changed = true, ret = false;
  while (Changed) {
    Changed = false;
    for (Function &F : *M)
      Changed |= runOnFunction(&F);
    ret |= Changed;
  }
  return ret;
}

void CallGraphPass::dumpCallees() {
}

