Add simplify to fix docs

Author: saulshanabrook

docs/how-to-guides.md

@@ -6,13 +6,18 @@ file_format: mystnb
 
 ## Parsing and running program strings
 
-You can provide your program in a special DSL language and run it with {meth}`egg_smol.bindings.EGraph.parse_and_run_program`:
+You can provide your program in a special DSL language. You can parse this with {meth}`egg_smol.bindings.EGraph.parse_program` and then run the result with You can parse this with {meth}`egg_smol.bindings.EGraph.run_program`::
 
 ```{code-cell}
 from egg_smol.bindings import EGraph
 
 egraph = EGraph()
-egraph.parse_and_run_program("(check (= (+ 1 2) 3))")
+commands = egraph.parse_program("(check (= (+ 1 2) 3))")
+commands
+```
+
+```{code-cell}
+egraph.run_program(*commands)
 ```
 
 ## Developing this package

docs/tutorials/getting-started.ipynb

@@ -133,7 +133,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 3,
+      "execution_count": 8,
       "id": "b6424530",
       "metadata": {},
       "outputs": [
@@ -143,7 +143,7 @@
               "(Dim.named(\"x\") * Dim(10)) * Dim(10)"
             ]
           },
-          "execution_count": 3,
+          "execution_count": 8,
           "metadata": {},
           "output_type": "execute_result"
         }
@@ -195,6 +195,7 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "id": "167722d1-60b8-452a-ae54-6a8df4db5b00",
       "metadata": {},
@@ -221,33 +222,12 @@
       "source": [
         "### Testing\n",
         "Going back to the notebook, we can test out the that the rewrites are working.\n",
-        "\n",
-        "First, we have to add our expression to the egraph. We can do this by defining an empty `let` whcih uses this expression:"
+        "We can run some number of iterations and extract out the lowest cost expression which is equivalent to our variable:"
       ]
     },
     {
       "cell_type": "code",
       "execution_count": 5,
-      "id": "29c6d0c1-3249-4597-9d01-ec1fb29dd13f",
-      "metadata": {
-        "tags": []
-      },
-      "outputs": [],
-      "source": [
-        "egraph.register(let(\"\", res))"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "id": "5dc24d6e-8145-4ab3-b114-024dbb53323f",
-      "metadata": {},
-      "source": [
-        "We can then run some number of iterations and extract out the lowest cost expression which is equivalent to our variable:"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": 6,
       "id": "31afa12e-da68-4398-91fa-14523f6c099a",
       "metadata": {
         "tags": []
@@ -259,48 +239,13 @@
               "Dim.named(\"x\") * Dim(100)"
             ]
           },
-          "execution_count": 6,
+          "execution_count": 5,
           "metadata": {},
           "output_type": "execute_result"
         }
       ],
       "source": [
-        "egraph.run(10)\n",
-        "egraph.extract(res)"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "id": "bd08d366-0ebe-4e43-b219-b57fbb13534d",
-      "metadata": {},
-      "source": [
-        "We can also extract a number of variants to see all the equivalent expresions, ordered by their cost:"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": 7,
-      "id": "17e39076-0257-4bcd-b1c4-04eb5a79791e",
-      "metadata": {
-        "tags": []
-      },
-      "outputs": [
-        {
-          "data": {
-            "text/plain": [
-              "[Dim(100) * Dim.named(\"x\"),\n",
-              " Dim.named(\"x\") * Dim(100),\n",
-              " Dim(10) * (Dim.named(\"x\") * Dim(10)),\n",
-              " (Dim.named(\"x\") * Dim(10)) * Dim(10)]"
-            ]
-          },
-          "execution_count": 7,
-          "metadata": {},
-          "output_type": "execute_result"
-        }
-      ],
-      "source": [
-        "egraph.extract_multiple(res, 10)"
+        "egraph.simplify(res, 10)"
       ]
     },
     {
@@ -316,7 +261,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 8,
+      "execution_count": 9,
       "id": "c5b96cfb",
       "metadata": {},
       "outputs": [],
@@ -378,7 +323,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 9,
+      "execution_count": 10,
       "id": "cb2b4fb8",
       "metadata": {},
       "outputs": [],
@@ -409,33 +354,26 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 10,
+      "execution_count": 13,
       "id": "8d18be2d",
       "metadata": {},
       "outputs": [
         {
-          "data": {
-            "text/plain": [
-              "(Dim.named(\"x\"), Dim.named(\"y\"))"
-            ]
-          },
-          "execution_count": 10,
-          "metadata": {},
-          "output_type": "execute_result"
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "Dim.named(\"y\")\n",
+            "Dim.named(\"x\")\n"
+          ]
         }
       ],
       "source": [
         "# If we multiply two identity matrices, we should be able to get the number of columns of the result\n",
         "x = Matrix.identity(Dim.named(\"x\"))\n",
         "y = Matrix.identity(Dim.named(\"y\"))\n",
         "x_mult_y = x @ y\n",
-        "x_mult_y_ncols = x_mult_y.ncols()\n",
-        "x_mult_y_nrows = x_mult_y.nrows()\n",
-        "\n",
-        "egraph.register(let(\"\", x_mult_y_ncols), let(\"\", x_mult_y_nrows))\n",
-        "\n",
-        "egraph.run(10)\n",
-        "egraph.extract(x_mult_y_nrows), egraph.extract(x_mult_y_ncols)"
+        "print(egraph.simplify(x_mult_y.ncols(), 10))\n",
+        "print(egraph.simplify(x_mult_y.nrows(), 10))"
       ]
     },
     {
@@ -451,7 +389,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 11,
+      "execution_count": 14,
       "id": "18a91684",
       "metadata": {},
       "outputs": [],
@@ -488,7 +426,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 12,
+      "execution_count": 15,
       "id": "303ce7f3",
       "metadata": {},
       "outputs": [],
@@ -522,18 +460,17 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 13,
+      "execution_count": 16,
       "id": "bb50ade6",
       "metadata": {},
       "outputs": [
         {
           "data": {
             "text/plain": [
-              "[kron(Matrix.identity(Dim.named(\"n\")), Matrix.named(\"B\")) @ kron(Matrix.named(\"A\"), Matrix.identity(Dim.named(\"m\"))),\n",
-              " kron(Matrix.named(\"A\"), Matrix.named(\"B\"))]"
+              "kron(Matrix.named(\"A\"), Matrix.named(\"B\"))"
             ]
           },
-          "execution_count": 13,
+          "execution_count": 16,
           "metadata": {},
           "output_type": "execute_result"
         }
@@ -554,12 +491,7 @@
         ")\n",
         "# Create an example which should equal the kronecker product of A and B\n",
         "ex1 = kron(Matrix.identity(n), B) @ kron(A, Matrix.identity(m))\n",
-        "egraph.register(let(\"\", ex1))\n",
-        "\n",
-        "egraph.run(20)\n",
-        "# Verify it matches the expected result\n",
-        "egraph.check(eq(ex1).to(kron(A, B)))\n",
-        "egraph.extract_multiple(ex1, 10)"
+        "egraph.simplify(ex1, 20)"
       ]
     },
     {
@@ -573,32 +505,24 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 14,
+      "execution_count": 17,
       "id": "d8dea199",
       "metadata": {},
       "outputs": [
         {
           "data": {
             "text/plain": [
-              "[(kron(Matrix.identity(Dim.named(\"p\")), Matrix.named(\"C\")) @ kron(Matrix.identity(Dim.named(\"n\")), Matrix.identity(Dim.named(\"m\")))) @ kron(\n",
-              "     Matrix.named(\"A\"), Matrix.identity(Dim.named(\"m\"))\n",
-              " ),\n",
-              " kron(Matrix.identity(Dim.named(\"p\")), Matrix.named(\"C\")) @ kron(Matrix.named(\"A\"), Matrix.identity(Dim.named(\"m\")))]"
+              "kron(Matrix.identity(Dim.named(\"p\")), Matrix.named(\"C\")) @ kron(Matrix.named(\"A\"), Matrix.identity(Dim.named(\"m\")))"
             ]
           },
-          "execution_count": 14,
+          "execution_count": 17,
           "metadata": {},
           "output_type": "execute_result"
         }
       ],
       "source": [
         "ex2 = kron(Matrix.identity(p), C) @ kron(A, Matrix.identity(m))\n",
-        "egraph.register(let(\"\", ex2))\n",
-        "\n",
-        "egraph.run(10)\n",
-        "# Verify it is not simplified\n",
-        "egraph.check(ex2 != kron(A, C))\n",
-        "egraph.extract_multiple(ex2, 10)"
+        "egraph.simplify(ex2, 20)"
       ]
     },
     {

python/egg_smol/egraph.py

@@ -116,6 +116,22 @@ def _get_egraph(self) -> bindings.EGraph:
             raise RuntimeError("Cannot get the e-graph")
         return self._egraph
 
+    def simplify(self, expr: EXPR, limit: int, *until: Fact) -> EXPR:
+        """
+        Simplifies the given expression.
+        """
+        return self._simplify(expr, limit, None, until)
+
+    def _simplify(self, expr: EXPR, limit: int, ruleset: Optional[Ruleset], until: tuple[Fact, ...]) -> EXPR:
+        tp, decl = expr_parts(expr)
+        egg_expr = decl.to_egg(self._decls)
+        self._run_program([bindings.Simplify(egg_expr, Config(limit, ruleset, until)._to_egg_config(self._decls))])
+        extract_report = self._get_egraph().take_extract_report()
+        if not extract_report:
+            raise ValueError("No extract report saved")
+        new_tp, new_decl = tp_and_expr_decl_from_egg(self._decls, extract_report.expr)
+        return cast(EXPR, RuntimeExpr(self._decls, new_tp, new_decl))
+
     def relation(self, name: str, *tps: Unpack[TS], egg_fn: Optional[str] = None) -> Callable[[Unpack[TS]], Unit]:
         """
         Defines a relation, which is the same as a function which returns unit.
@@ -705,6 +721,11 @@ def run(self, limit: int, *until: Fact) -> bindings.RunReport:
         """
         return self._egraph._run_schedule(config(limit, self, *until))
 
+    def simplify(self, expr: EXPR, limit: int, *until: Fact) -> EXPR:
+        """
+        Simplify the given expression with this ruleset.
+        """
+        return self._egraph._simplify(expr, limit, self, until)
 
 # We use these builders so that when creating these structures we can type check
 # if the arguments are the same type of expression
@@ -1037,13 +1058,14 @@ def __str__(self) -> str:
         args_str = ", ".join(map(str, [self.limit, self.ruleset, *self.until]))
         return f"config({args_str})"
 
-    def _to_egg(self, declerations: Declarations) -> bindings._Schedule:
-        return bindings.Run(
-            bindings.RunConfig(
-                self.ruleset.name if self.ruleset else "",
-                self.limit,
-                [fact_decl_to_egg(declerations, _fact_to_decl(fact)) for fact in self.until] if self.until else None,
-            )
+    def _to_egg(self, decls: Declarations) -> bindings._Schedule:
+        return bindings.Run(self._to_egg_config(decls))
+
+    def _to_egg_config(self, decls: Declarations) -> bindings.RunConfig:
+        return bindings.RunConfig(
+            self.ruleset.name if self.ruleset else "",
+            self.limit,
+            [fact_decl_to_egg(decls, _fact_to_decl(fact)) for fact in self.until] if self.until else None,
         )