#define DEBUG_TYPE "EscapeAnalysis"

#include "ArgumentHelpers.hpp"
#include "EscapeAnalysis.hpp"
#include "InterproceduralDefUse.hpp"
#include "Results.hpp"
#include "Utility.hpp"

#include "llvm/Argument.h"
#include "llvm/Constants.h"
#include "llvm/ADT/Statistic.h"
#include "llvm/Analysis/CallGraph.h"
#include "llvm/Support/CommandLine.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/InstIterator.h"
#include "llvm/Support/Streams.h"

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

#include <algorithm>
#include <iterator>
#include <iostream>
#include <numeric>

using namespace llvm;
using namespace boost;

extern cl::opt<bool> PermissiveFuncptr;

// STATISTIC(EscapePasses, "Number of iterations of the Escape analysis required.");
// STATISTIC(TotalEscapingVariables, "Total number of variables escaping from functions.");
// STATISTIC(EscapingArguments, "Number of function arguments escaping from their functions.");
namespace
{
  enum
  {
    Top,
    Weak,
    Strong
  };
  
  // The lattice for this analysis is the following:
  //     T
  //   /   \   //
  // Weak  |
  //   \   |
  //    \  |
  //     Strong
  //

  bool isaGlobalConstantExpr(Value* V)
  {
    ConstantExpr* CE = NULL;
    if(!(CE = dyn_cast<ConstantExpr>(V)))
      return false;

    // Add others as necessary
    switch(CE->getOpcode())
    {
    case Instruction::GetElementPtr:
    {
      Constant* C = CE->getOperand(0);
      return isa<GlobalValue>(C);
    }
    default:
      ;
    }

    return false;
  }
  
  // Return true if this value is a constituent of an argument.
  // This means it could be an argument itself, or some sub-part of
  // an array/struct argument (passed by value)
  bool tracesBackToArgumentOrGlobal(Value* dest, use_set& visited)
  {
    if(visited.find(dest) != visited.end())
      return false;

    visited.insert(dest);
    
    if(isa<Argument>(dest) || isa<GlobalValue>(dest) || isaGlobalConstantExpr(dest))
      return true;

    // Follow GEPs and Loads
    if(LoadInst* LI = dyn_cast<LoadInst>(dest))
    {
      Value* src = LI->getOperand(0);
      tracesBackToArgumentOrGlobal(src, visited);
    }
    else if(GetElementPtrInst* GEP = dyn_cast<GetElementPtrInst>(dest))
    {
      Value* src = GEP->getPointerOperand();
      return tracesBackToArgumentOrGlobal(src, visited);
    }
    else if(PHINode* PN = dyn_cast<PHINode>(dest))
    {
      for(unsigned int i = 0; i < PN->getNumIncomingValues(); ++i)
      {
        Value* iv = PN->getIncomingValue(i);
        if(tracesBackToArgumentOrGlobal(iv, visited))
          return true;
      }

      return false;
    }

    return false;
  }
  
  /*!
    A visitor to fill in the escape set for a given value
   */
  struct EscapeAnalyzer : public InstVisitor<EscapeAnalyzer>
  {
    EscapeAnalyzer(Value* V, use_set& escapes, InterproceduralDefUse& IDU, results_t* res)
    :_escapes(escapes), _V(V), _res(res), _IDU(IDU)
    {}
    
    void visitCallInst(CallInst& CI)
    {
      // For call instructions, consult _res to see if the relevant arguments (which
      // are recorded in by the TransitiveDefUse pass) escape in external functions.
      const aset* args_used = _IDU.getIDUArgSet(_V, &CI);
      // if(!args_used)
      if(!CI.getCalledFunction() && !PermissiveFuncptr)
      {
        // Called through a function pointer in non-permissive mode
        DOUT << format("[info - call] Escape through function pointer\n");
        _escapes.insert(&CI);
        return;
      }
      else if(!CI.getCalledFunction() || !args_used)
      {
        return;
      }

      for(aset::const_iterator it = args_used->begin(), ed = args_used->end(); it != ed; ++it)
      {
        const Argument* A = *it;
        if(!A)
        {
          // Vararg, ignore for any sane purposes.
          continue;
        }
        
        const Function* F = A->getParent();
        if(_res->argumentEscapesStrong(F, A))
        {
          DOUT << format("[info - call] escape via _results on argument %s\n") % A->getName();
          _escapes.insert(&CI);
        }
        else
        {
          DOUT << format("[info - call] No escape via argument %s\n") % A->getName();
        }
      }
    }
    
    void visitStoreInst(StoreInst& SI)
    {
      Value* dest = SI.getOperand(1);

      if(_V == dest)
        return;

      use_set visited;
      if(tracesBackToArgumentOrGlobal(dest, visited))
      {
        _escapes.insert(&SI);
      }
    }

    void visitReturnInst(ReturnInst& RI)
    {
      // These should be inserted unconditionally.  Analyses that want to treat these
      // specially can just parse through the use_set.
      _escapes.insert(&RI);
    }

    use_set& _escapes;
    Value* _V;
    results_t* _res;
    InterproceduralDefUse& _IDU;
  };
}

EscapeAnalysis::EscapeAnalysis(results_t * res)
 :CallGraphSCCPass((intptr_t)&ID),
 _res(res)
{

}

bool EscapeAnalysis::markEscape(const Argument & arg, short escapeType)
{
  const Function* F = arg.getParent();
  if(!escapeType)
    return false;
  cerr << format("  EA: Marking %s @ %s as escaping\n") % arg.getName() % F->getName();
  switch(escapeType)
  {
  case Weak:
    return _res->markWeakEscape(F, &arg);
    
  case Strong:
    return _res->markStrongEscape(F, &arg);

  default:
    // Top should never be requested, and any other values
    // are definitely out of range.
    cerr << "Invalid Escape: " << escapeType << std::endl;
    assert(false && "Invalid escape type");
  }

  return false;
}

void EscapeAnalysis::analyzeEscapes(Value* V)
{
  InterproceduralDefUse& IDU = getAnalysis<InterproceduralDefUse>();
  const use_set& uses = IDU.getIDUChain(V);
  
  // Examine all store destinations, function call instructions, and return statements
  EscapeAnalyzer EA(V, _escapes[V], IDU, _res);
  for(use_set::const_iterator it = uses.begin(), ed = uses.end(); it != ed; ++it)
  {
    if(Instruction* I = dyn_cast<Instruction>(*it))
    {
      EA.visit(I);
    }
  }
}

void EscapeAnalysis::visitCallInst(CallInst& CI)
{
  if(!isa<PointerType>(CI.getType()))
    return;

  analyzeEscapes(&CI);
}

void EscapeAnalysis::visitArgument(Argument& arg)
{
  if(!isa<PointerType>(arg.getType()) || arg.hasByValAttr())
    return;

  analyzeEscapes(&arg);
  if(_escapes[&arg].empty())
    return;

  // Remove self from the escape set - many instructions
  // allow the traversal to traverse back to the source
  // argument.
  _escapes[&arg].erase(&arg);

  // Also, if there is only one escape (a return node), that means
  // that the function allows its argument to pass through.  It is
  // undesirable to mark this as an escape for an argument if the
  // return is EXACTLY a return of the argument.
  ReturnInst* RI = NULL;
  if(_escapes[&arg].size() == 1 && (RI = dyn_cast<ReturnInst>(*_escapes[&arg].begin())))
  {
    if(RI->getOperand(0) == &arg)
    {
      _escapes[&arg].clear();
      return;
    }
  }
  
  markEscape(arg, Strong);
  cerr << format("  because of: \n");
  for(use_set::iterator it = _escapes[&arg].begin(), ed = _escapes[&arg].end(); it != ed; ++it)
  {
    Function* F = getParentFunction(*it);
    cerr << format("   %s @ %s\n") % F->getName() % valueToString(*it);
  }
}

void EscapeAnalysis::runOnFunction(Function* F)
{
  // Still not sure what an empty callgraph node is
  if(!F)
    return;

  cerr << format("[EscapeAnalysis] Analyzing %s\n") % F->getName();
  
  // Check each argument
  for(Function::arg_iterator it = firstArgument(*F), ed = F->arg_end(); it != ed; ++it)
  {
    visitArgument(*it);
  }

  // And every call instruction.
  visit(F);
}

bool EscapeAnalysis::runOnSCC(const std::vector<CallGraphNode*>& SCCM)
{
  do
  {
    _res->resetDirtyFlag();
    for(std::vector<CallGraphNode*>::const_iterator it = SCCM.begin(),
          ed = SCCM.end(); it != ed; ++it)
    {
      runOnFunction((*it)->getFunction());
    }
  }while(_res->isDirty());

  return false;
}

void EscapeAnalysis::getAnalysisUsage(llvm::AnalysisUsage &AU) const
{
  AU.setPreservesAll();

  AU.addRequired<TransitiveDefUse>();
  AU.addRequired<InterproceduralDefUse>();
}

const use_set& EscapeAnalysis::getEscapeInstructions(const Value* V) const
{
  assert(V);
  assert(isa<PointerType>(V->getType()));

  use_map::const_iterator it = _escapes.find(V);
  if(it == _escapes.end())
    throw 0;
  
  return it->second;
}

char EscapeAnalysis::ID = 0;
RegisterPass<EscapeAnalysis> EA("escape", "Analyzes pointers in functions to determine if they escape");


