move variant into rust

Author: saulshanabrook

python/egg_smol/bindings.pyi

@@ -3,7 +3,16 @@ from typing import Optional
 
 from typing_extensions import final
 
-from .bindings_py import Expr, Fact_, FunctionDecl, Rewrite, Variant
+from .bindings_py import Expr, Fact_, FunctionDecl, Rewrite
+
+@final
+class Variant:
+    def __init__(
+        self, name: str, types: list[str], cost: Optional[int] = None
+    ) -> None: ...
+    name: str
+    types: list[str]
+    cost: Optional[int]
 
 @final
 class EGraph:

python/egg_smol/bindings_py.py

@@ -5,13 +5,6 @@
 from typing import Optional, Union
 
 
-@dataclass(frozen=True)
-class Variant:
-    name: str
-    types: list[str]
-    cost: Optional[int] = None
-
-
 @dataclass(frozen=True)
 class FunctionDecl:
     name: str

python/tests/test.py

@@ -1,7 +1,7 @@
 import datetime
 
 import pytest
-from egg_smol.bindings import EggSmolError, EGraph
+from egg_smol.bindings import *
 from egg_smol.bindings_py import *
 
 
@@ -334,3 +334,15 @@ def test_check_fact(self):
                 ]
             )
         )
+
+    # def test_extract(self):
+    #     # Example from extraction-cost
+    #     egraph = EGraph()
+    #     egraph.declare_sort("Expr")
+    #     egraph.declare_constructor(Variant("Num", ["i64"], cost=5), "Expr")
+
+    #     egraph.define("x", Call("Num", [Lit(Int(1))]), cost=10)
+    #     egraph.define("y", Call("Num", [Lit(Int(2))]), cost=1)
+
+    #     assert egraph.extract("x") == Call("Num", [Lit(Int(1))])
+    #     assert egraph.extract("y") == Var("y")

src/conversions.rs

@@ -6,7 +6,7 @@ use std::time::Duration;
 // 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::{ffi::PyDateTime_Delta, prelude::*, types::PyDelta};
+use pyo3::prelude::*;
 
 // Execute the block and wrap the error in a type error
 fn wrap_error<T>(tp: &str, obj: &'_ PyAny, block: impl FnOnce() -> PyResult<T>) -> PyResult<T> {
@@ -18,24 +18,55 @@ fn wrap_error<T>(tp: &str, obj: &'_ PyAny, block: impl FnOnce() -> PyResult<T>)
     })
 }
 
-// Wrapped version of Variant
-pub struct WrappedVariant(egg_smol::ast::Variant);
+// A variant of a constructor
 
-impl FromPyObject<'_> for WrappedVariant {
-    fn extract(obj: &'_ PyAny) -> PyResult<Self> {
-        wrap_error("Variant", obj, || {
-            Ok(WrappedVariant(egg_smol::ast::Variant {
-                name: obj.getattr("name")?.extract::<String>()?.into(),
-                cost: obj.getattr("cost")?.extract()?,
-                types: obj
-                    .getattr("types")?
-                    .extract::<Vec<String>>()?
-                    .into_iter()
-                    .map(|x| x.into())
-                    .collect(),
-            }))
+#[pyclass(name = "Variant")]
+#[derive(Clone)]
+pub(crate) struct WrappedVariant(egg_smol::ast::Variant);
+
+#[pymethods]
+impl WrappedVariant {
+    #[new]
+    fn new(name: String, types: Vec<String>, cost: Option<usize>) -> Self {
+        Self(egg_smol::ast::Variant {
+            name: name.into(),
+            types: types.into_iter().map(|x| x.into()).collect(),
+            cost,
         })
     }
+    #[getter]
+    fn name(&self) -> &str {
+        self.0.name.into()
+    }
+    #[getter]
+    fn types(&self) -> Vec<String> {
+        self.0.types.iter().map(|x| x.to_string()).collect()
+    }
+    #[getter]
+    fn cost(&self) -> Option<usize> {
+        self.0.cost
+    }
+
+    fn __repr__(&self) -> String {
+        format!(
+            "Variant(name={}, types=[{}], cost={})",
+            self.0.name.to_string(),
+            self.0
+                .types
+                .iter()
+                .map(|x| x.to_string())
+                .collect::<Vec<_>>()
+                .join(", "),
+            match self.0.cost {
+                Some(x) => x.to_string(),
+                None => "None".to_string(),
+            }
+        )
+    }
+
+    fn __str__(&self) -> String {
+        format!("{:#?}", self.0)
+    }
 }
 
 impl From<WrappedVariant> for egg_smol::ast::Variant {
@@ -44,6 +75,12 @@ impl From<WrappedVariant> for egg_smol::ast::Variant {
     }
 }
 
+impl From<egg_smol::ast::Variant> for WrappedVariant {
+    fn from(other: egg_smol::ast::Variant) -> Self {
+        WrappedVariant(other)
+    }
+}
+
 // Wrapped version of FunctionDecl
 pub struct WrappedFunctionDecl(egg_smol::ast::FunctionDecl);
 impl FromPyObject<'_> for WrappedFunctionDecl {
@@ -140,6 +177,12 @@ impl From<WrappedExpr> for egg_smol::ast::Expr {
     }
 }
 
+impl From<egg_smol::ast::Expr> for WrappedExpr {
+    fn from(other: egg_smol::ast::Expr) -> Self {
+        WrappedExpr(other)
+    }
+}
+
 // Wrapped version of Literal
 pub struct WrappedLiteral(egg_smol::ast::Literal);
 
@@ -253,7 +296,7 @@ impl From<Duration> for WrappedDuration {
 impl IntoPy<PyObject> for WrappedDuration {
     fn into_py(self, py: Python<'_>) -> PyObject {
         let d = self.0;
-        PyDelta::new(
+        pyo3::types::PyDelta::new(
             py,
             0,
             0,

src/lib.rs

@@ -23,6 +23,26 @@ impl EGraph {
         }
     }
 
+    // fn __repr__(&mut self) -> PyResult<String> {
+    //     Ok(format!("{:#?}", self.egraph))
+    // }
+
+    /// Extract the best expression of a given value. Will also return
+    // variants number of additional options.
+    // #[pyo3(text_signature = "($self, value, variants=0)")]
+    // fn extract_expr(
+    //     &mut self,
+    //     value: WrappedExpr,
+    //     variants: usize,
+    // ) -> EggResult<(usize, WrappedExpr, Vec<WrappedExpr>)> {
+    //     let (cost, expr, exprs) = self.egraph.extract_expr(value.into(), variants)?;
+    //     Ok((
+    //         cost,
+    //         expr.into(),
+    //         exprs.into_iter().map(|x| x.into()).collect(),
+    //     ))
+    // }
+
     /// Check that a fact is true in the egraph.
     #[pyo3(text_signature = "($self, fact)")]
     fn check_fact(&mut self, fact: WrappedFact) -> EggResult<()> {
@@ -94,5 +114,6 @@ impl EGraph {
 fn bindings(_py: Python, m: &PyModule) -> PyResult<()> {
     m.add_class::<EGraph>()?;
     m.add_class::<EggSmolError>()?;
+    m.add_class::<WrappedVariant>()?;
     Ok(())
 }