#include <functional>

#include "AST.h"
#include "optimize.h"
#include "code-util.h"
#include "cfg.h"

// Basic copy-propagation optimization
// Author: Dan Bonachea <bonachea@cs.berkeley.edu>, 8/2001

static void propagateCopiesBblock(Bblock *b);
static void propagateCopiesFromAssign(TreeNode *assign);
TreeNode *getDefVar(TreeNode *t);
static int numReplacements;

// keep statistics about why candidates fail
static int stats_numLHSEligibleVAR = 0;   // number of copy-prop candidates where lhs is eligible to be propagated
static int stats_numLHSEligibleCONST = 0;   // number of copy-prop candidates where lhs is eligible to be propagated
static int stats_noRHSReachingAnalysis = 0; // number of candidates rejected because of lack of RHS reaching analysis
static int stats_numRHSOutOfScope = 0;  // number of candidates rejected because RHS out of scope
static int stats_numReplacementsVAR = 0; // number of candidates actually replaced
static int stats_numReplacementsCONST = 0; // number of candidates actually replaced

static void printStats() {
  cout << "Copy Propagation Statistics:" << endl;
  cout << " Total LHS Eligible for prop, varRHS:   " << stats_numLHSEligibleVAR << endl;
  cout << " Total LHS Eligible for prop, constRHS: " << stats_numLHSEligibleCONST << endl;
  cout << " Rejected varRHS RHSOutOfScope:         " << stats_numRHSOutOfScope << endl;
  cout << " Rejected varRHS noRHSReachingAnalysis: " << stats_noRHSReachingAnalysis << endl;
  cout << " Total Vars Accepted for propagation:   " << stats_numReplacementsVAR << endl;
  cout << " Total Consts Accepted for propagation: " << stats_numReplacementsCONST << endl;
  cout << endl;
}

// main entry point for copy propagation - 
// performs copy propagation on the enclosing method
// for best results, this optimization should be followed by dead-code elimination 
// assumes CFG and def-use analysis are valid
// returns the number of replacements made
int doCopyPropagation(TreeNode *t) {
  static bool firstTime = true;
  if (firstTime) {
    if (DEBUG_COPY) atexit(printStats);
    firstTime = false;
  }

  TreeNode *methoddeclnode = enclosingMethod(t);
  assert(isMethodDeclNode(methoddeclnode) || isConstructorDeclNode(methoddeclnode));

  // Methods without stmts don't have CFGs.
  if (!methoddeclnode->getBblockRoot()) {
    return 0;
  }

  numReplacements = 0;
  TraverseBblocks(methoddeclnode->getBblockRoot(), propagateCopiesBblock);
  return numReplacements;
}

// private helper for doCopyPropagation
// performs copy propagation on copy stmts in this bblock 
static void propagateCopiesBblock(Bblock *b) {
  // Traverse the CFG list for this block
  ListIterator<CfgNode *> cfgIt = b->cfgIter();
  for (; !cfgIt.isDone(); cfgIt.next()) {
    CfgNode *cfg = *cfgIt;
    TreeNode *t = cfg->astNode();
    if (isAssignNode(t) && !inDummyNode(t))
      propagateCopiesFromAssign(t);
  }  
}

// return true for nodes which form a lexical scope in the AST
bool isNodeWithScope(TreeNode *t) {
  return isBlockNode(t) || // local vars
    isMethodDeclNode(t) || isConstructorDeclNode(t) || // formal parameters
    isCatchNode(t) || // exn parameter
    isForEachStmtNode(t);// foreach point
}

// return true iff obj is in scope at the given context point
static bool inScopeAt(ObjectNode *obj, TreeNode *context) {
  TreeNode *objDeclNode = obj->decl()->source();
  TreeNode *scope = objDeclNode;

  while(scope && !isNodeWithScope(scope))  // find scope of obj
    scope = scope->parent();
  assert(scope != NULL);

  while (context && context != scope) { // check if it includes context
    context = context->parent();
  }

  bool result = (context == scope);
  if (!result) stats_numRHSOutOfScope++;
  return result;
}

// return true iff the value of obj at point 'from' is guaranteed 
// to reach point 'to' unchanged 
// on all paths leading from 'from' to 'to' that stay entirely
// within the scope of obj's declaration
static bool valueReaches(ObjectNode *obj, CfgNode *from, CfgNode *to) {
  
  // trivially true if obj has only one def point, but this is hard to check for

  // if both are in same extended bblock and ordered correctly, we can figure it out
  if (from->bblock == to->bblock) {
    // same bblock (we don't currently do extended bblock, although we could)
    CfgNode *c = from;
    while (true) {
      if (c == to) return true; 

      TreeNode *t = c->astNode();
      // the only way an ObjectNode can change is via an assignnode
      if (isAssignNode(t)) {
        TreeNode *lhs = t->child(0);

        if (isObjectNode(lhs) && lhs->decl() == obj->decl()) 
          return false; // found an intervening def
        else if (isIIFAN(lhs) && (lhs->object()->decl() == obj->decl()))
          return false; // found an intervening def of a field of this immutable obj
      }

      if (c->numSucc == 1) c = *c->succIter();
      else { // leaving bblock - 'to' must precede 'for' in bblock
        assert((*c->succIter())->bblock != from->bblock);
        return false;
      }
    }
  }
  
  // we don't currently have the analysis to answer this question across basic blocks
  stats_noRHSReachingAnalysis++;
  return false; // be conservative
}

// scan the subtree rooted at t and return a pointer to the ObjectNode matching decl d
// (or null if not found)
ObjectNode *findMatchingObjectNode(Decl *d, TreeNode *t) {
  if (isObjectNode(t) && t->decl() == d) return (ObjectNode*)t;
  for (int i=0; i < t->arity(); i++) {
    ObjectNode *o = findMatchingObjectNode(d, t->child(i));
    if (o) return o;
  }
  return NULL;
}


// return true iff obj is part of a larger expression that could result in 
// a def of the associated variable (e.g. on the lhs of an assignment expression)
bool isAssignTarget(ObjectNode *obj) {
  TreeNode *prev = obj;
  TreeNode *next = obj->parent();
  while (next) {
    if (isAssignNode(next) && (prev == next->child(0))) {  // lhs of assignment
      if (prev == obj) return true; // direct assignment
      if (isIIFAN(prev) && (prev->object() == obj)) return true; // IIFAN assignment
    }
    prev = next;
    next = next->parent();
  }
  return false;
}

// clone use lists in this subtree and update the corresponding defs accordingly
// subtree should not contain any defs
static void cloneDefUseInfo(TreeNode *t) {
  llist<TreeNode *>* defs = t->getDefs();
  t->setDefs(copylist(defs));

  for(;defs;defs = defs->tail()) {
    TreeNode *defvar = getDefVar(defs->front());
    assert(defvar != NULL);
    defvar->setUses(cons(t, defvar->getUses()));
  }

  assert(t->getUses() == NULL); 

  for (int i=0; i < t->arity(); i++) 
    cloneDefUseInfo(t->child(i));
}

// private helper for propagateCopiesBblock
// performs copy propagation on a particular assignment, if appropriate
// may change statements occuring later in this bblock or in any descendent bblocks
static void propagateCopiesFromAssign(TreeNode *assign) {
  assert(isAssignNode(assign));

  TreeNode *lhs = assign->opnd0();
  TreeNode *rhs = assign->opnd1();
  if (!(isObjectNode(lhs) && isLocalOrParam(lhs->decl()))) 
    return; // not a simple copy - nothing to do 

  if (!( (isObjectNode(rhs) && isLocalOrParam(rhs->decl()))
         || isConstant(rhs) // casts of constants and casts of null handled by constant folding
         || isNullPntrNode(rhs) 
       //  || isTrivialPointOrDomain(rhs)  // TODO - support this in some cases?
     )) return; // rhs is not a simple var or constant
  
  // assign is a simple copy of the form lhs <- rhs

  if (isObjectNode(rhs) && lhs->decl() == rhs->decl())
    return; // do not attempt to propagate dead assignments (x = x)

  // visit each use of lhs reached by this def
  llist<TreeNode *> *uselist = copylist(lhs->getUses()); // use a separate copy to prevent interference from updates
  uselist = destructive_filter(not1(ptr_fun(inDummyNode)), uselist);
  while (uselist) {
    TreeNode *usenode = uselist->front();

    // if the only reaching def(s) originate at our simple copy, 
    // then it's safe to propagate the copy to this use
    // (in the case of immutable var defs, the copy will generate defs for the var and all fields)
    llist<TreeNode *> *deflist = usenode->getDefs();
    while (deflist) {
      if (deflist->front() != assign) break;
      deflist = deflist->tail();
    }
    if (!deflist && // no mismatches - all defs come from our copy
        enclosingStmt(assign) != enclosingStmt(usenode)) { // don't propagate a copy to itself

      // drill down into usenode to find the matching ObjectNode to be replaced
      ObjectNode *useobj = findMatchingObjectNode(lhs->decl(), usenode);
      assert(useobj != NULL);

      // need to make sure use is not a LHS use
      if (!isAssignTarget(useobj)) {

        if (isObjectNode(rhs)) stats_numLHSEligibleVAR++;
        else stats_numLHSEligibleCONST++;

        // special and fn that evaluates both sides when debugging on (to collect stats)
        #define AND(x,y) (!!(DEBUG_COPY?((x)&(y)):((x)&&(y)))) 

        // now, we need to detect if the value of rhs at the copy reaches 
        // the target location for the propagation, and rhs is in scope

        if (!isObjectNode(rhs) || 
            (isObjectNode(rhs) && 
          AND(inScopeAt((ObjectNode*)rhs, usenode),  // check scoping
              valueReaches((ObjectNode*)rhs, 
                          enclosingStmt(assign)->getCfgExtent().exit, 
                          enclosingStmt(usenode)->getCfgExtent().entry)))
           ) {

          // everything OK - perform the replacement
          if (DEBUG_COPY)
            cout << "Propagating copy: " << pseudocode(enclosingStmt(assign)) <<
	      "\nto use: " << pseudocode(enclosingStmt(usenode)) << '\n';

          TreeNode *replacement = rhs->deepClone(); // create replacement

          cloneDefUseInfo(replacement); // give replacement code its own copy of def chains

          // remove the use entry from the lhs use list 
          llist<TreeNode*> *newuses = lhs->getUses();
          assert(find(usenode, newuses) != NULL);
          newuses = destructive_filter(bind1st(not_equal_to<TreeNode*>(), useobj), newuses);
          newuses = destructive_filter(bind1st(not_equal_to<TreeNode*>(), usenode), newuses);
          lhs->setUses(newuses);
          uselist = destructive_filter(bind1st(not_equal_to<TreeNode*>(), useobj), uselist); // update iteration list
          uselist = destructive_filter(bind1st(not_equal_to<TreeNode*>(), usenode), uselist); // update iteration list

          TreeNode *useparent = useobj->parent();
          int i;
          for (i=0; i < useparent->arity(); i++) {
            if (useparent->child(i) == useobj) {
	      cloneMIVEandInvarInfo(useobj, replacement);
              useparent->child(i, replacement);
              break;
            }
          }
          assert(i < useparent->arity());
          numReplacements++;

          if (isObjectNode(rhs)) stats_numReplacementsVAR++;
          else stats_numReplacementsCONST++;

          continue; // skip iteration increment - we already did it
        } // if rhs legal
      } // if not assign target
    } // if legal deflist

    uselist = uselist->free();
  } // while

}