#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 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();
		    } // if
		} // for
		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_mangleName( namer.newName( (*param)->get_mangleName() ) );
			    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 ) );
			} // if
			modVars.erase( var );
		    } // if
		} // for
	    } // if
	} // if
    }

    // this test is insufficient because it is possible for values to be modified by being passed to other polymorphic
    // routines (e.g., assignment operators) without having their addresses explicitly taken. Some thought is needed to
    // make sure that all of the correct cases are identified where copies are necessary.
    //
    // As a temporary measure, for correctness at the expense of performance, ignore the modVars list entirely and copy
    // every parameter of TypeInstType* when visiting the FunctionDecl.
    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() );
	    } // if
	} // if
    }
} // namespace GenPoly