import ast_nodes:*
import ast_transformation:*
import mem:*
import util:*
import vector:*
import stack:*
import string:*

fun make_this_noncached(object: *ast_node): *ast_node {
    return ast_identifier_ptr("this", object->type_def.self_type->clone_with_indirection(1), object)
}
fun possible_object_equality(lvalue: *ast_node, rvalue: *ast_node): *ast_node {
    var ltype = get_ast_type(lvalue)
    var rtype = get_ast_type(rvalue)
    if (ltype->indirection == 0 && (ltype->is_object() && has_method(ltype->type_def, "operator==", vector(rtype)))) {
        return make_method_call(lvalue, "operator==", vector(rvalue))
    } else if (ltype->is_object())
        // return false if object but no operator== (right now don't try for templated)
        return ast_value_ptr(string("false"), type_ptr(base_type::boolean()))
    return make_operator_call("==", vector(lvalue, rvalue))
}

// for now, needs source to already be in a variable for copy_constructing
fun assign_or_copy_construct_statement(lvalue: *ast_node, rvalue: *ast_node): *ast_node {
    var ltype = get_ast_type(lvalue)
    if (ltype->indirection == 0 && (ltype->is_object() && has_method(ltype->type_def, "copy_construct", vector(ltype->clone_with_increased_indirection()))))
        return make_method_call(lvalue, "copy_construct", vector(make_operator_call("&", vector(rvalue))))
    return ast_assignment_statement_ptr(lvalue, rvalue)
}

fun get_children_pointer(node: *ast_node): *vector<*ast_node> {
    var bc = null<vector<*ast_node>>()
    match(*node) {
        ast_node::translation_unit(backing) bc = &node->translation_unit.children
        ast_node::code_block(backing) bc = &node->code_block.children
    }
    return bc
}

fun remove_full_statement(node: *ast_node, parent_chain: *stack<*ast_node>): *ast_node {
    if (is_statement(node))
        return remove(node, parent_chain)
    if (is_statement(parent_chain->top()))
        return remove(parent_chain->top(), parent_chain->from_top(1))
    error(string("cannot remove full statement in ") + get_ast_name(parent_chain->top()))
}

fun remove(orig: *ast_node, in: *stack<*ast_node>): *ast_node
    return remove(orig, in->top())
fun remove(orig: *ast_node, in: *ast_node): *ast_node {
    var bc = get_children_pointer(in)
    if (bc) {
        var i = bc->find(orig)
        if (i >= 0) {
            var temp = bc->at(i)
            bc->remove(i)
            return temp
        }
    }
    error(string("cannot remove inside ") + get_ast_name(in))
}
fun replace_with_in(orig: *ast_node, new: *ast_node, in: *stack<*ast_node>)
    replace_with_in(orig, new, in->top())
fun replace_with_in(orig: *ast_node, new: *ast_node, in: *ast_node) {
    match (*in) {
        ast_node::statement(backing) { backing.child = new; return; }
        ast_node::return_statement(backing) { backing.return_value = new; return; }
        ast_node::assignment_statement(backing) {
            if (backing.to == orig) {
                backing.to = new
                return
            }
            if (backing.from == orig) {
                backing.from = new
                return
            }
        }
        ast_node::declaration_statement(backing) {
            if (backing.identifier == orig) {
                backing.identifier = new
                return
            }
            if (backing.expression == orig) {
                backing.expression = new
                return
            }
        }
        ast_node::if_statement(backing) {
            if (backing.condition == orig) {
                backing.condition = new
                return
            }
            if (backing.then_part == orig) {
                backing.then_part = new
                return
            }
            if (backing.else_part == orig) {
                backing.else_part = new
                return
            }
        }
        ast_node::for_loop(backing) {
            if (backing.init == orig) {
                backing.init = new
                return
            }
            if (backing.condition == orig) {
                backing.condition = new
                return
            }
            if (backing.update == orig) {
                backing.update = new
                return
            }
            if (backing.body == orig) {
                backing.body = new
                return
            }
        }
        ast_node::while_loop(backing) {
            if (backing.condition == orig) {
                backing.condition = new
                return
            }
            if (backing.statement == orig) {
                backing.statement = new
                return
            }
        }
        ast_node::function_call(backing) {
            if (backing.func == orig) {
                backing.func = new
                return
            }
            for (var i = 0; i < backing.parameters.size; i++;) {
                if (backing.parameters[i] == orig) {
                    backing.parameters[i] = new
                    return
                }
            }
        }
    }

    var bc = get_children_pointer(in)
    if (bc) {
        var i = bc->find(orig)
        if (i >= 0) {
            bc->set(i, new)
            return
        }
    }
    error(string("cannot replace_with_in inside ") + get_ast_name(in))
}

fun add_before_in(to_add: ref vector<*ast_node>, before: *ast_node, in: *stack<*ast_node>)
    to_add.for_each(fun(n: *ast_node) add_before_in(n, before, in);)
fun add_before_in(to_add: vector<*ast_node>, before: *ast_node, in: *ast_node)
    to_add.for_each(fun(n: *ast_node) add_before_in(n, before, in);)

fun add_before_in(to_add: *ast_node, before: *ast_node, in: *stack<*ast_node>)
    add_before_in(to_add, before, in->top())
fun add_before_in(to_add: *ast_node, before: *ast_node, in: *ast_node) {
    var bc = get_children_pointer(in)
    if (bc) {
        var i = bc->find(before)
        if (i >= 0) {
            bc->add(i, to_add)
            return
        }
    }
    error(string("cannot add_before_in to ") + get_ast_name(in))
}
fun add_after_in(to_add: *ast_node, before: *ast_node, in: *stack<*ast_node>)
    add_after_in(to_add, before, in->top())
fun add_after_in(to_add: *ast_node, before: *ast_node, in: *ast_node) {
    var bc = get_children_pointer(in)
    if (bc) {
        var i = bc->find(before)
        if (i >= 0) {
            bc->add(i+1, to_add)
            return
        }
    }
    error(string("cannot add_after_in to ") + get_ast_name(in))
}

fun empty_pass_first_half(node: *ast_node, parent_chain: *stack<*ast_node>): bool { return true; }
fun empty_pass_second_half(node: *ast_node, parent_chain: *stack<*ast_node>) {}
fun run_on_tree(func_before: fun(*ast_node,*stack<*ast_node>):void, func_after: fun(*ast_node,*stack<*ast_node>):void, tree: *ast_node)
    run_on_tree(fun(n: *ast_node, s: *stack<*ast_node>): bool {func_before(n, s);return true;}, func_after, tree)

fun run_on_tree(func_before: fun(*ast_node,*stack<*ast_node>):bool, func_after: fun(*ast_node,*stack<*ast_node>):void, tree: *ast_node) {
    var parent_stack = stack<*ast_node>()
    run_on_tree_helper(func_before, func_after, tree, &parent_stack, false)
}
fun run_on_tree_helper(func_before: fun(*ast_node,*stack<*ast_node>):bool, func_after: fun(*ast_node,*stack<*ast_node>):void, node: *ast_node, parent_chain: *stack<*ast_node>, do_func: bool) {
    if (!node || (!do_func && is_function(node))) return
    var do_children = func_before(node, parent_chain)
    parent_chain->push(node)
    if (do_children) {
        match(*node) {
            ast_node::translation_unit(backing) get_ast_children(node).for_each(fun(n: *ast_node) run_on_tree_helper(func_before, func_after, n, parent_chain, true);)
            ast_node::type_def(backing) get_ast_children(node).for_each(fun(n: *ast_node) run_on_tree_helper(func_before, func_after, n, parent_chain, true);)
            ast_node::adt_def(backing) get_ast_children(node).for_each(fun(n: *ast_node) run_on_tree_helper(func_before, func_after, n, parent_chain, true);)
            ast_node::function(backing) get_ast_children(node).for_each(fun(n: *ast_node) run_on_tree_helper(func_before, func_after, n, parent_chain, false);)
            ast_node::template(backing) get_ast_children(node).for_each(fun(n: *ast_node) run_on_tree_helper(func_before, func_after, n, parent_chain, true);)
            ast_node::code_block(backing) get_ast_children(node).for_each(fun(n: *ast_node) run_on_tree_helper(func_before, func_after, n, parent_chain, false);)
            ast_node::statement(backing) get_ast_children(node).for_each(fun(n: *ast_node) run_on_tree_helper(func_before, func_after, n, parent_chain, false);)
            ast_node::if_statement(backing) get_ast_children(node).for_each(fun(n: *ast_node) run_on_tree_helper(func_before, func_after, n, parent_chain, false);)
            ast_node::match_statement(backing) get_ast_children(node).for_each(fun(n: *ast_node) run_on_tree_helper(func_before, func_after, n, parent_chain, false);)
            ast_node::case_statement(backing) get_ast_children(node).for_each(fun(n: *ast_node) run_on_tree_helper(func_before, func_after, n, parent_chain, false);)
            ast_node::while_loop(backing) get_ast_children(node).for_each(fun(n: *ast_node) run_on_tree_helper(func_before, func_after, n, parent_chain, false);)
            ast_node::for_loop(backing) get_ast_children(node).for_each(fun(n: *ast_node) run_on_tree_helper(func_before, func_after, n, parent_chain, false);)
            ast_node::return_statement(backing) get_ast_children(node).for_each(fun(n: *ast_node) run_on_tree_helper(func_before, func_after, n, parent_chain, false);)
            ast_node::defer_statement(backing) get_ast_children(node).for_each(fun(n: *ast_node) run_on_tree_helper(func_before, func_after, n, parent_chain, false);)
            ast_node::assignment_statement(backing) get_ast_children(node).for_each(fun(n: *ast_node) run_on_tree_helper(func_before, func_after, n, parent_chain, false);)
            ast_node::declaration_statement(backing) get_ast_children(node).for_each(fun(n: *ast_node) run_on_tree_helper(func_before, func_after, n, parent_chain, false);)
            ast_node::if_comp(backing) get_ast_children(node).for_each(fun(n: *ast_node) run_on_tree_helper(func_before, func_after, n, parent_chain, true);)
            ast_node::function_call(backing) {
                if (!is_function(backing.func))
                    run_on_tree_helper(func_before, func_after, backing.func, parent_chain, false)
                node->function_call.parameters.for_each(fun(n: *ast_node) run_on_tree_helper(func_before, func_after, n, parent_chain, false);)
            }
            ast_node::cast(backing) get_ast_children(node).for_each(fun(n: *ast_node) run_on_tree_helper(func_before, func_after, n, parent_chain, false);)
        }
    }
    // function may have messed with the parent chain
    if (parent_chain->data.contains(node))
        while(parent_chain->pop() != node){}
    func_after(node, parent_chain)
}