/*
 * This file is part of the Cforall project
 *
 * $Id: CopyParams.cc,v 1.3 2005/08/29 20:14:13 rcbilson Exp $
 *
 */

#include <set>
#include <map>
#include <cassert>

#include "SynTree/Declaration.h"
#include "SynTree/Type.h"
#include "SynTree/Expression.h"
#include "SynTree/Statement.h"
#include "SynTree/Visitor.h"
#include "UniqueName.h"


namespace GenPoly {

class CopyParams : public Visitor
{
public:
  CopyParams();
  
  virtual void visit( FunctionDecl *funcDecl );
  virtual void visit( AddressExpr *addrExpr );

private:
  std::set< UniqueId > modVars;
  UniqueName namer;
};

void
copyParams( std::list< Declaration* > &translationUnit )
{
  CopyParams copier;
  acceptAll( translationUnit, copier );
}

CopyParams::CopyParams()
  : namer( "_cp" )
{
}

static const std::list< Label > noLabels;

void 
CopyParams::visit( FunctionDecl *funcDecl )
{
  if( funcDecl->get_statements() ) {
    funcDecl->get_statements()->accept( *this );
    
    if( !modVars.empty() ) {
      std::map< std::string, DeclarationWithType* > assignOps;
      // assume that the assignment operator is the first assert param after any "type" parameter
      for( std::list< TypeDecl* >::const_iterator tyVar = funcDecl->get_functionType()->get_forall().begin(); tyVar != funcDecl->get_functionType()->get_forall().end(); ++tyVar ) {
        if( (*tyVar)->get_kind() == TypeDecl::Any ) {
          assert( !(*tyVar)->get_assertions().empty() );
          assignOps[ (*tyVar)->get_name() ] = (*tyVar)->get_assertions().front();
        }
      }
      for( std::list< DeclarationWithType* >::iterator param = funcDecl->get_functionType()->get_parameters().begin(); param != funcDecl->get_functionType()->get_parameters().end(); ++param ) {
        std::set< UniqueId >::const_iterator var = modVars.find( (*param)->get_uniqueId() );
        if( var != modVars.end() ) {
          TypeInstType *typeInst = dynamic_cast< TypeInstType* >( (*param)->get_type() );
          assert( typeInst );
          std::map< std::string, DeclarationWithType* >::const_iterator assignOp = assignOps.find( typeInst->get_name() );
          if( assignOp != assignOps.end() ) {
            DeclarationWithType *oldParam = *param;
            *param = (*param)->clone();
            (*param)->set_name( namer.newName() );
            ApplicationExpr *assign = new ApplicationExpr( new VariableExpr( assignOp->second ) );
            assign->get_args().push_back( new VariableExpr( oldParam ) );
            assign->get_args().push_back( new VariableExpr( *param ) );
            funcDecl->get_statements()->get_kids().push_front( new ExprStmt( noLabels, assign ) );
            funcDecl->get_statements()->get_kids().push_front( new DeclStmt( noLabels, oldParam ) );
          }
          modVars.erase( var );
        }
      }
    }
  }
}

void 
CopyParams::visit( AddressExpr *addrExpr )
{
  if( VariableExpr *varExpr = dynamic_cast< VariableExpr* >( addrExpr->get_arg() ) ) {
    if( dynamic_cast< TypeInstType* >( varExpr->get_var()->get_type() ) ) {
      modVars.insert( varExpr->get_var()->get_uniqueId() );
    }
  }
}

} // namespace GenPoly