Add run command

Author: saulshanabrook

python/egg_smol/bindings.pyi

@@ -1,3 +1,4 @@
+from datetime import timedelta
 from typing import Optional
 
 from typing_extensions import final
@@ -12,6 +13,7 @@ class EGraph:
     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: ...
+    def run_rules(self, limit: int) -> tuple[timedelta, timedelta, timedelta]: ...
 
 @final
 class EggSmolError(Exception):

python/tests/test.py

@@ -1,3 +1,5 @@
+import datetime
+
 import pytest
 from egg_smol.bindings import EggSmolError, EGraph
 from egg_smol.bindings_py import *
@@ -92,3 +94,35 @@ def test_rewrite(self):
             )
         )
         assert isinstance(name, str)
+
+    def test_run_rules(self):
+        egraph = EGraph()
+        egraph.declare_sort("Math")
+        egraph.declare_constructor(Variant("Add", ["Math", "Math"]), "Math")
+        egraph.add_rewrite(
+            Rewrite(
+                Call(
+                    "Add",
+                    [
+                        Var("a"),
+                        Var("b"),
+                    ],
+                ),
+                Call(
+                    "Add",
+                    [
+                        Var("b"),
+                        Var("a"),
+                    ],
+                ),
+            )
+        )
+        start_time = datetime.datetime.now()
+        searching, applying, rebuilding = egraph.run_rules(10)
+        end_time = datetime.datetime.now()
+        assert isinstance(searching, datetime.timedelta)
+        assert isinstance(applying, datetime.timedelta)
+        assert isinstance(rebuilding, datetime.timedelta)
+        total_measured_time = searching + applying + rebuilding
+        # Verify  less than the total time (which includes time spent in Python).
+        assert total_measured_time < (end_time - start_time)

src/conversions.rs

@@ -1,10 +1,12 @@
+use std::time::Duration;
+
 // Create wrappers around input types so that convert from pyobjects to them
 // and then from them to the egg_smol types
 //
 // 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::*;
+use pyo3::{ffi::PyDateTime_Delta, prelude::*, types::PyDelta};
 
 // 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> {
@@ -237,3 +239,30 @@ impl From<WrappedFact> for egg_smol::ast::Fact {
         other.0
     }
 }
+
+// Wrapped version of Duration
+// Converts from a rust duration to a python timedelta
+pub struct WrappedDuration(Duration);
+
+impl From<Duration> for WrappedDuration {
+    fn from(other: Duration) -> Self {
+        WrappedDuration(other)
+    }
+}
+
+impl IntoPy<PyObject> for WrappedDuration {
+    fn into_py(self, py: Python<'_>) -> PyObject {
+        let d = self.0;
+        PyDelta::new(
+            py,
+            0,
+            0,
+            d.as_millis()
+                .try_into()
+                .expect("Failed to convert miliseconds to int32 when converting duration"),
+            true,
+        )
+        .expect("Failed to contruct timedelta")
+        .into()
+    }
+}

src/lib.rs

@@ -1,5 +1,7 @@
 mod conversions;
 mod error;
+use std::time::Duration;
+
 use conversions::*;
 use error::*;
 use pyo3::prelude::*;
@@ -22,6 +24,19 @@ impl EGraph {
         }
     }
 
+    /// Run the rules on the egraph until it reaches a fixpoint, specifying the max number of iterations.
+    /// Returns a tuple of the total time spen searching, applying, and rebuilding.
+    #[pyo3(text_signature = "($self, limit)")]
+    fn run_rules(
+        &mut self,
+        limit: usize,
+    ) -> EggResult<(WrappedDuration, WrappedDuration, WrappedDuration)> {
+        let [search, apply, rebuild] = self.egraph.run_rules(limit);
+        // Print all the timings
+        println!("Timings: {:?}, {:?}, {:?}", search, apply, rebuild);
+        Ok((search.into(), apply.into(), rebuild.into()))
+    }
+
     /// Define a rewrite rule, returning the name of the rule
     #[pyo3(text_signature = "($self, rewrite)")]
     fn add_rewrite(&mut self, rewrite: WrappedRewrite) -> EggResult<String> {