(datatype Var)
(datatype Term
    (App Term Term)
    (Lam Var Term)
    (TVar Var)
    (Let Var Term Term)
    (Add Term Term)
    (Num i64)
    (CaseSplit Term Term Term)
    (Cons Term Term))
(function NilConst () Term)
(let Nil (NilConst))

(function V (String) Var) 
(function From (Term) Var)

;; ==== FV ====
(sort StringSet (Set Var))
(function freer (Term) StringSet :merge (set-intersect old new))
(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 (TVar v)))
      ((set (freer e) (set-insert (set-empty) v))))
(rule ((= e (Let var e1 e2))
       (= (freer e1) fv1)
       (= (freer e2) fv2))
      ((set (freer e) (set-union fv1 (set-remove fv2 var)))))
(rule ((= e (Add e1 e2))
       (= (freer e1) fv1)
       (= (freer e2) fv2))
      ((set (freer e) (set-union fv1 fv2))))
(rule ((= e (Num v)))
      ((set (freer e) (set-empty))))
(rule ((= e (CaseSplit e1 e2 e3))
       (= (freer e1) fv1)
       (= (freer e2) fv2)
       (= (freer e3) fv3))
      ((set (freer e) (set-union (set-union fv1 fv2) fv3))))
(rule ((= e (Cons e1 e2))
       (= (freer e1) fv1)
       (= (freer e2) fv2))
      ((set (freer e) (set-union fv1 fv2))))
(rule ((= e Nil))
      ((set (freer e) (set-empty))))

;; ==== eval ====
; beta
(rewrite (App (Lam v body) e) (Let v e body))
; case-split-nil
(rewrite (CaseSplit Nil e1 e2) e1)
; case-split-cons
(rewrite (CaseSplit (Cons x xs) e1 e2) (App (App e2 x) xs))

; let-num
(rewrite (Let v e (Num n)) (Num n))
; let-nil
(rewrite (Let v e Nil) Nil)
; let-var-same
(rewrite (Let v1 e (TVar v1)) e)
; let-var-diff
(rewrite (Let v1 e (TVar v2)) (TVar v2) :when ((!= v1 v2)))

; let-lam-close
(rewrite (Let v1 e expr) expr :when ((set-not-contains (freer expr) v1)))      
; let-app
(rewrite (Let v e expr) (App (Let v e a) (Let v e b)) :when ((= expr (App a b)) (set-contains (freer expr) v)))
; let-add
(rewrite (Let v e expr) (Add (Let v e a) (Let v e b)) :when ((= expr (Add a b)) (set-contains (freer expr) v)))
; let-cons
(rewrite (Let v e expr) (Cons (Let v e x) (Let v e xs)) :when ((= expr (Cons x xs)) (set-contains (freer expr) v)))
; let-case-split
(rewrite (Let v e expr) 
         (CaseSplit (Let v e e1) (Let v e e2) (Let v e e3))
    :when ((= expr (CaseSplit e1 e2 e3))
           (set-contains (freer expr) v)))
; 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 ((set-contains (freer body) v1)
             (!= v1 v2)
             (= fvs (freer e))
             (set-not-contains fvs v2)))
(rule ((= expr (Let v1 e (Lam v2 body)))
       (set-contains (freer body) v1)
       (!= v1 v2)
       (= fvs (freer e))
       (set-contains fvs v2))
      ((union expr (Lam (From expr) (Let v1 e (Let v2 (TVar (From expr)) body))))))

(function pushdown (Term Term) Term :cost 10000)
(rewrite (App f (App (Lam x e) e2))
         (App (Lam x (pushdown f e)) e2))

(rewrite (pushdown f (CaseSplit e e1 (Lam x (Lam xs e2)))) 
         (CaseSplit e (App f e1) (Lam x (Lam xs (App f e2)))))

(relation is-tail (Term))
(rule ((= demand (pushdown f e)) (= e (App e1 e2))) ((is-tail e)))
(rule ((= demand (pushdown f e)) (= e (Lam x e))) ((is-tail e)))
(rule ((= demand (pushdown f e)) (= e (TVar x))) ((is-tail e)))
(rule ((= demand (pushdown f e)) (= e (Cons e1 e2))) ((is-tail e)))
(rule ((= demand (pushdown f e)) (= e Nil)) ((is-tail e)))
(rule ((= demand (pushdown f e)) (= e (Add e1 e2))) ((is-tail e)))
(rule ((= demand (pushdown f e)) (= e (Num n1))) ((is-tail e)))
(rewrite (pushdown f e) (App f e) :when ((is-tail e)))

;; ==== definition ====

(function sum () Term :cost 1000)
(function mapf () Term :cost 1000)
(function sum-o-mapf () Term)
(rewrite (App (sum) (App (mapf) x)) (App (sum-o-mapf) x))
(union (sum) (Lam (V "xs")
    (CaseSplit (TVar (V "xs")) 
        (Num 0)
        (Lam (V "x") (Lam (V "xs'") 
            (Add (TVar (V "x")) (App (sum) (TVar (V "xs'")))))))))

(union (mapf) (Lam (V "xs")
    (CaseSplit (TVar (V "xs"))
        Nil
        (Lam (V "x") (Lam (V "xs'")
            (Cons (Add (TVar (V "x")) (Num 1))
                  (App (mapf) (TVar (V "xs'")))))))))

(set (freer (sum)) (set-empty))
(set (freer (mapf)) (set-empty))

(let expr (App (sum) (App (mapf) (TVar (V "expr")))))

(run 100)

(query-extract (freer expr))


(let my-output
    (CaseSplit (TVar (V "expr")) (Num 0) 
           (Lam (V "x") (Lam (V "xs'") 
                (Add (Add (TVar (V "x")) (Num 1)) 
                     (App (sum-o-mapf) (TVar (V "xs'"))))))))

(check (= (App (sum-o-mapf) (TVar (V "expr")))
          (CaseSplit (TVar (V "expr")) (Num 0) 
                 (Lam (V "x") (Lam (V "xs'") 
                      (Add (Add (TVar (V "x")) (Num 1)) 
                           (App (sum-o-mapf) (TVar (V "xs'")))))))))