/*
 * This file is part of the Cforall project
 *
 * $Id: TypeSubstitution.h,v 1.9 2005/08/29 20:59:26 rcbilson Exp $
 *
 */

#ifndef SYNTREE_TYPESUBSTITUTION_H
#define SYNTREE_TYPESUBSTITUTION_H

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

#include "SynTree/Mutator.h"
#include "SynTree/Declaration.h"
#include "SynTree/Expression.h"


class TypeSubstitution : public Mutator
{
    typedef Mutator Parent;
    
public:
    TypeSubstitution();
    template< typename FormalIterator, typename ActualIterator >
    TypeSubstitution( FormalIterator formalBegin, FormalIterator formalEnd, ActualIterator actualBegin );
    TypeSubstitution( const TypeSubstitution &other );
    virtual ~TypeSubstitution();
    
    TypeSubstitution &operator=( const TypeSubstitution &other );
    
    template< typename SynTreeClass > int apply( SynTreeClass *&input );
    template< typename SynTreeClass > int applyFree( SynTreeClass *&input );
    
    void add( std::string formalType, Type *actualType );
    void add( const TypeSubstitution &other );
    void remove( std::string formalType );
    Type *lookup( std::string formalType ) const;
    bool empty() const;
    
    template< typename FormalIterator, typename ActualIterator >
    void add( FormalIterator formalBegin, FormalIterator formalEnd, ActualIterator actualBegin );
    
    template< typename TypeInstListIterator >
    void extract( TypeInstListIterator begin, TypeInstListIterator end, TypeSubstitution &result );
    
    void normalize();

    void print( std::ostream &os, int indent = 0 ) const;
    TypeSubstitution *clone() const { return new TypeSubstitution( *this ); }
    
private:
    virtual Type* mutate(TypeInstType *aggregateUseType);
    virtual Expression* mutate(NameExpr *nameExpr);
    
    template< typename TypeClass > Type *handleType( TypeClass *type );
    
    virtual Type* mutate(VoidType *basicType);
    virtual Type* mutate(BasicType *basicType);
    virtual Type* mutate(PointerType *pointerType);
    virtual Type* mutate(ArrayType *arrayType);
    virtual Type* mutate(FunctionType *functionType);
    virtual Type* mutate(StructInstType *aggregateUseType);
    virtual Type* mutate(UnionInstType *aggregateUseType);
    virtual Type* mutate(EnumInstType *aggregateUseType);
    virtual Type* mutate(ContextInstType *aggregateUseType);
    virtual Type* mutate(TupleType *tupleType);
    
    // TODO: worry about traversing into a forall-qualified function type or type decl with assertions
    
    void initialize( const TypeSubstitution &src, TypeSubstitution &dest );

    typedef std::map< std::string, Type* > TypeEnvType;
    typedef std::map< std::string, Expression* > VarEnvType;
    typedef std::set< std::string > BoundVarsType;
    TypeEnvType typeEnv;
    VarEnvType varEnv;
    BoundVarsType boundVars;
    int subCount;
    bool freeOnly;
};

template< typename FormalIterator, typename ActualIterator >
void 
TypeSubstitution::add( FormalIterator formalBegin, FormalIterator formalEnd, ActualIterator actualBegin )
{
    // FormalIterator points to a TypeDecl
    // ActualIterator points to a Type
    FormalIterator formalIt = formalBegin;
    ActualIterator actualIt = actualBegin;
    for( ; formalIt != formalEnd; ++formalIt, ++actualIt ) {
	if( TypeDecl *formal = dynamic_cast< TypeDecl* >( *formalIt ) ) {
	    if( TypeExpr *actual = dynamic_cast< TypeExpr* >( *actualIt ) ) {
		if( formal->get_name() != "" ) {
		    TypeEnvType::iterator i = typeEnv.find( formal->get_name() );
		    if( i != typeEnv.end() ) {
			delete i->second;
		    }
		    typeEnv[ formal->get_name() ] = actual->get_type()->clone();
		}
	    } else {
		throw SemanticError( "Attempt to provide non-type parameter for type parameter", formal );
	    }
	} else {
	    // TODO: type check the formal and actual parameters
	    if( (*formalIt)->get_name() != "" ) {
		varEnv[ (*formalIt)->get_name() ] = (*actualIt)->clone();
	    }
	}
    }
}

template< typename FormalIterator, typename ActualIterator >
TypeSubstitution::TypeSubstitution( FormalIterator formalBegin, FormalIterator formalEnd, ActualIterator actualBegin )
{
    add( formalBegin, formalEnd, actualBegin );
}

template< typename SynTreeClass >
int
TypeSubstitution::apply( SynTreeClass *&input )
{
    assert( input );
    subCount = 0;
    freeOnly = false;
    input = dynamic_cast< SynTreeClass *>( input->acceptMutator( *this ) );
    assert( input );
///	std::cout << "substitution result is: ";
///	newType->print( std::cout );
///	std::cout << std::endl;
    return subCount;
}
    
template< typename SynTreeClass >
int
TypeSubstitution::applyFree( SynTreeClass *&input )
{
    assert( input );
    subCount = 0;
    freeOnly = true;
    input = dynamic_cast< SynTreeClass *>( input->acceptMutator( *this ) );
    assert( input );
///	std::cout << "substitution result is: ";
///	newType->print( std::cout );
///	std::cout << std::endl;
    return subCount;
}
    
template< typename TypeInstListIterator >
void
TypeSubstitution::extract( TypeInstListIterator begin, TypeInstListIterator end, TypeSubstitution &result )
{
    while( begin != end ) {
	TypeEnvType::iterator cur = typeEnv.find( (*begin++)->get_name() );
	if( cur != typeEnv.end() ) {
	    result.typeEnv[ cur->first ] = cur->second;
	    typeEnv.erase( cur );
	}
    }
}

// helper function
template< typename FormalIterator, typename ActualIterator, typename MemberIterator, typename OutputIterator >
void
applySubstitution( FormalIterator formalBegin, FormalIterator formalEnd, ActualIterator actual, MemberIterator memberBegin, MemberIterator memberEnd, OutputIterator out )
{
    // Instantiate each member of the context given the actual parameters specified, and store the
    // instantiations for use by the indexer

    TypeSubstitution sub = TypeSubstitution( formalBegin, formalEnd, actual );
    for( std::list< Declaration* >::iterator i = memberBegin; i != memberEnd; ++i ) {
	Declaration *newdecl = (*i)->clone();
	sub.apply( newdecl );
	*out++ = newdecl;
    }
}


#endif /* #ifndef SYNTREE_TYPESUBSTITUTION_H */