/*!
  The goal of this analysis is to determine which parameters are used
  as arrays.

  The strategy is to record all values used as the base of a
  GetElementPtr instruction (the entry point for all indexing
  operations, array and structure).  This alone is not enough, though
  - not all pointer arithmetic leads to a dereference (though such
  code is odd, it is legal).  To catch these cases, all _uses_ of any
  value must be recorded.  In LLVM, a load or GEP indexing count as
  "uses", since LLVM is a Load/Store architecture.  In addition, the
  value resulting from any GEP instruction is stored and associated
  with the base pointer that produced it.

  When resolving whether or not a given parameter is used as an array,
  all that must be done is look through all of the arrayAccesses
  (obtained from the base of each GEP).  If the arrayAccess MustAlias
  a parameter, it is used in an array context IFF it was previously
  Used in the sense that it appears in the usedValues set.
*/
#define DEBUG_TYPE "ArrayAnalysis"

#include "ArrayAnalysis.hpp"
#include "AliasHelper.hpp"
#include "ArgumentHelpers.hpp"
#include "Results.hpp"
#include "StructArrayMemberAnalysis.hpp"
#include "Utility.hpp"
#include "ValueCounter.hpp"

#include "llvm/ADT/Statistic.h"
#include "llvm/Analysis/AliasAnalysis.h"
#include "llvm/Analysis/CallGraph.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/ref.hpp>

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

// Restrict the maximum array nesting depth to track for this
// analysis.  This restriction is arbitrary, but 8 should be
// more than enough for any reasonable use.
#define MAX_ARRAY_DEPTH 8

using namespace llvm;
using namespace boost;

// STATISTIC(ArrayParameters, "Total number of array parameters found.");
// STATISTIC(ArrayPasses, "Number of iterations required for the Array analysis.");
// STATISTIC(MaxDepth, "Deepest array nesting level.");

namespace
{
  /*!
    Debugging, print out a set of values (indented), one per line
  */
  void printValues(const value_set & vs)
  {
    for(value_set::const_iterator it = vs.begin(),
          ed = vs.end(); it != ed; ++it)
    {
      cerr << "  " << **it << std::endl;
    }
  }
}

struct ArrayAnalysisVisitor : public InstVisitor<ArrayAnalysisVisitor>
{
  ArrayAnalysisVisitor(AliasHelper * AH, results_t * res):_AH(AH), _res(res){}

  void visitBitCastInst(BitCastInst & BCI)
  {
    _casts[BCI.getOperand(0)] = &BCI;
  }
  
  void visitLoadInst(LoadInst &LI)
  {
    // NOTE: This may need a multimap or a map of values to sets.
    _loads[LI.getOperand(0)] = &LI;
  }
  
  /*
    Call instructions directly introduce array information (the
    operand used as an array argument has _at least_ the derived
    array levels).

    This information could, however, be augmented by other local
    information.

    Call instructions only need to be visited once.
   */
  void visitCallInst(CallInst &CI)
  {
    const Function * calledFunction = CI.getCalledFunction();
    if(!calledFunction) // Function pointer
      return;
  
    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);

      level_set levels = _res->arrayLevels(calledFunction, arg);
      if(levels.empty())
      {
        DOUT << format("IA: No levels available through call %s@%s\n")
          % calledFunction->getName() % arg->getName();
        continue;
      }

      value_set incomingValues = flattenPhiNode(op);
      for(value_set::iterator it = incomingValues.begin(),
            ed = incomingValues.end(); it != ed; ++it)
      {
        ArrayLevels[*it].resize(MAX_ARRAY_DEPTH);

        // The named value has at _least_ these array levels.
        for(level_set::iterator l = levels.begin(); l != levels.end(); ++l)
        {
          if(!ArrayLevels[*it].test(*l))
          {
            MakingProgress = true;
            ArrayLevels[*it].set(*l);
          }
        }
      }
    }
  }

  /*!
    Add new arrays to the growing array chain set.  This instantiates
    empty sets for the new elements (making them arrays of length 1).
   */
  void addNewValues(const value_set & discoveredValues)
  {
    for(value_set::iterator it = discoveredValues.begin(),
          ed = discoveredValues.end(); it != ed; ++it)
    {
      if(_arrayChains.find(*it) != _arrayChains.end())
        continue;

      // This is an array (at depth >= 0) that we never saw before.
      _arrayChains[*it];
      MakingProgress = true;
    }
  }

  /*!
    
   */
  value_set findKnownSubArrays(const value_set & discoveredValues)
  {
    value_set results;
    for(value_set::iterator it = discoveredValues.begin(),
          ed = discoveredValues.end(); it != ed; ++it)
    {
      results.insert(_arrayChains[*it].begin(),
                     _arrayChains[*it].end());
    }
    
    return results;
  }

  /*!
    GEP instructions build up an evolving view of instructions as
    arrays.

    These chains of array instructions (related by aliases relationships,
    but discriminated by type equality) need to be built up over many
    passes, then consolidated at the end into discrete array levels.
   */
  void visitGetElementPtrInst(GetElementPtrInst &I)
  {
    // We don't want to count structs here.
    if(I.getNumIndices() != 1)
      return;
    
    // The array
    const Value * discoveredArray = I.getPointerOperand();

    // The resulting array element
    const Value * arrayElement = &I;

    value_set discoveredValues = flattenPhiNode(discoveredArray);
    addNewValues(discoveredValues);

    value_set knownSubArrays = findKnownSubArrays(discoveredValues);
    // cerr << "Known Subarrays \n";
    // cerr << "----------------\n";
    // printValues(knownSubArrays);
    // cerr << "----------------\n";

    value_set aliasedArrays = findChainExtensions(arrayElement, knownSubArrays);
    // cerr << "Chain Extensions\n";
    // cerr << "----------------\n";
    // printValues(aliasedArrays);
    // cerr << "----------------\n";
    
    if(!aliasedArrays.empty())
    {
      MakingProgress = true;
      
      // If we know that this result is already an array, we can
      // grab that other known array (cleaning up its metadata)
      // and just use that in the chain for I.getPointerOperand()
      // without creating intermediate entries.

      for(value_set::iterator it = aliasedArrays.begin(), ed = aliasedArrays.end();
          it != ed; ++it)
      {
        addSubArray(*it, discoveredValues);
      }
      
      return;
    }
  }

  /*!
    Adds a new \var{subarray} to each array in \var{vs}; this essentially
    increases the depths of each array in \var{vs} by one.
   */
  void addSubArray(const Value * subarray, const value_set & vs)
  {
    for(value_set::const_iterator it = vs.begin(),
          ed = vs.end(); it != ed; ++it)
    {
      MakingProgress = true;
      _arrayChains[*it].insert(subarray);
    }
  }

  typedef std::map<const Value*, value_set> ChainMap;

  /*!
    \var{ArrayElement} is the result of a GEP instruction (as the name implies, that makes
    it an array element); this function checks to see if this element is itself an array
    (from among the set of known arrays).

    If \var{ArrayElement} _is_ an array, return the sub-array elements (these will later
    be used to extend the array associtaed with \var{ArrayElement}.

    The \var{except} parameter lists sub-arrays we already know about, so that we don't
    make phantom progress just hitting the same entry again and again.
   */
  value_set findChainExtensions(const Value * ArrayElement, const value_set & except)
  {
    value_set results;
    value_set ArrayElements = flattenPhiNode(ArrayElement);
    
    for(ChainMap::iterator it = _arrayChains.begin(),
          ed = _arrayChains.end(); it != ed; ++it)
    {
      // _arrayChains should have NO PHI nodes in it since we flatten
      // every node before inserting its incoming values.
      assert(!isa<PHINode>(it->first));

      // If we already have this subarray, don't include it again (or
      // we'll never stop iterating).
      if(except.find(it->first) != except.end())
        continue;

      for(value_set::const_iterator v_it = ArrayElements.begin(),
            v_ed = ArrayElements.end(); v_it != v_ed; ++v_it)
      {
        // An alias check here would be awesome, but the alias analysis
        // is a little too loose for that.
        if(_loads[*v_it] == it->first)
        {
          MakingProgress = true;
          results.insert(it->first);
        }
      }
    }

    return results;
  }

  /*!
    Recursively descend the chain of array accesses from v,
    returning the length of the chain
   */
  unsigned char calculateChainLength(const Value * v)
  {
    ChainMap::iterator it = _arrayChains.find(v);
    if(it == _arrayChains.end())
      return 0;

    unsigned char maxDepth = 0;
    for(value_set::iterator v_it = it->second.begin(),
          v_ed = it->second.end(); v_it != v_ed; ++v_it)
    {
      maxDepth = std::max(maxDepth, calculateChainLength(*v_it));
    }
    
    return 1 + maxDepth;
  }
  
  /*!
    Iterate over the chains built up by visiting GEP instructions
    and translate the results into discrete array levels (filling in
    ArrayLevels as it goes).
   */
  void calculateArrayLevels()
  {
    for(ChainMap::iterator it = _arrayChains.begin(),
          ed = _arrayChains.end(); it != ed; ++it)
    {
      unsigned char depth = calculateChainLength(it->first) - 1;
      fillInDepth(it->first, depth);
    }
  }

  /*!
    Set the bits in the bitset for \var{v} for levels zero
    through depth (inclusive).

    This doesn't really take into account arrays of pointers
    to arrays or anything fancy like that.  It could, though.
   */
  void fillInDepth(const Value * v, unsigned char depth)
  {
    assert(!isa<const PHINode>(v));
    ArrayLevels[v].resize(MAX_ARRAY_DEPTH);
    for(unsigned char i = 0; i <= depth; ++i)
    {
      ArrayLevels[v].set(i);
    }
  }

  /*!
    The external interface; returns the array depths for the
    given argument.  Should have called calculateArrayLevels first.

    If \var{arg} is not an array, returns an empty BitVector
   */
  const BitVector getArgumentDepths(const Argument & arg)
  {
    // If we directly found the argument, awesome.
    if(ArrayLevels.find(&arg) != ArrayLevels.end())
      return ArrayLevels[&arg];

    // Check to see if the argument aliases anything we found.
    // We only want aliases of the exact same type (modulo bitcasts).
    // This extra check is necessary due to some imprecision in the
    // alias analysis.
    for(std::map<const Value*, BitVector>::iterator it = ArrayLevels.begin(),
          ed = ArrayLevels.end(); it != ed; ++it)
    {
      bool sameType = arg.getType() == it->first->getType();
      // Check to see if we have a bitcast
      if(!sameType)
      {
        value_map::iterator c = _casts.find(&arg);
        if(c != _casts.end())
        {
          sameType = c->second->getType() == it->first->getType();
        }
      }
      
      //if(!sameType || _AH->alias(&arg, it->first, NoAlias))
      if(!sameType || &arg != normalizeValue(it->first))
        continue;

      return it->second;
    }

    // Check to see if the argument is an OUT array parameter (pointer
    // to an array).  This could be arbitrarily nested, but just check
    // one level for now.  FIXME


    // Otherwise, we just don't have it. (It is not an array)
    return BitVector();
  }
  
  AliasHelper * _AH;
  results_t * _res;
  ChainMap _arrayChains;
  value_map _loads;
  value_map _casts;
  
  std::map<const Value*, BitVector> ArrayLevels;
  bool MakingProgress;
};

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

void ArrayAnalysis::markArray(const BitVector & depths,
                              const Argument & arg)
{
  if(!isa<PointerType>(arg.getType()))
    return;

  int idx = depths.find_first();
  while(idx != -1)
  {
    _res->markArray(arg.getParent(), &arg, idx);

    idx = depths.find_next(idx);
  }
}

void ArrayAnalysis::checkArray(Function &F)
{
  ArrayAnalysisVisitor AAV(_AH.get(), _res);
  do
  {
    AAV.MakingProgress = false;
    AAV.visit(F);
  } while(AAV.MakingProgress);

  AAV.calculateArrayLevels();
  
  for(Function::const_arg_iterator it = firstArgument(F),
        ed = F.arg_end(); it != ed; ++it)
  {
    if(!isa<PointerType>(it->getType()) || it->hasByValAttr())
      continue;

    markArray(AAV.getArgumentDepths(*it), *it);
  }

  StructArrayMember& SAM = getAnalysis<StructArrayMember>();
  
  for(inst_iterator it = inst_begin(F), ed = inst_end(F); it != ed; ++it)
  {
    StoreInst* SI = NULL;
    if(!(SI = dyn_cast<StoreInst>(&*it)))
      continue;

    // Strip bitcasts from the src of the Store
    const Value* V = normalizeValue(SI->getOperand(0));
    const Argument* A = NULL;
    if(!(A = dyn_cast<const Argument>(V)))
      continue;

    const GetElementPtrInst* GEP = NULL;
    if(!(GEP = dyn_cast<const GetElementPtrInst>(SI->getOperand(1))))
      continue;

    int arrayDepth = SAM.getStructArrayDepth(GEP);
    if(arrayDepth == -1)
      continue;

    for(int i = 0; i < arrayDepth; ++i)
    {
      _res->markArray(&F, A, i);
    }
  }
}


bool ArrayAnalysis::runOnSCC(const std::vector<llvm::CallGraphNode *> &SCCM)
{
  if(SCCM.size() == 1)
  {
    Function * fn = SCCM.front()->getFunction();

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

    checkArray(*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;

      checkArray(*fn);
    }
  } while(_res->isDirty());

  return false;
}

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

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


char ArrayAnalysis::ID = 0;
RegisterPass<ArrayAnalysis> ArrayAnalysisRegistration("array", "Analyzes pointer arguments to determine if they are used as arrays");
