Add datatype support

Author: saulshanabrook

python/egg_smol/bindings_py.py

@@ -0,0 +1,12 @@
+# TODO: Figure out what these modules should be called
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Optional
+
+
+@dataclass
+class Variant:
+    name: str
+    types: list[str]
+    cost: Optional[int] = None

python/tests/test.py

@@ -1,5 +1,6 @@
 import pytest
 from egg_smol.bindings import EggSmolError, EGraph
+from egg_smol.bindings_py import Variant
 
 
 class TestEGraph:
@@ -18,3 +19,11 @@ def test_parse_and_run_program_exception(self):
             match='Check failed: Value { tag: "i64", bits: 5 } != Value { tag: "i64", bits: 4 }',
         ):
             egraph.parse_and_run_program(program)
+
+    def test_datatype(self):
+        egraph = EGraph()
+        egraph.declare_sort("Math")
+        egraph.declare_constructor(Variant("Num", ["i64"]), "Math")
+        egraph.declare_constructor(Variant("Var", ["String"]), "Math")
+        egraph.declare_constructor(Variant("Add", ["Math", "Math"]), "Math")
+        egraph.declare_constructor(Variant("Mul", ["Math", "Math"]), "Math")

src/lib.rs

@@ -25,6 +25,24 @@ impl EggSmolError {
     }
 }
 
+// Convert a Python Variant object into a rust variable, by getting the attributes
+fn get_variant(obj: &PyAny) -> PyResult<egg_smol::ast::Variant> {
+    // TODO: Is there a way to do this more automatically?
+    Ok(egg_smol::ast::Variant {
+        name: obj
+            .getattr(pyo3::intern!(obj.py(), "name"))?
+            .extract::<String>()?
+            .into(),
+        cost: obj.getattr(pyo3::intern!(obj.py(), "cost"))?.extract()?,
+        types: obj
+            .getattr(pyo3::intern!(obj.py(), "types"))?
+            .extract::<Vec<String>>()?
+            .into_iter()
+            .map(|x| x.into())
+            .collect(),
+    })
+}
+
 #[pymethods]
 impl EGraph {
     #[new]
@@ -34,6 +52,28 @@ impl EGraph {
         }
     }
 
+    /// declare_sort($self, name)
+    /// --
+    ///
+    /// Declare a new sort with the given name.
+    fn declare_sort(&mut self, name: &str) -> PyResult<()> {
+        // TODO: Should the name be a symbol? If so, how should we expose that
+        // to Python?
+        self.egraph
+            .declare_sort(name)
+            .map_err(|e| PyErr::new::<EggSmolError, _>(e.to_string()))
+    }
+
+    fn declare_constructor(
+        &mut self,
+        #[pyo3(from_py_with = "get_variant")] variant: egg_smol::ast::Variant,
+        sort: &str,
+    ) -> PyResult<()> {
+        self.egraph
+            .declare_constructor(variant, sort)
+            .map_err(|e| PyErr::new::<EggSmolError, _>(e.to_string()))
+    }
+
     /// parse_and_run_program($self, input)
     /// --
     ///