/* This file is for dominators and related information.  Dominators are
   computed over the CfgNodes, but being able to ask if CfgNode x
   dominates CfgNode y turns out to be not what was most useful.  So
   we convert that dominator info to a form that allows us to ask what
   each TreeNode dominates.
*/

#include "AST.h"
#include "cfg.h"
#include "optimize.h"
#include "OrderlySet.h"

static bool debug_dom, debug_domv;

typedef OrderlySet<TreeNode *> TreeSetFactory;
typedef TreeSetFactory::set OrderlyTreeSet;
typedef map< TreeNode *, OrderlyTreeSet > map_tree_to_OrderlyTreeSet;
typedef map< TreeNode *, map_tree_to_OrderlyTreeSet > 
    map_tree_tree_to_OrderlyTreeSet;
typedef map< TreeNode *, TreeSetFactory * > map_loop_to_Factory;

typedef OrderlySet<CfgNode *> CfgSetFactory;
typedef CfgSetFactory::set CfgSet;
typedef map< CfgNode *, CfgSet > map_cfgnode_to_cfgset;
typedef map< CfgSet, OrderlyTreeSet > map_cfgset_to_OrderlyTreeSet;

/* dominators[loop][t] = set of nodes that are dominators of t */
static map_tree_tree_to_OrderlyTreeSet dominators;
static map_tree_to_OrderlyTreeSet loopExits, everyIter;
static map_loop_to_Factory o;

static map_tree_to_OrderlyTreeSet *
computeDominators(TreeSetFactory &F, llist<CfgNode *> *l, CfgNode *n0,
		  bool ignoreNSE, bool *ignored = 0);

void freeDomInfo(TreeNode *l)
{
  /* cout << "Freeing dom info for foreach at " <<
     l->position().asString() << endl; */

  delete o[l];

  dominators.erase(l);
  loopExits.erase(l);
  everyIter.erase(l);
  o.erase(l);

  /*  cout << "done freeing dom info for foreach at " <<
      l->position().asString() << endl; */
}

static void printOrderlyTreeSet(TreeSetFactory &F,
				OrderlyTreeSet s, ostream &os, int m)
{
  if (F.isEmpty(s)) {
    os << "empty" << endl;
  } else
    for (TreeSetFactory::set_iterator i = F.begin(s); i != F.end(s); ++i) {
      indent(os, m);
      os << PLCODE2(os, *i) << endl;
    }
}

static OrderlyTreeSet 
convert_cfgset(TreeSetFactory &F, CfgSetFactory &G, CfgSet s,
	       map_cfgset_to_OrderlyTreeSet &memo)
{
  // cout << "convert_cfgset()" << endl;
  OrderlyTreeSet &r = memo[s], tmp;
  if (r) {
    if (DEBUG_ORDERLY)
      cout << "memo hit" << endl;
    return r;
  }
  if (s != G.empty() && (tmp = memo[G.tail(s)])) {
    if (DEBUG_ORDERLY)
      cout << "memo tail hit" << endl;
    return (r = F.adjoin(G.head(s)->astNode(), tmp));
  }
  TreeSetFactory::lazyset result;
  for (CfgSetFactory::set_iterator i = G.begin(s);
       i != G.end(s);
       ++i)
    result.adjoin((*i)->astNode());
  return (r = F.makeset(result));
}

static OrderlyTreeSet
convert_treeset_to_OrderlyTreeSet(TreeSetFactory &F, const treeSet *s)
{
  OrderlyTreeSet result = F.empty();
  for (treeSet::const_iterator i = s->begin(); i != s->end(); ++i)
    result = F.adjoin((TreeNode *) *i, result);
  return result;
}

void assertAppearsOnEveryIter(const TreeNode *ee, const TreeNode *ll)
{
  TreeNode *e = const_cast<TreeNode *>(ee);
  TreeNode *l = const_cast<TreeNode *>(ll);
  o[l]->adjoin(e, everyIter[l]);
}

bool appearsOnEveryIter(const TreeNode *ee, const TreeNode *ll)
{
  TreeNode *e = const_cast<TreeNode *>(ee);
  TreeNode *l = const_cast<TreeNode *>(ll);
  return o[l]->contains(everyIter[l], e);
}

// is everything in s dominated by t?
bool dominated(OrderlyTreeSet s, TreeNode *t, TreeNode *WRTloop)
{
  map_tree_to_OrderlyTreeSet &d = dominators[WRTloop];
  TreeSetFactory &F = *o[WRTloop];
  for (TreeSetFactory::set_iterator i = F.begin(s); i != F.end(s); ++i)
    if (!F.contains(d[*i], t)) {
      if (debug_domv) {
	cout << PLCODE(t); cout << " does not dominate "; cout << PLCODE(*i) << endl;
      }
      return false;
    }

  return true;
}

bool dominated(treeSet *s, TreeNode *t, TreeNode *WRTloop)
{
  return dominated(convert_treeset_to_OrderlyTreeSet(*o[WRTloop], s),
		   t, WRTloop);
}

static llist<CfgNode *> * current_list;
#define ADD_TO_LIST() ((current_list ?					\
			(free_all(current_list), current_list = NULL) :	\
			NULL),						\
		       addToList)

static void addToList(CfgNode *x)
{
  if (debug_domv)
    cout << "loop includes " << unique((void *) x) << endl;
  current_list = cons(x, current_list);
}

static map_tree_to_OrderlyTreeSet *
convert_cfgmap_to_astmap(TreeSetFactory &F, CfgSetFactory &G,
			 map_cfgnode_to_cfgset * r)
{
  // cout << "convert_cfgmap_to_astmap()" << endl;
  map_cfgset_to_OrderlyTreeSet memo;
  map_tree_to_OrderlyTreeSet *result = new map_tree_to_OrderlyTreeSet;
  for (map_cfgnode_to_cfgset::iterator i = r->begin(); i != r->end(); ++i) {
    OrderlyTreeSet &s = (*result)[(*i).first->astNode()],   
      t = convert_cfgset(F, G, (*i).second, memo);
#if 0
    size_t old_s_size = s->size();
#endif
    s = F.merge(s, t);
#if 0
    if (true || old_s_size != 0)
      cout << "Merge sets of size " << old_s_size << " and " << t->size() <<
	": " << s->size() << endl;
#endif
  }
  delete r;
  return result;
}

CfgSet list_to_CfgSet(CfgSetFactory &U, const llist<CfgNode *> *l)
{
  CfgSet s = U.empty();
  foreach_const (p, llist<CfgNode *>, *l)
    s = U.adjoin(*p, s);
  return s;
}

// Uses the algorithm from the Dragon book (Aho, Sethi, and Ullman) to
// compute dominators of a loop whose nodes are in l, assuming start
// node n0.  Returns a mapping from each node in l to a set of nodes
// in l that dominate it.  If ignoreNSE is set then edges marked NSE
// in the CFG are ignored.  This code could be improved by using bit
// vectors to represent sets, by leaking less memory, and by switching
// to a better algorithm (e.g., Tarjan and Lengauer's).
static map_tree_to_OrderlyTreeSet *
computeDominators(TreeSetFactory &F, llist<CfgNode *> *l, const CFGset *ls,
		  CfgNode *n0, bool ignoreNSE, bool *ignored)
{
  // cout << "computeDominators()" << endl;
  CfgSetFactory U;
  map_cfgnode_to_cfgset & result = *(new map_cfgnode_to_cfgset);

  // Initialization:
  //   D(n0) = n0
  //   D(n) = N   for n != n0, where N is the set of all nodes (handled below)
  result[n0] = U.singleton(n0);
  
  CfgSet everything = U.adjoin(n0, list_to_CfgSet(U, l));

  // Iteration:
  //  D(n) = n + intersection over predecessors p of n of D(p)
  int k = 0, count = 0, oldcount = -1;
  while (count != oldcount) {
    if (debug_dom) {
      cout << "Dom loop " << ++k << ", count = " << count <<
	  ", ignoreNSE = " << ignoreNSE << endl;
      if (k == 1 && DEBUG_DOM_VERBOSE) {
	cout << endl << "summary of CFG:" << endl << endl;
	summarizeCFG(l, cout);
	cout << endl << "end of summary of CFG." << endl << endl;
      }
    }
    oldcount = count;
    count = 0;
    int p = 0, cfgsize = l->size();
    foreach (n, llist<CfgNode *>, *l) {
      CfgNode *a = *n;
      if (a == n0) {
	if (debug_domv) {
	  cout << "(" << ++p << " of " << cfgsize << ")"
	    " dominators for " << QCFG(n0) << ": self" << endl;
	}
      } else {
	if (debug_domv) {
	  cout << "(" << ++p << " of " << cfgsize << ")"
	    " working on dominators for " << QCFG(a) << endl;
	}
	ListIterator<CfgNode *> m = a->predIter(ignoreNSE, ls, ignored);
	if (m.isDone())
	  continue;
	CfgSet u = result[*m];
	if (u == NULL) {
	  u = result[*m] = everything;
	  if (debug_domv)
	    cout << "  doms are EVERYTHING" << endl;
	} else if (debug_domv) {
	  cout << "  doms are ";
	  // shortCfgSet(u, cout);
	  cout << endl;
	}
	m.next();
	for ( ; !m.isDone(); m.next()) {
	  if (result[*m] != NULL)
	    u = U.intersect(u, result[*m]);
	  if (debug_domv) {
	    cout << "  doms are ";
	    // shortCfgSet(u, cout);
	    cout << endl;
	  }
	}
	u = U.adjoin(a, u);
	result[a] = u;
	int size = u->size();
	if (debug_domv)
	  cout << int2string(size) << " dominators for " << QCFG(a) <<
	    " at " << a->astNode()->position().asString() << endl;
	count += size;
      }
    }
  }

  delete ls;
  return convert_cfgmap_to_astmap(F, U, &result);
}

llist<TreeNode *> * cfglist_to_treelist(const llist<CfgNode *> *l)
{
  if (l == NULL)
    return NULL;
  llist<TreeNode *> * result = NULL;
  l = dreverse((llist<CfgNode *> *) l);
  foreach_const (p, llist<CfgNode *>, *l)
    result = cons((*p)->astNode(), result);
  l = dreverse((llist<CfgNode *> *) l);
  return result;
}

/* Return a set of nodes that are loop exits.  t is the loop.  l is
   the loop contents.  l plus t->stmt() is the list of nodes we should
   check.  For each element of l, if any successors are outside of l,
   then include that element in the result. */
static OrderlyTreeSet computeLoopExits(TreeSetFactory &F, llist<CfgNode *> *l,
				       TreeNode *t, bool ignoreNSE,
				       bool *ignored)
{
  OrderlyTreeSet lc = F.makeset(cfglist_to_treelist(l));
  // Assume t->stmt() is in the result.  The cfg is always built that way.
  OrderlyTreeSet result = F.singleton(t->stmt());

  ListIterator<CfgNode *> w;

  foreach (p, llist<CfgNode *>, *l)
    if ((*p)->astNode() != t->stmt() && (*p)->astNode() != t)
      for (w = (*p)->succIter(); !w.isDone(); w.next()) {
	if (debug_domv)
	  cout << "(" << unique((void *) *p) << "'s succ " <<
	    unique((void *) *w) << ") checking if loop contains " <<
	    pseudocode((*w)->astNode()) << ": " <<
	    (F.contains(lc, (*w)->astNode()) ? "yes" : "no") << endl; 
	if (!F.contains(lc, (*w)->astNode()))
	  if (ignoreNSE && (*w)->isNSE(*p)) {
	    *ignored = true;
	    if (debug_domv)
	      cout << " (ignored)\n";
	  } else {
	    result = F.adjoin((*p)->astNode(), result);
	    if (debug_domv) {
	      cout << "Found loop exit: ";
	      (*p)->print(cout);
	      cout << " which is dominated by:" << endl;
	      printOrderlyTreeSet(F, dominators[t][(*p)->astNode()], cout, 2);
	    }
	    break;
	  }
      }
  return result;
}

static bool exprnode(TreeNode *t)
{
  return isExprNode(t);
}

/* A node in l appears on every iteration if it dominates every exit
   from the loop. */
OrderlyTreeSet computeEveryIter(TreeSetFactory &F,
				llist<CfgNode *> *l,
				OrderlyTreeSet exits,
				map_tree_to_OrderlyTreeSet &dom,
				const string &where)
{
  llist<TreeNode *> *exprnodes = filter(exprnode, cfglist_to_treelist(l));
  OrderlyTreeSet result = F.makeset(exprnodes);

  for (TreeSetFactory::set_iterator i = F.begin(exits); i != F.end(exits); ++i)
    result = F.intersect(result, dom[(TreeNode *) *i]);
  
  if (debug_dom) {
    cout << "Every Iter (loop at " << where << "):" << endl;
    printOrderlyTreeSet(F, result, cout, 2);
    cout << endl;
  }

  return result;
}

void ForEachStmtNode::computeForeachDominators(bool ignoreNSE)
{
  debug_domv = DEBUG_DOM_VERBOSE;
  debug_dom = debug_domv || DEBUG_DOM;

  // Erase previously computed dominator info, if any.
  if (o[this] != NULL)
    freeDomInfo(this);

  TreeSetFactory &F = *(o[this] = new TreeSetFactory);

  llist<CfgNode *> *l;
  CfgExtent extent = getCfgExtent();
  CfgNode *from = extent.entry;
  CfgNode *to = extent.exit;

  TraverseCfgScope(this, ADD_TO_LIST());
  l = dreverse(current_list);

  if (debug_domv) {
    cout << "Computing dominators for foreach at " <<
      position().asString() << " (cfg size=" << l->size() << ")" << endl;
    cout << "Extents:" << endl;
    from->print(cout);
    to->print(cout);
    cout << endl << "Entire loop:" << endl << pseudocode(stmt()) << endl;
    int i = 0;
    foreach (p, llist<CfgNode *>, *l) {
      cout << i++ << ':' << endl;
      (*p)->print(cout);
    }
    cout << endl;
  } else if (debug_dom)
    cout << "Computing dominators(" << ignoreNSE << ") for" << endl <<
      pseudocode(this) << endl << endl;
  
  if (ignoreNSE)
    _partialDomain = false;
  const CFGset *ls = list_to_CFGset(l);
  dominators[this] = *computeDominators(F, l, ls, from, ignoreNSE,
					&_partialDomain);
#if 0
  if (ignoreNSE && !_partialDomain) {
    /* Look for successor edges from nodes in the loop that are marked NSE. */
    foreach (p, llist<CfgNode *>, *l)
      for (ListIterator<CfgNode *> w = (*p)->succIter(); !w.isDone(); w.next()) {
	if (debug_dom)
	  cout << "Checking for NSE from " << unique((void *) *p) << " to " <<
	    unique((void *) *w) << endl;
	(*w)->predIter(true, ls, &_partialDomain);
	if (_partialDomain) {
	  if (debug_dom)
	    cout << "Found NSE from " << unique((void *) *p) << " to " <<
	      unique((void *) *w) << endl;
	  goto done_partialDomain;
	}
      }
  done_partialDomain: ;
  }
#endif
  loopExits[this] = computeLoopExits(F, l, this, ignoreNSE, &_partialDomain);
  everyIter[this] = computeEveryIter(F, l, loopExits[this], dominators[this],
				     position().asString());
  TreeNode::computeForeachDominators(ignoreNSE);
}

void TreeNode::computeForeachDominators(bool ignoreNSE)
{
  for (int i = arity(); i-- > 0; )
    child(i)->computeForeachDominators(ignoreNSE);
}