;; NOTE: This file contains several unsafe operations
(datatype Value (Num i64))
(function TrueConst () Value)
(let True (TrueConst))
(function FalseConst () Value)
(let False (FalseConst))
(datatype VarType)
(datatype Term
    (Val Value)
    (Var VarType)
    (Add Term Term)
    (Eq Term Term)
    (App Term Term)
    (Lam VarType Term)
    (Let VarType Term Term)
    (Fix VarType Term)
    (If Term Term Term))
(function V (String) VarType) 
(function From (Term) VarType)

; All free variables are free,
; but some free variables are more free than others
; Freer variables will only contain variables
; that will affect the evaluation
; of the corresponding term
; e.g., x is free in x - x, but not freer in x - x
(sort StringSet (Set VarType))
(function freer (Term) StringSet :merge (set-intersect old new))
(rule ((= e (Val v)))
      ((set (freer e) (set-empty))))
(rule ((= e (Var v)))
      ((set (freer e) (set-insert (set-empty) v))))
(rule ((= e (Add e1 e2))
       (= (freer e1) fv1)
       (= (freer e2) fv2))
      ((set (freer e) (set-union fv1 fv2))))
(rule ((= e (Eq e1 e2))
       (= (freer e1) fv1)
       (= (freer e2) fv2))
      ((set (freer e) (set-union fv1 fv2))))
(rule ((= e (App e1 e2))
       (= (freer e1) fv1)
       (= (freer e2) fv2))
      ((set (freer e) (set-union fv1 fv2))))
(rule ((= e (Lam var body))
       (= (freer body) fv))
      ((set (freer e) (set-remove fv var))))
(rule ((= e (Let var e1 e2))
       (= (freer e1) fv1)
       (= (freer e2) fv2))
      ((set (freer e) (set-union fv1 (set-remove fv2 var)))))
(rule ((= e (Fix var body))
       (= (freer body) fv))
      ((set (freer e) (set-remove fv var))))
(rule ((= e (If c e1 e2))
       (= (freer c) fv1)
       (= (freer e1) fv2)
       (= (freer e2) fv3))
      ((set (freer e) (set-union fv1 (set-union fv2 fv3)))))

;; START evals-to
(function evals-to (Term) Value)

(rule ((= e (Val val)))
      ((union (evals-to e) val)))
(rule ((= e (Add a b))
       (= (Num va) (evals-to a))
       (= (Num vb) (evals-to b)))
      ((union (evals-to e) (Num (+ va vb)))))
(rule ((= e (Eq a b))
       (= (evals-to b) (evals-to a)))
      ((union (evals-to e) True)))
(rule ((= e (Eq a b))
       (= va (evals-to a))
       (= vb (evals-to b))
       (!= va vb)) ; TODO: if true and false are non-mergeable datatype, 
                   ; we should be able to do != over it
      ((union (evals-to e) False)))
(rule ((= v (evals-to e)))
      ((union e (Val v))))
;; END evals-to

; if-true
(rewrite (If (Val True) then else) then)
; if-false
(rewrite (If (Val False) then else) else)
; if-elim
(rule ((= term (If (Eq (Var x) e) then else)))
      ((Let x e then)
       (Let x e else)))
(rewrite (If (Eq (Var x) e) then else) else
    :when ((= (Let x e then) (Let x e else))))

; add-comm
(rewrite (Add a b) (Add b a))
; add-assoc
(rewrite (Add (Add a b) c) (Add a (Add b c)))
; eq-comm
(rewrite (Eq a b) (Eq b a))

; fix
(rewrite (Fix v e) (Let v (Fix v e) e))
; beta
(rewrite (App (Lam v body) e) (Let v e body))
; let-app
(rewrite (Let v e (App a b)) (App (Let v e a) (Let v e b)))
; let-add
(rewrite (Let v e (Add a b)) (Add (Let v e a) (Let v e b)))
; let-eq
(rewrite (Let v e (Eq a b)) (Eq (Let v e a) (Let v e b)))
; let-const
(rewrite (Let v e c) c :when ((= const (evals-to c))))
; let-if
(rewrite (Let v e (If cond then else)) 
         (If (Let v e cond) (Let v e then) (Let v e else)))
; let-var-same
(rewrite (Let v1 e (Var v1)) e)
; let-var-diff
(rewrite (Let v1 e (Var v2)) (Var v2) :when ((!= v1 v2)))
; let-lam-same
(rewrite (Let v1 e (Lam v1 body)) (Lam v1 body))
; let-lam-diff
(rewrite (Let v1 e (Lam v2 body)) (Lam v2 (Let v1 e body))
      :when ((!= v1 v2)
             (= fvs (freer e))
             (set-not-contains fvs v2)))
(rule ((= expr (Let v1 e (Lam v2 body)))
       (!= v1 v2)
       (= fvs (freer e))
       (set-contains fvs v2))
      ((union expr (Lam (From expr) (Let v1 e (Let v2 (Var (From expr)) body))))))

;; lambda_under
(push)
(let e 
        (Lam (V "x") 
             (Add (Val (Num 4))
                  (App (Lam (V "y") (Var (V "y"))) (Val (Num 4))))))
(run 10)
(check (= e (Lam (V "x") (Val (Num 8)))))
(pop)

;; lambda_if_elim
(push)
(let e2 (If (Eq (Var (V "a")) (Var (V "b")))
              (Add (Var (V "a")) (Var (V "a")))
              (Add (Var (V "a")) (Var (V "b")))))
(run 10)
(check (= e2 (Add (Var (V "a")) (Var (V "b")))))
(pop)

;; lambda_let_simple
(push)
(let e3 (Let (V "x") (Val (Num 0))
          (Let (V "y") (Val (Num 1))
          (Add (Var (V "x")) (Var (V "y"))))))
(run 10)
(check (= e3 (Val (Num 1))))
(pop)

;; lambda_capture
(push)
(let e4 (Let (V "x") (Val (Num 1)) 
          (Lam (V "x") (Var (V "x")))))
(run 10)
(fail (check (= e4 (Lam (V "x") (Val (Num 1))))))
(pop)

;; lambda_capture_free
(push)
(let e5  (Let (V "y") (Add (Var (V "x")) (Var (V "x"))) 
           (Lam (V "x") (Var (V "y")))))
(run 10)
(check (set-contains (freer (Lam (V "x") (Var (V "y")))) (V "y")))
(fail (check (= e5 (Lam (V "x") (Add (Var (V "x")) (Var (V "x")))))))
(pop)

;; lambda_closure_not_seven
(push)
(let e6
     (Let (V "five") (Val (Num 5))
     (Let (V "add-five") (Lam (V "x") (Add (Var (V "x")) (Var (V "five"))))
     (Let (V "five") (Val (Num 6))
     (App (Var (V "add-five")) (Val (Num 1)))))))
(run 10)
(fail (check (= e6 (Val (Num 7)))))
(check (= e6 (Val (Num 6))))
(pop)


;; lambda_compose
(push)
(let e7
    (Let (V "compose") (Lam (V "f") 
                       (Lam (V "g") 
                       (Lam (V "x") (App (Var (V "f"))
                                         (App (Var (V "g")) 
                                              (Var (V "x")))))))
    (Let (V "add1") (Lam (V "y") (Add (Var (V "y")) (Val (Num 1))))
    (App (App (Var (V "compose")) (Var (V "add1"))) (Var (V "add1"))))))
(run 20)
(check (= e7 (Lam x (Add (Var x) (Val (Num 2))))))
(pop)

;; lambda_if_simple
(push)
(let e10 (If (Eq (Val (Num 1)) (Val (Num 1))) (Val (Num 7)) (Val (Num 9))))
(run 4)
(check (= e10 (Val (Num 7))))
(pop)

;; lambda_compose_many
(push)
(let e11
    (Let (V "compose") (Lam (V "f") (Lam (V "g") (Lam (V "x") (App (Var (V "f"))
                                                              (App (Var (V "g")) (Var (V "x")))))))
    (Let (V "add1") (Lam (V "y") (Add (Var (V "y")) (Val (Num 1))))
    (App (App (Var (V "compose")) (Var (V "add1")))
         (App (App (Var (V "compose")) (Var (V "add1")))
              (App (App (Var (V "compose")) (Var (V "add1")))
                   (App (App (Var (V "compose")) (Var (V "add1")))
                        (App (App (Var (V "compose")) (Var (V "add1")))
                             (App (App (Var (V "compose")) (Var (V "add1")))
                                  (Var (V "add1")))))))))))

(run 30)
(check (= e11 (Lam x (Add (Var x) (Val (Num 7))))))
(pop)

;; lambda_if
(push)
(let e8
    (Let (V "zeroone") (Lam (V "x")
         (If (Eq (Var (V "x")) (Val (Num 0)))
             (Val (Num 0))
             (Val (Num 1))))
         (Add (App (Var (V "zeroone")) (Val (Num 0)))
              (App (Var (V "zeroone")) (Val (Num 10)))))
)
(run 20)
(check (= e8 (Val (Num 1))))
(pop)