#include "AST.h"
#include "code-util.h"
#include "Worklist.h"
#include "optimize.h"
#include "code-MIVE.h"
#include "pseudocode.h"

/* An "invariant" is a loop invariant with respect to (WRT) some loop. */

/* Utilities for keeping track of invariants */

/* invar[l] is the set of invariants for loop l. */
static map_tree_to_treeSet invar, nInvar, lInvar;

void initInvar(TreeNode *l)
{
  l->invariants() = invar[l] = new treeSet;
  l->lifted_invariants() = lInvar[l] = new treeSet;
  nInvar[l] = new treeSet;
}

/* Free the data structures allocated in initInvar. */
void freeInvarInfo(const TreeNode *ll)
{
  TreeNode *l = const_cast<TreeNode *>(ll);
  delete invar[l];
  invar.erase(l);
  delete nInvar[l];
  nInvar.erase(l);
  delete lInvar[l];
  lInvar.erase(l);
}

treeSet *& TreeNode::invariants()
{
  static treeSet *bogus = NULL;
  undefined("invariants");
  return bogus;
}

treeSet *& ForEachStmtNode::invariants()
{
  return _invariants;
}

treeSet *& TreeNode::lifted_invariants()
{
  static treeSet *bogus = NULL;
  undefined("lifted_invariants");
  return bogus;
}

treeSet *& ForEachStmtNode::lifted_invariants()
{
  return _lifted_invariants;
}

/* Return whether t appears on the list of invariants for WRTloop. */
bool isInvar(const TreeNode *t, const TreeNode *WRTloop)
{
  TreeNode *l = const_cast<TreeNode *>(WRTloop);
  return (invar[l]->find(t) != invar[l]->end());
}

/* Return whether t appears on the list of lifted_invariants for WRTloop. */
bool isLiftedInvar(const TreeNode *t, const TreeNode *WRTloop)
{
  TreeNode *l = const_cast<TreeNode *>(WRTloop);
  return (lInvar[l]->find(t) != lInvar[l]->end());
}

/* Return whether t appears on the list of non-invariants for WRTloop. */
static bool isNotInvar(const TreeNode *t, const TreeNode *WRTloop)
{
  TreeNode *l = const_cast<TreeNode *>(WRTloop);
  return (nInvar[l]->find(t) != nInvar[l]->end());
}

/* Record that t is not invariant with respect to given loop. */
static void notInvar(const TreeNode *t, const TreeNode *WRTloop)
{
  TreeNode *l = const_cast<TreeNode *>(WRTloop);
  nInvar[l]->insert(t);
}

/* Add t to the set of lifted_invariants for WRTloop. */
void liftedInvar(const TreeNode *t, const TreeNode *WRTloop)
{
  TreeNode *l = const_cast<TreeNode *>(WRTloop);
  lInvar[l]->insert(t);
}

void foundInvariant(const TreeNode *inv, TreeNode *WRTloop)
{
  assert(isForEachStmtNode(WRTloop));
  assert(!isInvar(inv, WRTloop)); /* no repeats */

  invar[WRTloop]->insert(inv);

  if (DEBUG_LOOP_INVAR) {
    cout << "found invariant " << inv->oper_name() <<
      " (WRTloop " << WRTloop->position().asString() << "): ";
    inv->pseudoprint(cout, 0);
    cout << endl;
  }
}

/* Mark dest as invar with respect to all loops for which src is invar. */
void cloneMIVEandInvarInfo(const TreeNode *src, const TreeNode *dest)
{
  if (DEBUG_LOOP_INVAR)
    cout << "cloneMIVEandInvarInfo: " <<
      pseudocode1(src) << " to " << pseudocode1(dest) << endl;
  for (TreeNode *ancestor = src->parent(); ancestor != NULL;
       ancestor = ancestor->parent())
    if (isForEachStmtNode(ancestor)) {
      if (isInvar(src, ancestor) && !isInvar(dest, ancestor))
	foundInvariant(dest, ancestor);
      MIVE *m;
      if ((m = const_cast<TreeNode *>(src)->getMIVE(ancestor)) != NULL)
	const_cast<TreeNode *>(dest)->setMIVE(ancestor, m);
    }
}


/////////////////////////////////////////////////////////////////////////////

static bool allInvariant(TreeNode *expr, TreeNode *loop, treeSet *loopContents,
			 void action(TreeNode *, TreeNode *))
{
  foriter (p, expr->allChildren(), TreeNode::ChildIter)
    if (!((*p)->loopInvariant(loop, loopContents, action)))
      return false;

  return true;
}

/////////////////////////////////////////////////////////////////////////////
// the loopInvariant() method and associated helpers.
/////////////////////////////////////////////////////////////////////////////

/* If we know the answer from before, put it in *answer and return true.
   Otherwise return false. */
bool memo_loopInvariant(TreeNode *t, TreeNode *WRTloop, bool *answer)
{
  if (isInvar(t, WRTloop)) {
    *answer = true;
    return true;
  } else if (isNotInvar(t, WRTloop)) {
    *answer = false;
    return true;
  }
  return false;
}

bool TreeNode::loopInvariant(TreeNode *WRTloop, treeSet *loopContents,
			     void action(TreeNode *, TreeNode *))
{
  bool b;
  if (memo_loopInvariant(this, WRTloop, &b)) return b;
  foriter (p, allChildren(), ChildIter) {
    (*p)->loopInvariant(WRTloop, loopContents, action);
  }
  notInvar(this, WRTloop); return false;
}

#define INVAR_IF_CHILD0_IS_INVAR(NODE)					\
bool NODE::loopInvariant(TreeNode *WRTloop, treeSet *loopContents,	\
			 void action(TreeNode *, TreeNode *))		\
{									\
  bool b;								\
  if (memo_loopInvariant(this, WRTloop, &b)) return b;			\
  if (child(0)->loopInvariant(WRTloop, loopContents, action)) {		\
    action(this, WRTloop);						\
    return true;							\
  }									\
  notInvar(this, WRTloop); return false;				\
}

#define INVAR_IF_CHILDREN_ARE_INVAR(NODE)				\
bool NODE::loopInvariant(TreeNode *WRTloop, treeSet *loopContents,	\
			 void action(TreeNode *, TreeNode *))		\
{									\
  bool b;								\
  if (memo_loopInvariant(this, WRTloop, &b)) return b;			\
  bool result = true;							\
  int i, a = arity();							\
  for (i = 0; i < a; i++) {						\
    if (!child(i)->loopInvariant(WRTloop, loopContents, action))	\
      result = false;							\
  }									\
  if (result) {								\
    action(this, WRTloop);						\
    return true;							\
  }									\
									\
  notInvar(this, WRTloop); return result;				\
}

INVAR_IF_CHILD0_IS_INVAR(DynamicTypeNode)

INVAR_IF_CHILDREN_ARE_INVAR(UnaryArithNode)
INVAR_IF_CHILDREN_ARE_INVAR(BinaryArithNode)
INVAR_IF_CHILDREN_ARE_INVAR(ComplementNode)
INVAR_IF_CHILDREN_ARE_INVAR(NotNode)
INVAR_IF_CHILDREN_ARE_INVAR(ShiftNode)
INVAR_IF_CHILDREN_ARE_INVAR(RelationNode)
INVAR_IF_CHILDREN_ARE_INVAR(EqualityNode)
INVAR_IF_CHILDREN_ARE_INVAR(BitwiseNode)
INVAR_IF_CHILDREN_ARE_INVAR(LogCondNode)

bool AssignNode::loopInvariant(TreeNode *WRTloop, treeSet *loopContents, void action(TreeNode *, TreeNode *))
{
  bool b;
  if (DEBUG_LOOP_INVAR)
    cout << "entering AssignNode::loopInvariant(): " << pseudocode1(this) <<
      "\n";
  if (memo_loopInvariant(this, WRTloop, &b))
    return b;
  if (opnd1()->loopInvariant(WRTloop, loopContents, action)) {
    if (DEBUG_LOOP_INVAR)
      cout << "AssignNode::loopInvariant(): rhs of " << pseudocode1(this) <<
	" is invar\n";
    action(this, WRTloop);
    action(opnd0(), WRTloop);
    return true;
  }
  notInvar(opnd0(), WRTloop);
  notInvar(this, WRTloop);
  if (DEBUG_LOOP_INVAR)
    cout << "exiting AssignNode::loopInvariant(): " << pseudocode1(this) <<
      " is not invar\n";
  return false;
}

/* True iff all DEFs to t from s have the same value and that
   value is loop invariant with respect to the given loop. */
static bool matchingInvarDefsFromSet(TreeNode *t, TreeNode *WRTloop,
				     treeSet *s)
{
  MIVE *p = NULL;

  llist<TreeNode *> *defs = t->getDefs();

  if (DEBUG_LOOP_INVAR) {
    cout << "  Defs of {" << t->oper_name() << "} ";
    cout << PCODE(t); cout << ": ";
    if (defs == NULL)
      cout << "None";
    else {
      bool first = true;
      foreach (u, llist<TreeNode *>, *defs) {
	if (first)
	  first = false;
	else
	  cout << "; ";
	(*u)->pseudoprint(cout, 2);
      }
    }
    cout << endl;
    flush(cout);
  }

  foreach (u, llist<TreeNode *>, *defs)
    if (s->count(*u) > 0) {
      TreeNode *source = valueTransferred((*u), t);
      if (source == NULL || !isInvar(source, WRTloop))
	return false;
      if (p && !p->equals(source->getMIVE(WRTloop)))
	return false;
      else p = source->getMIVE(WRTloop);
    }

  return true;
}


static bool methodTakesLoopInvariantsToLoopInvariant(TreeNode *t)
{
  bool pure = t->isPureFunction();
  if (DEBUG_LOOP_INVAR)
    cout << "methodTakesLoopInvariantsToLoopInvariant(" <<
      pseudocode(t) << "): " << pure << endl;
  return pure;
}

bool MethodCallNode::loopInvariant(TreeNode *WRTloop, treeSet *loopContents,
				   void action(TreeNode *, TreeNode *))
{
  bool b;
  if (memo_loopInvariant(this, WRTloop, &b)) return b;
  if (arithCall(this)) {
    TreeNode *left = child(0)->child(0);
    TreeNode *right = child(1)->child(0);
    bool lefti = left->loopInvariant(WRTloop, loopContents, action);
    bool righti = right->loopInvariant(WRTloop, loopContents, action);
    if (lefti && righti &&
	isTiOrNumericType(left) && isTiOrNumericType(right)) {
      action(this, WRTloop);
      return true;
    }
    if (DEBUG_LOOP_INVAR) {
      cout << "non-invariant arithCall (lefti=" << lefti
	   << " righti=" << righti << "): ";
      pseudoprint(cout, 12);
      cout << '\n';
    }
  } else if (methodTakesLoopInvariantsToLoopInvariant(this)
	     && (isTypeFAN(child(0))
		 || isOFAN(child(0))
		 && child(0)->child(0)->loopInvariant(WRTloop, loopContents,
						      action))) {
    bool b = true;
    foriter (p, args()->allChildren(), ChildIter) {
      if (!(*p)->loopInvariant(WRTloop, loopContents, action)) b = false;
    }
    if (b) {
      action(this, WRTloop);
      return true;
    }
  }
  
  foriter (p, allChildren(), ChildIter) {
    (*p)->loopInvariant(WRTloop, loopContents, action);
  }
  notInvar(this, WRTloop); return false;
}

bool PrimitiveLitNode::loopInvariant(TreeNode *WRTloop, treeSet *loopContents,
				     void action(TreeNode *, TreeNode *))
{
  action(this, WRTloop);
  return true;
}

bool NullPntrNode::loopInvariant(TreeNode *WRTloop, treeSet *loopContents,
				 void action(TreeNode *, TreeNode *))
{
  action(this, WRTloop);
  return true;
}

bool ArrayAccessNode::loopInvariant(TreeNode *WRTloop, treeSet *loopContents,
				    void action(TreeNode *, TreeNode *))
{
  bool b;
  if (memo_loopInvariant(this, WRTloop, &b)) return b;
  if (DEBUG_LOOP_INVAR) {
    cout << "checking if array access is loop invariant: ";
    pseudoprint(cout, 0);
    cout << "\n";
  }
  if (matchingInvarDefsFromSet(this, WRTloop, loopContents) &&
      allInvariant(this, WRTloop, loopContents, action)) {
    action(this, WRTloop);
    return true;
  } 
  notInvar(this, WRTloop); return false;
}

bool subExprsAreInvar(TreeNode *t,
		      TreeNode *WRTloop, treeSet *loopContents,
		      void action(TreeNode *, TreeNode *))
{
  if (t->isExprNode())
    return t->loopInvariant(WRTloop, loopContents, action);
  else {
    bool result = true;
    for (int i = t->arity(); i-- > 0; )
      if (!subExprsAreInvar(t->child(i), WRTloop, loopContents, action))
	result = false;
    return result;
  }
}

bool PointNode::loopInvariant(TreeNode *WRTloop, treeSet *loopContents,
			      void action(TreeNode *, TreeNode *))
{
  bool b;
  if (memo_loopInvariant(this, WRTloop, &b)) return b;
  if (subExprsAreInvar(child(0), WRTloop, loopContents, action)) {
    action(this, WRTloop);
    return true;
  }
  notInvar(this, WRTloop); return false;
}

bool DomainNode::loopInvariant(TreeNode *WRTloop, treeSet *loopContents,
			       void action(TreeNode *, TreeNode *))
{
  bool b;
  if (memo_loopInvariant(this, WRTloop, &b)) return b;
  if (subExprsAreInvar(child(0), WRTloop, loopContents, action)) {
    action(this, WRTloop);
    return true;
  }
  notInvar(this, WRTloop); return false;
}

/*
bool FieldAccessNode::loopInvariantField(TypeNode *type,
					 const string *field,
					 TreeNode *WRTloop)
{
  bool b;
  if (memo_loopInvariant(this, WRTloop, &b)) return b;
  if (WRTloop->defs()->fieldNotModified(type, *field)) {
    action(this, WRTloop);
    return true;
  }
  notInvar(this, WRTloop); return false;
}
*/

bool ObjectFieldAccessNode::loopInvariant(TreeNode *WRTloop,
					  treeSet *loopContents,
					  void action(TreeNode *, TreeNode *))
{
  bool b;
  if (memo_loopInvariant(this, WRTloop, &b)) return b;
  if (object()->loopInvariant(WRTloop, loopContents, action)
      && matchingInvarDefsFromSet(this, WRTloop, loopContents)) {
    action(this, WRTloop);
    return true;
  }
  notInvar(this, WRTloop); return false;
}

bool TypeFieldAccessNode::loopInvariant(TreeNode *WRTloop,
					treeSet *loopContents,
					void action(TreeNode *, TreeNode *))
{
  bool b;
  if (memo_loopInvariant(this, WRTloop, &b)) return b;
  if (matchingInvarDefsFromSet(this, WRTloop, loopContents)) {
    action(this, WRTloop);
    return true;
  }
  notInvar(this, WRTloop); return false;
}

/* A ThisNode should probably be treated like an ObjectNode (i.e.,
   like any other object).  However, for now, we'll conservatively
   assume that this is not loop invariant. */
bool ThisNode::loopInvariant(TreeNode *WRTloop, treeSet *loopContents,
			     void action(TreeNode *, TreeNode *))
{
  bool b;
  if (memo_loopInvariant(this, WRTloop, &b)) return b;
  notInvar(this, WRTloop);
  return false;
}

bool ObjectNode::loopInvariant(TreeNode *WRTloop, treeSet *loopContents,
			       void action(TreeNode *, TreeNode *))
{
  bool b;
  if (memo_loopInvariant(this, WRTloop, &b)) return b;
  if (name()->loopInvariant(WRTloop, loopContents, action)) {
    action(this, WRTloop);
    return true;
  }
  notInvar(this, WRTloop); return false;
}

bool NameNode::loopInvariant(TreeNode *WRTloop, treeSet *loopContents,
			     void action(TreeNode *, TreeNode *))
{
  bool b;
  if (memo_loopInvariant(this, WRTloop, &b)) return b;
  if (isInvar(this, WRTloop)) {
    if (DEBUG_LOOP_INVAR) {
      cout << "namenode ";
      pseudoprint(cout, 0);
      cout << " is loop invariant\n";
    }
    return true;
  }
  
  if (DEBUG_LOOP_INVAR) {
    cout << "checking if namenode ";
    pseudoprint(cout, 0);
    cout << " at " << position().asString() << " is loop invariant" << endl;
    cout << "Its decl's source is" << endl;
    PLCODE(parent()->decl()->source());
    cout << endl;
  }
  
  if (isObjectNode(parent()) &&
      parent()->decl()->source() != WRTloop->vars()->child(0) &&
      matchingInvarDefsFromSet(parent(), WRTloop, loopContents)) {
    action(this, WRTloop);
    return true;
  }
  if (DEBUG_LOOP_INVAR) {
    cout << "checked  if namenode ";
    pseudoprint(cout, 0);
    cout << " is loop invariant: No\n";
  }
  notInvar(this, WRTloop);
  return false;
}