#include "StructArrayMemberAnalysis.hpp"

#include "llvm/Constants.h"
#include "llvm/Function.h"
#include "llvm/Instructions.h"
#include "llvm/Module.h"
#include "llvm/Support/InstIterator.h"

#include <boost/format.hpp>

using namespace boost;
using namespace llvm;

StructArrayMember::StructArrayMember()
 :ModulePass((intptr_t)&ID)
{
  
}

static GetElementPtrInst* getArrayAccessParent(Value* V)
{
  if(!V)
    return NULL;
  
  // Strip off any bitcasts
  while(BitCastInst* BCI = dyn_cast<BitCastInst>(V))
    V = BCI->getOperand(0);

  // Now we need to see a load of a GEP (which is the parent we
  // want to return)
  LoadInst* LI = NULL;
  if(!(LI = dyn_cast<LoadInst>(V)))
    return NULL;

  return dyn_cast<GetElementPtrInst>(LI->getOperand(0));
}

void StructArrayMember::analyzeGEP(GetElementPtrInst* GEP)
{
  // First, see if this is an array access or a struct access
  // There are actually more complex forms of the GEP instruction (with up to 4 indices)
  // which could _also_ be arrays (if the first index != 0).  We ignore those for now...
  // I don't think that llvm-gcc actually generates these.
  if(GEP->getNumIndices() != 1)
    return;

  int arrayLevels = 0;
  int fieldNo = -1;
  // Next, see if the pointer operand (essentially the parent) is
  // a struct (another GEP instruction). (Do walk through bitcasts,
  // just in case the array was stored in a void*).
  Value* V = GEP->getPointerOperand();
  while(GetElementPtrInst* parent = getArrayAccessParent(V))
  {
    ++arrayLevels;
    V = parent->getPointerOperand();
    if(parent->getNumIndices() > 1)
    {
      Value* idxOp = parent->getOperand(2);
      ConstantInt* idxVal = NULL;
      if(!(idxVal = dyn_cast<ConstantInt>(idxOp)))
      {
        // Not sure what this is, but we can't use it in this analysis.
        return;
        // assert(false && "Struct index is not a constant int");
      }

      // This must be signed, so take the zext rather than sign extended value
      fieldNo = idxVal->getZExtValue();
      break;
    }
  }

  if(arrayLevels == 0)
  {
    cerr << format("  [info] No array levels\n");
    return;
  }

  // Not a recognized pattern, just ignore it.
  if(fieldNo < 0)
    return;
  
  const Type* T = V->getType();

  array_depth_key descriptor(std::make_pair(T, fieldNo));
  if(_arrayDepths.find(descriptor) != _arrayDepths.end() &&
     arrayLevels <= _arrayDepths[descriptor])
  {
    // This array finding is less-specific than a previous one
    return;
  }
  
  std::string structName = (format("%s / %s") % V->getName() % V->getNameStr()).str();
  cerr << format("  [info] Found an array in struct %s at field %d (%d levels)\n")
    % structName % fieldNo % arrayLevels;
  _arrayDepths[descriptor] = arrayLevels;
}

void StructArrayMember::analyzeFunction(Function* F)
{
  for(inst_iterator it = inst_begin(F), ed = inst_end(F); it != ed; ++it)
  {
    GetElementPtrInst* GEP = NULL;
    if(!(GEP = dyn_cast<GetElementPtrInst>(&*it)))
      continue;

    analyzeGEP(GEP);
  }
}

bool StructArrayMember::runOnModule(Module& M)
{
  for(Module::iterator F = M.begin(), E = M.end(); F != E; ++F)
  {
    cerr << format("[StructArrayMember] Analyzing %s\n") % F->getName();
    analyzeFunction(F);
  }
  
  return false;
}

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

short StructArrayMember::getStructArrayDepth(const GetElementPtrInst* GEP) const
{
  // This GEP is itself an array access and this analysis has no information
  if(GEP->getNumIndices() == 1)
    return -1;

  const Value* parent = GEP->getPointerOperand();
  const Type* T = parent->getType();

  Value* idxOp = GEP->getOperand(2);
  ConstantInt* idxVal = NULL;
  if(!(idxVal = dyn_cast<ConstantInt>(idxOp)))
  {
    // This could be a complex access of a global array; don't assert here,
    // just indicate that it is not the kind of array dealt with by this
    // analysis pass.
    return -1;
  }

  // This must be signed, so take the zext rather than sign extended value
  short fieldNo = idxVal->getZExtValue();

  array_depth_key descriptor(std::make_pair(T, fieldNo));

  array_depth_map::const_iterator it = _arrayDepths.find(descriptor);
  if(it == _arrayDepths.end())
    return -1;

  return it->second;
}

char StructArrayMember::ID = 0;
RegisterPass<StructArrayMember> SAMReg("structArrayMember", "Determines which struct fields are used as arrays");

