#define DEBUG_TYPE "AllocatorAnalysis"

#include "AllocatorAnalysis.hpp"
#include "AliasHelper.hpp"
#include "EscapeAnalysis.hpp"
#include "Results.hpp"
#include "Utility.hpp"

#include "llvm/Constants.h"
#include "llvm/ADT/Statistic.h"
#include "llvm/Analysis/AliasAnalysis.h"
#include "llvm/Analysis/CallGraph.h"
#include "llvm/Support/CommandLine.h"
#include "llvm/Support/Debug.h"

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

#include <algorithm>
#include <iterator>
#include <iostream>
#include <set>

using namespace llvm;
using namespace boost;

STATISTIC(AllocatorCount, "Total number of allocation functions.");
// STATISTIC(AllocatorPasses, "Number of iterations required for the Allocator analysis.");

extern cl::opt<bool> NoNullConstructors;

namespace
{
  /*!

   */
  bool isAMustAlloc(AliasHelper * DSA,
                    results_t * res,
                    const value_set & vs,
                    const Value * v,
                    std::set<const Function*> & allocators)
  {
    // We want to know early if this is a phi node; the results
    // of this part of the analysis would be incorrect if
    // performed on one.
    assert(!isa<PHINode>(v));
  
    // The NULL pointer is an allowable return value from an allocator,
    // as it can safely be freed.
    if(isa<ConstantPointerNull>(v))
    {
      DOUT << "CheckIfAllocator: Null Pointer (OK)\n";

      // NOTE:  If we want to test much of an effect allowing NULL
      // returns gives, we can just comment out this statement.
      return true && !NoNullConstructors;
    }
  
    for(value_set::const_iterator it = vs.begin(), ed = vs.end();
        it != ed; ++it)
    {
      const CallInst* ci = cast<CallInst>(*it);
      const Function* calledFunction = ci->getCalledFunction();

      // If the calledFunction is NULL, it is a function pointer.
      // We can't make any strong statements about function pointers,
      // so we skip this particular call (assuming it is not an
      // allocator).  We allow the loop to continue running, though,
      // as another aliased call could possibly be an allocator (although
      // probably not, due to the must_alloc restriction).
      if(!calledFunction)
      {
        continue;
      }

      DOUT << format("CheckIfAllocator: Checking %s\n")
        % calledFunction->getName();

      const Value * v1 = v;
      const Value * v2 = *it;

      if(DSA->alias(v1, v2, MayAlias))
      {
        DOUT << format("CheckIfAllocator: May alias for %s<->%s\n")
          % *v1 % *v2;
      }

      if(!res->isAllocator(calledFunction))
      {
        DOUT << format("Callee[%s] is not an allocator\n")
          % calledFunction->getName();
      }

      if(DSA->alias(v1, v2, MustAlias) &&
         res->isAllocator(calledFunction))
      {
        DOUT << "CheckIfAllocator: Must alias an allocator (OK)\n";
        if(!calledFunction->isDeclaration())
          allocators.insert(calledFunction);
        return true;
      }
    }
  
    // If we make it through all of the called functions
    // without already returning, this value is not provably
    // the result of an allocator.
    DOUT << "CheckIfAllocator: Does not alias an allocator (Fail)\n";
    return false;
  }

  bool isNullPointer(const Value * v)
  {
    return isa<ConstantPointerNull>(v);
  }
}

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

void AllocatorAnalysis::visitCallInst(CallInst &CI)
{
  DOUT << format("Found a function call %s\n") % CI;
  functionResults.insert(&CI);
}

void AllocatorAnalysis::visitReturnInst(ReturnInst &RI)
{
  returnValue = RI.getReturnValue();
  if(returnValue)
  {
    DOUT << format("Found a return value: %s\n") % *returnValue;
  }
}

static bool onlyContainsReturns(const use_set& escapes)
{
  for(use_set::const_iterator it = escapes.begin(), ed = escapes.end(); it != ed; ++it)
  {
    DOUT << format("  [info - onlyContainsReturns]: %s\n") % valueToString(*it);

    if(isa<const ReturnInst>(*it))
      continue;

    return false;
  }

  return true;
}

void AllocatorAnalysis::checkAllocator(Function &F)
{
  // We assume the function is a deallocator by default; we do not
  // need to take any action when it proves this to us.
  if(!returnValue) 
    return;

  cerr << format("AllocatorAnalysis: %s\n") % F.getName();

  // Check this separately; if returnValue is a PHINode
  // and one of the incoming values is a constant zero,
  // we still want to count that as a heap allocation.
  std::set<const Function*> allocation_functions;
  if(!isHeapAllocation(returnValue, allocation_functions))
    return;

  value_set nodes = flattenPhiNode(returnValue);
  
  for(value_set::const_iterator it = nodes.begin(), ed = nodes.end();
      it != ed; ++it)
  {
    if(isa<ConstantPointerNull>(*it))
      continue;
    
    // NULL is allowed to escape (and only NULL), since it
    // will likely be used in many contexts.
    cerr << format(" Found %d allocators\n") % allocation_functions.size();
    try
    {
      const use_set& escapes = _E->getEscapeInstructions(*it);
      if(!onlyContainsReturns(escapes) && !isa<ConstantPointerNull>(*it))
      {
        cerr << "  return escape strong\n";
        return;
      }
    }
    catch(...)
    {
      cerr << "  return escape strong due to exception (probably pointer arithmetic)\n";
      return;
    }
  }

  ++AllocatorCount;
  cerr << format("  Found allocator!\n");
  _res->markAllocator(&F);
}

/*
  Technically, a heap allocated value need not originate as a return
  value of a function; one could be introduced through an Out parameter.
  This is probably unlikely, so ignore it for now.
*/
bool AllocatorAnalysis::isHeapAllocation(const Value * v, std::set<const Function*> & allocators) const
{
  value_set must_alias = flattenPhiNode(v);

  size_t nulls = std::count_if(must_alias.begin(),
                               must_alias.end(),
                               isNullPointer);

  bool allMustAlloc = true;
  for(value_set::iterator it = must_alias.begin(),
        ed = must_alias.end(); it != ed; ++it)
  {
    allMustAlloc = allMustAlloc &&
      isAMustAlloc(_AH.get(), _res, functionResults, *it, allocators);
  }

  // This is an allocator if NO possible aliases were NOT also
  // allocators (all aliases are allocators).
  //
  // Also, ensure that the function does not return JUST NULL.
  // That means that we require it to at least be able to
  // return a malloced value sometimes.
  return allMustAlloc && (nulls < must_alias.size());
}

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

void AllocatorAnalysis::resetStateVariables()
{
  returnValue = NULL;
  functionResults.clear();
}

bool AllocatorAnalysis::runOnSCC(const std::vector<CallGraphNode *> &SCCM)
{
  AliasAnalysis & aa = getAnalysis<AliasAnalysis>();
  _AH.reset(new AliasHelper(aa));

  _E = &getAnalysis<EscapeAnalysis>();

  if(SCCM.size() == 1)
  {
    Function * fn = SCCM.front()->getFunction();
    if(!fn || fn->isDeclaration())
      return false;

    resetStateVariables();

    visit(*fn);
    checkAllocator(*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();

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

  return false;
}

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

  AU.addRequired<AliasAnalysis>();
  AU.addRequired<EscapeAnalysis>();
}

char AllocatorAnalysis::ID = 0;
RegisterPass<AllocatorAnalysis> AllocatorAnalysisRegistration("allocators", "Analyzes function return values to determine which are allocated via malloc, and escape");
