#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 ( labeledBlock ) {
      assert( ! enclosingBlocks.empty() );
      if( ! enclosingBlocks.back().get_breakExit().empty() )
	set_breakLabel( enclosingBlocks.back().get_breakExit() );
      enclosingBlocks.pop_back();
    }

    //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 switch/choose
      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 not necessary)

    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 ); }
      break;
    case BranchStmt::Continue:
      if ( check->get_contExit() != "" )
	newLabel = check->get_contExit();
      else
	{ newLabel = generator->newLabel(); check->set_contExit( newLabel ); }
      break;

    default:
      // shouldn't be here
      return 0;
    }

    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 );
    }

    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 ( 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 ) );
      }
    }

    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