#include <cassert>
#include <algorithm>

#include "MLEMutator.h"
#include "SynTree/Statement.h"


namespace ControlStruct {
    MLEMutator::~MLEMutator() {
	delete targetTable;
	targetTable = 0;
    }

    CompoundStmt* MLEMutator::mutate( CompoundStmt *cmpndStmt ) {
	bool labeledBlock = false;
	if ( !((cmpndStmt->get_labels()).empty()) ) {
	    labeledBlock = true;
	    enclosingBlocks.push_back( Entry( cmpndStmt ) );
	}

	std::list< Statement * > &kids = cmpndStmt->get_kids();
	for ( std::list< Statement * >::iterator k = kids.begin(); k != kids.end(); k++ ) {
	    *k = (*k)->acceptMutator(*this);

	    if ( !get_breakLabel().empty() ) {
		std::list< Statement * >::iterator next = k; next++;
		if ( next == kids.end() ) {
		    std::list<Label> ls; ls.push_back( get_breakLabel() );
		    kids.push_back( new NullStmt( ls ) );
		} else
		    (*next)->get_labels().push_back( get_breakLabel() );

		set_breakLabel("");
	    } // if
	} // for

	if ( labeledBlock ) {
	    assert( ! enclosingBlocks.empty() );
	    if ( ! enclosingBlocks.back().get_breakExit().empty() )
		set_breakLabel( enclosingBlocks.back().get_breakExit() );
	    enclosingBlocks.pop_back();
	} // if

	//mutateAll( cmpndStmt->get_kids(), *this );
	return cmpndStmt;
    }

    Statement *MLEMutator::mutate( WhileStmt *whileStmt ) {
	enclosingLoops.push_back( Entry( whileStmt ) );
	whileStmt->set_body ( whileStmt->get_body()->acceptMutator( *this ) );

	Entry &e = enclosingLoops.back();
	assert ( e == whileStmt );
	whileStmt->set_body( mutateLoop( whileStmt->get_body(), e ) );
	enclosingLoops.pop_back();

	return whileStmt;
    }

    Statement *MLEMutator::mutate( ForStmt *forStmt ) {
	enclosingLoops.push_back( Entry( forStmt ) );
	maybeMutate( forStmt->get_body(), *this );

	Entry &e = enclosingLoops.back();
	assert ( e == forStmt );
	forStmt->set_body( mutateLoop( forStmt->get_body(), e ) );
	enclosingLoops.pop_back();

	return forStmt;
    }

    Statement *MLEMutator::mutate( BranchStmt *branchStmt ) throw ( SemanticError ) {
	if ( branchStmt->get_type() == BranchStmt::Goto )
	    return branchStmt;

	// test if continue target is a loop
	if ( branchStmt->get_type() == BranchStmt::Continue && enclosingLoops.empty() )
	    throw SemanticError( "'continue' outside a loop" );

	if ( branchStmt->get_type() == BranchStmt::Break && (enclosingLoops.empty() && enclosingSwitches.empty() && enclosingBlocks.empty() ) )
	    throw SemanticError( "'break' outside a loop or switch" );

	if ( branchStmt->get_target() == "" ) return branchStmt;

	if ( targetTable->find( branchStmt->get_target() ) == targetTable->end() )
	    throw SemanticError("The label defined in the exit loop statement does not exist." );  // shouldn't happen (since that's already checked)

	std::list< Entry >::iterator check;
	if ( ( check = std::find( enclosingLoops.begin(), enclosingLoops.end(), (*targetTable)[branchStmt->get_target()] ) ) == enclosingLoops.end() )
	    // not in loop, checking if in block
	    if ( (check = std::find( enclosingBlocks.begin(), enclosingBlocks.end(), (*targetTable)[branchStmt->get_target()] )) == enclosingBlocks.end() )
		// neither in loop nor in block, checking if in switch/choose
		if ( (check = std::find( enclosingSwitches.begin(), enclosingSwitches.end(), (*targetTable)[branchStmt->get_target()] )) == enclosingSwitches.end() )
		    throw SemanticError("The target specified in the exit loop statement does not correspond to an enclosing loop.");

	if ( enclosingLoops.back() == (*check) )
	    return branchStmt;				// exit the innermost loop (labels unnecessary)

	Label newLabel;
	switch( branchStmt->get_type() ) {
	  case BranchStmt::Break:
	    if ( check->get_breakExit() != "" )
		newLabel = check->get_breakExit();
	    else {
		newLabel = generator->newLabel();
		check->set_breakExit( newLabel );
	    } // if
	    break;
	  case BranchStmt::Continue:
	    if ( check->get_contExit() != "" )
		newLabel = check->get_contExit();
	    else {
		newLabel = generator->newLabel();
		check->set_contExit( newLabel );
	    } // if
	    break;
	  default:
	    return 0;					// shouldn't be here
	} // switch

	return new BranchStmt( std::list<Label>(), newLabel, BranchStmt::Goto );
    }


    Statement *MLEMutator::mutate( SwitchStmt *switchStmt ) {
	Label brkLabel = generator->newLabel();
	enclosingSwitches.push_back( Entry(switchStmt, "", brkLabel) );
	mutateAll( switchStmt->get_branches(), *this ); {
	    // check if this is necessary (if there is a break to this point, otherwise do not generate
	    std::list<Label> temp; temp.push_back( brkLabel );
	    switchStmt->get_branches().push_back( new BranchStmt( temp, Label(""), BranchStmt::Break ) );
	}
	assert ( enclosingSwitches.back() == switchStmt );
	enclosingSwitches.pop_back();
	return switchStmt;
    }

    Statement *MLEMutator::mutate( ChooseStmt *switchStmt ) {
	Label brkLabel = generator->newLabel();
	enclosingSwitches.push_back( Entry(switchStmt,"", brkLabel) );
	mutateAll( switchStmt->get_branches(), *this ); {
	    // check if this is necessary (if there is a break to this point, otherwise do not generate
	    std::list<Label> temp; temp.push_back( brkLabel );
	    switchStmt->get_branches().push_back( new BranchStmt( temp, Label(""), BranchStmt::Break ) );
	}
	assert ( enclosingSwitches.back() == switchStmt );
	enclosingSwitches.pop_back();
	return switchStmt;
    }

    Statement *MLEMutator::mutateLoop( Statement *bodyLoop, Entry &e ) {
	CompoundStmt *newBody;
	if ( ! (newBody = dynamic_cast<CompoundStmt *>( bodyLoop )) ) {
	    newBody = new CompoundStmt( std::list< Label >() );
	    newBody->get_kids().push_back( bodyLoop );
	} // if

	Label endLabel = e.get_contExit();

	if ( e.get_breakExit() != "" ) {
	    if ( endLabel == "" ) endLabel = generator->newLabel();
	    // check for whether this label is used or not, so as to not generate extraneous gotos
	    if (e.breakExitUsed)
		newBody->get_kids().push_back( new BranchStmt( std::list< Label >(), endLabel, BranchStmt::Goto ) );
	    // xxx
	    //std::list< Label > ls; ls.push_back( e.get_breakExit() );
	    set_breakLabel( e.get_breakExit() );
	    //newBody->get_kids().push_back( new BranchStmt( ls, "", BranchStmt::Break ) );
	} // if

	if ( e.get_breakExit() != "" || e.get_contExit() != "" ){
	    if (dynamic_cast< NullStmt *>( newBody->get_kids().back() ))
		newBody->get_kids().back()->get_labels().push_back( endLabel );
	    else {
		std::list < Label > ls; ls.push_back( endLabel );
		newBody->get_kids().push_back( new NullStmt( ls ) );
	    } // if
	} // if

	return newBody;
    }

    //*** Entry's methods
    void MLEMutator::Entry::set_contExit( Label l ) {
	assert ( contExit == "" || contExit == l );
	contExit = l;
    }

    void MLEMutator::Entry::set_breakExit( Label l ) {
	assert ( breakExit == "" || breakExit == l );
	breakExit = l;
    }
} // namespace ControlStruct