#include "AST.h"
#include "ArrayAccessSet.h"
#include "code-util.h"
#include "optimize.h"
#include "MIVEset.h"

/* Given an array arr, return the source for its VarDecl. */
TreeNode *arrayVarDecl(const TreeNode *arr)
{
  Decl *d = arr->decl();
  if (!d->hasSource()) return NULL;
  return d->source();
}

MIVEset * ArrayAccessSet::allMIVEs(const TreeNode *arr) const
{
  const MIVEset *s = MIVEs(arr);
  MIVEset *result = new MIVEset;
  for (MIVEset::const_iterator j = s->begin(); j != s->end(); j++)
    result = ::adjoin(result, *j);
  arr = arrayVarDecl(arr);
  for (osr_iterator i = begin(); !i.isDone(); i.next())
    if (i.array() == arr)
      result = ::adjoin(result, (MIVE *) i.dest());
  return result;
}

MIVEset * & ArrayAccessSet::MIVEs(const TreeNode *arr)
{
  return m[arrayVarDecl(arr)];
}

const MIVEset * ArrayAccessSet::MIVEs(const TreeNode *arr) const
{
  return m[arrayVarDecl(arr)];
}

string ArrayAccessSet::OSRTempname(TreeNode *WRTloop, TreeNode *arr,
				   const MIVE *m)
{
  TreeNode *d = arrayVarDecl(arr);
  int i = find_or_append_MIVE(m, q[d]);
  return string("OSR_") + int2alpha(uniqueInt((void *) WRTloop)) + "_" +
    int2alpha(uniqueInt((void *) arr->decl()->source())) + "_" +
    int2alpha(i / 2);
}

const MIVE *ArrayAccessSet::baseMIVEForOSR(TreeNode *arr, const MIVE *m) const
{
  TreeNode *d = arrayVarDecl(arr);
  if (d == NULL)
    return NULL;
  llist<const MIVE *> *l = q[d];
  if (l == NULL)
    q.erase(d);
  while (l != NULL)
    if (l->tail()->front()->equals(m))
      return l->front();
    else
      l = l->tail()->tail();
  return NULL;
}

/* x is the base, and we're going to do y off it. */
void ArrayAccessSet::adjoinOSR(TreeNode *t, TreeNode *WRTloop,
			       const MIVE *x, const MIVE *y)
{
  if (debug_sr) {
    cout << "AAS::adjoinOSR ";
    cout << PLCODE(t);
    cout << endl;
  }
  llist<const MIVE *> * & l = q[arrayVarDecl(t->array())];
  l = cons(x, cons(y, l));
}

void ArrayAccessSet::adjoinConstant(TreeNode *t, TreeNode *WRTloop)
{
  TreeNode *d = arrayVarDecl(t->array());
  TreeNode *indexdecl = t->index()->decl()->source();
  if (d != NULL && indexdecl != NULL) {
    treeSet *& s = c[d];
    if (s == NULL)
      s = new treeSet;
    ::adjoin(s, indexdecl);
    use(d, indexdecl) = t;
    if (debug_sr) {
      cout << "AAS::adjoinConstant ";
      cout << PLCODE(t);
      cout << endl;
    }
  }
}

void ArrayAccessSet::adjoin(TreeNode *t, TreeNode *WRTloop)
{
  MIVE *m = t->index()->getMIVE(WRTloop);
  if (m == NULL)
    adjoinConstant(t, WRTloop);
  else {
    if (debug_sr) {
      cout << "AAS::adjoin ";
      cout << PLCODE(t);
      cout << endl;
    }
    MIVEset * & s = MIVEs(t->array());
    if (s == NULL)
      s = new MIVEset;
    ::adjoin(s, m);
    use(arrayVarDecl(t->array())) = t;
  }
}

bool ArrayAccessSet::containsOSR(TreeNode *t, TreeNode *WRTloop) const
{
  return (baseMIVEForOSR(t->array(), t->index()->getMIVE(WRTloop)) != NULL);
}

bool ArrayAccessSet::containsConstant(TreeNode *t, TreeNode *WRTloop) const
{
  TreeNode *d = arrayVarDecl(t->array());
  TreeNode *indexdecl = t->index()->decl()->source();
  bool result =
    (d != NULL && indexdecl != NULL &&
     (c[d] != NULL || (c.erase(d), false)) &&
     ::contains(c[d], indexdecl));
  if (debug_sr) {
    cout << "AAS::containsConstant ";
    cout << PLCODE(t);
    cout << ": " << result << endl;
  }
  return result;
}

bool ArrayAccessSet::contains(TreeNode *t, TreeNode *WRTloop) const
{
  TreeNode *d = arrayVarDecl(t->array());
  if (d == NULL)
    return false;
  const MIVE *im = t->index()->getMIVE(WRTloop);
  bool result =
    (im == NULL && containsConstant(t, WRTloop) ||
     im != NULL &&
     (m[d] != NULL || (m.erase(d), false)) &&
     ::contains(m[d], t->index()->getMIVE(WRTloop)));
  if (im != NULL && debug_sr) {
    cout << "AAS::contains ";
    cout << PLCODE(t);
    cout << ": " << result << endl;
  }
  return result;
}

/* Return a new set that includes the var decl for each array used in
   strength reduced index expression. */
treeSet *ArrayAccessSet::SRarrays() const
{
  treeSet *s = new treeSet;
  for (map_tree_to_MIVEset::const_iterator i = m.begin();
       i != m.end();
       i++) {
    s = ::adjoin(s, (*i).first);
  }
  return s;
}

/* A certain array is accesses with index expressions represented by q
   and r.  m[i] is the largest stride that the array might have in
   dimension i (0 means we don't know).  For each dimension of the
   array, update m[i] to be the smaller if possible, using the difference
   between q and r as evidence. */
static void seekUnitStride(map_int_to_int &m, const MIVE *q, const MIVE *r)
{
  if (q == NULL || r == NULL)
    return;
  int n = q->arity();
  assert(r->arity() == n);
  const MIVE *diff = q->add(r->negate());
  if (diff == NULL)
    return;
  for (int i = 0; i < n; i++) {
    const Poly *p = diff->p[i];
    if (p != NULL && p->constantp()) {
      int &g = m[i], d = p->absoluteValue();
      if (d != 0)
	g = (g == 0) ? d : gcd(g, d);
    }
  }
}

/* Search through all array accesses and try to use the differences
   between them to prove that certain arrays must have unit stride in
   certain dimensions.  E.g., if we have A[e + [1, 0]] and A[e] then we
   know A must have unit stride in dimension 0.  This would be recorded
   by adding A to provablyUnitStride[0]. 

   Works by taking pairwise differences of all LIPs and pairwise
   differences of all MIVEs for each dimension of each array.
   The gcd of all such differences is the maximum possible stride.
   We ignore differences that are not known integers at compile time.
*/
void ArrayAccessSet::seekUnitStride(map_int_to_treeSet &provablyUnitStride,
				    TreeNode *l)
{
  if (DEBUG_BC)
    cout << endl << "Seeking unit stride arrays in loop at " <<
      l->position().asString() << endl << "LIP:" << endl;
    
  // y[A][dim] is the minimum that the stride for A might be in dimension dim 
  map_tree_int_to_int y;
  treeSet *s = new treeSet;

  {
    map_tree_to_cMIVElist x;

    for (lip_iterator i = lip_begin(); !i.isDone(); i.next()) {
      TreeNode *acc = use(i.array(), i.index());
      TreeNode *a = arrayVarDecl(acc->array()), *index = acc->index();
      const MIVE *z;
      if (a == NULL || index == NULL || (z = index->getMIVE(l)) == NULL)
	continue;

      if (DEBUG_BC) {
	cout << "array " << a << ": ";
	a->pseudoprint(cout, 2);
	cout << endl << "MIVE of index: ";
	SHOWMIVE(z);
	cout << endl;
      }

      x[a] = cons(z, x[a]);
      s = ::adjoin(s, a);
    }
    for (treeSet::iterator i = s->begin(); i != s->end(); ++i) {
      TreeNode *a = (TreeNode *) *i;
      llist<const MIVE *> *l = x[a];
      int n;
      if (l == NULL || (n = l->size()) < 2)
	continue;
      for (int j = 0; j < n; j++)
	for (int k = j + 1; k < n; k++)
	  ::seekUnitStride(y[a], (*l)[j], (*l)[k]);
    }
  }

  if (DEBUG_BC)
    cout << "OSR:" << endl;

  {
    map_tree_to_MIVEset x;
    for (osr_iterator i = begin(); !i.isDone(); i.next()) {
      TreeNode *a = i.array();
      const MIVE *b = i.base(), *d = i.dest();
      if (a == NULL || b == NULL || d == NULL)
	continue;

      if (DEBUG_BC) {
	cout << "array " << a << ": ";
	a->pseudoprint(cout, 0);
	cout << endl;
	if (d) { cout << "d: "; SHOWMIVE(d); cout << endl; }
	else { cout << "d: NULL" << endl; }
	if (b) { cout << "b: "; SHOWMIVE(b); cout << endl; }
	else { cout << "b: NULL" << endl; }
      }

      if (x[a] == NULL)
	x[a] = new MIVEset;
      x[a] = ::adjoin(::adjoin(x[a], (MIVE *) b), (MIVE *) d); 
      s = ::adjoin(s, a);
    }
    for (treeSet::iterator i = s->begin(); i != s->end(); ++i) {
      TreeNode *a = (TreeNode *) *i;
      llist<const MIVE *> *l = MIVEset_to_list(x[a]);
      int n;
      if (l == NULL || (n = l->size()) < 2)
	continue;
      for (int j = 0; j < n; j++)
	for (int k = j + 1; k < n; k++)
	  ::seekUnitStride(y[a], (*l)[j], (*l)[k]);
      int arity = l->front()->arity();
      for (int i = 0; i < arity; i++)
	if (y[a][i] == 1) {
	  treeSet * & p = provablyUnitStride[i];
	  if (p == NULL)
	    p = new treeSet;
	  p = ::adjoin(p, a);
	  if (DEBUG_BC) {
	    cout << "provablyUnitStride: ";
	    a->pseudoprint(cout, 3);
	    cout << " in dim " << i << endl;
	  }
	}
    }
  }
}