#include <list>
#include <vector>
#include <cassert>
#include <algorithm>

#include "MultRet.h"
#include "SynTree/Statement.h"
#include "SynTree/Expression.h"
#include "SynTree/Declaration.h"
#include "SynTree/Type.h"

namespace Tuples {

  MVRMutator::MVRMutator() : newVars( 0 ), newCode( 0 )
  {
  }

  MVRMutator::~MVRMutator()
  {
  }

  int MVRMutator::curVal = 0;

  Statement *MVRMutator::mutate( ExprStmt *exprStmt ){
    MVRMutator toplevel;
    exprStmt->set_expr( maybeMutate( exprStmt->get_expr(), toplevel ) );

    if ( toplevel.hasCode() ) {
      assert ( toplevel.getVars() != 0 );

      typedef std::list<Statement *> Statements;

      CompoundStmt *code = new CompoundStmt( std::list< Label >() );

      // copy variables
      Statements &vars = toplevel.getVars()->get_kids();
      for( Statements::iterator i = vars.begin(); i != vars.end(); i++ )
	code->get_kids().push_back( *i );

      // copy statements
      Statements &block = toplevel.getCode()->get_kids();
      for( Statements::iterator i = block.begin(); i != block.end(); i++ )
	code->get_kids().push_back( *i );

      return code;
    } else
      return exprStmt;
  }


  Expression *MVRMutator::mutate( ApplicationExpr *appExpr )
  {
    // appExpr->set_function(  maybeMutate( appExpr->get_function(), *this ) );
    bool mulretp = false;
    VariableExpr *funname;
    if ( (funname = dynamic_cast<VariableExpr *>(appExpr->get_function())) == 0 ) goto DoArgs;

    FunctionDecl *fundecl;
    if ((fundecl = dynamic_cast<FunctionDecl *>(funname->get_var())) == 0) goto DoArgs;

    {
      typedef std::list<DeclarationWithType*> RetType;

      RetType &rets = fundecl->get_functionType()->get_returnVals();
      if ( rets.size() <= 1 ) goto DoArgs;
      mulretp = true;

      if( newVars == 0 )
	newVars = new CompoundStmt( std::list<Label>(0) );

      for (RetType::iterator i = rets.begin() ; i != rets.end(); i++) {
	DeclStmt *arg = newVar( *i );
	newVars->get_kids().push_back( arg );
	add_pending( arg->get_decl() );
      }
    }

  DoArgs:
     // mutate the argument list
    typedef std::list< Expression *> Exprs;
    Exprs &args = appExpr->get_args();
    std::list< Expression * > newArgs;
    for( Exprs::iterator i = args.begin(); i != args.end(); i++ ) {
      MVRMutator next;
      Expression *mutated = (*i)->acceptMutator( next );

      if ( next.hasCode() ) {
	// merge new vars and bodies
	typedef std::list< Statement * > Stmts;
	Stmts &vars = next.getVars()->get_kids();
	Stmts &block = next.getCode()->get_kids();

	if (newVars == 0)
	  newVars = new CompoundStmt( std::list< Label >() );
	for( Stmts::iterator i = vars.begin(); i != vars.end(); i++ )  // std::splice? -- need to append lists
	  newVars->get_kids().push_back( *i );

	if (newCode == 0)
	  newCode = new CompoundStmt( std::list< Label >() );
	for( Stmts::iterator i = block.begin(); i != block.end(); i++ )
	  newCode->get_kids().push_back( *i );

      }

      if ( next.hasResults() ) {
	Exprs &res = next.get_results();
	for( Exprs::iterator i = res.begin(); i != res.end(); i++ )
	  newArgs.push_back( *i );
      } else
	newArgs.push_back( mutated );
    }

    appExpr->get_args() = newArgs;  // new argument list


    if ( mulretp ) {
      // add 'out' parameters
      if ( ! argsToAdd.empty() )
	for(std::list< Expression *>::iterator i = argsToAdd.begin(); i != argsToAdd.end(); i++)
	  (appExpr->get_args()).push_back( *i );
      // clear 'out' parameters ( so that the list can be reused -- substitute by auto_ptr later? )

      if (newCode == 0)
	newCode = new CompoundStmt( std::list<Label>(0) );
    }

    // add to block of code
    if ( newCode != 0 )
      newCode->get_kids().push_back( new ExprStmt( std::list<Label>(), appExpr ) );

    return appExpr;
  }

  // Auxiliary function to generate new names for the `output' parameters
  DeclStmt *MVRMutator::newVar( DeclarationWithType *reqDecl ) {
    // std::ostrstream os;
    // os << "__" << curVal++ << "__";// << std::ends;
    // os.freeze( false );

    ObjectDecl *decl;
    if ((decl = dynamic_cast<ObjectDecl *>( reqDecl )) != 0)
      // return new DeclStmt( new ObjectDecl( std::string (os.str(), os.pcount()), );
      return new DeclStmt( std::list<Label>(), decl );
    else
      return 0;
  }

  void MVRMutator::add_pending( Declaration *decl ) {
    ObjectDecl *obj;
    if ( (obj = dynamic_cast< ObjectDecl * >( decl )) == 0 ) return;

    VariableExpr *var = new VariableExpr(obj, 0 );
    results.push_back( var ); // probably change this name to newResults or something
    argsToAdd.push_back( new AddressExpr( var ) );
    return;
  }
}