#define DEBUG_TYPE "InterproceduralDefUse"

#include "ArgumentHelpers.hpp"
#include "InterproceduralDefUse.hpp"
#include "TransitiveDefUse.hpp"
#include "Utility.hpp"

#include "llvm/ADT/Statistic.h"
#include "llvm/Analysis/CallGraph.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/InstIterator.h"

#include <boost/format.hpp>

using namespace boost;
using namespace llvm;

InterproceduralDefUse::InterproceduralDefUse()
 :CallGraphSCCPass((intptr_t)&ID), _TDU(NULL)
{
}

bool InterproceduralDefUse::doInitialization(CallGraph & CG)
{
  return false;
}

static void printUse(const Value* V)
{
  DOUT << format("    %s\n") % valueToString(V);
}

bool InterproceduralDefUse::mergeValueTDU(Value* V)
{
  assert(V);
  assert(isa<PointerType>(V->getType()));

  size_t startingSize = _iuses[V].size();

  // Copy over the local uses of V
  const use_set& uses = _TDU->getTDUChain(V);
  _iuses[V].insert(uses.begin(), uses.end());
  DOUT << format("   %d uses (base)\n") % uses.size();
  std::for_each(_iuses[V].begin(), _iuses[V].end(), printUse);

  // Make a copy so that we don't invalidate our iterators
  use_set copy = _iuses[V];
  for(use_set::iterator it = copy.begin(), ed = copy.end(); it != ed; ++it)
  {
    CallInst* CI = NULL;
    if(!(CI = dyn_cast<CallInst>(*it)))
      continue;

    const aset* argset = _TDU->getTDUArgSet(V, CI);
    if(!CI->getCalledFunction())
    {
      DOUT << format("Use disappearing into function pointer\n");
      continue;
    }
    else if(!argset)
    {
      // Nothing..?
      continue;
    }

    // For each argument we find that is not null, merge _iuses[arg] with _iuses[V]
    for(aset::const_iterator arg = argset->begin(), e_arg = argset->end(); arg != e_arg; ++arg)
    {
      Argument* A = *arg;
      // TODO: Something to indicate varargs, possibly, but we largely assume that
      // varargs are benign so this should be sufficient
      if(!A)
      {
        DOUT << format("[info] Skipping vararg");
        continue;
      }

      _iuses[V].insert(_iuses[A].begin(), _iuses[A].end());
      DOUT << format("[INFO(mergeValueTDU)] Argument [%s] adding to %s\n")
        % A->getName() % valueToString(V);
      
      std::for_each(_iuses[A].begin(), _iuses[A].end(), printUse);
    }
  }

  return _iuses[V].size() != startingSize;
}

void InterproceduralDefUse::addLocalTDUArgSet(Value* V)
{
  const use_set& uses = _TDU->getTDUChain(V);
  Function* F = getParentFunction(V);

  for(use_set::const_iterator it = uses.begin(), ed = uses.end(); it != ed; ++it)
  {
    Function* instParent = getParentFunction(*it);
    CallInst* CI = NULL;
    // We only want to add CallInstructions from the _current_ Function
    if(!(CI = dyn_cast<CallInst>(*it)) || (F != instParent))
      continue;

    addArgumentTDUSets(F, CI);
  }
}

void InterproceduralDefUse::addArgumentTDUSets(Function* F, CallInst* CI)
{
  for(Function::arg_iterator it = firstArgument(*F), ed = F->arg_end(); it != ed; ++it)
  {
    if(!isa<PointerType>(it->getType()) || it->hasByValAttr())
      continue;

    const aset* argset = _TDU->getTDUArgSet(&*it, CI);
    if(!argset || argset->empty())
    {
      DOUT << format("  [info - addArgumentTDUSets] No argset in TDU for %s / %s\n")
        % F->getName() % valueToString(it);
      continue;
    }
    
    _iargs[std::make_pair(&*it, CI)] = *argset;
  }
}

bool InterproceduralDefUse::mergeValueTDUArgSet(Value* V)
{
  assert(V);
  assert(isa<PointerType>(V->getType()));
  DOUT << format("  [info - mergeValueTDUArgSet] %s @ %s\n")
    % getParentFunction(V)->getName() % valueToString(V);
  
  Argument* A = NULL;
  if(!(A = dyn_cast<Argument>(V)))
    return false;
  
  const use_set& uses = _TDU->getTDUChain(V);

  // 0) Add the TDUArgSet for the function containing V.
  addLocalTDUArgSet(V);
  
  // 1) For each call C (call to function F) in the IDU of V that has an entry in the IDUArgSet:
  bool makingProgress;
  bool madeAnyProgress = false;
  do
  {
    makingProgress = false;
    for(use_set::const_iterator it = uses.begin(), ed = uses.end(); it != ed; ++it)
    {
      // Skip entries that are not callinsts with entries in the _iargs table.
      CallInst* CI = NULL;
      if(!(CI = dyn_cast<CallInst>(*it)))
        continue;

      amap::iterator argit = _iargs.find(std::make_pair(V, CI));
      if(argit == _iargs.end())
        continue;
      
      // If there is such an entry, we can look at the function call "chained" off of it
      // and make a new entry for that.

      // If there are no actual uses of this parameter, don't descend farther.
      if(argit->second.empty())
        continue;

      
      CallInst* chainedCall = argit->first.second;
      // Function* chainedFunction = chainedCall->getFunction();
      // if(!chainedFunction)
      // {
      //   // FuncPtr call
      // }
      makingProgress = descendTDUArgSet(V, chainedCall, argit->second, uses) || makingProgress;
      madeAnyProgress = makingProgress || madeAnyProgress;
    }
  }while(makingProgress);
  // 2) Note the argument(s) in F corresponding to this call: A_j
  // 3) Make a new entry for this with the local entry/arg and the looked up value of A_j.
  // 4) Repeat until no changes

  return madeAnyProgress;
}

bool InterproceduralDefUse::descendTDUArgSet(Value* V, CallInst* CI, const aset& arguses, const use_set& uses)
{
  assert(V);
  assert(isa<PointerType>(V->getType()));
  assert(CI);
  
  Function* calledFunction = CI->getCalledFunction();
  if(!calledFunction)
  {
    DOUT << "  [info - descendTDUArgSet] Call through function pointer\n";
    // Function pointer
    return false;
  }

  bool madeProgress = false;
  // Now, if calledFunction is the parent of some call in the uses set, we can transitively
  // add the entries (arguses) of this CI for V to the interprocedural argument set.
  for(use_set::const_iterator it = uses.begin(), ed = uses.end(); it != ed; ++it)
  {
    CallInst* callee = NULL;
    if(!(callee = dyn_cast<CallInst>(*it)) ||
       (callee->getCalledFunction() != calledFunction) ||
       (callee == CI))
      continue;

    madeProgress = _iargs.find(std::make_pair(V, callee)) == _iargs.end() || madeProgress;
    _iargs[std::make_pair(V, callee)] = arguses;
  }

  return madeProgress;
}

bool InterproceduralDefUse::mergeFunctionTDU(Function* F)
{
  bool madeProgress = false;
  DOUT << format("mergeFunctionTDU: %s\n") % F->getName();  
  for(Function::arg_iterator it = firstArgument(*F),
        ed = F->arg_end(); it != ed; ++it)
  {
    if(!isa<PointerType>(it->getType()) || it->hasByValAttr())
      continue;
    DOUT << format("  [arg]: %s\n") % it->getName();
    madeProgress = mergeValueTDU(it) || madeProgress;
  }

  // Also, callinsts
  for(inst_iterator I = inst_begin(F), E = inst_end(F); I != E; ++I)
  {
    if(!isa<CallInst>(&*I))
      continue;

    if(!isa<PointerType>(I->getType()))
      continue;

    madeProgress = mergeValueTDU(&*I) || madeProgress;
  }

  return madeProgress;
}

bool InterproceduralDefUse::mergeFunctionTDUArgSet(Function* F)
{
  bool madeProgress = false;
  for(Function::arg_iterator it = firstArgument(*F),
        ed = F->arg_end(); it != ed; ++it)
  {
    if(!isa<PointerType>(it->getType()) || it->hasByValAttr())
      continue;

    madeProgress = mergeValueTDUArgSet(&*it) || madeProgress;
  }

  return madeProgress;
}

/*
  For each function, compute a transitive closure of def-use information:
  1) Constructing an alias set for each argument, merging def-use information for
     each potential alias until reaching a fixed-point
  2) [Optional] Track return values for all call instructions.
  3) For each Use of an argument (taking in to account the transitive-closure over
     aliasing) occuring in the context of a call instruction, merge the interprocedural
     def-use information from the called function.

  To handle external function calls, any Use in the context of an external
  function call needs to have some metadata attached to it indicating which
  function and which parameter, so that clients can take appropriate action
  (either checking against provided data or by making conservative assumptions).

  Similar steps need to be taken for Uses that are passed to function pointers
  (and virtual method calls later on).
  
  Possible implementation strategy: boost::variant
  Alternative: side-by-side mapping of "external dependencies" (probably better)
 */

bool InterproceduralDefUse::runOnSCC(const std::vector<CallGraphNode *> &SCCM)
{
  if(!_TDU)
  {
    _TDU = &getAnalysis<TransitiveDefUse>();
  }
  bool makingProgress = false;
  do
  {
    makingProgress = false;
    for(std::vector<CallGraphNode*>::const_iterator it = SCCM.begin(),
          ed = SCCM.end(); it != ed; ++it)
    {
      Function* F = (*it)->getFunction();
      if(!F)
        continue;

      cerr << format("[InterproceduralDefUse] Analyzing %s\n") % F->getName();
      
      makingProgress = mergeFunctionTDU(F) || makingProgress;
      makingProgress = mergeFunctionTDUArgSet(F) || makingProgress;
    }
    // SCCM.size() == 1 => we only need one pass
  } while((SCCM.size() != 1) && makingProgress);
  
  return false;
}

void InterproceduralDefUse::getAnalysisUsage(AnalysisUsage &AU) const
{
  AU.setPreservesAll();
  AU.addRequired<TransitiveDefUse>();
}

static void printUses(const Value* V, const use_set& uses)
{
  DOUT << format("[ITDU]: %s\n") % valueToString(V);

  if(const Argument* arg = dyn_cast<Argument>(V))
  {
    DOUT << format("(%s)") % arg->getParent()->getName();
  }
  DOUT << std::endl;
  for(use_set::const_iterator it = uses.begin(), ed = uses.end(); it != ed; ++it)
  {
    DOUT << format(" %s\n") % valueToString(*it);
  }
}

bool InterproceduralDefUse::doFinalization(CallGraph &CG)
{
  for(use_map::iterator it = _iuses.begin(), ed = _iuses.end(); it != ed; ++it)
  {
    printUses(it->first, it->second);
  }
  
  return false;
}

const use_set& InterproceduralDefUse::getIDUChain(Value* V) const
{
  assert(V);
  assert(isa<PointerType>(V->getType()));
  assert(isa<Argument>(V) || isa<CallInst>(V));

  use_map::const_iterator it = _iuses.find(V);
  assert(it != _iuses.end());

  return it->second;
}

const aset* InterproceduralDefUse::getIDUArgSet(Value* V, CallInst* CI) const
{
  assert(V);
  assert(isa<PointerType>(V->getType()));
  assert(CI);

  amap::const_iterator it = _iargs.find(std::make_pair(V, CI));
  // Function pointer
  if(it == _iargs.end())
    return NULL;

  return &it->second;
}

char InterproceduralDefUse::ID = 0;
RegisterPass<InterproceduralDefUse> IDUReg("interDefUse", "Builds interprocedural def-use chains as in [Harrold and Soffa 94]");
