Add rewrite

Author: saulshanabrook

python/egg_smol/bindings.pyi

@@ -2,7 +2,7 @@ from typing import Optional
 
 from typing_extensions import final
 
-from .bindings_py import Expr, FunctionDecl, Variant
+from .bindings_py import Expr, FunctionDecl, Rewrite, Variant
 
 @final
 class EGraph:
@@ -11,6 +11,7 @@ class EGraph:
     def declare_sort(self, name: str) -> None: ...
     def declare_function(self, decl: FunctionDecl) -> None: ...
     def define(self, name: str, expr: Expr, cost: Optional[int] = None) -> None: ...
+    def add_rewrite(self, rewrite: Rewrite) -> str: ...
 
 @final
 class EggSmolError(Exception):

python/egg_smol/bindings_py.py

@@ -1,7 +1,7 @@
 # TODO: Figure out what these modules should be called
 from __future__ import annotations
 
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from typing import Optional, Union
 
 
@@ -62,3 +62,23 @@ class Unit:
 
 
 Literal = Union[Int, String, Unit]
+
+
+@dataclass(frozen=True)
+class Rewrite:
+    lhs: Expr
+    rhs: Expr
+    conditions: list[Fact_] = field(default_factory=list)
+
+
+@dataclass(frozen=True)
+class Fact:
+    expr: Expr
+
+
+@dataclass
+class Eq:
+    exprs: list[Expr]
+
+
+Fact_ = Union[Fact, Eq]

python/tests/test.py

@@ -68,3 +68,27 @@ def test_define(self):
                 ],
             ),
         )
+
+    def test_rewrite(self):
+        egraph = EGraph()
+        egraph.declare_sort("Math")
+        egraph.declare_constructor(Variant("Add", ["Math", "Math"]), "Math")
+        name = egraph.add_rewrite(
+            Rewrite(
+                Call(
+                    "Add",
+                    [
+                        Var("a"),
+                        Var("b"),
+                    ],
+                ),
+                Call(
+                    "Add",
+                    [
+                        Var("b"),
+                        Var("a"),
+                    ],
+                ),
+            )
+        )
+        assert isinstance(name, str)

src/conversions.rs

@@ -3,6 +3,7 @@
 //
 // Converts from Python classes we define in pure python so we can use dataclasses
 // to represent the input types
+// TODO: Copy strings of these from egg-smol... Maybe actually wrap those isntead.
 use pyo3::prelude::*;
 
 // Execute the block and wrap the error in a type error
@@ -175,3 +176,64 @@ impl From<WrappedLiteral> for egg_smol::ast::Literal {
         other.0
     }
 }
+
+// Wrapped version of Rewrite
+pub struct WrappedRewrite(egg_smol::ast::Rewrite);
+
+impl FromPyObject<'_> for WrappedRewrite {
+    fn extract(obj: &'_ PyAny) -> PyResult<Self> {
+        wrap_error("Rewrite", obj, || {
+            Ok(WrappedRewrite(egg_smol::ast::Rewrite {
+                lhs: obj.getattr("lhs")?.extract::<WrappedExpr>()?.into(),
+                rhs: obj.getattr("rhs")?.extract::<WrappedExpr>()?.into(),
+                conditions: obj
+                    .getattr("conditions")?
+                    .extract::<Vec<WrappedFact>>()?
+                    .into_iter()
+                    .map(|x| x.into())
+                    .collect(),
+            }))
+        })
+    }
+}
+
+impl From<WrappedRewrite> for egg_smol::ast::Rewrite {
+    fn from(other: WrappedRewrite) -> Self {
+        other.0
+    }
+}
+
+// Wrapped version of Fact
+pub struct WrappedFact(egg_smol::ast::Fact);
+
+impl FromPyObject<'_> for WrappedFact {
+    fn extract(obj: &'_ PyAny) -> PyResult<Self> {
+        wrap_error("Fact", obj, || {
+            extract_fact_eq(obj)
+                .or_else(|_| extract_fact_fact(obj))
+                .map(WrappedFact)
+        })
+    }
+}
+
+fn extract_fact_eq(obj: &PyAny) -> PyResult<egg_smol::ast::Fact> {
+    Ok(egg_smol::ast::Fact::Eq(
+        obj.getattr("exprs")?
+            .extract::<Vec<WrappedExpr>>()?
+            .into_iter()
+            .map(|x| x.into())
+            .collect(),
+    ))
+}
+
+fn extract_fact_fact(obj: &PyAny) -> PyResult<egg_smol::ast::Fact> {
+    Ok(egg_smol::ast::Fact::Fact(
+        obj.getattr("expr")?.extract::<WrappedExpr>()?.into(),
+    ))
+}
+
+impl From<WrappedFact> for egg_smol::ast::Fact {
+    fn from(other: WrappedFact) -> Self {
+        other.0
+    }
+}

src/lib.rs

@@ -22,6 +22,13 @@ impl EGraph {
         }
     }
 
+    /// Define a rewrite rule, returning the name of the rule
+    #[pyo3(text_signature = "($self, rewrite)")]
+    fn add_rewrite(&mut self, rewrite: WrappedRewrite) -> EggResult<String> {
+        let res = self.egraph.add_rewrite(rewrite.into())?;
+        Ok(res.to_string())
+    }
+
     /// Define a new named value.
     #[pyo3(
         text_signature = "($self, name, expr, cost=None)",