(datatype Math
    (Diff Math Math)
    (Integral Math Math)
    
    (Add Math Math)
    (Sub Math Math)
    (Mul Math Math)
    (Div Math Math)
    (Pow Math Math)
    (RShift Math Math)
    (LShift Math Math)
    (Mod Math Math)
    (Not Math)
    
    (Const i64)
    (Var String))

(relation MathU (Math))
(rule ((= e (Diff x y))) ((MathU e)))
(rule ((= e (Integral x y))) ((MathU e)))
(rule ((= e (Add x y))) ((MathU e)))
(rule ((= e (Sub x y))) ((MathU e)))
(rule ((= e (Mul x y))) ((MathU e)))
(rule ((= e (Div x y))) ((MathU e)))
(rule ((= e (Pow x y))) ((MathU e)))
(rule ((= e (Const x))) ((MathU e)))
(rule ((= e (Var x))) ((MathU e)))
(rule ((= e (RShift x y))) ((MathU e)))
(rule ((= e (LShift x y))) ((MathU e)))
(rule ((= e (Mod x y))) ((MathU e)))
(rule ((= e (Not x))) ((MathU e)))

(relation evals-to (Math i64))
(rule ((evals-to x vx)) ((union x (Const vx))))
(rule ((= e (Const c))) ((evals-to e c)))

(relation is-not-zero (Math))
(rule ((MathU a) (!= a (Const 0))) ((is-not-zero a)))

;; Evaluation
(rewrite (Add (Const a) (Const b))
         (Const (+ a b)))
(rewrite (Sub (Const a) (Const b))
            (Const (- a b)))
(rewrite (Mul (Const a) (Const b)) (Const (* a b)))
(rewrite (Div (Const a) (Const b)) (Const (/ a b)) :when ((!= 0 b)))
(rewrite (RShift (Const a) (Const b)) (Const (>> a b)))
(rewrite (LShift (Const a) (Const b)) (Const (<< a b)))
(rewrite (Not (Const a)) (Const (not-i64 a)))

;; Properties
(rewrite (Add a b) (Add b a))
(rewrite (Mul a b) (Mul b a))
(rewrite (Add a (Add b c)) (Add (Add a b) c))
(rewrite (Mul a (Mul b c)) (Mul (Mul a b) c))

(rewrite (Sub a b) (Add a (Mul (Const -1) b)))

(rewrite (Add a (Const 0)) a)
(rewrite (Mul a (Const 0)) (Const 0))
(rewrite (Mul a (Const 1)) a)

(rule ((MathU a) (!= a (Const 0))) ((union a (Add a (Const 0)))))
(rule ((MathU a) (!= a (Const 1))) ((union a (Mul a (Const 1)))))

(rewrite (Sub a a) (Const 0))
(rewrite (Div a a) (Const 1) :when ((is-not-zero a)))

(rewrite (Mul a (Add b c)) (Add (Mul a b) (Mul a c)))
(rewrite (Add (Mul a b) (Mul a c)) (Mul a (Add b c)))

; This rule doesn't work when pow is negative - consider 2^-1 * 2^1, which is 0, but 2^0 = 1
(rewrite (Mul (Pow a b) (Pow a c)) (Pow a (Add b c)) :when ((is-not-zero b) (is-not-zero c)))

(rewrite (Pow x (Const 0)) (Const 1) :when ((is-not-zero x)))
(rewrite (Pow x (Const 1 )) x)
(rewrite (Pow x (Const 2)) (Mul x x))

(rewrite (Pow x (Const -1)) (Div (Const 1) x) :when ((is-not-zero x)))

(rewrite (Mul x (Pow (Const 2) y)) (LShift x y))
(rewrite (Div x (Pow (Const 2) y)) (RShift x y))

(rewrite (Not (Not x)) x)


;; Tests
(let start-expr (Div (
                        Mul (Var "a") (Pow (Const 2) (Const 3))
                    ) (
                        Add (Var "c") (
                            Sub (Mul (Var "b") (Const 2)) (Mul (Var "b") (Const 2))
                        )
                    )))

(let equiv-expr (Div (
                    LShift (Var "a") (Const 3)
                ) (
                    Mul (Var "c") (Not (Not (Const 1)))
                )
                ))

(run 4)

(check (= start-expr equiv-expr))