egglog 11.2.0__cp314-cp314-manylinux_2_17_x86_64.manylinux2014_x86_64.whl

egglog/examples/bignum.py

@@ -0,0 +1,32 @@
+# mypy: disable-error-code="empty-body"
+"""
+BigNum/BigRat Example
+=====================
+"""
+
+from __future__ import annotations
+
+from egglog import *
+
+x = BigInt(-1234)
+y = BigInt.from_string("2")
+z = BigRat(x, y)
+
+egraph = EGraph()
+
+assert egraph.extract(z.numer.to_string()).value == "-617"
+
+
+@function
+def bignums(x: BigInt, y: BigInt) -> BigRat: ...
+
+
+egraph.register(set_(bignums(x, y)).to(z))
+
+c = var("c", BigRat)
+a, b = vars_("a b", BigInt)
+egraph.check(
+    bignums(a, b) == c,
+    c.numer == a >> 1,
+    c.denom == b >> 1,
+)

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,44 @@
+# 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)
+
+egraph = EGraph()
+egraph.register(expr1, expr2)
+
+egraph.run(
+    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
+)
+
+egraph.check(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,42 @@
+# mypy: disable-error-code="empty-body"
+"""
+Higher Order Functions
+======================
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+
+from egglog import *
+
+
+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/jointree.py

@@ -0,0 +1,67 @@
+# mypy: disable-error-code="empty-body"
+
+"""
+Join Tree (custom costs)
+========================
+
+Example of using custom cost functions for jointree.
+
+From https://egraphs.zulipchat.com/#narrow/stream/328972-general/topic/How.20can.20I.20find.20the.20tree.20associated.20with.20an.20extraction.3F
+"""
+
+from __future__ import annotations
+
+from egglog import *
+
+
+class JoinTree(Expr):
+    def __init__(self, name: StringLike) -> None: ...
+
+    def join(self, other: JoinTree) -> JoinTree: ...
+
+    @method(merge=lambda old, new: old.min(new))  # type:ignore[prop-decorator]
+    @property
+    def size(self) -> i64: ...
+
+
+ra = JoinTree("a")
+rb = JoinTree("b")
+rc = JoinTree("c")
+rd = JoinTree("d")
+re = JoinTree("e")
+rf = JoinTree("f")
+
+query = ra.join(rb).join(rc).join(rd).join(re).join(rf)
+
+egraph = EGraph()
+egraph.register(
+    set_(ra.size).to(50),
+    set_(rb.size).to(200),
+    set_(rc.size).to(10),
+    set_(rd.size).to(123),
+    set_(re.size).to(10000),
+    set_(rf.size).to(1),
+)
+
+
+@egraph.register
+def _rules(s: String, a: JoinTree, b: JoinTree, c: JoinTree, asize: i64, bsize: i64):
+    # cost of relation is its size minus 1, since the string arg will have a cost of 1 as well
+    yield rule(JoinTree(s).size == asize).then(set_cost(JoinTree(s), asize - 1))
+    # cost/size of join is product of sizes
+    yield rule(a.join(b), a.size == asize, b.size == bsize).then(
+        set_(a.join(b).size).to(asize * bsize), set_cost(a.join(b), asize * bsize)
+    )
+    # associativity
+    yield rewrite(a.join(b)).to(b.join(a))
+    # commutativity
+    yield rewrite(a.join(b).join(c)).to(a.join(b.join(c)))
+
+
+egraph.register(query)
+egraph.run(1000)
+print(egraph.extract(query))
+print(egraph.extract(query.size))
+
+
+egraph

egglog/examples/lambda_.py

@@ -0,0 +1,287 @@
+# mypy: disable-error-code="empty-body"
+
+"""
+Lambda Calculus
+===============
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from typing import ClassVar
+
+from egglog import *
+from egglog import Expr
+
+
+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)),
+)