optimizer2

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 str(var)
        case Val(val):
            return str(val)


def optimize_step(e: Node) -> Node:
    match e:
        case Op(op, Val(v), Val(w)):
            match op:
                case "*":
                    return Val(v * w)
                case "+":
                    return Val(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) -> 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")))
    assert optimize_step(Op('*', Val(5), Val(3))) == Val(15)
    assert optimize_step(Op('+', Var('x'), Var('x'))) == Op('*', Val(2), Var('x'))
    assert optimize_step(Op('+', Var('x'), Var('y'))) == None
    e = Op('*', Val(3), Var('x'))
    assert optimize_step(Op('+', e, e)) == Op('*', Val(2), e)
    e1 = Op('*', Val(2), Val(3))
    e2 = Op('*', Val(3), Val(2))
    assert optimize_step(Op('+', e1, e2)) == Op('+', Val(6), e2)
    assert optimize_step(Op('+', Val(6), e2)) == Op('+', Val(6), Val(6))
    assert optimize_step(Op('+', Val(6), Val(6))) == Val(12)
    assert optimize_step(Val(12)) == None