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

// Dead-variable removal optimization
// Author: Dan Bonachea <bonachea@cs.berkeley.edu>, 8/2001

static set<LocalVarDecl *> possiblyDeadDecls;
static void discoverDeadVars(TreeNode *t);
static TreeNode *removeDeadVars(TreeNode *t);

// main entry point for dead-variable removal
// remove any unused VariableDeclNodes in the given AST subtree
// to avoid cluttering up the stack frame we present to the C compiler with dead vars
// for best results, this optimization should follow dead-code elimination 
// returns modified tree
TreeNode *doDeadVariableRemoval(TreeNode *t) {
  if (!underPragma(Pragma::noOpt, t)) {
    if (DEBUG_DEADVAR)
      cout << "Dead-variable elimination..." << endl;

    possiblyDeadDecls.clear();
    discoverDeadVars(t);
    t = removeDeadVars(t);
    possiblyDeadDecls.clear();
  }
  return t;
}

static void discoverDeadVars(TreeNode *t) {
  if (t->isPragma(Pragma::noOpt))
    return;

  if (isVarDeclNode(t)) {
    // start by assuming all VarDeclNodes are dead
    LocalVarDecl *lvd = dynamic_cast<LocalVarDecl *>(t->decl());
    assert(lvd != NULL);
    possiblyDeadDecls.insert(lvd);
  }
  else if (isObjectNode(t) && 
	   (t->decl()->category() & Decl::LocalVar)) {
    // record a use
    LocalVarDecl *lvd = dynamic_cast<LocalVarDecl *>(t->decl());
    assert(lvd != NULL);
    possiblyDeadDecls.erase(lvd);
  }
   
  for (int i=0; i < t->arity(); i++) 
    discoverDeadVars(t->child(i));
}

static TreeNode *removeDeadVars(TreeNode *t) {
  int cnt = 0;
  treeSet declNodesToRemove;
  // collect the dead VarDeclNodes
  for (set<LocalVarDecl *>::iterator it = possiblyDeadDecls.begin();
       it != possiblyDeadDecls.end();
       it++) {
         LocalVarDecl *lvd = *(it);
         VarDeclNode *vdn = dynamic_cast<VarDeclNode *>(lvd->source());
         assert(vdn != NULL);
         if (DEBUG_DEADVAR)
           cout << "  removing dead VarDeclNode: " << *vdn->simpName()->ident() << endl;
         declNodesToRemove.insert(vdn);
         cnt++;
       }
  // delete them
  t = deleteDeclsFromBlocks(t, &declNodesToRemove);
  if (DEBUG_DEADVAR)
    cout << "Dead-variable elimination removed " << cnt << " dead VarDeclNodes." << endl;

  return t;
}