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

import pass_common:*

fun defer_lower(name_ast_map: *map<str, pair<*tree<symbol>,*ast_node>>, ast_to_syntax: *map<*ast_node, *tree<symbol>>) {
    var visited = hash_set<*ast_node>()
    name_ast_map->for_each(fun(name: str, syntax_ast_pair: pair<*tree<symbol>,*ast_node>) {
        var defer_triple_stack = stack<stack<stack<*ast_node>>>()
        var loop_stack = stack(-1)
        var helper_before = fun(node: *ast_node, parent_chain: *stack<*ast_node>) {
            match(*node) {
                ast_node::defer_statement(backing) {
                    if (is_code_block(parent_chain->top())) {
                        remove(node, parent_chain)
                        defer_triple_stack.top().top().push(backing.statement)
                    } else {
                        replace_with_in(node, backing.statement, parent_chain)
                    }
                }
                ast_node::code_block(backing) {
                    defer_triple_stack.top().push(stack<*ast_node>())
                }
                ast_node::for_loop(backing) {
                    loop_stack.push(defer_triple_stack.top().size())
                }
                ast_node::while_loop(backing) {
                    loop_stack.push(defer_triple_stack.top().size())
                }
                ast_node::function(backing) {
                    defer_triple_stack.push(stack<stack<*ast_node>>())
                }
            }
        }
        var helper_after = fun(node: *ast_node, parent_chain: *stack<*ast_node>) {
            match(*node) {
                ast_node::branching_statement(backing) {
                    var block = _code_block()
                    add_to_scope("~enclosing_scope", parent_chain->item_from_top_satisfying(fun(i: *ast_node): bool return is_code_block(i) || is_function(i);), block)
                    replace_with_in(node, block, parent_chain)
                    for (var i = 0; i < defer_triple_stack.top().size() - loop_stack.top(); i++;)
                        block->code_block.children.add_all(defer_triple_stack.top().from_top(i).reverse_vector())
                    block->code_block.children.add(node)
                }
                ast_node::return_statement(backing) {
                    var block = parent_chain->top()
                    if (!is_code_block(block))
                        error("defer doesn't have block - it should from obj lower")
                    for (var i = 0; i < defer_triple_stack.top().size(); i++;) {
                        defer_triple_stack.top().from_top(i).reverse_vector().for_each(fun(c: *ast_node) {
                            add_before_in(c, node, block)
                        })
                    }
                }
                ast_node::code_block(backing) {
                    node->code_block.children.add_all(defer_triple_stack.top().pop().reverse_vector())
                }
                ast_node::for_loop(backing) {
                    loop_stack.pop()
                }
                ast_node::while_loop(backing) {
                    loop_stack.pop()
                }
                ast_node::function(backing) {
                    defer_triple_stack.pop()
                }
            }
        }
        run_on_tree(helper_before, helper_after, syntax_ast_pair.second, &visited)
    })
}