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 Dim (Times Dim Dim) (NamedDim String) (Lit i64))

(rewrite (Times a (Times b c)) (Times (Times a b) c))
(rewrite (Times (Times a b) c) (Times a (Times b c)) )
(rewrite (Times (Lit i) (Lit j)) (Lit (* i j)))
(rewrite (Times a b) (Times b a))

(datatype MExpr
    (MMul MExpr MExpr)
    (Kron MExpr MExpr)
    (NamedMat String)
    (Id Dim)
    ; DSum
    ; HStack
    ; VStack
    ; Transpose
    ; Inverse
    ; Zero Math Math
    ; ScalarMul
)

; alternative encoding (type A) = (Matrix n m) may be more useful for "large story example"
(constructor nrows (MExpr) Dim)
(constructor ncols (MExpr) Dim)

(rewrite (nrows (Kron A B)) (Times (nrows A) (nrows B)))
(rewrite (ncols (Kron A B)) (Times (ncols A) (ncols B)))

(rewrite (nrows (MMul A B)) (nrows A))
(rewrite (ncols (MMul A B)) (ncols B))

(rewrite (nrows (Id n)) n)
(rewrite (ncols (Id n)) n)

(rewrite (MMul (Id n) A) A)
(rewrite (MMul A (Id n)) A)

(rewrite (MMul A (MMul B C)) (MMul (MMul A B) C))
(rewrite (MMul (MMul A B) C) (MMul A (MMul B C)))

(rewrite (Kron A (Kron B C)) (Kron (Kron A B) C))
(rewrite (Kron (Kron A B) C) (Kron A (Kron B C)))

(rewrite (Kron (MMul A C) (MMul B D))  (MMul (Kron A B) (Kron C D)))


(rewrite (MMul (Kron A B) (Kron C D))
    (Kron (MMul A C) (MMul B D))
    :when
        ((= (ncols A) (nrows C))
        (= (ncols B) (nrows D)))
)

; demand
(rule ((= e (MMul A B)))
((ncols A)
(nrows A)
(ncols B)
(nrows B))
)

(rule ((= e (Kron A B)))
((ncols A)
(nrows A)
(ncols B)
(nrows B))
)


(let $n (NamedDim "n"))
(let $m (NamedDim "m"))
(let $p (NamedDim "p"))

(let $A (NamedMat "A"))
(let $B (NamedMat "B"))
(let $C (NamedMat "C"))

(union (nrows $A) $n)
(union (ncols $A) $n)
(union (nrows $B) $m)
(union (ncols $B) $m)
(union (nrows $C) $p)
(union (ncols $C) $p)
(let $ex1 (MMul (Kron (Id $n) $B) (Kron $A (Id $m))))
(let $rows1 (nrows $ex1))
(let $cols1 (ncols $ex1))

(run 20)

(check (= (nrows $B) $m))
(check (= (nrows (Kron (Id $n) $B)) (Times $n $m)))
(let $simple_ex1 (Kron $A $B))
(check (= $ex1 $simple_ex1))

(let $ex2 (MMul (Kron (Id $p) $C) (Kron $A (Id $m))))
(run 10)
(fail (check (= $ex2 (Kron $A $C))))