egglog 2.0.0

egglog is a language that combines the benefits of equality saturation and datalog. It can be used for analysis, optimization, and synthesis of programs. It is the successor to the popular rust library egg.
Documentation
(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))
(constructor NilConst () Term)
(let $Nil (NilConst))

(constructor V (String) Var) 
(constructor 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))))))

(constructor 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 ====

(constructor sum () Term :cost 1000)
(constructor mapf () Term :cost 1000)
(constructor 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)

(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'")))))))))