#include "AST.h"
#include "clone.h"
#include "code-MIVE.h"
#include "code-util.h"
#include "dominator.h"
#include "free_map.h"
#include "optimize.h"

/***************************************************************************
  Lifting of loop invariant expressions out of foreach loops
  ----------------------------------------------------------

  Idea: for each assignment that appears on every iteration and is
  invariant, decide whether it is legal and worth lifting, and if
  so put it on a list.  Move everything on the list to the loop
  preheader, respecting ordering constraints.

  If a foreach contains a foreach in its body we do not lift code from
  the inner foreach all the way out because the inner foreach body may
  not execute (the iteration domain could be empty).  This follows from
  the "appears on every iteration" stipulation.

  The way we determine what is loop invariant is different, but our
  legality checks for invariant code-motion are basically the same as
  those in the Dragon Book, pages 639--641.
***************************************************************************/

bool debug_lift2 = false;

bool debug_lift;

/* Let d = the VarDecl for the variable v.  If d appears in container,
   return d.  Else return NULL. */
static TreeNode *findVarDecl(TreeNode *v, TreeNode *container)
{
 top:
  if (isDynamicTypeNode(v))
    { v = v->opnd0(); goto top; }
  if (v->isArrayAccessNode())
    { v = v->array(); goto top; }
  if (isOFAN(v))
    { v = v->object(); goto top; }

  return isObjectNode(v) ? isAncestor(container, v->decl()->source()) : NULL;
}

static void printUses(TreeNode *t)
{
  llist<TreeNode *> *uses = t->getUses();

  cout << "  Uses of {" << t->oper_name() << "} " << pseudocode(t) << ": ";
  if (uses == NULL)
    cout << "None";
  else {
    bool first = true;
    foreach (u, llist<TreeNode *>, *uses) {
      if (first)
	first = false;
      else
	cout << "; ";
      PLCODE(*u);
    }
  }
  cout << endl;
  flush(cout);
}

static bool noOtherAssignmentsWithSameLHS(TreeNode *e, treeSet *loopContents)
{
  treeSet *s = destructively_intersect(LvalueUses(e->child(0)),
				       loopContents);
  int size = s->size();
  bool r = (size == 0) || (size == 1) && ((*(s->begin()))->parent() == e);
  if (debug_lift) {
    cout << "Lvalue uses: ";
    printTreeSet(s, cout, 1);
    cout << '\n';
  }
  delete s;
  return r;
}

/* Assumes E is an assignment to a variable x.  True iff E is the only
   assignment to x in LOOPCONTENTS and E dominates all uses of x in
   LOOPCONTENTS.  See Dragon Book for why we'd want to check these
   things (if it returns true and several other conditions hold then
   it is legal to lift E out of the loop.) 
*/

static bool 
uniquelyDominates (TreeNode *e, treeSet *loopContents,
		   TreeNode *WRTloop)
{
  if (debug_lift) {
    cout << "Checking legality of lifting " << pseudocode(e) << endl;
    printUses(e->child(0));
  }
 
  treeSet *s = destructively_intersect(RvalueUses(e->child(0)),loopContents);
  bool r =
    dominated(s,
	      e,
	      WRTloop) &&
    noOtherAssignmentsWithSameLHS(e, loopContents);
  
  if (debug_lift)
    cout << "         legality of lifting " << pseudocode(e) << ": " << r << endl;

  delete s;
  return r;
}

bool isSequentiallyMobile(const TreeNode *e)
{
  const bool rigid = sequential_consistency
    && isAssignNode(e)
    && (e->opnd0()->isSharedAccess()
	|| e->opnd1()->isSharedAccess());
  
  return !rigid;
}

treeSet *AssignableNodes(TreeNode *e);

static bool notVarDeclNode(const TreeNode *t)
{
  return !isVarDeclNode(t);
}

static treeSet *mustMoveFirst(TreeNode *e, treeSet *loopContents)
{
  treeSet *an = AssignableNodes(e);
  treeSet *result = destructively_intersect(allDefs(an), loopContents);
  delete an;
  result = destructively_filter(notVarDeclNode, result);
  result->erase(e);
  if (debug_lift) {
    cout << "mustMoveFirst(" << pseudocode(e) << "):" << endl;
    printTreeSet(result, cout, 1);
  }
  return result;
}

static bool sameSimpleNode(TreeNode *x, TreeNode *y)
{
  bool result =
    isThisNode(x) && isThisNode(y) ||
    isNameNode(x) && isNameNode(y) && x->decl() == y->decl();
  /*
  (x->qualifier()->absent() && y->qualifier->absent() ||
     !x->qualifier()->absent() && !y->qualifier->absent() &&
     canReuse(x->qualifier(), y->qualifier())
  */
  return result;
}

static bool sameLit(TreeNode *x, TreeNode *y)
{
  return
    isPrimitiveLitNode(y) &&
    isPrimitiveLitNode(x) &&
    y->literal().eq(x->literal()).boolValue();
}

static bool canReuse(TreeNode *x, TreeNode *y)
{
  bool result =
    (isObjectNode(y) && isObjectNode(x) &&
     sameSimpleNode(x->name(), y->name())) ||
    (isThisNode(y) && isThisNode(x)) ||
    sameLit(x, y) ||
    (isDynamicTypeNode(y) && isDynamicTypeNode(x) &&
     x->dtype() == y->dtype() && canReuse(x->opnd0(), y->opnd0())) ||
    (isOFAN(y) && isOFAN(x) &&
     canReuse(x->object(), y->object()) &&
     sameSimpleNode(x->simpName(), y->simpName()));
  if (debug_lift) {
    cout << "Checking if we can reuse " << pseudocode(x) << " for " <<
      pseudocode(y) << ": " << (result ? "yes" : "no") << endl;
  }
  return result;
}

/* e is an invariant expression.  If there is an assignment in moved
   that has an equivalent RHS and its LHS is an ObjectNode then
   return the LHS. */
static TreeNode *findReusableInvar(treeSet &moved, TreeNode *e)
{
  for (treeSet::iterator i = moved.begin(); i != moved.end(); ++i) {
    TreeNode *assignment = (TreeNode *) *i;
    TreeNode *lhs = assignment->opnd0();
    TreeNode *rhs = assignment->opnd1();
    if (canReuse(rhs, e) && isObjectNode(lhs))
      return lhs;
  }    
  return NULL;
}

/////////////////////////////////////////////////////////////////////////////
// Patches: maps from VarDeclNodes to ObjectNodes
//  When an ObjectNode that refers to the object named by the
//  VarDeclNode is used, it may be replaced by a copy of the
//  corresponding ObjectNode.
/////////////////////////////////////////////////////////////////////////////

static TreeNode *patchOf(map_tree_to_tree &patches, TreeNode *t)
{
  return isObjectNode(t) ? patches[t->name()->decl()->source()] : NULL;
}

static void addPatch(map_tree_to_tree &patches, TreeNode *x,
		     TreeNode *replacement)
{
  if (x == NULL)
    return;

  TreeNode *z;
  while ((z = patchOf(patches, replacement)) != NULL)
    replacement = z;
  patches[x] = replacement;

  if (debug_lift)
    cout << "Added patch: uses of " << pseudocode(x) << " -> " <<
      pseudocode(replacement) << endl;
}

static TreeNode *CloneWithUseDefInfo(TreeNode *t)
{
  return CloneTree(t, Duplicate);
}

static TreeNode *applyPatches(map_tree_to_tree &patches, TreeNode *t,
			      llist<TreeNode *> *enclosingLoops)
{
  TreeNode *z;
  if (isObjectNode(t) && !isLHSofAssignNode(t) &&
      ((z = patchOf(patches, t)) != NULL)) {
    TreeNode *result = CloneWithUseDefInfo(z);
    if (debug_lift)
      cout << "Patching " << pseudocode(t) << "; checking " <<
	enclosingLoops->size() << " enclosing loops" << endl;
    foreach (l, llist<TreeNode *>, *enclosingLoops) {
      TreeNode *WRTloop = *l;
      if (appearsOnEveryIter(t, WRTloop))
	assertAppearsOnEveryIter(result, WRTloop);
      if (isInvar(t, WRTloop))
	foundInvariant(result, WRTloop);
    }
    return result;
  }
  else {
    if (isForEachStmtNode(t))
      enclosingLoops = cons(t, enclosingLoops);
    for (int i = t->arity(); i-- > 0; )
      t->child(i, applyPatches(patches, t->child(i), enclosingLoops));
    if (isForEachStmtNode(t))
      enclosingLoops->free();
    return t;
  }
}

/* Create (and return) a BlockNode that contains code for
   invariants moved out of the loop t, followed by t itself.
   Destructively modifies t.  Exception: if opt_lift is false, do nothing.
*/
static StatementNode *lifted(ForEachStmtNode *t)
{
  assert(t->numLifted() == 0);
  if (!opt_lift) return t;
  llist<TreeNode *> *l = NULL;
  Worklist w(t);
  treeSet *loopContents = list_to_set(cons(t->vars()->child(0), w.asList()));
  while (!w.empty()) {
    TreeNode *e = w.pop();
    if (isAssignNode(e) &&
	isInvar(e, t) &&
	appearsOnEveryIter(e, t) &&
	uniquelyDominates(e, loopContents, t))
      l = cons(e, l);
  }
  if (l == NULL) {
    delete loopContents;
    return t;
  }

  l = dreverse(l);
  treeSet moved;
  llist<TreeNode *> *result = NULL;
  treeSet deleted;
  bool progress;
  map_tree_to_treeSet pre;
  foreach (c, llist<TreeNode *>, *l)
    pre[*c] = mustMoveFirst(*c, loopContents);

  map_tree_to_tree patches;
  do {
    progress = false;
    foreach (e, llist<TreeNode *>, *l) {
      TreeNode *c = *e;
      if (!treeSetContains(moved, c) &&
	  isSubset(*(pre[c]), moved)) {
	TreeNode *previous =
	  isObjectNode(c->opnd0()) ?
	  findReusableInvar(moved, c->opnd1()) :
	  NULL;
	moved.insert(c);
	adjoin(&deleted, c);
	if (previous == NULL)
	  result = cons((TreeNode *)
			(new ExpressionStmtNode(c, c->position())),
			result);
	else {
	  // We've previously lifted an equivalent expression and stored it
	  // it in previous, so make a note to replace uses of c->opnd0()
	  // with previous.
	  addPatch(patches, c->opnd0()->name()->decl()->source(), previous);

	  // Add an assignment to c->opnd0() just in case something after
	  // the loop happens to care.  Most of the time this assignment is
	  // not necessary and, in that case, UAE will remove it later.
	  previous = CloneWithUseDefInfo(previous);
	  result =
	    cons((TreeNode *)
		 (new ExpressionStmtNode(new AssignNode(c->opnd0(),
							previous,
							c->position()),
					 c->position())),
		 result);
	}

	// If a decl we need is in the loop, lift it to the preheader too.
	TreeNode *decl = findVarDecl(c->opnd0(), t);
	if (decl != NULL && !treeSetContains(deleted, decl)) {
	  result = cons(decl, result);
	  adjoin(&deleted, decl);
	}
	decl = findVarDecl(c->opnd1(), t);
	if (decl != NULL && !treeSetContains(deleted, decl)) {
	  result = cons(decl, result);
	  adjoin(&deleted, decl);
	}

	progress = true;
	if (debug_lift) {
	  cout << "Lift " << pseudocode(c->opnd0()) << " = " <<
	    pseudocode(c->opnd1()) << endl;
	  cout << "all uses of LHS:" << endl;
	  printTreeSet(allUses(c->opnd0()), cout, 1);
	  cout << endl;
	}
	++(t->numLifted());
	liftedInvar(c, t);
      }
    }
  } while (progress);

  free_map(pre);

  if (result == NULL)
    return t;

  destructively_filter( isSequentiallyMobile, &deleted );

  foreach (lifted, llist<TreeNode *>, *result)
    if (!isVarDeclNode(*lifted)) {
      *lifted = CloneWithUseDefInfo(*lifted);
      if (debug_lift) {
	TreeNode *c = (*lifted)->child(0), *lhs = c->opnd0();
	cout << "In lifted clone, " << pseudocode(lhs) << " = " <<
	  pseudocode(c->opnd1()) << ", all uses of LHS:" << endl;
	printTreeSet(allUses(lhs), cout, 1);
	cout << endl;
      }
    }
  
  t->deleteStmts(&deleted);
  result = dreverse(cons((TreeNode *) t, result));
  return (StatementNode *)
    applyPatches(patches,
		 SortVarDeclsInBlock(new BlockNode(result, NULL, t->position())),
		 NULL);
}

TreeNode *TreeNode::liftInvariantExprs()
{
  for (int i = arity(); i-- > 0; )
    child(i, child(i)->liftInvariantExprs());
  return this;
}

/* If this is a foreach+ already then just do lifting.  Otherwise, 
   rewrite "foreach (p in D) { ... }" to
   "if (D.isNotNull()) { L; foreach+ (p in D) { ... } }" where L is a loop
   preheader that contains lifted computations.  Destructively modifies this.
   foreach+ is a foreach node that is known to have a positive number of
   iterations.  We rely on the fact (see lowering pass) that D is
   a local variable that does not have any defs inside the loop.
*/
TreeNode *ForEachStmtNode::liftInvariantExprs()
{
  static const string * isNotNull_string = intern("isNotNull");
  TreeNode *l = lifted(this);
  stmt(stmt()->liftInvariantExprs());
  if (_cannotBeEmpty == NULL) {
    TreeNode *dom = vars()->child(0)->initExpr();
    assert(dom != NULL && !dom->absent());
    bool rectangular = dom->type()->isRectDomainType();
    int tiArity = dom->type()->tiArity();
    Decl * isNotNull_decl =
      rectangular ?
      getRectDomainMethodDecl(tiArity, isNotNull_string) :
      getDomainMethodDecl(tiArity, isNotNull_string);
    SourcePosn pos = position();
    assert(isInvar(dom, this));
    TreeNode *isnotnull =
      new ObjectFieldAccessNode(CloneWithUseDefInfo(dom),
				new NameNode(TreeNode::omitted,
					     isNotNull_string,
					     isNotNull_decl, pos),
				pos);
    TreeNode *not_empty = new MethodCallNode(isnotnull, NULL, pos);
    // TreeNode *l = lifted(this);
    // stmt(stmt()->liftInvariantExprs());
    l = new IfStmtNode(not_empty, l, TreeNode::omitted, pos);
    _cannotBeEmpty = l; /* foreach -> foreach+ */
  }
  if (debug_lift) {
    cout << endl << "Foreach loop after lifting:" << endl << endl;
    PCODE(l);
    cout << endl;
  }
  return l;
}