Benkex

optimizer2

Dec 1st, 2021
667
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
  1. from tree import Node, Op, Var, Val
  2. from parser import parse
  3.  
  4. def node_to_str(node: Node):
  5.     match node:
  6.         case Op(op, left, right):
  7.             return "(" + node_to_str(left) + f" {op} " + node_to_str(right) + ")"
  8.         case Var(var):
  9.             return str(var)
  10.         case Val(val):
  11.             return str(val)
  12.            
  13.  
  14. def optimize_step(e: Node) -> Node:
  15.     match e:           
  16.         case Op(op, Val(v), Val(w)):
  17.             match op:
  18.                 case "*":
  19.                     return Val(v * w)
  20.                 case "+":
  21.                     return Val(v + w)
  22.         case Op("+", X, Y) if X == Y:
  23.             return Op("*", Val(2), X)
  24.         case Op("+", X, Op("+", Y, Z)):
  25.             return Op("+", Op("+", X, Y), Z)
  26.         case Op(op, X, Y):
  27.             if (opt_x := optimize_step(X)) is not None:
  28.                 return Op(op, opt_x, Y)
  29.             elif (opt_y := optimize_step(Y)) is not None:
  30.                 return Op(op, X, opt_y)
  31.             else:
  32.                 return None
  33.         case _:
  34.             return None
  35.  
  36.  
  37. def optimize_step_if(e: Node) -> Node:
  38.     match e:           
  39.         case Op(op, Val(v), Val(w)):
  40.             match op:
  41.                 case "*":
  42.                     return v * w
  43.                 case "+":
  44.                     return v + w
  45.         case Op("+", X, Y) if X == Y:
  46.             return Op("*", Val(2), X)
  47.         case Op("+", X, Op("+", Y, Z)):
  48.             return Op("+", Op("+", X, Y), Z)
  49.         case Op(op, X, Y):
  50.             if opt_x := optimize_step(X) is not None:
  51.                 return Op(op, opt_x, Y)
  52.             elif opt_y := optimize_step(Y) is not None:
  53.                 return Op(op, X, opt_y)
  54.             else:
  55.                 return None
  56.         case _:
  57.             return None
  58.  
  59.            
  60. if __name__ == "__main__":
  61.     print(node_to_str(parse("2 * x + y")))
  62.     assert optimize_step(Op('*', Val(5), Val(3))) == Val(15)
  63.     assert optimize_step(Op('+', Var('x'), Var('x'))) == Op('*', Val(2), Var('x'))
  64.     assert optimize_step(Op('+', Var('x'), Var('y'))) == None
  65.     e = Op('*', Val(3), Var('x'))
  66.     assert optimize_step(Op('+', e, e)) == Op('*', Val(2), e)
  67.     e1 = Op('*', Val(2), Val(3))
  68.     e2 = Op('*', Val(3), Val(2))
  69.     assert optimize_step(Op('+', e1, e2)) == Op('+', Val(6), e2)
  70.     assert optimize_step(Op('+', Val(6), e2)) == Op('+', Val(6), Val(6))
  71.     assert optimize_step(Op('+', Val(6), Val(6))) == Val(12)
  72.     assert optimize_step(Val(12)) == None
RAW Paste Data