// This is a simple example that illustrates the use of an 
// iterator or a walker to report the number of direct and
// indirect calls in each procedure.

#include <suifnodes/suif.h>
#include <suifkernel/suif_env.h>
#include <suifkernel/module_subsystem.h>
#include <suifkernel/utilities.h>
#include <suifkernel/group_walker.h>
#include <suifpasses/suifpasses.h>
#include <suifkernel/visitor_map.h>


// IsDirect returns 1 if the address of a callee procedure
// is a SymbolAddress, 0 otherwise.

int
IsDirect(SuifEnv *env, Expression * addr) {
  return (is_kind_of<SymbolAddressExpression>(addr));
}

// An iterator iterates over the objects (transitively) owned
// by a starting object.  You can restrict the objects that are returned 
// by the iterator by specifying that they must be subtypes of a specified
// type (in the example below, subtypes of CallStatement).  If you want
// to iterate over all the objects, then you can specify "SuifObject"
// instead of the "CallStatement" in the example code.  Note that you
// should not modify the objects you are iterating over.

class CountCallsWithIteratorPass : public PipelinablePass {
  SuifEnv *_env;
public:
  CountCallsWithIteratorPass(SuifEnv *env, const LString &name): PipelinablePass(env, name){
    _env = env;
  }

  virtual ~CountCallsWithIteratorPass(void)  { }

  Module *clone() const { return(Module*)this;} 

  void do_procedure_definition(ProcedureDefinition* proc_def) {
    unsigned direct_calls = 0, indirect_calls = 0;
    Iter<CallStatement> it = object_iterator<CallStatement>(proc_def);
    for (; it.is_valid(); it.next()) {
      CallStatement& cur = it.current();
      if (IsDirect(_env, cur.get_callee_address()))
         direct_calls++; else indirect_calls++;
    }
    printf("(Iterator) Caller: %s  Direct: %d, Indirect: %d\n", 
	   proc_def->get_procedure_symbol()->get_name().c_str(),
	   direct_calls, indirect_calls);
  }
};


// The following code does the same as the above but uses a walker.  

// A walker is a powerful pattern for traversing the tree rooted at a node.
// (Actually, a walker can also be used to walk non-ownership links 
// but then it is your responsibility to ensure termination...).  
// It is more powerful than an iterator though the following simple example 
// does not illustrate its full power.  
// First, a walker gives finer control over traversal than an iterator (via
// the is_walkable function and the set_pre_order).  Second, using a walker
// you can actually rewrite the tree you are traversing over (this action
// is specified in the return code of the operator()).  Rewriting trees
// as you are walking over them is however tricky so be careful.
// There are several flavors of walkers.  The one below is a subclass of
// SelectiveWalker which walks only over the types of nodes specified
// by the parameter to its constructor (in this case subclass of CallStatement).

class CountCallsWalker: public SelectiveWalker {
  unsigned indirect_calls, direct_calls;
  SuifEnv *_env;

public:

  CountCallsWalker(SuifEnv *env): 
    SelectiveWalker(env, CallStatement::get_class_name()) { 
      _env = env; 
      indirect_calls = direct_calls = 0;
    }

  Walker::ApplyStatus operator() (SuifObject *x) {

    // Operator is invoked only on CallStatements, so the 
    // following cast to CallStatement will always succeed.

    CallStatement *cs = to<CallStatement>(x);
    if (IsDirect(_env, cs->get_callee_address())) direct_calls++; 
    else indirect_calls++;

    // Continue walking.  This return code could be used for several
    // purposes, such as pruning the walk or replacing part of the tree
    // being walked over.  See header file for alternatives.
    return Continue;
  }
  unsigned get_indirect_calls() { return indirect_calls; }
  unsigned get_direct_calls() { return direct_calls; }
};

// This pass creates an instance of the CountIndirectCalls_Walker, invokes
// it on the current procedure and prints out the results.

class CountCallsWithWalkerPass : public PipelinablePass {
  SuifEnv *_env;

public:
  CountCallsWithWalkerPass(SuifEnv *env, const LString &name): PipelinablePass(env, name){
    _env = env;
  }
  virtual ~CountCallsWithWalkerPass(void)  { }
  Module *clone() const { return(Module*)this;}

  void do_procedure_definition(ProcedureDefinition *proc_def) {

    CountCallsWalker count_calls_walker(_env);
    proc_def->walk(count_calls_walker);
    printf("(Walker) Caller: %s  Direct: %d, Indirect: %d\n", 
	   proc_def->get_procedure_symbol()->get_name().c_str(),
	   count_calls_walker.get_direct_calls(), 
	   count_calls_walker.get_indirect_calls());
  }
};


extern "C" void init_ex(SuifEnv *suif_env) {
  ModuleSubSystem *ms = suif_env->get_module_subsystem();
  ms->register_module(
     new CountCallsWithIteratorPass (suif_env, "count_calls_iteratively"));
  ms->register_module(
     new CountCallsWithWalkerPass (suif_env, "count_calls_with_a_walker"));
}