#define DEBUG_TYPE "OutputParameterAnalysis"

#include "AliasHelper.hpp"
#include "ArgumentHelpers.hpp"
#include "ArrayAnalysis.hpp"
#include "DataFlow.hpp"
#include "OutputParameterAnalysis.hpp"
#include "Results.hpp"
#include "Utility.hpp"
#include "ValueCounter.hpp"

#include "llvm/Constants.h"
#include "llvm/DerivedTypes.h"
#include "llvm/Type.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/bind.hpp>
#include <boost/format.hpp>

#include <algorithm>

using namespace llvm;
using namespace boost;

// STATISTIC(InOutPasses, "Number of passes required for the InOut analysis.");
STATISTIC(TotalOutParams, "Total number of OUT parameters identified.");
STATISTIC(TotalInParams, "Total number of IN parameters identified.");
STATISTIC(TotalInOutParams, "Total number of parameters that are INOUT identified.  Included in both IN and OUT counts.");

extern cl::opt<bool> PermissiveFuncptr;


typedef std::map<const llvm::BasicBlock *, var_states> block_map;


std::ostream & operator << (std::ostream & os, const ArgType & at)
{
  switch(at._type)
  {
  case Opaque:
    os << "OPAQUE";
    break;
  case In:
    os << "IN";
    break;
  case Out:
    os << "OUT";
    break;
  case InOut:
    os << "INOUT";
    break;
  default:
    assert(0 && "Unhandled ArgType");
  }

  return os;
}

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

size_t OutputParameterAnalysis::valueIndex(const Value * v)
{
  if(_mapping.find(v) != _mapping.end())
    return _mapping[v];

  _mapping[v] = _indexSource++;
  return _mapping[v];
}

void OutputParameterAnalysis::visitCallInst(CallInst & CI)
{
  calls.insert(&CI);
}

void OutputParameterAnalysis::visitBitCastInst(BitCastInst & BCI)
{
  assert(casts.find(&BCI) == casts.end() ||
         casts[&BCI] == BCI.getOperand(0));
  
  // We need to worry about bitcasts to catch casts
  // from pointer types to void* (among other things).
  casts.insert(std::make_pair(&BCI, BCI.getOperand(0)));
}

namespace
{
  /*!
    The meet operation that defines movement along the InOut
    lattice.
   */
  ArgType meet(const ArgType & at1, const ArgType & at2)
  {
    if((at1.isIn() && (at2.isOut() || at2.isInOut())) ||
       (at1.isOut() && (at2.isIn() || at2.isInOut())))
      return ArgType(InOut);

    if(at1 < at2)
      return at2;

    return at1;
  }
}

void OutputParameterAnalysis::printState(const var_states & s)
{
  cerr << format("InOutStates\n");
  for(std::tr1::unordered_map<const llvm::Value*, unsigned>::iterator it = _mapping.begin(),
        ed = _mapping.end(); it != ed; ++it)
  {
    if(_interestingIndices.find(it->second) == _interestingIndices.end())
      continue;

    cerr << format("  %s -> %d\n") % it->first->getName() % s[it->second];
  }

  cerr << "\n\n";
}

void OutputParameterAnalysis::checkInOut(Function & F)
{
  cerr << format("OutputParameterAnalysis: %s\n") % F.getName();
  
  ValueCounter& vc = getAnalysis<ValueCounter>();
  _totalValues = vc.getValueCount(F);
  
  var_states s = inOutAnalyzer(F);

  DEBUG(printState(s));
  
  std::for_each(firstArgument(F), F.arg_end(),
                bind(&OutputParameterAnalysis::markInOut, this, cref(F), s, _1));
}


/*!
  Walk the chain of values formed by related bitcasts
  to give a canonical representation of @var{target}.
  
  This keeps all information about a given logical value
  assigned to a single actual value.
*/
const Value * OutputParameterAnalysis::resolveCasts(const Value * target)
{
  value_map::const_iterator it = casts.find(target);
  if(it == casts.end())
    return target;
  
  return resolveCasts(it->second);
}
/*!
  A small stub to run the forward data flow analysis and
  return only the information we are interested in (the
  values at the end of the function).
  
  Resolver is a function used to provide a canonical reference
  to each logical value encountered.  This is important because
  we need to track references and stores for each value, but
  casts (and possibly some other operations) create "new" LLVM
  values that really refer to the same logical value.
*/
var_states OutputParameterAnalysis::inOutAnalyzer(Function & F)
{
  block_map states = forward_flow(F);
  
  const BasicBlock * exitBlock = &F.back();    
  return states[exitBlock];
}

/*!
  Examine one basic block, recording InOut data for each
  variable at this block.  Combines locally gathered information
  with the input states given to it (@var{inputStates}).

  This only needs to examine a few types of instruction:
  - GetElementPtr (a kind of load)
  - Load
  - Store
  - Call (for the interprocedural analysis).

  It may be worth it to convert this to an InstructionVisitor
  later.
*/
var_states OutputParameterAnalysis::transfer(var_states & inputStates,
                                   const BasicBlock * bb)
{
  var_states localState(_totalValues);
  if(inputStates.size() != _totalValues)
    inputStates.resize(_totalValues);
  localState = inputStates;

  for(BasicBlock::const_iterator it = (*bb).begin(), ed = (*bb).end();
      it != ed; ++it)
  {
    if(const GetElementPtrInst * gep =
       dyn_cast<const GetElementPtrInst>(it))
    {
      const Value * target = gep->getPointerOperand();
      DOUT << format("  IN(gep): %s\n") % target->getName();

      registerValue(localState, resolveCasts(target), In);
    }
    else if(const LoadInst * li =
            dyn_cast<const LoadInst>(it))
    {
      const Value * target = li->getOperand(0);
      DOUT << format("  IN: %s\n") % target->getName();

      registerValue(localState, resolveCasts(target), In);
    }
    else if(const StoreInst * si =
            dyn_cast<const StoreInst>(it))
    {
      const Value * target = si->getOperand(1);
      // Need to ignore things of the form:
      //  store i8* %str, i8** %str_addr
      // These are artificial and from LLVM
      if(si->getOperand(1)->getName() ==
         si->getOperand(0)->getName() + "_addr")
        continue;
      
      DOUT << format("  OUT: %s\n") % target->getName();

      registerValue(localState, resolveCasts(target), Out);
    }
    else if(const CallInst * ci =
            dyn_cast<const CallInst>(it))
    {
      const Function * calledFunction = ci->getCalledFunction();

      // Function pointer call - mark all pointers to primitive types
      // as inout
      if(!calledFunction && !PermissiveFuncptr)
      {
        for(const Use * opit = operandStartIter(ci), *oped = ci->op_end();
            opit != oped; ++opit)
        {
          const Value * v = opit->get();
          if(!isa<PointerType>(v->getType()))
            continue;

          registerValue(localState, v, InOut);
        }
        
        continue;
      }

      if(!calledFunction)
        continue;
      
      Function::const_arg_iterator arg = firstArgument(*calledFunction);
        
      for(unsigned i = operandStart(ci);
          i != ci->getNumOperands() &&
            arg != calledFunction->arg_end();
          ++i, ++arg)
      {
        const Value * op = ci->getOperand(i);
        op = resolveCasts(op);
          
        // See if op has been bitcast from a more useful type.
        if(_res->isIn(calledFunction, arg))
          registerValue(localState, op, In);

        if(_res->isOut(calledFunction, arg))
          registerValue(localState, op, Out);
      }
    }
  }

  return localState;
}


/*!
  Merge new entries into the return map, update existing
  entries with a meet operator.
*/
var_states* OutputParameterAnalysis::join(var_states * accum, var_states * vs)
{
  if(accum->size() != _totalValues)
    accum->resize(_totalValues);

  if(vs->size() != _totalValues)
  {
    // If the vs (the state of one predecessor) was not the
    // correct size, it could not have contained any useful
    // information (since it had not yet been resized).  We
    // don't need to perform the merge on it.
    vs->resize(_totalValues);
    return accum;
  }

  var_states & state = *accum;
  
  for(std::set<size_t>::const_iterator it = _interestingIndices.begin(),
        ed = _interestingIndices.end(); it != ed; ++it)
  {
    state[*it] = meet(state[*it], (*vs)[*it]);
  }

  return accum;
}


/*!
  Update the given value state set with new usage information @var{t}
  for value @var{target}.  This will flatten target if it is a PHI node
  and register all incoming values.
*/
void OutputParameterAnalysis::registerValue(var_states & vs,
                                  const Value * target,
                                  ArgType_tag t)
{
  assert(vs.size() == _totalValues);
  
  value_set incomingValues = flattenPhiNode(target);
  for(value_set::iterator it = incomingValues.begin(),
        ed = incomingValues.end(); it != ed; ++it)
  {
    size_t idx = valueIndex(*it);
    _interestingIndices.insert(idx);
    vs[idx] = meet(vs[idx], ArgType(t));
  }
}


/*!
  Use the final set of state information gathered by the
  analysis to mark the (pointer) arguments of @var{F} with
  their correct classifications.

  Both local and interprocedural information is available
  here, so no work should be required after this.
*/
void OutputParameterAnalysis::markInOut(const Function & F,
                              const var_states & s,
                              const Argument & arg)
{
  assert(s.size() == _totalValues);
  
  const Value * v = &arg;
  if(!isa<const PointerType>(v->getType()) || arg.hasByValAttr())
    return;

  const ArgType & at = s[valueIndex(v)];
  
  if(at.isIn())
  {
    ++TotalInParams;
    _res->markIn(&F, &arg);
    return;
  }

  // A parameter that is an array at level 0 (e.g. int*a; is
  // used as a[5]) cannot be an OUT parameter.  They look
  // somewhat similar in use, but the "arrayness" is the
  // overriding characteristic.
  level_set ls = _res->arrayLevels(&F, &arg);
  if(at.isOut() && (ls.find(0) == ls.end()))
  {
    ++TotalOutParams;
    _res->markOut(&F, &arg);
    return;
  }
    
  if(at.isInOut())
  {
    ++TotalInOutParams;
    _res->markIn(&F, &arg);
    // Again, array0 parameters are not OUT.
    if(ls.find(0) == ls.end())
      _res->markOut(&F, &arg);
  }
}

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

void OutputParameterAnalysis::resetStateVariables()
{
  calls.clear();
  casts.clear();
  _mapping.clear();
  _interestingIndices.clear();
  _indexSource = 0;
}

bool OutputParameterAnalysis::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();

    visit(*fn);
    checkInOut(*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);
      checkInOut(*fn);
    }
  } while(_res->isDirty());
  
  return false;
}

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

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


char OutputParameterAnalysis::ID = 0;
RegisterPass<OutputParameterAnalysis> IORegistration("OutputParameter", "Analyzes pointer arguments to determine if they are in, out, or both");


