#include "ArgumentHelpers.hpp"
#include "Results.hpp"
#include "Scheme.hpp"

#include <llvm/Constants.h>
#include <llvm/Function.h>
#include <llvm/Module.h>
#include <llvm/Support/Debug.h>

#include <boost/format.hpp>
#include <boost/regex.hpp>
#include <boost/tuple/tuple.hpp>

// This is the prefix that LLVM uses before all struct names.
#define STRUCT_NAME_PREFIX "struct."

using namespace boost;
using namespace llvm;

namespace
{
  enum
  {
    DEALLOCATES,
    IN,
    OUT,
    ESCAPES_STRONG,
    ESCAPES_WEAK,
    // This is a sentinal value to record
    // the number of required bits
    STATE_COUNT
  };

  const Argument * argFromIndex(const Function * F, short index)
  {
    short current = 0;
    boost::match_results<std::string::const_iterator> what;
    boost::regex expr("\\.[[:digit:]]+");
  
    for(Function::const_arg_iterator it = firstArgument(*F),
          ed = F->arg_end(); it != ed; ++it)
    {
      if(current == index)
        return &*it;
      
      // the StructRetAttr stuff is taken care of in firstArgument
      
      if(it->getName().find(".0") != std::string::npos)
      {
        ++current;
        continue;
      }
      
      if(boost::regex_search(it->getName(), what, expr))
        continue;
      
      ++current;
    }
    
    return NULL;
  }
}

argument_t::argument_t(const llvm::Function * fn,
                       const llvm::Argument * arg)
:_fn(fn), _arg(arg)
{
}


bool argument_t::isVarg() const
{
  return _arg == NULL;
}

bool argument_t::operator < (const argument_t &a) const
{
  return _arg < a._arg;
}

state_t::state_t()
 :_State(llvm::BitVector(STATE_COUNT))
{
  // We assume that all (known) functions are
  // deallocators by default.
  _State.set(DEALLOCATES);
}

void state_t::setBit(unsigned i)
{
  _State.set(i);
}

void state_t::unsetBit(unsigned i)
{
  _State.reset(i);
}

bool state_t::queryBit(unsigned i) const
{
  return _State.test(i);
}

results_t::results_t(const Module* M)
:_dirty(false), _M(M){}

void results_t::resetDirtyFlag()
{
  _dirty = false;
}

bool results_t::isDirty() const
{
  return _dirty;
}

void results_t::markAllocator(const Function * fn)
{
  // Don't reinsert if the function is already known to be
  // an allocator - toggling the dirty flag in this case will
  // lead to infinite loops.
  if(_allocators.find(fn) != _allocators.end())
    return;

  _dirty = true;
  _allocators.insert(fn);
}

void results_t::markNotDeallocator(const Function * fn,
                                   const Argument * arg)
{
  argument_t a(fn, arg);
  if(_pinnedDeallocators.find(a) != _pinnedDeallocators.end())
  {
    // We need to explictly instantiate a result object for this
    // argument so that the deallocator annotation is added
    // to the output.
    _results[a];
    return;
  }
  
  if(_results.find(a) != _results.end() &&
     !_results[a].queryBit(DEALLOCATES))
    return;
    
  _dirty = true;
  _results[a].unsetBit(DEALLOCATES);
}

void results_t::markArray(const Function * fn,
                          const Argument * arg,
                          short depth)
{
  const argument_t a(fn, arg);
  assert(isa<PointerType>(arg->getType()));
         
  std::pair<array_map::iterator, array_map::iterator> values =
    _arrays.equal_range(a);
  for(array_map::iterator it = values.first; it != values.second; ++it)
  {
    if(it->second == depth)
      return;
  }

  _dirty = true;
  _arrays.insert(std::make_pair(a, depth));
}

bool results_t::markStrongEscape(const Function * fn, const Argument * arg)
{
  argument_t a(fn, arg);
  
  if(_results.find(a) != _results.end() &&
     _results[a].queryBit(ESCAPES_STRONG))
    return false;

  _dirty = true;
  _results[a].setBit(ESCAPES_STRONG);
  return true;
}

bool results_t::markWeakEscape(const Function * fn, const Argument * arg)
{
  argument_t a(fn, arg);
  if(_results.find(a) != _results.end() &&
     _results[a].queryBit(ESCAPES_WEAK))
    return false;

  _dirty = true;
  _results[a].setBit(ESCAPES_WEAK);
  return true;
}

void results_t::markIn(const Function * fn, const Argument * arg)
{
  argument_t a(fn, arg);
  if(_results.find(a) != _results.end() &&
     _results[a].queryBit(IN))
    return;
  
  _dirty = true;
  _results[a].setBit(IN);
}

void results_t::markOut(const Function * fn, const Argument * arg)
{
  argument_t a(fn, arg);
  if(_results.find(a) != _results.end() &&
     _results[a].queryBit(OUT))
    return;
  
  _dirty = true;
  _results[a].setBit(OUT);
}

void results_t::registerWeakField(const std::string & sname, short index)
{
  std::string internalName = STRUCT_NAME_PREFIX + sname;
  _weak.insert(std::make_pair(internalName, index));
}


void results_t::pinDeallocator(const std::string & function_name, short index)
{
  const Function * finalizer = _M->getFunction(function_name);
  const Argument * arg = argFromIndex(finalizer, index);

  argument_t a(finalizer, arg);
  _pinnedDeallocators.insert(a);
}

void results_t::registerFinalizer(const std::string & allocator,
                                  const std::string & finalizer)
{
  const Function * alloc = _M->getFunction(allocator);
  const Function * dealloc = _M->getFunction(finalizer);
  _pairings[alloc] = dealloc;
  _allocators.insert(alloc);

  // Look up the first (and only) argument of finalizer
  // and register it as a deallocator.  Pin it so that
  // the analysis can't disprove it later.
  Function::const_arg_iterator it = firstArgument(*dealloc);
  if(it == dealloc->arg_end())
  {
    cerr << format("WARNING: Finalizers must have exactly one argument (%s)\n")
      % finalizer;
    return;
  }

  const Argument * arg1 = &*it;
  ++it;
  if(it != dealloc->arg_end())
  {
    cerr << format("WARNING: Finalizers must have exactly one argument (%s)\n")
      % finalizer;
    return;
  }

  argument_t a(dealloc, arg1);
  _pinnedDeallocators.insert(a);
}

bool results_t::isWeakField(const std::string & structName, short index)
{
  weak_map::const_iterator it,ed;
  tie(it, ed) = _weak.equal_range(structName);

  for(; it != ed; ++it)
  {
    if(it->second == index)
      return true;
  }

  return false;
}

bool results_t::isAllocator(const Function * fn)
{
  argument_t a(fn, NULL);
  ensureAllocator(a);
  return _allocators.find(fn) != _allocators.end();
}

level_set results_t::arrayLevels(const Function * fn, const Argument * arg)
{
  level_set ret;

  argument_t a(fn, arg);
  ensureDependencyMet(a);

  std::pair<array_map::iterator, array_map::iterator> values =
    _arrays.equal_range(a);
  for(array_map::iterator it = values.first; it != values.second; ++it)
  {
    ret.insert(it->second);
  }

  
  return ret;
}

bool results_t::deallocatesArgument(const Function * fn, const Argument *arg)
{
  argument_t a(fn, arg);
  ensureDependencyMet(a);
  return _results[a].queryBit(DEALLOCATES);
}

bool results_t::argumentEscapesStrong(const Function * fn, const Argument * arg)
{
  argument_t a(fn, arg);
  ensureDependencyMet(a);
  return _results[a].queryBit(ESCAPES_STRONG);
}

bool results_t::argumentEscapesWeak(const Function * fn, const Argument * arg)
{
  argument_t a(fn, arg);
  ensureDependencyMet(a);
  return _results[a].queryBit(ESCAPES_WEAK);
}

bool results_t::isIn(const Function *fn, const Argument *arg)
{
  argument_t a(fn, arg);
  ensureDependencyMet(a);
  return _results[a].queryBit(IN);
}

bool results_t::isOut(const Function *fn, const Argument *arg)
{
  argument_t a(fn, arg);
  ensureDependencyMet(a);
  return _results[a].queryBit(OUT);
}

bool results_t::entryInSystem(argument_t arg)
{
  if(!arg._fn->isDeclaration())
    return true;

  if(hasExternalFunction(arg))
    return true;

  return false;
}

void results_t::ensureDependencyMet(argument_t arg)
{
  if(arg._fn->isDeclaration() && !hasExternalFunction(arg))
    guileFetchDependency(const_cast<results_t*>(this), arg);
}

void results_t::ensureAllocator(argument_t arg)
{
  if(arg._fn->isDeclaration() && !hasExternalFunction(arg))
    guileFetchAllocator(const_cast<results_t*>(this), arg);
}

bool results_t::hasExternalFunction(argument_t arg)
{
  return _results.find(arg) != _results.end() ||
    _allocators.find(arg._fn) != _allocators.end();
}

void results_t::removeResultsFor(const Function * F)
{
  // Remove F from _allocators, if it is there
  // Remove all of F's arguments from _results and _arrays
  _allocators.erase(F);
  for(Function::const_arg_iterator it = F->arg_begin(),
        ed = F->arg_end(); it != ed; ++it)
  {
    argument_t arg(F, &*it);
    _results.erase(arg);
    _arrays.erase(arg);
  }

  _hidden_functions.insert(F);
}

void results_t::addAlias(const Function * F, const GlobalAlias* A)
{
  _aliases.insert(std::make_pair(F, A->getName()));
}

