#define DEBUG_TYPE "FinalizerAnalysis"

#include "AliasHelper.hpp"
#include "ArgumentHelpers.hpp"
#include "DataFlow.hpp"
#include "FinalizerAnalysis.hpp"
#include "Results.hpp"
#include "Utility.hpp"
#include "ValueCounter.hpp"

#include "llvm/Function.h"
#include "llvm/Constants.h"
#include "llvm/ADT/Statistic.h"
#include "llvm/Analysis/AliasAnalysis.h"
#include "llvm/Analysis/CallGraph.h"
#include "llvm/Support/CFG.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/InstIterator.h"

#include <boost/format.hpp>
#include <boost/bind.hpp>
#include <list>
#include <set>
#include <algorithm>
#include <numeric>
#include <functional>

using namespace llvm;
using namespace boost;

// STATISTIC(DeallocatorPasses, "Number of iterations required for the Deallocator analysis.");
STATISTIC(DeallocatorsRuledOut, "Total number of potential deallocators eliminated.");

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

void FinalizerAnalysis::visitCallInst(CallInst &CI)
{
  DOUT << format("Call Instruction: %s\n") % CI;
  calls.insert(&CI);
}

namespace
{  
  std::set<int> argumentAliasedBy(AliasHelper * AH,
                                  const CallInst * ci,
                                  const Value * v)
  {
    std::set<int> ret;
    
    for(unsigned i = operandStart(ci); i < ci->getNumOperands(); ++i)
    {
      value_set incomingValues = flattenPhiNode(ci->getOperand(i));
      for(value_set::iterator it = incomingValues.begin(),
            ed = incomingValues.end(); it != ed; ++it)
      {
        // Maybe change this to no alias?  This is probably only
        // used by the deallocator analysis.
        if(!AH->alias(v, *it, MustAlias) ||
           isa<ConstantExpr>(*it))
          break;
        
        // There would be other problems if anyone used 2 billion
        // arguments, but best to be safe?
        assert(i < static_cast<unsigned long>(std::numeric_limits<int>::max()));
        
        ret.insert(static_cast<int>(i));
      }
    }
    
    return ret;
  }
  
  bool functionDeallocatesAnArgument(const Function * fn, results_t & res)
  {
    for(Function::const_arg_iterator it = fn->arg_begin(), ed = fn->arg_end();
        it != ed; ++it)
    {
      if(res.deallocatesArgument(fn, &*it))
        return true;
    }

    return false;
  }
  
  bool operator != (const deallocator_info & LHS,
                    const deallocator_info & RHS)
  {
    return LHS.deallocatedOrNull != RHS.deallocatedOrNull ||
      LHS.nullValues != RHS.nullValues ||
      LHS.deallocatedValues != RHS.deallocatedValues;
  }
  
  std::string argSetString(const deallocator_info & args)
  {
    std::string ret = "FN[";

    for(value_set::const_iterator it = args.deallocatedOrNull.begin(),
          ed = args.deallocatedOrNull.end(); it != ed; ++it)
    {
      ret += (*it)->getName() + ", ";
    }

    ret += "]; N[";

    for(value_set::const_iterator it = args.nullValues.begin(),
          ed = args.nullValues.end(); it != ed; ++it)
    {
      ret += (*it)->getName() + ", ";
    }
    
    ret += "]";
    return ret;
  }


  /*!
    If cmp compares a pointer against zero (either directly
    or one step "back" to check the negation of that),
    AND
    the successor is the Basicblock branched to when that
    value is NOT NULL,
    THEN
    return that value.
    OTHERWISE
    return NULL.
   */
  const Value * valueComparedAgainstZero(const ICmpInst * cmp,
                                         const BranchInst * br,
                                         const BasicBlock * successor)
  {
    const Value * v1 = cmp->getOperand(0);
    const Value * v2 = cmp->getOperand(1);

    if(!isa<const ConstantPointerNull>(v2))
      return NULL;

    if(!isa<const PointerType>(v1->getType()))
      return NULL;

    unsigned nonZeroBranch = 0;

    if(cmp->getUnsignedPredicate() == ICmpInst::ICMP_NE)
    {
      /*
        if(x != 0) {useX(x)}

        The THEN branch is operand 1 to the branch
        instruction (operand 0 is the predicate)
       */
      nonZeroBranch = 1;
      DOUT << "  NEQ\n";
    }
    else if(cmp->getUnsignedPredicate() == ICmpInst::ICMP_EQ)
    {
      /*
        if(x == 0) { something } else {useX(x)}

        The ELSE branch is operand 2
       */
      nonZeroBranch = 0;
      DOUT << "  EQ\n";
    }
    else
    {
      return NULL;
    }

    assert(br->getNumSuccessors() == 2);
    // Need a case for when the successor branches to
    // a return, since then we have pred == EQ but we
    // need nonZeroBranch = 1
    DOUT << format("  (%d) Successor: %s  && op: %s\n")
      % nonZeroBranch % successor->getName() %
      br->getSuccessor(nonZeroBranch)->getName();

    if(br->getSuccessor(nonZeroBranch) == successor)
      return v1;
    
    return NULL;
  }

  const Value * predecessorNull(const BasicBlock * bb,
                                const BasicBlock * successor)
  {
    const TerminatorInst * lastInstruction = bb->getTerminator();
    
    // Is the last instruction a branch?
    if(!isa<const BranchInst>(lastInstruction))
      return NULL;

    const BranchInst * br = cast<const BranchInst>(lastInstruction);
    if(br->isUnconditional())
      return NULL;
    
    // Is the target (op0) an icmp?
    const Value * condition = br->getCondition();
    if(!isa<const ICmpInst>(condition))
      return NULL;

    const ICmpInst * cmp = cast<const ICmpInst>(condition);    
    
    const Value * zeroAtSuccessor =
      valueComparedAgainstZero(cmp, br, successor);

    // Either null or the correct value.
    return zeroAtSuccessor;
  }
  

  bool argumentAliasesSomeValue(AliasHelper * AH, const Argument & arg,
                                const value_set & areDeallocated)
  {
    unsigned aliases = 0;
    for(value_set::iterator it = areDeallocated.begin(),
          ed = areDeallocated.end(); it != ed; ++it)
    {
      if(AH->alias(&arg, flattenPhiNode(*it), MustAlias))
        ++aliases;
    }

    return aliases != 0;
  }
  
  void markIfNotDeallocator(AliasHelper * AH, const Function * F,
                            results_t * res,
                            const value_set & areDeallocated,
                            const Argument & arg)
  {
    if(argumentAliasesSomeValue(AH, arg, areDeallocated))
      return;

    ++DeallocatorsRuledOut;
    res->markNotDeallocator(F, &arg);
  }
}

void FinalizerAnalysis::checkDeallocator(Function &F)
{
  cerr << format("DeallocatorAnalysis: %s\n") % F.getName();

  value_set deallocated = deallocatedArguments(F);
  std::for_each(F.arg_begin(), F.arg_end(),
                bind(&markIfNotDeallocator, _AH.get(), &F,
                     _res, deallocated, _1));
}

value_set FinalizerAnalysis::deallocatedArguments(Function &F)
{
  block_map states = forward_flow(F);

  const BasicBlock * exitBlock = &F.back();

  for(block_map::iterator it = states.begin(),
        ed = states.end(); it != ed; ++it)
  {
    DOUT << format("  %s => %s\n")
      % it->first->getName()
      % argSetString(it->second);
  }

  return states[exitBlock].deallocatedOrNull;
}

deallocator_info FinalizerAnalysis::transfer(const block_map & m,
                                        deallocator_info & inputStates,
                                        const BasicBlock * bb)
{
  // Only one new value can become NULL at the start of
  // a basic block (since the immediate predecessor could
  // only check one pointer against zero at a time.
  // Multiple values can become non-NULL at one join
  // point, though.
  const Value * newNullValue = newlyNull(m, bb);
  
  deallocator_info ret;
  ret.merge(inputStates);
    
  if(newNullValue)
  {
    ret.addNull(newNullValue);
  }
    
  for(BasicBlock::const_iterator it = (*bb).begin(), ed = (*bb).end();
      it != ed; ++it)
  {
    if(!isa<const CallInst>(it))
      continue;
    
    const CallInst * ci = cast<const CallInst>(it);
    const Function * calledFunction = ci->getCalledFunction();

    // If there is no function, this is a call through a function
    // pointer and we won't analyze it.
    // TODO: collect stats about how many times this trips;
    //       this could give some information on how much
    //       unannotated function pointers are impacting results.
    if(!calledFunction)
      continue;
    
    // For each operand in ci that is deallocated,
    // insert it into ret.
    for(CallInst::const_op_iterator oit = operandStartIter(ci),
          oed = ci->op_end(); oit != oed; ++oit)
    {
      checkArgument(ret, ci, *oit);
    }

    DOUT << format("    (In Loop) Basic Block: %s -> %s\n") %
      bb->getName() %
      argSetString(ret);
  }
    
  DOUT << format("  Basic Block: %s -> %s\n") %
    bb->getName() %
    argSetString(ret);
        
  // The exit state of this node is the union of whatever was
  // calculated here AND the exit states of the predecessors
  // of this node.
  return ret;
}

deallocator_info* FinalizerAnalysis::join(deallocator_info * accum,
                                     deallocator_info * s2)
{
  value_set tmp;

  std::set_intersection(accum->nullValues.begin(),
                        accum->nullValues.end(),
                        s2->nullValues.begin(),
                        s2->nullValues.end(),
                        std::inserter(tmp, tmp.begin()));

  accum->nullValues.swap(tmp);
  tmp.clear();
  
  std::set_intersection(accum->deallocatedValues.begin(),
                        accum->deallocatedValues.end(),
                        s2->deallocatedValues.begin(),
                        s2->deallocatedValues.end(),
                        std::inserter(tmp, tmp.begin()));

  accum->deallocatedValues.swap(tmp);
  tmp.clear();
  
  std::set_intersection(accum->deallocatedOrNull.begin(),
                        accum->deallocatedOrNull.end(),
                        s2->deallocatedOrNull.begin(),
                        s2->deallocatedOrNull.end(),
                        std::inserter(tmp, tmp.begin()));

  accum->deallocatedOrNull.swap(tmp);
  
  return accum;
}


void FinalizerAnalysis::checkArgument(deallocator_info & accum,
                                 const CallInst * ci,
                                 const Value * arg)
{
  // Also, don't bother checking if this is being used
  // in a vararg position.
  if(argAliasesVararg(arg, ci))
  {
    DOUT << format("\t%s aliases a vararg\n") % arg->getName();
    return;
  }
    
  const Function * calledFunction = ci->getCalledFunction();
  // Called function will be valid here, since the caller
  // would have bailed out if ci was using a function pointer.
  assert(calledFunction &&
         "Checking deallocator status through function pointer.");

  DOUT << format("Checking if argument %s is used in %s\n")
    % arg->getName()
    % calledFunction->getName();

  value_set incomingValues = flattenPhiNode(arg);
  for(value_set::iterator vit = incomingValues.begin(),
        ved = incomingValues.end(); vit != ved; ++vit)
  {
    std::set<int> argumentUses = argumentAliasedBy(_AH.get(), ci, *vit);

    DOUT << format("Used %d times\n") % argumentUses.size();
    
    if(argumentUses.empty())
      continue;
    
    for(std::set<int>::iterator it = argumentUses.begin(),
          ed = argumentUses.end(); it != ed; ++it)
    {
      Function::const_arg_iterator calleeArg =
        calledFunction->arg_begin();
          
      unsigned long umax =
        static_cast<unsigned long>(std::numeric_limits<int>::max());
      assert(calledFunction->arg_size() < umax);

      DOUT << format("Argument used at index %d/%d\n")
        % *it % calledFunction->arg_size();

      if(*it <= static_cast<int>(calledFunction->arg_size()))
      {
        // The - 1 offset is to account for the first operand of
        // all CallInstructions - the function being called.
        std::advance(calleeArg, *it - 1);
        if(_res->deallocatesArgument(calledFunction, calleeArg))
        {
          DOUT << format("%s frees its %dst argument\n")
            % calledFunction->getName() % (*it - 1);

          accum.addDeallocated(arg);
        }
      }
    }
  }
}

bool FinalizerAnalysis::argAliasesVararg(const Value * arg,
                                    const CallInst * ci)
{
  const Function * fn = ci->getCalledFunction();

  if(!fn || !fn->isVarArg())
    return false;

  unsigned numStaticArgs = realNumArgs(fn);

  // If there are the same number of operands (except for the
  // ones we must skip - the function and sret parameters),
  // this cannot possibly be a vararg situation.
  if(ci->getNumOperands()-operandStart(ci) == numStaticArgs)
    return false;

  /*
    No required arguments, 1 actual:  0 < 1 + 1 - 1 GOOD (true - vararg)
    One required, 1 actual: 1 < 1 + 1 - 1 GOOD (false - not vararg)
    One required, 2 actual: 1 < 2 + 1 - 1 GOOD (false no vararg)
    SRET, no required, 1 actual: 0 < 2 + 1 - 2 GOOD (true - vararg)
  */
  for(unsigned i = operandStart(ci); i < ci->getNumOperands(); ++i)
  {
    if(isa<ConstantExpr>(ci->getOperand(i)))
      continue;

    DOUT << format("Static Args(%s): %d\nDynamic Arg: %d\n") %
      fn->getName() % numStaticArgs % i;

    if((numStaticArgs < (i + 1 - operandStart(ci))) &&
       !_AH->alias(flattenPhiNode(arg),
                   flattenPhiNode(ci->getOperand(i)), NoAlias))
    {
      DOUT << format("Vararg alias(%s)\n") % *arg;
      return true;
    }
  }

  return false;
}

bool FinalizerAnalysis::isAliasOfVFreedInBB(const block_map & m,
                                       const Value * v,
                                       const BasicBlock * bb)
{
  block_map::const_iterator di = m.find(bb);
  assert(di != m.end());
    
  for(value_set::iterator it = di->second.deallocatedValues.begin(),
        ed = di->second.deallocatedValues.end(); it != ed; ++it)
  {
    if(_AH->alias(flattenPhiNode(*it),
                  flattenPhiNode(v), MustAlias))
      return true;
  }

  return false;
}
  
/*!
  Examine the predecessors of @var{bb} and, if they all conditionally
  branch to @var{bb} based on the NULLness of a particular value,
  return that value.  Otherwise, return NULL.
*/
const Value * FinalizerAnalysis::newlyNull(const block_map & m,
                                      const BasicBlock * bb)
{
  const Value * nn = NULL;

  // Find at least one value that is NULL at the start
  // of this block (also at the end of the predecessor
  // block).
  for(pred_const_iterator it = pred_begin(bb),
        ed = pred_end(bb); it != ed; ++it)
  {
    nn = predecessorNull(*it, bb);
    if(nn)
      break;
  }

  // If there was no such value, no new values are NULL
  if(!nn)
    return NULL;
    
  for(pred_const_iterator it = pred_begin(bb),
        ed = pred_end(bb); it != ed; ++it)
  {
    const Value * next = predecessorNull(*it, bb);
    
    if(next)
    {
      DOUT << format("  NULL: %s is null at the end of %s\n")
        % next->getName()
        % (*it)->getName();
    }

    if(isAliasOfVFreedInBB(m, nn, *it))
    {
      DOUT << format("  NULL: %s is deallocated in %s\n")
        % nn->getName()
        % (*it)->getName();
    }
      
    if(nn != next && !isAliasOfVFreedInBB(m, nn, *it))
      // AND NOT (an alias of) nn is deallocated in *it
      return NULL;
  }

  return nn;
}

bool FinalizerAnalysis::onInitialize(CallGraph & CG)
{
  return false;
}

void FinalizerAnalysis::resetStateVariables()
{
  _Entry = NULL;
  calls.clear();
}

bool FinalizerAnalysis::runOnSCC(const std::vector<CallGraphNode *> &SCCM)
{
  AliasAnalysis & aa = getAnalysis<AliasAnalysis>();
  _AH.reset(new AliasHelper(aa));
  
  if(SCCM.size() == 1)
  {
    Function * fn = SCCM.front()->getFunction();
    if(!fn || fn->isDeclaration())
      return false;

    resetStateVariables();
    _Entry = &fn->getEntryBlock();
    
    visit(*fn);
    checkDeallocator(*fn);
    
    return false;
  }

  do
  {
    _res->resetDirtyFlag();
    for(size_t i = 0; i < SCCM.size(); ++i)
    {
      Function * fn = SCCM[i]->getFunction();
      if(!fn || fn->isDeclaration())
        continue;
      
      resetStateVariables();
      _Entry = &fn->getEntryBlock();

      visit(*fn);
      checkDeallocator(*fn);
    }
  } while(_res->isDirty());

  return false;
}

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

  AU.addRequired<ValueCounter>();
  AU.addRequiredTransitive<AliasAnalysis>();
}


char FinalizerAnalysis::ID = 0;
RegisterPass<FinalizerAnalysis> FinalizerAnalysisRegistration("finalizers", "Analyzes pointer arguments to determine which are deallocated");


