#include "ArgumentHelpers.hpp"

#include <llvm/Instructions.h>

#include <boost/regex.hpp>

using namespace boost;
using namespace llvm;

unsigned realArgno(const Argument * arg)
{
  const Function * fn = arg->getParent();

  size_t realArg = 0;
  boost::match_results<std::string::const_iterator> what;
  boost::regex expr("\\.[[:digit:]]+");
  for(Function::const_arg_iterator it = fn->arg_begin(), ed = fn->arg_end();
      it != ed; ++it)
  {
    if(&*it == arg)
      break;

    if(it->hasStructRetAttr())
    {
      continue;
    }
    // We need to count the _first_ element of by-value aggregates
    // that get expanded by llvm, but not subsequent ones.
    else if(it->getName().find(".0") != std::string::npos)
    {
      ++realArg;
      continue;
    }
    else if(boost::regex_search(it->getName(), what, expr))
    {
      continue;
    }
    
    ++realArg;
  }

  return realArg;
}

unsigned realNumArgs(const Function * fn)
{
  if(fn->hasStructRetAttr())
    return fn->arg_size() - 1;
  
  return fn->arg_size();
}

Function::const_arg_iterator firstArgument(const Function & F)
{
  Function::const_arg_iterator it = F.arg_begin();
  if(F.hasStructRetAttr())
    ++it;
  
  return it;
}

Function::arg_iterator firstArgument(Function & F)
{
  Function::arg_iterator it = F.arg_begin();
  if(F.hasStructRetAttr())
    ++it;
  
  return it;
}

unsigned operandStart(const CallInst * ci)
{
  const Function * fn = ci->getCalledFunction();

  if(!fn)
  {
    const Value * calledValue = ci->getCalledValue();
    const PointerType * fpt = cast<const PointerType>(calledValue->getType());
    const FunctionType * ft = cast<const FunctionType>(fpt->getContainedType(0));
    
    if(ft->getReturnType()->isAggregateType())
      return 2;

    return 1;
  }

  // If this is a void function that was transformed
  // then need to skip the function (op0) and the
  // OUT parameter (op1).
  if(fn->hasStructRetAttr())
    return 2;

  // Otherwise, we only ever need to skip one argument,
  // the function.
  return 1;
}

const Use * operandStartIter(const CallInst * ci)
{
  assert(ci != NULL);
  return operandStartIter(const_cast<CallInst*>(ci));
}

Use* operandStartIter(CallInst * ci)
{
  Function * fn = ci->getCalledFunction();
  Use * ret = ci->op_begin() + 1;

  if(!fn)
  {
    const Value * calledValue = ci->getCalledValue();
    const PointerType * fpt = cast<const PointerType>(calledValue->getType());
    const FunctionType * ft = cast<const FunctionType>(fpt->getContainedType(0));
    
    if(ft->getReturnType()->isAggregateType())
    {
      ++ret;
      return ret;
    }
    
    return ret;
  }

  // If this is a void function that was transformed
  // then need to skip the function (op0) and the
  // OUT parameter (op1).
  if(fn->hasStructRetAttr())
  {
    ++ret;
    return ret;
  }
  
  // Otherwise, we only ever need to skip one argument,
  // the function.
  return ret;
}
