#define DEBUG_TYPE "TransitiveDefUse"

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

#include "llvm/Instructions.h"
#include "llvm/Module.h"
#include "llvm/Type.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/InstIterator.h"
#include "llvm/Support/InstVisitor.h"

#include <boost/bind.hpp>
#include <boost/format.hpp>
#include <boost/function.hpp>

#include <deque>
#include <iostream>

using namespace boost;
using namespace llvm;

namespace
{
  void printDUChain(Value* v)
  {
    DOUT << "[DU-Chain]\n";
    for(Value::use_iterator it = v->use_begin(), ed = v->use_end(); it != ed; ++it)
    {
      DOUT << valueToString(*it) << "\n\n";
    }
  }

  void printTDUChain(const use_set& us)
  {
    DOUT << "[TDU-Chain]\n";
    for(use_set::const_iterator it = us.begin(), ed = us.end();
        it != ed; ++it)
    {
      DOUT << valueToString(*it) << "\n\n";
    }
  }
  
  void mergeUseSets(use_set& base, const use_set& extra)
  {
    base.insert(extra.begin(), extra.end());
  }

  void printUseSet(const std::string & indent, const use_set& uses)
  {
    for(use_set::const_iterator it = uses.begin(), ed = uses.end(); it != ed; ++it)
    {
      DOUT << format("%s%s\n") % indent % valueToString(*it);
    }
  }
  struct TransitiveClosureVisitor : public InstVisitor<TransitiveClosureVisitor>
  {
    TransitiveClosureVisitor(Value* V, use_map& tuses, use_map& rdeps, amap& targs,
                             worklist_t& worklist, function<void (Value*)> calculateValueTDU)
    :_V(V),
     _tuses(tuses),
     _rdeps(rdeps),
     _targs(targs),
     _worklist(worklist),
     _calculateValueTDU(calculateValueTDU),
     _phase_size(0)
    {}

    void startPhase()
    {
      _phase_size = _tuses[_V].size();
    }

    bool makingProgress() const
    {
      return _phase_size != _tuses[_V].size();
    }
    
    // Any use of a PHI node could be a use of the root value with which it is associated (_V)
    void visitPHINode(PHINode &PN)
    {
      use_set& value_uses = _tuses[_V];
      
      // The eventual result of V depends on this phi node
      _rdeps[&PN].insert(_V);
      
      // Not yet analyzed
      if(_tuses.find(&PN) == _tuses.end())
      {
        // Might as well just bail out - the PHI node will just
        // add this node to the worklist again anyway.
        _worklist.push_back(&PN);
        return;
      }
              
      mergeUseSets(value_uses, _tuses[&PN]);
    }

    // This is largely similar to PHI nodes, but even simpler
    void visitBitCastInst(BitCastInst& BCI)
    {
      use_set& value_uses = _tuses[_V];
      _rdeps[&BCI].insert(_V);
      
      if(_tuses.find(&BCI) == _tuses.end())
      {
        _worklist.push_back(&BCI);
        return;
      }

      mergeUseSets(value_uses, _tuses[&BCI]);
    }
  
    void visitGetElementPtrInst(GetElementPtrInst& I)
    {
      DOUT << format("[info - visitor] Visiting GEP: %s\n") % valueToString(&I);
      // The GEP itself is in the value_uses set already;
      // just add the src (operand 0) to the value_uses set.
      use_set& value_uses = _tuses[_V];
      Value* src = I.getPointerOperand();
      
      _rdeps[src].insert(_V);
      _rdeps[&I].insert(_V);

      if(_tuses.find(src) == _tuses.end())
      {
        _worklist.push_back(src);
        _worklist.push_back(&I);
        return;
      }

      mergeUseSets(value_uses, _tuses[src]);
      value_uses.insert(src);
    }

    void visitLoadInst(LoadInst& LI)
    {
      use_set& value_uses = _tuses[_V];
      
      // Operand 0 is the source (and only operand for load)
      Value* src = LI.getOperand(0);
      
      // Not sure if this case can occur or even what it means, but it
      // is best to not get bitten later
      if(_V != src && value_uses.find(src) == value_uses.end())
        return;

      // Standard propagation of loads UP TO the same "indirection depth"
      // as _V.  Further loads would be /dereferences/ of _V, which we don't
      // really care about. (Do we?)  Under our assumption that libraries are
      // internally consistent and well-behaved, this type of escape is permissible.
      _rdeps[&LI].insert(_V);

      if(_tuses.find(&LI) == _tuses.end())
      {
        _worklist.push_back(&LI);
        return;
      }

      if(!typeMoreIndirectThan(_V->getType(), LI.getType()))
      {
        return;
      }
      
      mergeUseSets(value_uses, _tuses[&LI]);
      value_uses.insert(&LI);
    }

    bool typeMoreIndirectThan(const Type* target, const Type* use)
    {
      if(!isa<PointerType>(target) && !isa<PointerType>(use))
        return true;

      if(!isa<PointerType>(target))
        return true;
      
      if(!isa<PointerType>(use))
        return false;

      const PointerType* pTarget = cast<PointerType>(target);
      const PointerType* pUse = cast<PointerType>(use);
      assert(pTarget->getNumContainedTypes() == 1);
      assert(pUse->getNumContainedTypes() == 1);
      
      return typeMoreIndirectThan(pTarget->getContainedType(0),
                                  pUse->getContainedType(0));
    }

    void visitStoreInst(StoreInst& SI)
    {
      DOUT << format("[info visitor]: Visiting: %s\n") % valueToString(&SI);
      
      use_set& value_uses = _tuses[_V];
      
      // Operand 0 is the source
      Value* src = SI.getOperand(0);
      Value* dst = SI.getOperand(1);

      // If src is _V or in value_uses, add dst to value_uses.
      // If not, then dst should be _v or in value_uses, but this case doesn't matter to us (but it
      // would to a real alias analysis.
      if(_V != src && value_uses.find(src) == value_uses.end())
        return;
      
      _rdeps[dst].insert(_V);
      _rdeps[&SI].insert(_V);

      if(_tuses.find(dst) == _tuses.end())
      {
        DOUT << format("  Adding dst and SI to worklist\n");
        _worklist.push_back(dst);
        _worklist.push_back(&SI);
        return;
      }

      mergeUseSets(value_uses, _tuses[dst]);
      value_uses.insert(dst);
    }

    /*
      A set with NULL entries indicates that the value was passed as a vararg
      The lack of an entry indicates that the call was through a function pointer
     */
    void visitCallInst(CallInst& CI)
    {
      use_set& value_uses = _tuses[_V];

      // Check each operand to see which one(s) triggered
      // this use dependence.  Add them to the argument mapping,
      // which is consulted during the interprocedural phase.
      // NULL in the arg_map denotes a call through a function pointer.
      Function* F = CI.getCalledFunction();
      if(!F)
      {
        return;
      }

      // Forcibly instantiate an empty set.  This is essentially a sentinel so that
      // we can always find a table entry for this value/call pair as long as there
      // is not a call through a function pointer.
      (void)_targs[std::make_pair(_V, &CI)];
      
      DOUT << format(" [info - TDU] call to function %s, tracking value %s\n")
        % F->getName() % valueToString(_V);
      // printUseSet("   VSE: ", value_uses);
        
      bool isVararg = false;
      Function::arg_iterator arg = firstArgument(*F);
      for(Instruction::op_iterator it = operandStartIter(&CI), ed = CI.op_end(); it != ed; ++it)
      {
        isVararg = isVararg || arg == F->arg_end();

        // We want to make sure that arguments which have not been used anywhere
        // else are still represented interprocedurally when they pass through.
        // if(value_uses.count(it->get()) == 0 && !isa<Argument>(it->get()))
        if(value_uses.find(it->get()) == value_uses.end())
        {
          // Remember to advance the parallel argument iterator
          if(!isVararg) ++arg;
          
          continue;
        }
                
        DOUT << format("  Operand: %s\n") % valueToString(*it);

        if(!isVararg)
        {
          DOUT << format("  Argument: %s\n") % valueToString(arg);
          DOUT << format(" [info - visitCallInst/TDU] Adding %s to argset for %s\n")
            % arg->getName() % valueToString(_V);
          DOUT << format("    due to op %s\n") % valueToString(it->get());

          _targs[std::make_pair(_V, &CI)].insert(arg);

          // Need to advance our parallel argument iterator
          ++arg;
        }
        else
        {
          DOUT << format("  Argument: VARARG\n");
          
          _targs[std::make_pair(_V, &CI)].insert(NULL);
        }
      }
    }
    
  private:
    Value* _V;
    use_map& _tuses;
    use_map& _rdeps;
    amap& _targs;
    worklist_t& _worklist;
    function<void (Value*)> _calculateValueTDU;

    size_t _phase_size;
  };
  
  void propagateRDepsToParents(Value* V, use_map& rdeps)
  {
    use_set& ruses = rdeps[V];
    for(use_set::iterator it = ruses.begin(), ed = ruses.end(); it != ed; ++it)
    {
      rdeps[*it].insert(ruses.begin(), ruses.end());
    }
  }
}

TransitiveDefUse::TransitiveDefUse()
 :ModulePass((intptr_t)&ID)
{
}

void TransitiveDefUse::calculateValueTDU(Value* V)
{
  use_set& value_uses = _tuses[V];
  
  size_t pre_alias_size = value_uses.size();  
  
  // Add the normal DU-chain calculated by LLVM
  for(Value::use_iterator it = V->use_begin(), ed = V->use_end(); it != ed; ++it)
  {
    // We don't want to introduce stores INTO our variable;
    // they mess up later stages of the analysis and are not
    // interesting
    if(StoreInst* SI = dyn_cast<StoreInst>(V))
    {
      Value* dst = SI->getOperand(1);
      if(dst == *it)
        continue;
    }

    value_uses.insert(*it);
  }

  value_uses.insert(V);

  TransitiveClosureVisitor TCV(V, _tuses, _rdeps, _targs, _worklist,
                               bind(&TransitiveDefUse::calculateValueTDU, const_cast<TransitiveDefUse*>(this), _1));
  // Reach a fixed-point on this instruction set
  do
  {
    // Make a copy so that we don't invalidate our iterators below
    TCV.startPhase();
    use_set local_value_uses = _tuses[V];
    for(use_set::iterator it = local_value_uses.begin(),
          ed = local_value_uses.end(); it != ed; ++it)
    {
      Value* useV = *it;
      if(Instruction* I = dyn_cast<Instruction>(useV))
      {
        TCV.visit(I);
      }
      else
      {
        DOUT << format("[WARNING] Non-instruction in worklist processing");
      }
      // Only other entities are Arguments, BasicBlocks, and Constants.
      // BasicBlocks should never be here, Constants are not a concern.
      // Arguments are trivial.
    }
  }while(TCV.makingProgress());
  
  // Merge DU chains derived from alias relations.  If value_uses grows at all
  // from any alias, add that value; if the alias has no entry in transitive_uses,
  // also add it (on first visit, it will have nothing).  Also, re-add self
  // AliasSet * 	getAliasSetForPointerIfExists (Value *P, unsigned Size)

  // A visitor would be nice here, but it is somewhat useful to be able to visit
  // all of one class of use first (in particular, all GEPs before considering any loads).
  
  // Transitively and recursively follow PHI nodes
  // phiTransitiveClosure(V, value_uses);


  // Also transitively follow stores (either p = q, p = &q, p.x = q, or p.x = &q)
  // The first two cases are the same (the second just requires two steps) due to LLVM
  //   normalizations.
  // The second two require the same two steps, plus one of following the store target if
  //   it is a GEP.
  // storeTransitiveClosure(V, value_uses);


  // Handle bitcasts of the value
  // bitcastTransitiveClosure(V, value_uses);

  // And handle GEP instructions (in a field sensitive manner for structs, insensitively for
  // arrays).
  // gepTransitiveClosure(V, value_uses);

  // Finally, handle Load instructions; these only matter for GEP instructions, indicating
  // USES of GEP instructions.  As an example, this denotes retrieval of an item from an array
  // or struct.
  
  // Instead of getting recursive, it would be best to use a worklist here.  Just add all
  // aliases and PHI nodes to a worklist and retry this until no changes occur.
  if(pre_alias_size != value_uses.size())
  {
    propagateRDepsToParents(V, _rdeps);
    _worklist.insert(_worklist.end(), _rdeps[V].begin(), _rdeps[V].end());
  }
}

void TransitiveDefUse::calculateFunctionTDU(Function* F) 
{
  for(Function::arg_iterator it = F->arg_begin(), ed = F->arg_end(); it != ed; ++it)
  {
    if(!isa<PointerType>(it->getType()) || it->hasByValAttr())
      continue;

    DOUT << format("==Argument@%s: %s\n") % F->getName() % valueToString(it);
    // printDUChain(it);

    _worklist.push_back(it);

    while(!_worklist.empty())
    {
      Value* v = _worklist.front();
      _worklist.pop_front();
      calculateValueTDU(v);//, AST);
    }

    // printTDUChain(_tuses[&*it]);
  }

  // Also, take the return values of all function calls (callinstrs)
  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;

    DOUT << format("==CallInst: %s\n") % valueToString(&*I);
    // printDUChain(&*I);
    
    _worklist.push_back(&*I);
    while(!_worklist.empty())
    {
      Value* v = _worklist.front();
      _worklist.pop_front();
      calculateValueTDU(v);
    }
    // printTDUChain(_tuses[&*I]);
  }
}

bool TransitiveDefUse::runOnModule(Module& M)
{
  for(Module::iterator F = M.begin(), E = M.end(); F != E; ++F)
  {
    cerr << format("[TransitiveDefUse] Analyzing %s\n") % F->getName();
    calculateFunctionTDU(F);
  }

  // Release a lot of memory
  use_map tmp;
  _rdeps.swap(tmp);
  
  return false;
}

void TransitiveDefUse::getAnalysisUsage(AnalysisUsage &AU) const
{
  AU.setPreservesAll();
}


const use_set& TransitiveDefUse::getTDUChain(Value* V) const
{
  assert(V);
  assert(isa<PointerType>(V->getType()));
  // Either argument or return value
  assert(isa<Argument>(V) || isa<CallInst>(V));

  use_map::const_iterator it = _tuses.find(V);
  if(it == _tuses.end())
    assert(false && "Unknown failure, value missing");

  return it->second;
}

const aset* TransitiveDefUse::getTDUArgSet(Value* V, CallInst* CI) const
{
  assert(CI);
  amap::const_iterator it = _targs.find(std::make_pair(V, CI));
  if(it == _targs.end())
    return NULL;

  return &it->second;
}

char TransitiveDefUse::ID = 0;
RegisterPass<TransitiveDefUse> TDUReg("transDefUse", "Augments local def-use chains with _may_ uses through alais information and PHI node tracking");
