/* Oort
 * Copyright 2007, Soren Sandmann Pedersen
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.
 */

#include "ast.h"

#if 0
#define TRUE (1)
#define FALSE ({ g_print ("return at %d\n", __LINE__); 0; })
#endif

/*
 * Various helpers
 */
static ast_type_spec_t *
new_int32 (void)
{
    return ast_type_spec_new (AST_INT32_TYPE);
}

static ast_type_spec_t *
new_double (void)
{
    return ast_type_spec_new (AST_DOUBLE_TYPE);
}

static ast_type_spec_t *
new_null (void)
{
    return ast_type_spec_new (AST_NULL_TYPE);
}

static ast_type_spec_t *
new_bool (void)
{
    return ast_type_spec_new (AST_BOOL_TYPE);
}

static ast_type_spec_t *
new_string (void)
{
    return ast_type_spec_new (AST_STRING_TYPE);
}

static gboolean
is_bool (ast_type_spec_t *type)
{
    return type->common.type == AST_BOOL_TYPE;
}

static gboolean
is_null (ast_type_spec_t *type)
{
    return type->common.type == AST_NULL_TYPE;
}

static gboolean
is_int32 (ast_type_spec_t *type)
{
    return type->common.type == AST_INT32_TYPE;
}

static gboolean
is_void (ast_type_spec_t *type)
{
    return type->common.type == AST_VOID_TYPE;
}

static gboolean
is_function (ast_type_spec_t *type)
{
    return type->common.type == AST_FUNCTION_TYPE;
}

static gboolean
is_double (ast_type_spec_t *type)
{
    return type->common.type == AST_DOUBLE_TYPE;
}

static gboolean
is_numeric (ast_type_spec_t *type)
{
    if (is_int32 (type) || is_double (type))
	return TRUE;

    return FALSE;
}

static ast_type_spec_t *
promote_numeric (const ast_type_spec_t *left,
		 const ast_type_spec_t *right)
{
    ast_type_spec_type_t new_type;

    if (left->common.type == AST_DOUBLE_TYPE	||
	right->common.type == AST_DOUBLE_TYPE)
    {
	new_type = AST_DOUBLE_TYPE;
    }
    else
    {
	new_type = AST_INT32_TYPE;
    }

    return ast_type_spec_new (new_type);
}

static ast_type_spec_t *
get_type (ast_expression_t *expr)
{
    return expr->common.type_spec;
}

static gboolean assignable (ast_type_spec_t *left,
			    ast_type_spec_t *right);

static gboolean
compatible_functions (ast_function_type_spec_t *left,
		      ast_function_type_spec_t *right)
{
    int n_left = array_length ((gpointer *)left->args);
    int n_right = array_length ((gpointer *)right->args);
    int i;

    if (n_left != n_right)
    {
	return report_error ("incompatible functions "
			     "(left: %d args, right: %d args)\n",
			     n_left, n_right);
    }

    for (i = 0; i < n_left; ++i)
    {
	ast_type_spec_t *left_arg = left->args[i];
	ast_type_spec_t *right_arg = right->args[i];

	if (!assignable (right_arg, left_arg))
	    return report_error ("Incompatible argument types\n");
    }

    if (!assignable (left->return_, right->return_)		&&
	!is_void (left->return_))
    {
	return FALSE;
    }

    return TRUE;
}

static gboolean
is_class (ast_type_spec_t *type)
{
    g_return_val_if_fail (type != NULL, FALSE);

    return type->common.type == AST_OBJECT_TYPE;
}

static gboolean
is_array (ast_type_spec_t *type)
{
    return type->common.type == AST_ARRAY_TYPE;
}

static gboolean
is_label (ast_type_spec_t *type)
{
    return type->common.type == AST_LABEL_TYPE;
}

static gboolean
is_string (ast_type_spec_t *type)
{
    return type->common.type == AST_STRING_TYPE;
}

static gboolean
assignable (ast_type_spec_t *left,
	    ast_type_spec_t *right)
{
    g_return_val_if_fail (left != NULL, FALSE);
    g_return_val_if_fail (right != NULL, FALSE);

    if ((is_numeric (left) && is_numeric (right)) ||
	(is_bool (left) && is_bool (right)))
    {
	return TRUE;
    }

    if ((is_label (left) && is_label (right)))
	return TRUE;

    if ((is_label (left) && is_null (right)))
	return TRUE;
    
    if ((is_function (left) && is_function (right)))
    {
	return compatible_functions (&left->function, &right->function);
    }

    if (is_class (left) && is_class (right))
    {
	/* FIXME, check that right is a subclass of left */
	return TRUE;
    }

    if (is_array (left) && is_array (right))
    {
	/* FIXME: figure out what to do about covariance (and genericity
	 * in general).
	 */
	return
	    assignable (left->array.child_type, right->array.child_type) &&
	    assignable (right->array.child_type, left->array.child_type);
    }

    if (is_class (left) && is_null (right))
	return TRUE;

    if (is_array (left) && is_null (right))
	return TRUE;

    if (is_string (left) && is_string (right))
	return TRUE;
    
    return FALSE;
}

/*
 * First pass to gather the types of functions
 */
static void
type_check_pass1 (ast_t *ast)
{
    GList *list;
    ast_function_definition_t *function = NULL;

    if (ast_is (ast, AST_DEFINITION, AST_FUNCTION_DEFINITION))
    {
	function = &ast->definition.function;
    }
    else if (ast_is (ast, AST_EXPRESSION, AST_LAMBDA_EXPRESSION))
    {
	function = ast->expression.lambda.function;
    }

    if (function)
    {
	GPtrArray *type_specs;
	ast_type_spec_t **parameter_types;
	int i;

	/* Parameter types */
	type_specs = g_ptr_array_new ();
	for (i = 0; function->parameters[i] != NULL; ++i)
	    g_ptr_array_add (type_specs, function->parameters[i]->type_spec);
	g_ptr_array_add (type_specs, NULL);

	parameter_types = (ast_type_spec_t **)g_ptr_array_free (type_specs, FALSE);

	/* Function type */
	g_assert (function->return_type);
	function->type_spec = ast_type_spec_new_function (
	    function->return_type, parameter_types);
    }

    if (ast_is (ast, AST_DEFINITION, AST_CLASS_DEFINITION))
    {
	ast_type_spec_t *type = ast_type_spec_new_object (&ast->definition.class);

	ast->definition.class.type_spec = type;
    }

    for (list = ast->common.children->head; list; list = list->next)
	type_check_pass1 (list->data);
}

/* Pass 1.5 to resolve identifiers */
static void
type_check_resolve_identifiers (ast_t *ast)
{
    GList *list;

    for (list = ast->common.children->head; list != NULL; list = list->next)
	type_check_resolve_identifiers (list->data);

    while (ast_is (ast, AST_TYPE_SPEC, AST_IDENTIFIER_TYPE))
    {
	ast_identifier_type_spec_t *id = &ast->type_spec.identifier;
	ast_type_spec_t *spec = &ast->type_spec;

	if (id->definition->common.type == AST_CLASS_DEFINITION)
	    *spec = *id->definition->class.type_spec;
	else if (id->definition->common.type == AST_TYPE_DEFINITION)
	    *spec = *id->definition->type.type_spec;
	else
	    g_assert_not_reached();
    }
}

/*
 * Second pass to do the actual checking
 */
static gboolean type_check_pass2 (ast_t *ast);
static gboolean type_check_expression (ast_expression_t *expr);

static gboolean
type_check_binary_expression (ast_binary_expression_t *expr)
{
    ast_type_spec_t *left_type;
    ast_type_spec_t *right_type;
    gboolean check_assign = FALSE;

    /* We must check the right type before the left type, since
     * the type of the left expression may be inferred from
     * the type of the right one
     */
    if (!type_check_expression (expr->right))
	return FALSE;

    if (!type_check_expression (expr->left))
	return FALSE;

    right_type = get_type (expr->right);
    left_type = get_type (expr->left);

    switch (expr->operator)
    {
    case AST_PLUS:
    case AST_MINUS:
    case AST_DIVIDE:
    case AST_TIMES:
	if (!is_numeric (left_type) || !is_numeric (right_type))
	    return report_error ("Wrong type arguments to +, -, *, or /\n");

	expr->common.type_spec = promote_numeric (left_type, right_type);
	break;

    case AST_MOD:
	if (!is_int32 (left_type) || !is_int32 (right_type))
	    return report_error ("Type error to %% expression\n");

	expr->common.type_spec = new_int32();
	break;
    case AST_BOR:
    case AST_BXOR:
    case AST_BAND:
	if (!is_int32 (left_type) || !is_int32 (right_type))
	    return report_error ("Type error in |, ^, or & expression\n");

	expr->common.type_spec = new_int32();
	break;

    case AST_LSHIFT:
    case AST_RSHIFT:
	if (!is_int32 (left_type) && !is_int32 (right_type))
	    return report_error ("Type error in shift expression\n");

	expr->common.type_spec = new_int32();
	break;

    case AST_ASSIGN:
	check_assign = TRUE;
	break;

    case AST_LT:
    case AST_GT:
    case AST_LTE:
    case AST_GTE:
	if ((!is_numeric (left_type) || !is_numeric (right_type)))
	    return report_error ("Type error in relational expression\n");

	expr->common.type_spec = new_bool();
	break;

    case AST_EQUAL:
    case AST_NOT_EQUAL:
	if (!((is_numeric (left_type) && is_numeric (right_type))	||
	      (is_bool (left_type) && is_bool (right_type))		||
	      (is_function (left_type) && is_function (right_type))	||
	      (is_null (left_type) && is_class (right_type))		||
	      (is_class (left_type) && is_null (right_type))		||
	      (is_null (left_type) && is_null (right_type))		||
	      (is_class (left_type) && is_class (right_type))		||
	      (is_array (left_type) && is_null (right_type))		||
	      (is_null (right_type) && is_array (left_type))))
	{
	    return report_error (
		"Type error in == expression\n");
	}

	if (is_function (left_type) || is_function (right_type))
	{
	    /* FIXME: The interpreter actually does support comparing
	     * closures, so we could just remove this error.
	     *
	     * However, it's not clear that just comparing the address
	     * of the functions is right. Strictly speaking, the
	     * *closures* are not necessarily equal just because
	     * the functions are.
	     *
	     * Another possibility would be to also compare the
	     * environment, but most likely most people want is just
	     * comparing the function addresses.
	     */
	    return report_error (
		"Comparing closures is not supported for now\n");
	}

	expr->common.type_spec = new_bool();
	break;

    case AST_OR:
    case AST_AND:
	if (!is_bool (left_type) || !is_bool (right_type))
	    return report_error ("Type error in || or && expression\n");

	expr->common.type_spec = new_bool();
	break;
    }

    if (check_assign)
    {
	if (!assignable (left_type, right_type))
	    return report_error ("Type error in assignment\n");

	expr->common.type_spec = left_type;
    }

    return TRUE;
}

static gboolean
type_check_unary_expression (ast_unary_expression_t *unary)
{
    if (!type_check_expression (unary->expr))
	return FALSE;

    switch (unary->operator)
    {
    case AST_PREFIX_INC:
    case AST_PREFIX_DEC:
    case AST_POSTFIX_DEC:
    case AST_POSTFIX_INC:
	if (!is_int32 (get_type (unary->expr)))
	{
	    return report_error (
		"Argument to inc/dec expression must be an integer\n");
	}

	unary->common.type_spec = new_int32 ();
	break;

    case AST_UNARY_MINUS:
    case AST_UNARY_PLUS:
	if (!is_numeric (unary->expr->common.type_spec))
	    return report_error ("Argument to unary + or - must be numeric\n");

	unary->common.type_spec = get_type (unary->expr);
	break;

    case AST_UNARY_BNEG:
	if (!is_int32 (get_type (unary->expr)))
	    return report_error ("Argument to ~ must be integer\n");

	unary->common.type_spec = get_type (unary->expr);
	break;

    case AST_UNARY_NOT:
	if (!is_bool (get_type (unary->expr)))
	    return report_error ("Argument to ! must be bool\n");

	unary->common.type_spec = new_bool();
	break;
    }

    return TRUE;
}

static gboolean
type_check_arguments (ast_type_spec_t **types,
		      ast_expression_t **arguments)
{
    int n_types = array_length ((gpointer *)types);
    int n_args  = array_length ((gpointer *)arguments);
    int i;

    if (n_types < n_args)
	return report_error ("Too many arguments\n");

    if (n_types > n_args)
	return report_error ("Too few arguments\n");

    for (i = 0; types[i] != NULL; ++i)
    {
	if (!assignable (types[i], arguments[i]->common.type_spec))
	{
	    g_print ("type %d vs %d\n", arguments[i]->common.type_spec->common.type, types[i]->common.type);
	    return report_error ("Type error in argument list\n");
	}
    }

    return TRUE;
}

static gboolean
type_check_expression (ast_expression_t *expr)
{
    ast_type_spec_t *then_type;
    ast_type_spec_t *else_type;
    ast_identifier_expression_t *id_expr;
    ast_type_spec_t *callee_type;
    ast_type_spec_t *left_type;
    ast_definition_t *definition;
    ast_class_definition_t *class;
    int i;

    switch (expr->common.type)
    {
    case AST_INT_LITERAL_EXPRESSION:
	expr->common.type_spec = new_int32();
	break;

    case AST_FLOAT_LITERAL_EXPRESSION:
	expr->common.type_spec = new_double();
	break;
	
    case AST_BOOL_LITERAL_EXPRESSION:
	expr->common.type_spec = new_bool();
	break;

    case AST_STRING_LITERAL_EXPRESSION:
	expr->common.type_spec = new_string();
	break;

    case AST_NULL_EXPRESSION:
	expr->common.type_spec = new_null();
	break;

    case AST_THIS_EXPRESSION:
	class = ast_enclosing_class ((ast_t *)expr);

	if (!class)
	    return report_error ("'this' outside of class\n");

	expr->common.type_spec = ast_type_spec_new_object (class);
	break;

    case AST_IDENTIFIER_EXPRESSION:
	id_expr = &expr->identifier;

	switch (id_expr->definition->common.type)
	{
	case AST_VARIABLE_DEFINITION:
	    expr->common.type_spec = id_expr->definition->variable.type_spec;

	    if (expr->common.type_spec->common.type == AST_INFERRED_TYPE)
	    {
		return report_error (
		    "Type of %s cannot be determined at this point\n",
		    id_expr->definition->variable.name);
	    }
	    break;

	case AST_LABEL_DEFINITION:
	    expr->common.type_spec = ast_type_spec_new (AST_LABEL_TYPE);
	    break;

	case AST_FUNCTION_DEFINITION:
	    expr->common.type_spec = id_expr->definition->function.type_spec;
	    g_assert (expr->common.type_spec);
	    g_assert (expr->common.type_spec->common.type == AST_FUNCTION_TYPE);
	    break;

	case AST_CLASS_DEFINITION:
	    return report_error ("Name %s is a class\n", expr->identifier.name);
	    break;

	case AST_TYPE_DEFINITION:
	    return report_error ("Name %s is a type\n", expr->identifier.name);
	    break;
	}
	break;

    case AST_UNARY_EXPRESSION:
	if (!type_check_unary_expression (&expr->unary))
	    return FALSE;
	break;

    case AST_BINARY_EXPRESSION:
	if (!type_check_binary_expression (&expr->binary))
	    return FALSE;
	break;

    case AST_TERNARY_EXPRESSION:
	if (!type_check_expression (expr->ternary.cond))
	    return FALSE;
	if (!type_check_expression (expr->ternary.then_expr))
	    return FALSE;
	if (!type_check_expression (expr->ternary.else_expr))
	    return FALSE;

	if (!is_bool (get_type (expr->ternary.cond)))
	    return report_error ("Condition in ?: operator must be bool\n");

	then_type = get_type (expr->ternary.then_expr);
	else_type = get_type (expr->ternary.else_expr);

	if (is_numeric (then_type) && is_numeric (else_type))
	{
	    expr->common.type_spec =
		promote_numeric (then_type, else_type);
	}
	else if (is_bool (then_type) && is_bool (else_type))
	{
	    expr->common.type_spec = new_bool ();
	}
	else
	{
	    return report_error ("Type error in ?: expression\n");
	}
	break;

    case AST_LAMBDA_EXPRESSION:
	if (!type_check_pass2 ((ast_t *)expr->lambda.function->body))
	    return FALSE;

	expr->common.type_spec = expr->lambda.function->type_spec;
	g_assert (expr->common.type_spec);
	break;

    case AST_CALL_EXPRESSION:
	if (!type_check_expression (expr->call.callee))
	    return FALSE;

	for (i = 0; expr->call.args[i] != NULL; ++i)
	{
	    if (!type_check_expression (expr->call.args[i]))
		return FALSE;
	}

	callee_type = get_type (expr->call.callee);

	if (!is_function (callee_type))
	    return report_error ("Called object is not a function\n");

	if (!type_check_arguments (callee_type->function.args, expr->call.args))
	    return FALSE;

	expr->common.type_spec = get_type (expr->call.callee)->function.return_;
	break;

    case AST_STATEMENT_EXPRESSION:
	if (!type_check_pass2 ((ast_t *)expr->statement.statement))
	    return FALSE;

	if (!type_check_expression (expr->statement.expression))
	    return FALSE;

	expr->common.type_spec = expr->statement.expression->common.type_spec;
	break;

    case AST_NEW_EXPRESSION:
	for (i = 0; expr->new.args[i] != NULL; ++i)
	{
	    if (!type_check_expression (expr->new.args[i]))
		return FALSE;
	}

	if (!is_class (expr->new.type) && !is_array (expr->new.type))
	    return report_error ("Type is not class or array\n");

	if (is_array (expr->new.type))
	{
	    if (array_length ((void **)expr->new.args) != 1	||
		!expr->new.args[0]				||
		!is_int32 (get_type (expr->new.args[0])))
	    {
		return report_error (
		    "Array constructions must have exactly one "
		    "integer argument");
	    }
	}
	else if (expr->new.args[0])
	{
	    /* FIXME: check that the parameters match the constructor */

	    return report_error ("FIXME: Arguments to new not implemented\n");
	}

	expr->common.type_spec = expr->new.type;
	break;

    case AST_BLOCK_EXPRESSION:
	if (!type_check_expression (expr->block.body))
	    return FALSE;

	expr->common.type_spec = expr->block.body->common.type_spec;
	break;

    case AST_DEFINITION_EXPRESSION:
	if (!type_check_expression (expr->definition.initializer))
	    return FALSE;

	definition = expr->definition.definition;

	g_assert (definition->common.type == AST_VARIABLE_DEFINITION);

	if (definition->variable.type_spec->common.type == AST_INFERRED_TYPE)
	{
	    return report_error (
		"Can't infer type of %s\n", definition->variable.name);
	}

	if (!assignable (definition->variable.type_spec,
			 expr->definition.initializer->common.type_spec))
	{
	    return report_error ("Type error in initializer\n");
	}

	expr->common.type_spec = definition->variable.type_spec;
	break;

    case AST_DOT_EXPRESSION:
	if (!type_check_expression (expr->dot.left))
	    return FALSE;

	left_type = get_type (expr->dot.left);
	if (left_type->common.type != AST_OBJECT_TYPE)
	{
	    return report_error (
		"Left-hand side of dot expression is not an object\n");
	}

	definition = ast_class_definition_lookup_var (left_type->object.class,
						      expr->dot.name);

	if (!definition)
	{
	    return report_error ("Name %s is not defined in class %s\n",
				 expr->dot.name, left_type->object.class->name);
	}

	expr->dot.definition = definition;

	switch (definition->common.type)
	{
	case AST_VARIABLE_DEFINITION:
	    expr->common.type_spec = definition->variable.type_spec;
	    break;

	case AST_FUNCTION_DEFINITION:
	    expr->common.type_spec = definition->function.type_spec;
	    break;

	case AST_LABEL_DEFINITION:
	    return report_error ("Name %s is a label\n", expr->dot.name);
	    break;

	case AST_CLASS_DEFINITION:
	    return report_error ("Name %s is an inner class\n", expr->dot.name);
	    break;

	case AST_TYPE_DEFINITION:
	    return report_error ("Name %s is a type\n", expr->dot.name);
	    break;
	}
	break;

    case AST_INDEX_EXPRESSION:
	if (!type_check_expression (expr->index.left))
	    return FALSE;

	if (!type_check_expression (expr->index.right))
	    return FALSE;

	if (is_array (get_type (expr->index.left))		&&
	    is_int32 (get_type (expr->index.right)))
	{
	    expr->common.type_spec =
		get_type (expr->index.left)->array.child_type;
	}
	break;

    case AST_VOID_EXPRESSION:
	expr->common.type_spec = ast_type_spec_new (AST_VOID_TYPE);
	break;
    }

    g_assert (expr->common.type_spec);

    if (expr->common.type_inferrer)
	expr->common.type_inferrer->type_spec = expr->common.type_spec;

    return TRUE;
}

static gboolean
type_check_return (ast_return_statement_t *return_)
{
    ast_function_definition_t *function;
    ast_type_spec_t *type;

    function = ast_enclosing_function ((ast_t *)return_);

    g_assert (function->return_type);

    type = get_type (return_->expr);

    if (is_void (function->return_type) && !is_void (type))
	return report_error ("return with value in void function\n");

    if (is_void (type) && !is_void (function->return_type))
	return report_error ("return without value in non-void function\n");

    if (!is_void (function->return_type) || !is_void (type))
    {
	if (!assignable (function->return_type, type))
	    return report_error ("Type error in return\n");
    }

    return TRUE;
}

static gboolean
type_check_pass2 (ast_t *ast)
{
    GList *list;

    if (ast->common.type == AST_EXPRESSION)
	return type_check_expression (&ast->expression);

    for (list = ast->common.children->head; list; list = list->next)
    {
	if (!type_check_pass2 (list->data))
	    return FALSE;
    }

    if (ast_is (ast, AST_STATEMENT, AST_LOOP_STATEMENT))
    {
	if (!is_bool (get_type (ast->statement.loop.condition)))
	{
	    return report_error (
		"Condition in loop must be boolean\n");
	}
    }
    else if (ast_is (ast, AST_STATEMENT, AST_GOTO_STATEMENT))
    {
	if (!is_label (get_type (ast->statement.goto_.target)))
	{
	    return report_error (
		"Target of goto must be a label\n");
	}
    }
    else if (ast_is (ast, AST_STATEMENT, AST_IF_STATEMENT))
    {
	if (!is_bool (get_type (ast->statement.if_.condition)))
	    return report_error ("Condition in if statement must be boolean\n");
    }
    else if (ast_is (ast, AST_STATEMENT, AST_SWITCH_STATEMENT))
    {
	ast_type_spec_t *cond_type = get_type (ast->statement.switch_.condition);
	ast_case_t **cases = ast->statement.switch_.cases;
	int i;

	for (i = 0; cases[i] != NULL; ++i)
	{
	    switch (cases[i]->common.type)
	    {
	    case AST_EXPRESSION_CASE:
		if (!assignable (
			cond_type, get_type (cases[i]->expression.expr)))
		{
		    return report_error ("Type error in case\n");
		}
		break;

	    case AST_DEFAULT_CASE:
		break;
	    }
	}
    }
    else if (ast_is (ast, AST_STATEMENT, AST_RETURN_STATEMENT))
    {
	return type_check_return (&ast->statement.return_);
    }

    return TRUE;
}

gboolean
type_check (ast_t *ast)
{
    type_check_pass1 (ast);
    type_check_resolve_identifiers (ast);

    if (!type_check_pass2 (ast))
	return FALSE;

    return TRUE;
}