optimizer

from tree import Node, Op, Var, Val
from parser import parse

def node_to_str(node: Node):
    match node:
        case Op(op, left, right):
            return "(" + node_to_str(left) + f" {op} " + node_to_str(right) + ")"
        case Var(var):
            return var
        case Val(val):
            return val


def optimize_step(e: Node):
    match e:
        case Op(op, Val(v), Val(w)):
            match op:
                case "*":
                    return v * w
                case "+":
                    return v + w
        case Op("+", X, Y) if X == Y:
            return Op("*", Val(2), X)
        case Op("+", X, Op("+", Y, Z)):
            return Op("+", Op("+", X, Y), Z)
        case Op(op, X, Y):
            if opt_x := optimize_step(X) is not None:
                return Op(op, opt_x, Y)
            elif opt_y := optimize_step(Y) is not None:
                return Op(op, X, opt_y)
            else:
                return None
        case _:
            return None


def optimize_step_if(e: Node):
    match e:
        case Op(op, Val(v), Val(w)):
            match op:
                case "*":
                    return v * w
                case "+":
                    return v + w
        case Op("+", X, Y) if X == Y:
            return Op("*", Val(2), X)
        case Op("+", X, Op("+", Y, Z)):
            return Op("+", Op("+", X, Y), Z)
        case Op(op, X, Y):
            if opt_x := optimize_step(X) is not None:
                return Op(op, opt_x, Y)
            elif opt_y := optimize_step(Y) is not None:
                return Op(op, X, opt_y)
            else:
                return None
        case _:
            return None


if __name__ == "__main__":
    print(node_to_str(parse("2 * x + y")))