from __future__ import annotations
from optimizations.Optimization import Optimization
from optimizations.DeadAssignmentElimination import DeadAssignmentElimination
from syntax import EXPRESSION

# Traverses the AST and applies the given optimizations
def __apply_optimizations(node: EXPRESSION, optimizations: list[Optimization]) -> EXPRESSION:
    for key, value in list(node.__dict__.items()):
        if key == "pp":
            continue
        if isinstance(value, EXPRESSION):
            new_child = __apply_optimizations(value, optimizations)
            if new_child is not value:
                setattr(node, key, new_child)
        elif isinstance(value, list):
            for i, elem in enumerate(value):
                if isinstance(elem, EXPRESSION):
                    new_elem = __apply_optimizations(elem, optimizations)
                    if new_elem is not elem:
                        value[i] = new_elem

    result: EXPRESSION = node
    for opt in optimizations:
        result = opt.apply(result)
    return result

# Stores the active optimizations
ACTIVE_OPTIMIZATIONS: list[Optimization] = [
    DeadAssignmentElimination(),
]

# Apply optimizations to the given AST node
def optimize(ast: EXPRESSION) -> EXPRESSION:
    for opt in ACTIVE_OPTIMIZATIONS:
        opt.setup(ast)
    return __apply_optimizations(ast, ACTIVE_OPTIMIZATIONS)