Add egglog tutorials

Author: saulshanabrook

docs/tutorials.md

@@ -5,4 +5,8 @@ auto_examples/index
 tutorials/getting-started
 tutorials/sklearn
 tutorials/tut_1_basics
+tutorials/tut_2_datalog
+tutorials/tut_3_analysis
+tutorials/tut_4_scheduling
+tutorials/tut_5_extraction
 ```

docs/tutorials/__init__.py

@@ -0,0 +1,3 @@
+"""
+Tutorials
+"""

docs/tutorials/tut_1_basics.py

@@ -10,6 +10,7 @@
 # mypy: disable-error-code="empty-body"
 from __future__ import annotations
 from typing import TypeAlias
+from collections.abc import Iterable
 from egglog import *
 
 
@@ -19,67 +20,94 @@ def __init__(self, value: i64Like) -> None: ...
     @classmethod
     def var(cls, name: StringLike) -> Num: ...
 
-    def __add__(self, other: ExprLike) -> Num: ...
-    def __mul__(self, other: ExprLike) -> Num: ...
+    def __add__(self, other: NumLike) -> Num: ...
+    def __mul__(self, other: NumLike) -> Num: ...
 
     # Support inverse operations for convenience
     # they will be translated to non-reversed ones
-    def __radd__(self, other: ExprLike) -> Num: ...
-    def __rmul__(self, other: ExprLike) -> Num: ...
+    def __radd__(self, other: NumLike) -> Num: ...
+    def __rmul__(self, other: NumLike) -> Num: ...
 
 
-ExprLike: TypeAlias = Num | StringLike | i64Like
-converter(i64, Num, Num)
-converter(str, Num, Num.var)
+NumLike: TypeAlias = Num | StringLike | i64Like
 # -
 
+
+# The signature here takes `NumLike` not `Num` so that you can write `Num(1) + 2` instead of
+# `Num(1) + Num(2)`. This is helpful for ease of use and also for compatibility when you are trying to
+# create expressions that act like Python objects which perform upcasting.
+#
+# To support this, you must define conversions between primitive types and your expression types.
+# When a value is passed into a function, it will find the type it should be converted to and
+# transitively apply the conversions you have defined:
+
+converter(i64, Num, Num)
+converter(String, Num, Num.var)
+
 # Now, let's define some simple expressions.
 
+egraph = EGraph()
 x = Num.var("x")
-expr1 = 2 * (x * 3)
-expr2 = 6 * x
+expr1 = egraph.let("expr1", 2 * (x * 3))
+expr2 = egraph.let("expr2", 6 * x)
 
 # You should see an e-graph with two expressions.
 
-egraph = EGraph()
-egraph.register(expr1, expr2)
 egraph
 
-# We can print the values of the expressions as well to see their fully expanded forms.
+# We can `.extract` the values of the expressions as well to see their fully expanded forms.
 
-String("Hello, world!")
+egraph.extract(String("Hello, world!"))
 
-i64(42)
+egraph.extract(i64(42))
 
-expr1
+egraph.extract(expr1)
 
-expr2
+egraph.extract(expr2)
 
 # We can use the `check` commands to check properties of our e-graph.
 
 x, y = vars_("x y", Num)
-assert egraph.check_bool(expr1 == x * y)
+egraph.check(expr1 == x * y)
 
 # This checks if `expr1` is equivalent to some expression `x * y`, where `x` and `y` are
 # variables that can be mapped to any `Num` expression in the e-graph.
 #
 # Checks can fail. For example the following check fails because `expr1` is not equivalent to
 # `x + y` for any `x` and `y` in the e-graph.
 
-assert not egraph.check_bool(expr1 == x + y)
+egraph.check_fail(expr1 == x + y)
 
 # Let us define some rewrite rules over our small DSL.
 
-egraph.register(rewrite(x + y).to(y + x))
+
+@egraph.register
+def _add_comm(x: Num, y: Num):
+    yield rewrite(x + y).to(y + x)
+
+
+# This could also been written like:
+#
+# ```python
+# x, y = vars_("x y", Num)
+# egraph.register(rewrite(x + y).to(y + x))
+# ```
+#
+# In this tutorial we will use the function form to define rewrites and rules, because then then we only
+# have to write the variable names once as arguments and they are not leaked to the outer scope.
+
 
 # This rule asserts that addition is commutative. More concretely, this rules says, if the e-graph
 # contains expressions of the form `x + y`, then the e-graph should also contain the
 # expression `y + x`, and they should be equivalent.
 #
 # Similarly, we can define the associativity rule for addition.
 
-z = var("z", Num)
-egraph.register(rewrite(x + (y + z)).to((x + y) + z))
+
+@egraph.register
+def _add_assoc(x: Num, y: Num, z: Num) -> Iterable[RewriteOrRule]:
+    yield rewrite(x + (y + z)).to((x + y) + z)
+
 
 # This rule says, if the e-graph contains expressions of the form `x + (y + z)`, then the e-graph should also contain
 # the expression `(x + y) + z`, and they should be equivalent.
@@ -89,20 +117,22 @@ def __rmul__(self, other: ExprLike) -> Num: ...
 # 1. Defining a rule is different from running it. The following check would fail at this point
 #    because the commutativity rule has not been run (we've inserted `x + 3` but not yet derived `3 + x`).
 
-assert not egraph.check_bool((x + 3) == (3 + x))
+egraph.check_fail((x + 3) == (3 + x))
 
 # 2. Rules are not instantiated for every possible term; they are only instantiated for terms that are
 #    in the e-graph. For instance, even if we ran the commutativity rule above, the following check would
 #    still fail because the e-graph does not contain either of the terms `Num(-2) + Num(2)` or `Num(2) + Num(-2)`.
 
-assert not egraph.check_bool(Num(-2) + 2 == Num(2) + -2)
+egraph.check_fail(Num(-2) + 2 == Num(2) + -2)
 
 # Let's also define commutativity and associativity for multiplication.
 
-egraph.register(
-    rewrite(x * y).to(y * x),
-    rewrite(x * (y * z)).to((x * y) * z),
-)
+
+@egraph.register
+def _mul(x: Num, y: Num, z: Num) -> Iterable[RewriteOrRule]:
+    yield rewrite(x * y).to(y * x)
+    yield rewrite(x * (y * z)).to((x * y) * z)
+
 
 # `egglog` also defines a set of built-in functions over primitive types, such as `+` and `*`,
 # and supports operator overloading, so the same operator can be used with different types.
@@ -116,11 +146,12 @@ def __rmul__(self, other: ExprLike) -> Num: ...
 # With primitives, we can define rewrite rules that talk about the semantics of operators.
 # The following rules show constant folding over addition and multiplication.
 
-a, b = vars_("a b", i64)
-egraph.register(
-    rewrite(Num(a) + Num(b)).to(Num(a + b)),
-    rewrite(Num(a) * Num(b)).to(Num(a * b)),
-)
+
+@egraph.register
+def _const_fold(a: i64, b: i64) -> Iterable[RewriteOrRule]:
+    yield rewrite(Num(a) + Num(b)).to(Num(a + b))
+    yield rewrite(Num(a) * Num(b)).to(Num(a * b))
+
 
 # While we have defined several rules, the e-graph has not changed since we inserted the two
 # expressions. To run rules we have defined so far, we can use `run`.