#include <list>
#include <cassert>

#include "LabelFixer.h"
#include "MLEMutator.h"
#include "SynTree/Statement.h"
#include "SynTree/Declaration.h"
#include "utility.h"

namespace ControlStruct {
    LabelFixer::Entry::Entry( Statement *to, Statement *from ) : definition ( to ) {
	if ( from != 0 )
	    usage.push_back( from );
    }

    bool LabelFixer::Entry::insideLoop() {
	return ( dynamic_cast< ForStmt * > ( definition ) ||
		 dynamic_cast< WhileStmt * > ( definition )  );
    }

    LabelFixer::LabelFixer( LabelGenerator *gen ) : generator ( gen ) {
	if ( generator == 0 )
	    generator = LabelGenerator::getGenerator();
    }

    void LabelFixer::visit( FunctionDecl *functionDecl ) {
	if ( functionDecl->get_statements() != 0 )
	    functionDecl->get_statements()->accept( *this );

	MLEMutator mlemut( resolveJumps(), generator );
	functionDecl->acceptMutator( mlemut );
    }

    void LabelFixer::visit( Statement *stmt ) {
	std::list< Label > &labels = stmt->get_labels();

	if ( ! labels.empty() ) {
	    Label current = setLabelsDef( labels, stmt );
	    labels.clear();
	    labels.push_front( current );
	} // if
    }

    void LabelFixer::visit( BranchStmt *branchStmt ) {
	visit ( ( Statement * )branchStmt );  // the labels this statement might have

	Label target;
	if ( (target = branchStmt->get_target()) != "" ) {
	    setLabelsUsg( target, branchStmt );
	} //else       /* computed goto or normal exit-loop statements */
    }

    Label LabelFixer::setLabelsDef( std::list< Label > &llabel, Statement *definition ) {
	assert( definition != 0 );
	Entry *entry = new Entry( definition );
	bool used = false;

	for ( std::list< Label >::iterator i = llabel.begin(); i != llabel.end(); i++ )
	    if ( labelTable.find( *i ) == labelTable.end() )
		{ used = true; labelTable[ *i ] = entry; } // undefined and unused
	    else
		if ( labelTable[ *i ]->defined() )
		    throw SemanticError( "Duplicate definition of label: " + *i );
		else
		    labelTable[ *i ]->set_definition( definition );

	if (! used ) delete entry;

	return labelTable[ llabel.front() ]->get_label();  // this *must* exist
    }

    Label LabelFixer::setLabelsUsg( Label orgValue, Statement *use ) {
	assert( use != 0 );

	if ( labelTable.find( orgValue ) != labelTable.end() )
	    labelTable[ orgValue ]->add_use( use );         // the label has been defined or used before
	else
	    labelTable[ orgValue ] = new Entry( 0, use );

	return labelTable[ orgValue ]->get_label();
    }

    std::map<Label, Statement * > *LabelFixer::resolveJumps() throw ( SemanticError ) {
	std::map< Statement *, Entry * > def_us;

	for ( std::map< Label, Entry *>::iterator i = labelTable.begin(); i != labelTable.end(); i++ ) {
	    Entry *e = i->second;

	    if ( def_us.find ( e->get_definition() ) == def_us.end() )
		def_us[ e->get_definition() ] = e;
	    else
		if ( e->used() )
		    def_us[ e->get_definition() ]->add_uses( e->get_uses() );
	}

	// get rid of labelTable
	for ( std::map< Statement *, Entry * >::iterator i = def_us.begin(); i != def_us.end(); i++ ) {
	    Statement *to = (*i).first;
	    std::list< Statement *> &from = (*i).second->get_uses();
	    Label finalLabel = generator->newLabel();
	    (*i).second->set_label( finalLabel );

	    if ( to == 0 ) {
		BranchStmt *first_use = dynamic_cast<BranchStmt *>(from.back());
		Label undef("");
		if ( first_use != 0 )
		    undef = first_use->get_target();
		throw SemanticError ( "'" + undef + "' label not defined");
	    }

	    to->get_labels().clear();
	    to->get_labels().push_back( finalLabel );

	    for ( std::list< Statement *>::iterator j = from.begin(); j != from.end(); j++ ) {
		BranchStmt *jumpTo = dynamic_cast< BranchStmt * > ( *j );
		assert( jumpTo != 0 );
		jumpTo->set_target( finalLabel );
	    } // for
	} // for

	// reverse table
	std::map< Label, Statement * > *ret = new std::map< Label, Statement * >();
	for ( std::map< Statement *, Entry * >::iterator i = def_us.begin(); i != def_us.end(); i++ ) 
	    (*ret)[ (*i).second->get_label() ] = (*i).first;

	return ret;
    }
}  // namespace ControlStruct