#include "symbol_checking.hh"
#include "type_checking.hh"
#include "visitor.hh"
#include "options.hh"

using namespace std;
using namespace ast;
using namespace symbol_checking;
using namespace type_checking;

// global scopes
static symbol_map_t globals_map;
static symbol_map_t parameters_map;
static symbol_map_t locals_map;

// exceptions
void
Redefinition::print_error_msg (std::ostream &os)
{
    std::cerr << "Symbol " << _symbol
	      << " redefining previously declared variable."
	      << std::endl;
}

// symbol definitions
namespace symbol_checking {

    class SymbolDefinitionVisitor : public Visitor {
	SymbolTable &_scope;

	void define_globals (DProg &dprog);
	void define_parameters (DProg &dprog);
	void define_locals (DProg &dprog);

    public:
	SymbolDefinitionVisitor(SymbolTable &scope)
	    : _scope(scope)
	{};

	virtual void visit (RExpr 	&rexpr);

	virtual void visit (WhereExpr 	&where_expr);

	virtual void visit (Range 	&range);
	virtual void visit (Values 	&values);
	virtual void visit (SimpleFun 	&simple_fun);

	virtual void visit (FunCallExpr &fun_call_expr);
	virtual void visit (MatrixExpr  &matrix_expr);
	virtual void visit (BinOpExpr   &binop_expr);


	virtual void visit (DProg 	&dprog);
    };

    class GlobalsDefine : public Visitor {
	SymbolTable &_scope;

    public:
	GlobalsDefine(SymbolTable &scope)
	    : _scope(scope)
	{};

	virtual void visit (ValDeclaration 	&val_def);
	virtual void visit (FunDeclaration 	&fun_def);
	virtual void visit (MatrixDeclaration 	&matrix_def);
    };

    class ParametersDefine : public Visitor {
	SymbolTable &_scope;

    public:
	ParametersDefine(SymbolTable &scope)
	    : _scope(scope)
	{};

	virtual void visit (ValDeclaration 	&val_def);
	virtual void visit (FunDeclaration 	&fun_def);
	virtual void visit (MatrixDeclaration 	&matrix_def);

	virtual void visit (RExpr     		&rexpr);
	virtual void visit (IDExpr              &id_expr);
	virtual void visit (Update		&update);
    };

    class LocalsDefine : public Visitor {
	SymbolTable &_scope;

    public:
	LocalsDefine(SymbolTable &scope)
	    : _scope(scope)
	{};

	virtual void visit (ValDeclaration 	&val_def);
	virtual void visit (FunDeclaration 	&fun_def);
	virtual void visit (MatrixDeclaration 	&matrix_def);
    };

};

void
SymbolDefinitionVisitor::define_globals (DProg &dprog)
{
    GlobalsDefine gdef(_scope);

    const list<Declaration*> *globals = dprog.globals();
    list<Declaration*>::const_iterator g_itr;
    for (g_itr = globals->begin(); g_itr != globals->end(); ++g_itr)
	(*g_itr)->accept(gdef);
}

void
SymbolDefinitionVisitor::define_parameters (DProg &dprog)
{
    ParametersDefine pdef(_scope);

    const list<Declaration*> *parameters = dprog.parameters();
    list<Declaration*>::const_iterator p_itr;
    for (p_itr = parameters->begin(); p_itr != parameters->end(); ++p_itr)
	(*p_itr)->accept(pdef);

    // matrices being updated are assumed parameters (unless
    // explicitly declared)
    const list<Update*> *update_list = dprog.update_list();
    list<Update*>::const_iterator u_itr;
    for (u_itr = update_list->begin(); u_itr != update_list->end(); ++u_itr)
	(*u_itr)->accept(pdef);

    // all the outer 'where' expression bounds should be parameters
    const list<RExpr*> *rexpr_list = dprog.rexpr_list();
    list<RExpr*>::const_iterator re_itr;
    for (re_itr = rexpr_list->begin(); re_itr != rexpr_list->end(); ++re_itr)
	(*re_itr)->accept(pdef);
}

void
SymbolDefinitionVisitor::define_locals (DProg &dprog)
{
    LocalsDefine ldef(_scope);

    const list<Declaration*> *locals = dprog.locals();
    list<Declaration*>::const_iterator itr;
    for (itr = locals->begin(); itr != locals->end(); ++itr)
	(*itr)->accept(ldef);
}

void
SymbolDefinitionVisitor::visit (RExpr &rexpr)
{
    if (_scope.lookup_this_scope(rexpr.index()) != 0)
	throw new Redefinition(rexpr, rexpr.index());
    _scope.define_symbol(rexpr.index(), new IndexType(rexpr.index()));
}

void
SymbolDefinitionVisitor::visit (WhereExpr &where_expr)
{
    const list<RExpr*> *rexpr_list = where_expr.rexpr_list();
    list<RExpr*>::const_iterator re_itr;
    for (re_itr = rexpr_list->begin(); re_itr != rexpr_list->end(); ++re_itr)
	(*re_itr)->accept(*this);
}

void
SymbolDefinitionVisitor::visit (Range &range)
{
    // new indices defined in where_expr in range
    range.where_expr()->accept(*this);
}

void
SymbolDefinitionVisitor::visit (Values &values)
{
    // nothing new defined in Values
}

void
SymbolDefinitionVisitor::visit (MatrixExpr &matrix_expr)
{
    // matrix and functions are implicitly defined unless strict_decl
    if (options::strict_decl) return;

    if (_scope.lookup(matrix_expr.id()) == 0)
	{
	    parameters_map[matrix_expr.id()] =
		new MatrixType(matrix_expr.id(),
			       matrix_expr.expr_list()->size());
	}
}

void
SymbolDefinitionVisitor::visit (FunCallExpr &fun_call_expr)
{
    if (_scope.lookup(fun_call_expr.id()) == 0)
	{
	    globals_map[fun_call_expr.id()] =
		new FunctionType(fun_call_expr.id(),
				 fun_call_expr.expr_list()->size());
	}
}

void
SymbolDefinitionVisitor::visit (BinOpExpr &binop_expr)
{
    binop_expr.left()->accept(*this);
    binop_expr.right()->accept(*this);
}

void
SymbolDefinitionVisitor::visit (SimpleFun &simple_fun)
{
    // traverse to get implicit definitions unless strict
    if (options::strict_decl) return;
    simple_fun.expr()->accept(*this);
}

void
SymbolDefinitionVisitor::visit (DProg &dprog)
{
    define_globals(dprog);
    define_locals(dprog);	// we need to define locals before parameters,
                                // since parameters are also implicitly defined
    define_parameters(dprog);

    const list<RExpr*> *rexpr_list = dprog.rexpr_list();
    list<RExpr*>::const_iterator re_itr;
    for (re_itr = rexpr_list->begin(); re_itr != rexpr_list->end(); ++re_itr)
	(*re_itr)->accept(*this);
}


void
GlobalsDefine::visit (ValDeclaration &val_def)
{
    if (_scope.lookup(val_def.name()) != 0)
	throw new Redefinition(val_def, val_def.name());
    globals_map[val_def.name()] =
	new ValType(val_def.name(), val_def.type());
}

void
GlobalsDefine::visit (FunDeclaration &fun_def)
{
    if (_scope.lookup(fun_def.name()) != 0)
	throw new Redefinition(fun_def, fun_def.name());
    globals_map[fun_def.name()] =
	new FunctionType(fun_def.name(), fun_def.parameter_types(),
			 fun_def.return_type());
}

void
GlobalsDefine::visit (MatrixDeclaration &matrix_def)
{
    if (_scope.lookup(matrix_def.name()) != 0)
	throw new Redefinition(matrix_def, matrix_def.name());
    globals_map[matrix_def.name()] =
	new MatrixType(matrix_def.name(), matrix_def.dim_list(),
		       matrix_def.type());
}




void
ParametersDefine::visit (ValDeclaration &val_def)
{
    if (_scope.lookup(val_def.name()) != 0)
	throw new Redefinition(val_def, val_def.name());
    parameters_map[val_def.name()] =
	new ValType(val_def.name(), val_def.type());
}

void
ParametersDefine::visit (FunDeclaration &fun_def)
{
    std::cerr << "Parameter functions is currently not supported.\n"
	      << "Sorry.\n"
	      << std::endl;
    exit(4);
}

void
ParametersDefine::visit (MatrixDeclaration &matrix_def)
{
    if (_scope.lookup(matrix_def.name()) != 0)
	throw new Redefinition(matrix_def, matrix_def.name());
    parameters_map[matrix_def.name()] =
	new MatrixType(matrix_def.name(), matrix_def.dim_list(),
		       matrix_def.type());
}


void
ParametersDefine::visit (RExpr &rexpr)
{
    rexpr.begin()->accept(*this);
    rexpr.end()->accept(*this);
}

void
ParametersDefine::visit (IDExpr &id_expr)
{
    // implicit definition if non-strict
    if (options::strict_decl) return;
    if (_scope.lookup(id_expr.id()) == 0)
	{
	    parameters_map[id_expr.id()]
		= new ValType(id_expr.id(), TYPE_INT);
	}
}

void
ParametersDefine::visit (Update &update)
{
    // implicit definition if non-strict
    if (options::strict_upd_decl) return;
    if (_scope.lookup(update.name()) == 0)
	{
	    parameters_map[update.name()] =
		new MatrixType(update.name(),
			       update.indices()->size(),
			       update.type());
	}
}


void
LocalsDefine::visit (ValDeclaration &val_def)
{
    std::cerr << "Local values are currently not supported.\n"
	      << "Sorry.\n"
	      << std::endl;
    exit(4);
}

void
LocalsDefine::visit (FunDeclaration &fun_def)
{
    std::cerr << "Local functions are currently not supported.\n"
	      << "Sorry.\n"
	      << std::endl;
    exit(4);
}

void
LocalsDefine::visit (MatrixDeclaration &matrix_def)
{
    if (_scope.lookup(matrix_def.name()) != 0)
	throw new Redefinition(matrix_def, matrix_def.name());
    locals_map[matrix_def.name()] =
	new MatrixType(matrix_def.name(), matrix_def.dim_list(),
		       matrix_def.type());
}



// symbol table
SymbolTable::SymbolTable(Ast &ast, SymbolTable *outer_scope)
    : _outer_scope(outer_scope)
{
    SymbolDefinitionVisitor sdv(*this);
    ast.accept(sdv);
}

const tinfo_t*
SymbolTable::lookup_this_scope (const char *symbol) const
{
    symbol_map_t::const_iterator i = _this_scope.find(symbol);
    if (i != _this_scope.end()) return i->second;
    else                        return 0;
}

const tinfo_t*
SymbolTable::lookup_local_scopes (const char *symbol) const
{
    const tinfo_t *tinfo = lookup_this_scope(symbol);
    if (tinfo != 0)             return tinfo;
    else if (_outer_scope != 0) return _outer_scope->lookup(symbol);
    else                        return 0;
}

const tinfo_t*
SymbolTable::lookup_global_scopes (const char *symbol)
{
    symbol_map_t::const_iterator i = globals_map.find(symbol);
    if (i != globals_map.end()) return i->second;

    i = parameters_map.find(symbol);
    if (i != parameters_map.end()) return i->second;

    i = locals_map.find(symbol);
    if (i != locals_map.end()) return i->second;

    return 0;
}

const tinfo_t*
SymbolTable::lookup (const char *symbol) const
{
    const tinfo_t *tinfo = lookup_local_scopes(symbol);
    if (tinfo != 0) return tinfo;
    else            return lookup_global_scopes(symbol);
}

void
SymbolTable::define_symbol (const char *symbol, tinfo_t *tinfo)
{
    _this_scope[symbol] = tinfo;
}



// access to global scopes.
set<tinfo_t*> *
symbol_checking::globals ()
{
    static set<tinfo_t*> *symbols = 0;
    if (!symbols)
	{
	    symbols = new set<tinfo_t*>;
	    symbol_map_t::const_iterator i;
	    for (i = globals_map.begin(); i != globals_map.end(); ++i)
		symbols->insert(i->second);
	}
    return symbols;
}

set<tinfo_t*> *
symbol_checking::parameters ()
{
    static set<tinfo_t*> *symbols = 0;
    if (!symbols)
	{
	    symbols = new set<tinfo_t*>;
	    symbol_map_t::const_iterator i;
	    for (i = parameters_map.begin(); i != parameters_map.end(); ++i)
		symbols->insert(i->second);
	}
    return symbols;
}

set<tinfo_t*> *
symbol_checking::locals ()
{
    static set<tinfo_t*> *symbols = 0;
    if (!symbols)
	{
	    symbols = new set<tinfo_t*>;
	    symbol_map_t::const_iterator i;
	    for (i = locals_map.begin(); i != locals_map.end(); ++i)
		symbols->insert(i->second);
	}
    return symbols;
}