#ifndef __DATAFLOW__HPP
#define __DATAFLOW__HPP

#include "llvm/Support/CFG.h"
#include <boost/bind.hpp>
#include <boost/function.hpp>
#include <vector>
#include <deque>
#include <map>
#include <algorithm>
#include <numeric>

template<typename T>
const T * addrof(const T & item)
{
  return &item;
}

template<typename T>
const T* mapToTPtr(std::map<const llvm::BasicBlock*, T> & m,
             const llvm::BasicBlock * bb)
{
  const T* ret = &(m[bb]);
  assert(ret);
  return &m[bb];
}

/*!
  Essentially the same as the normal DataFlow module,
  except the transfer function is called with intermediate
  results available.
 */
template<typename T>
class DataFlowIntermediate
{
public:
  virtual ~DataFlowIntermediate(){}
  std::map<const llvm::BasicBlock*, T>
    forward_flow(const llvm::Function & F)
  {
    std::map<const llvm::BasicBlock*, T> ret;
  
    std::deque<const llvm::BasicBlock *> worklist;
    std::transform(F.begin(), F.end(),
                   std::back_inserter(worklist),
                   addrof<llvm::BasicBlock>);

    T predecessorStates;
    while(!worklist.empty())
    {
      predecessorStates.clear();
      
      const llvm::BasicBlock * current = worklist.front();
      worklist.pop_front();
    
      // First, join the exit states of all predecessors
      llvm::pred_const_iterator pi = pred_begin(current);
      llvm::pred_const_iterator ed = pred_end(current);
      
      // Initialize the exit state, in case we are performing
      // an intersection.
      if(pi != ed)
      {
        typedef std::vector<T*> MergeSet;
        MergeSet merge_sets; //(16);
        {
          llvm::pred_const_iterator it = pi;
          for(;it != ed; ++it)
          {
            merge_sets.push_back(&ret[*it]);
          }
        }

        // NOTE: May need to enforce all ret[x] be initialized
        predecessorStates = ret[*pi];
        T * accum = &predecessorStates;
        for(typename MergeSet::iterator mit = merge_sets.begin(),
              med = merge_sets.end(); mit != med; ++mit)
        {
          accum = join(accum, *mit);
        }
      }
      
      // Union (or intersect) with current node result f(item)
      // Part of the transfer function
      T blockResults = transfer(ret, predecessorStates, current);

      // If this iteration changed the exit state for this
      // node, we need to recalculate the exit states for
      // each successor.
      if(blockResults != ret[current])
      {
        worklist.insert(worklist.end(),
                        succ_begin(current),
                        succ_end(current));
        ret[current] = blockResults;
      }
    }

    return ret;
  }
  
  virtual T* join(T * accum, T * item) = 0;
  virtual T transfer(const std::map<const llvm::BasicBlock*, T> & current,
                     T & inputStates, const llvm::BasicBlock* bb) = 0;
};

template<typename T>
class DataFlow
{
public:
  virtual ~DataFlow(){}
  std::map<const llvm::BasicBlock*, T>
  forward_flow(const llvm::Function & F)
  {
    std::map<const llvm::BasicBlock*, T> ret;
  
    std::deque<const llvm::BasicBlock *> worklist;
    std::transform(F.begin(), F.end(),
                   std::back_inserter(worklist),
                   addrof<llvm::BasicBlock>);

    while(!worklist.empty())
    {
      T predecessorStates;
      
      const llvm::BasicBlock * current = worklist.front();
      worklist.pop_front();
    
      // First, join the exit states of all predecessors
      llvm::pred_const_iterator pi = pred_begin(current);
      llvm::pred_const_iterator ed = pred_end(current);
      
      // Initialize the exit state, in case we are performing
      // an intersection.
      if(pi != ed)
      {
        typedef std::vector<T*> MergeSet;
        MergeSet merge_sets; //(16);
        {
          llvm::pred_const_iterator it = pi;
          for(;it != ed; ++it)
          {
            merge_sets.push_back(&ret[*it]);
          }
        }

        // NOTE: May need to enforce all ret[x] be initialized
        predecessorStates = ret[*pi];
        T * accum = &predecessorStates;
        for(typename MergeSet::iterator mit = merge_sets.begin(),
              med = merge_sets.end(); mit != med; ++mit)
        {
          assert(*mit);
          assert(accum);
          accum = join(accum, *mit);
        }
      }
      
      // Union (or intersect) with current node result f(item)
      // Part of the transfer function
      T blockResults = transfer(predecessorStates, current);

      // If this iteration changed the exit state for this
      // node, we need to recalculate the exit states for
      // each successor.
      if(blockResults != ret[current])
      {
        worklist.insert(worklist.end(),
                        succ_begin(current),
                        succ_end(current));
        ret[current] = blockResults;
      }
    }

    return ret;
  }

  virtual T* join(T * accum, T * item) = 0;
  virtual T transfer(T & inputStates, const llvm::BasicBlock* bb) = 0;
};

#endif // __DATAFLOW__HPP
