import symbol:*
import tree:*
import map:*
import vec:*
import set:*
import hash_set:*
import util:*
import str:*
import mem:*
import io:*
import ast_nodes:*
import ast_transformation:*

import pass_common:*

fun has_ref_inside(t: *type): bool {
    if (t->is_ref)
        return true
    if (t->is_function()) {
        for (var i = 0; i < t->parameter_types.size; i++;)
            if (has_ref_inside(t->parameter_types[i]))
                return true
        if (has_ref_inside(t->return_type))
            return true
    }
    return false
}
fun remove_ref(t: *type) {
    if (t->is_ref) {
        t->is_ref = false
        t->indirection++
     }
    if (t->is_function()) {
        for (var i = 0; i < t->parameter_types.size; i++;)
            remove_ref(t->parameter_types[i])
        remove_ref(t->return_type)
    }
}

fun ref_lower(name_ast_map: *map<str, pair<*tree<symbol>,*ast_node>>, ast_to_syntax: *map<*ast_node, *tree<symbol>>) {
    var remove_ref_type_set = set<pair<str,*type>>()
    var modify_reference_use_set = set<pair<*ast_node, *ast_node>>()
    var modify_return_set = set<*ast_node>()
    var visited = hash_set<*ast_node>()
    name_ast_map->for_each(fun(name: str, syntax_ast_pair: pair<*tree<symbol>,*ast_node>) {
        var helper_before = fun(node: *ast_node, parent_chain: *stack<*ast_node>) {
            match(*node) {
                ast_node::identifier(backing) {
                    if (backing.type->is_ref)
                        modify_reference_use_set.add(make_pair(node, parent_chain->top()))
                    if (has_ref_inside(backing.type)) {
                        remove_ref_type_set.add(make_pair("identifier: " + backing.name, backing.type))
                    }
                }
                ast_node::function(backing) {
                    var full_name = get_fully_scoped_name(node) + " - (" + backing.name + ")"
                    if (has_ref_inside(backing.type))
                        remove_ref_type_set.add(make_pair(full_name, backing.type))
                }
                ast_node::function_call(backing) {
                    // check to see if it's taking a ref, if so add in the &
                    var func_type_params = get_ast_type(backing.func)->parameter_types
                    for (var i = 0; i < func_type_params.size; i++;)
                        if (func_type_params[i]->is_ref)
                            backing.parameters[i] = make_operator_call("&", vec(backing.parameters[i]))
                    // add the function call to the modify_reference_use set if the function returns a ref
                    if (get_ast_type(backing.func)->return_type->is_ref)
                        modify_reference_use_set.add(make_pair(node, parent_chain->top()))
                }
                ast_node::return_statement(backing) {
                    // check to see if it's returning a ref, if so add in the &
                    if (parent_chain->item_from_top_satisfying(is_function)->function.type->return_type->is_ref)
                        modify_return_set.add(node)
                }
            }
        }
        run_on_tree(helper_before, empty_pass_second_half(), syntax_ast_pair.second, &visited)
    })
    remove_ref_type_set.for_each(fun(p: pair<str, *type>) {
        var t = p.second
        remove_ref(t)
    })
    modify_reference_use_set.for_each(fun(p: pair<*ast_node, *ast_node>) {
        // if we haven't modified it's indirection yet
        if (is_identifier(p.first) && p.first->identifier.type->is_ref) {
            // remove ref, add 1 to indirection
            p.first->identifier.type = p.first->identifier.type->clone_with_increased_indirection(1, false);
        }
        // note that we definitly want to replace the type for unused parameters, but we don't want to add the * for paramters
        // in function declarations or the new identifier in declaration statements (but we do for expressions in declaration statements)
        if (!is_identifier(p.first) || (!is_function(p.second) && (!is_declaration_statement(p.second) || p.second->declaration_statement.identifier != p.first)))
            replace_with_in(p.first, make_operator_call("*", vec(p.first)), p.second);
    })
    modify_return_set.for_each(fun(r: *ast_node) {
        r->return_statement.return_value = make_operator_call("&", vec(r->return_statement.return_value));
    })
}