egglog 9.0.0__cp311-cp311-win_amd64.whl

egglog/examples/bool.py

@@ -0,0 +1,38 @@
+# mypy: disable-error-code="empty-body"
+"""
+Boolean data type example and test
+==================================
+"""
+
+from __future__ import annotations
+
+from egglog import *
+
+T = Bool(True)
+F = Bool(False)
+check(eq(T & T).to(T))
+check(eq(T & F).to(F))
+check(eq(T | F).to(T))
+check(ne(T | F).to(F))
+
+check(eq(i64(1).bool_lt(2)).to(T))
+check(eq(i64(2).bool_lt(1)).to(F))
+check(eq(i64(1).bool_lt(1)).to(F))
+
+check(eq(i64(1).bool_le(2)).to(T))
+check(eq(i64(2).bool_le(1)).to(F))
+check(eq(i64(1).bool_le(1)).to(T))
+
+R = relation("R", i64)
+
+
+@function
+def f(i: i64Like) -> Bool: ...
+
+
+i = var("i", i64)
+check(
+    eq(f(0)).to(T),
+    ruleset(rule(R(i)).then(set_(f(i)).to(T))) * 3,
+    R(i64(0)),
+)

egglog/examples/eqsat_basic.py

@@ -0,0 +1,45 @@
+# mypy: disable-error-code="empty-body"
+"""
+Basic equality saturation example.
+==================================
+"""
+
+from __future__ import annotations
+
+from egglog import *
+
+egraph = EGraph()
+
+
+class Num(Expr):
+    def __init__(self, value: i64Like) -> None: ...
+
+    @classmethod
+    def var(cls, name: StringLike) -> Num: ...
+
+    def __add__(self, other: Num) -> Num: ...
+
+    def __mul__(self, other: Num) -> Num: ...
+
+
+expr1 = Num(2) * (Num.var("x") + Num(3))
+expr2 = Num(6) + Num(2) * Num.var("x")
+
+a, b, c = vars_("a b c", Num)
+i, j = vars_("i j", i64)
+
+check(
+    # Check that these expressions are equal
+    eq(expr1).to(expr2),
+    # After running these rules, up to ten times
+    ruleset(
+        rewrite(a + b).to(b + a),
+        rewrite(a * (b + c)).to((a * b) + (a * c)),
+        rewrite(Num(i) + Num(j)).to(Num(i + j)),
+        rewrite(Num(i) * Num(j)).to(Num(i * j)),
+    )
+    * 10,
+    # On these two initial expressions
+    expr1,
+    expr2,
+)

egglog/examples/fib.py

@@ -0,0 +1,28 @@
+# mypy: disable-error-code="empty-body"
+"""
+Fibonacci numbers example
+=========================
+"""
+
+from __future__ import annotations
+
+from egglog import *
+
+
+@function
+def fib(x: i64Like) -> i64: ...
+
+
+f0, f1, x = vars_("f0 f1 x", i64)
+check(
+    eq(fib(i64(7))).to(i64(21)),
+    ruleset(
+        rule(
+            eq(f0).to(fib(x)),
+            eq(f1).to(fib(x + 1)),
+        ).then(set_(fib(x + 2)).to(f0 + f1)),
+    )
+    * 7,
+    set_(fib(0)).to(i64(1)),
+    set_(fib(1)).to(i64(1)),
+)

egglog/examples/higher_order_functions.py

@@ -0,0 +1,45 @@
+# mypy: disable-error-code="empty-body"
+"""
+Higher Order Functions
+======================
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from egglog import *
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+
+class Math(Expr):
+    def __init__(self, i: i64Like) -> None: ...
+    def __add__(self, other: Math) -> Math: ...
+
+
+class MathList(Expr):
+    def __init__(self) -> None: ...
+    def append(self, i: Math) -> MathList: ...
+    def map(self, f: Callable[[Math], Math]) -> MathList: ...
+
+
+@ruleset
+def math_ruleset(i: i64, j: i64, xs: MathList, x: Math, f: Callable[[Math], Math]):
+    yield rewrite(Math(i) + Math(j)).to(Math(i + j))
+    yield rewrite(xs.append(x).map(f)).to(xs.map(f).append(f(x)))
+    yield rewrite(MathList().map(f)).to(MathList())
+
+
+@function(ruleset=math_ruleset)
+def incr_list(xs: MathList) -> MathList:
+    return xs.map(lambda x: x + Math(1))
+
+
+egraph = EGraph()
+y = egraph.let("y", incr_list(MathList().append(Math(1)).append(Math(2))))
+egraph.run(math_ruleset.saturate())
+egraph.check(eq(y).to(MathList().append(Math(2)).append(Math(3))))
+
+egraph

egglog/examples/lambda_.py

@@ -0,0 +1,288 @@
+# mypy: disable-error-code="empty-body"
+
+"""
+Lambda Calculus
+===============
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, ClassVar
+
+from egglog import *
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+
+class Val(Expr):
+    """
+    A value is a number or a boolean.
+    """
+
+    TRUE: ClassVar[Val]
+    FALSE: ClassVar[Val]
+
+    def __init__(self, v: i64Like) -> None: ...
+
+
+class Var(Expr):
+    def __init__(self, v: StringLike) -> None: ...
+
+
+class Term(Expr):
+    @classmethod
+    def val(cls, v: Val) -> Term: ...
+
+    @classmethod
+    def var(cls, v: Var) -> Term: ...
+
+    def __add__(self, other: Term) -> Term: ...
+
+    def __eq__(self, other: Term) -> Term:  # type: ignore[override]
+        ...
+
+    def __call__(self, other: Term) -> Term: ...
+
+    def eval(self) -> Val: ...
+
+    def v(self) -> Var: ...
+
+
+@function
+def lam(x: Var, t: Term) -> Term: ...
+
+
+@function
+def let_(x: Var, t: Term, b: Term) -> Term: ...
+
+
+@function
+def fix(x: Var, t: Term) -> Term: ...
+
+
+@function
+def if_(c: Term, t: Term, f: Term) -> Term: ...
+
+
+StringSet = Set[Var]
+
+
+@function(merge=lambda old, new: old & new)
+def freer(t: Term) -> StringSet: ...
+
+
+(v, v1, v2) = vars_("v v1 v2", Val)
+(t, t1, t2, t3, t4) = vars_("t t1 t2 t3 t4", Term)
+(x, y) = vars_("x y", Var)
+fv, fv1, fv2, fv3 = vars_("fv fv1 fv2 fv3", StringSet)
+i1, i2 = vars_("i1 i2", i64)
+lamdba_ruleset = ruleset(
+    # freer
+    rule(eq(t).to(Term.val(v))).then(set_(freer(t)).to(StringSet.empty())),
+    rule(eq(t).to(Term.var(x))).then(set_(freer(t)).to(StringSet.empty().insert(x))),
+    rule(eq(t).to(t1 + t2), eq(freer(t1)).to(fv1), eq(freer(t2)).to(fv2)).then(set_(freer(t)).to(fv1 | fv2)),
+    rule(eq(t).to(t1 == t2), eq(freer(t1)).to(fv1), eq(freer(t2)).to(fv2)).then(set_(freer(t)).to(fv1 | fv2)),
+    rule(eq(t).to(t1(t2)), eq(freer(t1)).to(fv1), eq(freer(t2)).to(fv2)).then(set_(freer(t)).to(fv1 | fv2)),
+    rule(eq(t).to(lam(x, t1)), eq(freer(t1)).to(fv)).then(set_(freer(t)).to(fv.remove(x))),
+    rule(eq(t).to(let_(x, t1, t2)), eq(freer(t1)).to(fv1), eq(freer(t2)).to(fv2)).then(
+        set_(freer(t)).to(fv1.remove(x) | fv2)
+    ),
+    rule(eq(t).to(fix(x, t1)), eq(freer(t1)).to(fv)).then(set_(freer(t)).to(fv.remove(x))),
+    rule(eq(t).to(if_(t1, t2, t3)), eq(freer(t1)).to(fv1), eq(freer(t2)).to(fv2), eq(freer(t3)).to(fv3)).then(
+        set_(freer(t)).to(fv1 | fv2 | fv3)
+    ),
+    # eval
+    rule(eq(t).to(Term.val(v))).then(union(t.eval()).with_(v)),
+    rule(eq(t).to(t1 + t2), eq(Val(i1)).to(t1.eval()), eq(Val(i2)).to(t2.eval())).then(
+        union(t.eval()).with_(Val(i1 + i2))
+    ),
+    rule(eq(t).to(t1 == t2), eq(t1.eval()).to(t2.eval())).then(union(t.eval()).with_(Val.TRUE)),
+    rule(eq(t).to(t1 == t2), eq(t1.eval()).to(v1), eq(t2.eval()).to(v2), ne(v1).to(v2)).then(
+        union(t.eval()).with_(Val.FALSE)
+    ),
+    rule(eq(v).to(t.eval())).then(union(t).with_(Term.val(v))),
+    # if
+    rewrite(if_(Term.val(Val.TRUE), t1, t2)).to(t1),
+    rewrite(if_(Term.val(Val.FALSE), t1, t2)).to(t2),
+    # if-elim
+    # Adds let rules so next one can match on them
+    rule(eq(t).to(if_(Term.var(x) == t1, t2, t3))).then(let_(x, t1, t2), let_(x, t1, t3)),
+    rewrite(if_(Term.var(x) == t1, t2, t3)).to(
+        t3,
+        eq(let_(x, t1, t2)).to(let_(x, t1, t3)),
+    ),
+    # add-comm
+    rewrite(t1 + t2).to(t2 + t1),
+    # add-assoc
+    rewrite((t1 + t2) + t3).to(t1 + (t2 + t3)),
+    # eq-comm
+    rewrite(t1 == t2).to(t2 == t1),
+    # Fix
+    rewrite(fix(x, t)).to(let_(x, fix(x, t), t)),
+    # beta reduction
+    rewrite(lam(x, t)(t1)).to(let_(x, t1, t)),
+    # let-app
+    rewrite(let_(x, t, t1(t2))).to(let_(x, t, t1)(let_(x, t, t2))),
+    # let-add
+    rewrite(let_(x, t, t1 + t2)).to(let_(x, t, t1) + let_(x, t, t2)),
+    # let-eq
+    rewrite(let_(x, t, t1 == t2)).to(let_(x, t, t1) == let_(x, t, t2)),
+    # let-const
+    rewrite(let_(x, t, Term.val(v))).to(Term.val(v)),
+    # let-if
+    rewrite(let_(x, t, if_(t1, t2, t3))).to(if_(let_(x, t, t1), let_(x, t, t2), let_(x, t, t3))),
+    # let-var-same
+    rewrite(let_(x, t, Term.var(x))).to(t),
+    # let-var-diff
+    rewrite(let_(x, t, Term.var(y))).to(Term.var(y), ne(x).to(y)),
+    # let-lam-same
+    rewrite(let_(x, t, lam(x, t1))).to(lam(x, t1)),
+    # let-lam-diff
+    rewrite(let_(x, t, lam(y, t1))).to(lam(y, let_(x, t, t1)), ne(x).to(y), eq(fv).to(freer(t)), fv.not_contains(y)),
+    rule(eq(t).to(let_(x, t1, lam(y, t2))), ne(x).to(y), eq(fv).to(freer(t1)), fv.contains(y)).then(
+        union(t).with_(lam(t.v(), let_(x, t1, let_(y, Term.var(t.v()), t2))))
+    ),
+)
+
+result = relation("result")
+
+
+def l(fn: Callable[[Term], Term]) -> Term:  # noqa: E743
+    """
+    Create a lambda term from a function
+    """
+    # Use first var name from fn
+    x = fn.__code__.co_varnames[0]
+    return lam(Var(x), fn(Term.var(Var(x))))
+
+
+def assert_simplifies(left: Expr, right: Expr) -> None:
+    """
+    Simplify and print
+    """
+    print(f"{left} ➡  {right}")
+    check(eq(left).to(right), lamdba_ruleset * 30, left)
+
+
+assert_simplifies((Term.val(Val(1))).eval(), Val(1))
+assert_simplifies((Term.val(Val(1)) + Term.val(Val(2))).eval(), Val(3))
+
+
+# lambda under
+assert_simplifies(
+    l(lambda x: Term.val(Val(4)) + l(lambda y: y)(Term.val(Val(4)))),
+    l(lambda x: Term.val(Val(8))),
+)
+# lambda if elim
+a = Term.var(Var("a"))
+b = Term.var(Var("b"))
+assert_simplifies(if_(a == b, a + a, a + b), a + b)
+
+# lambda let simple
+x = Var("x")
+y = Var("y")
+assert_simplifies(
+    let_(x, Term.val(Val(0)), let_(y, Term.val(Val(1)), Term.var(x) + Term.var(y))),
+    Term.val(Val(1)),
+)
+# lambda capture
+assert_simplifies(
+    let_(x, Term.val(Val(1)), l(lambda x: x)),
+    l(lambda x: x),
+)
+# lambda capture free
+egraph = EGraph()
+e5 = egraph.let("e5", let_(y, Term.var(x) + Term.var(x), l(lambda x: Term.var(y))))
+egraph.run(lamdba_ruleset * 10)
+egraph.check(freer(l(lambda x: Term.var(y))).contains(y))
+egraph.check_fail(eq(e5).to(l(lambda x: x + x)))
+
+# lambda_closure_not_seven
+egraph = EGraph()
+e6 = egraph.let(
+    "e6",
+    let_(
+        Var("five"),
+        Term.val(Val(5)),
+        let_(
+            Var("add-five"),
+            l(lambda x: x + Term.var(Var("five"))),
+            let_(Var("five"), Term.val(Val(6)), Term.var(Var("add-five"))(Term.val(Val(1)))),
+        ),
+    ),
+)
+egraph.run(lamdba_ruleset * 10)
+egraph.check_fail(eq(e6).to(Term.val(Val(7))))
+egraph.check(eq(e6).to(Term.val(Val(6))))
+
+
+# lambda_compose
+egraph = EGraph()
+compose = Var("compose")
+add1 = Var("add1")
+e7 = egraph.let(
+    "e7",
+    let_(
+        compose,
+        l(
+            lambda f: l(
+                lambda g: l(
+                    lambda x: f(g(x)),
+                ),
+            ),
+        ),
+        let_(
+            add1,
+            l(lambda y: y + Term.val(Val(1))),
+            Term.var(compose)(Term.var(add1))(Term.var(add1)),
+        ),
+    ),
+)
+egraph.run(lamdba_ruleset * 20)
+egraph.register(
+    rule(
+        eq(t1).to(l(lambda x: Term.val(Val(1)) + l(lambda y: Term.val(Val(1)) + y)(x))),
+        eq(t2).to(l(lambda x: x + Term.val(Val(2)))),
+    ).then(result())
+)
+egraph.run(1)
+egraph.check(result())
+
+
+# lambda_if_simple
+assert_simplifies(if_(Term.val(Val(1)) == Term.val(Val(1)), Term.val(Val(7)), Term.val(Val(9))), Term.val(Val(7)))
+
+
+# # lambda_compose_many
+assert_simplifies(
+    let_(
+        compose,
+        l(lambda f: l(lambda g: l(lambda x: f(g(x))))),
+        let_(
+            add1,
+            l(lambda y: y + Term.val(Val(1))),
+            Term.var(compose)(Term.var(add1))(
+                Term.var(compose)(Term.var(add1))(
+                    Term.var(compose)(Term.var(add1))(
+                        Term.var(compose)(Term.var(add1))(
+                            Term.var(compose)(Term.var(add1))(Term.var(compose)(Term.var(add1))(Term.var(add1)))
+                        )
+                    )
+                )
+            ),
+        ),
+    ),
+    l(lambda x: x + Term.val(Val(7))),
+)
+
+# lambda_if
+zeroone = Var("zeroone")
+assert_simplifies(
+    let_(
+        zeroone,
+        l(lambda x: if_(x == Term.val(Val(0)), Term.val(Val(0)), Term.val(Val(1)))),
+        Term.var(zeroone)(Term.val(Val(0))) + Term.var(zeroone)(Term.val(Val(10))),
+    ),
+    Term.val(Val(1)),
+)

egglog/examples/matrix.py

@@ -0,0 +1,175 @@
+"""
+Matrix multiplication and Kronecker product.
+============================================
+"""
+
+from __future__ import annotations
+
+from egglog import *
+
+egraph = EGraph()
+
+
+class Dim(Expr):
+    """
+    A dimension of a matix.
+
+    >>> Dim(3) * Dim.named("n")
+    Dim(3) * Dim.named("n")
+    """
+
+    @method(egg_fn="Lit")
+    def __init__(self, value: i64Like) -> None: ...
+
+    @method(egg_fn="NamedDim")
+    @classmethod
+    def named(cls, name: StringLike) -> Dim:  # type: ignore[empty-body]
+        ...
+
+    @method(egg_fn="Times")
+    def __mul__(self, other: Dim) -> Dim:  # type: ignore[empty-body]
+        ...
+
+
+a, b, c, n = vars_("a b c n", Dim)
+i, j = vars_("i j", i64)
+egraph.register(
+    rewrite(a * (b * c)).to((a * b) * c),
+    rewrite((a * b) * c).to(a * (b * c)),
+    rewrite(Dim(i) * Dim(j)).to(Dim(i * j)),
+    rewrite(a * b).to(b * a),
+)
+
+
+class Matrix(Expr, egg_sort="MExpr"):
+    @method(egg_fn="Id")
+    @classmethod
+    def identity(cls, dim: Dim) -> Matrix:  # type: ignore[empty-body]
+        """
+        Create an identity matrix of the given dimension.
+        """
+
+    @method(egg_fn="NamedMat")
+    @classmethod
+    def named(cls, name: StringLike) -> Matrix:  # type: ignore[empty-body]
+        """
+        Create a named matrix.
+        """
+
+    @method(egg_fn="MMul")
+    def __matmul__(self, other: Matrix) -> Matrix:  # type: ignore[empty-body]
+        """
+        Matrix multiplication.
+        """
+
+    @method(egg_fn="nrows")
+    def nrows(self) -> Dim:  # type: ignore[empty-body]
+        """
+        Number of rows in the matrix.
+        """
+
+    @method(egg_fn="ncols")
+    def ncols(self) -> Dim:  # type: ignore[empty-body]
+        """
+        Number of columns in the matrix.
+        """
+
+
+@function(egg_fn="Kron")
+def kron(a: Matrix, b: Matrix) -> Matrix:  # type: ignore[empty-body]
+    """
+    Kronecker product of two matrices.
+
+    https://en.wikipedia.org/wiki/Kronecker_product#Definition
+    """
+
+
+A, B, C, D = vars_("A B C D", Matrix)
+egraph.register(
+    # The dimensions of a kronecker product are the product of the dimensions
+    rewrite(kron(A, B).nrows()).to(A.nrows() * B.nrows()),
+    rewrite(kron(A, B).ncols()).to(A.ncols() * B.ncols()),
+    # The dimensions of a matrix multiplication are the number of rows of the first
+    # matrix and the number of columns of the second matrix.
+    rewrite((A @ B).nrows()).to(A.nrows()),
+    rewrite((A @ B).ncols()).to(B.ncols()),
+    # The dimensions of an identity matrix are the input dimension
+    rewrite(Matrix.identity(n).nrows()).to(n),
+    rewrite(Matrix.identity(n).ncols()).to(n),
+)
+egraph.register(
+    # Multiplication by an identity matrix is the same as the other matrix
+    rewrite(Matrix.identity(n) @ A).to(A),
+    rewrite(A @ Matrix.identity(n)).to(A),
+    # Matrix multiplication is associative
+    rewrite(A @ (B @ C)).to((A @ B) @ C),
+    rewrite((A @ B) @ C).to(A @ (B @ C)),
+    # Kronecker product is associative
+    rewrite(kron(A, kron(B, C))).to(kron(kron(A, B), C)),
+    rewrite(kron(kron(A, B), C)).to(kron(A, kron(B, C))),
+    # Kronecker product distributes over matrix multiplication
+    rewrite(kron(A @ C, B @ D)).to(kron(A, B) @ kron(C, D)),
+    rewrite(kron(A, B) @ kron(C, D)).to(
+        kron(A @ C, B @ D),
+        # Only when the dimensions match
+        eq(A.ncols()).to(C.nrows()),
+        eq(B.ncols()).to(D.nrows()),
+    ),
+)
+egraph.register(
+    # demand rows and columns when we multiply matrices
+    rule(eq(C).to(A @ B)).then(
+        A.ncols(),
+        A.nrows(),
+        B.ncols(),
+        B.nrows(),
+    ),
+    # demand rows and columns when we take the kronecker product
+    rule(eq(C).to(kron(A, B))).then(
+        A.ncols(),
+        A.nrows(),
+        B.ncols(),
+        B.nrows(),
+    ),
+)
+
+
+# Define a number of dimensions
+n = egraph.let("n", Dim.named("n"))
+m = egraph.let("m", Dim.named("m"))
+p = egraph.let("p", Dim.named("p"))
+
+# Define a number of matrices
+A = egraph.let("A", Matrix.named("A"))
+B = egraph.let("B", Matrix.named("B"))
+C = egraph.let("C", Matrix.named("C"))
+
+# Set each to be a square matrix of the given dimension
+egraph.register(
+    union(A.nrows()).with_(n),
+    union(A.ncols()).with_(n),
+    union(B.nrows()).with_(m),
+    union(B.ncols()).with_(m),
+    union(C.nrows()).with_(p),
+    union(C.ncols()).with_(p),
+)
+# Create an example which should equal the kronecker product of A and B
+ex1 = egraph.let("ex1", kron(Matrix.identity(n), B) @ kron(A, Matrix.identity(m)))
+rows = egraph.let("rows", ex1.nrows())
+cols = egraph.let("cols", ex1.ncols())
+
+egraph.run(20)
+
+egraph.check(eq(B.nrows()).to(m))
+egraph.check(eq(kron(Matrix.identity(n), B).nrows()).to(n * m))
+
+# Verify it matches the expected result
+simple_ex1 = egraph.let("simple_ex1", kron(A, B))
+egraph.check(eq(ex1).to(simple_ex1))
+
+ex2 = egraph.let("ex2", kron(Matrix.identity(p), C) @ kron(A, Matrix.identity(m)))
+
+egraph.run(10)
+# Verify it is not simplified
+egraph.check_fail(eq(ex2).to(kron(A, C)))
+egraph

egglog/examples/multiset.py

@@ -0,0 +1,61 @@
+# mypy: disable-error-code="empty-body"
+"""
+Multiset example based off of egglog version
+============================================
+"""
+
+from __future__ import annotations
+
+from collections import Counter
+
+from egglog import *
+
+
+class Math(Expr):
+    def __init__(self, x: i64Like) -> None: ...
+
+
+@function
+def square(x: Math) -> Math: ...
+
+
+@ruleset
+def math_ruleset(i: i64):
+    yield rewrite(square(Math(i))).to(Math(i * i))
+
+
+egraph = EGraph()
+
+xs = MultiSet(Math(1), Math(2), Math(3))
+egraph.register(xs)
+
+with egraph.set_current():
+    assert xs == MultiSet(Math(1), Math(3), Math(2))
+    assert xs != MultiSet(Math(1), Math(1), Math(2), Math(3))
+
+    assert Counter(xs) == Counter({Math(1): 1, Math(2): 1, Math(3): 1})
+
+    inserted = MultiSet(Math(1), Math(2), Math(3), Math(4))
+    egraph.register(inserted)
+    assert xs.insert(Math(4)) == inserted
+
+    assert xs.contains(Math(1))
+    assert xs.not_contains(Math(4))
+    assert Math(1) in xs
+    assert Math(4) not in xs
+
+    assert xs.remove(Math(1)) == MultiSet(Math(2), Math(3))
+
+    assert xs.length() == i64(3)
+    assert len(xs) == 3
+
+    assert MultiSet(Math(1), Math(1)).length() == i64(2)
+
+    assert MultiSet(Math(1)).pick() == Math(1)
+
+    mapped = xs.map(square)
+    egraph.register(mapped)
+    egraph.run(math_ruleset)
+    assert mapped == MultiSet(Math(1), Math(4), Math(9))
+
+    assert xs + xs == MultiSet(Math(1), Math(2), Math(3), Math(1), Math(2), Math(3))