From 2b969c7d798c897ba3131d7124161af68a215839 Mon Sep 17 00:00:00 2001
From: Saul Shanabrook <s.shanabrook@gmail.com>
Date: Tue, 24 Oct 2023 16:35:16 -0400
Subject: [PATCH] Add support for matching on rationals

---
 README.md            |  1 +
 src/sort/rational.rs | 11 +++++++----
 tests/rational.egg   | 14 ++++++++++++++
 3 files changed, 22 insertions(+), 4 deletions(-)
 create mode 100644 tests/rational.egg

diff --git a/README.md b/README.md
index ed9841103..bab1ce495 100644
--- a/README.md
+++ b/README.md
@@ -128,6 +128,7 @@ Rational numbers (fractions) with 64-bit precision for numerator and denominator
 + - * /         ; arithmetic
 min max neg abs floor ceil round
 rational        ; construct from a numerator and denominator
+numer denom     ; get numerator and denominator
 pow log sqrt
 < > <= >=       ; comparisons
 ```
diff --git a/src/sort/rational.rs b/src/sort/rational.rs
index 86ed67a83..e3b14fddd 100644
--- a/src/sort/rational.rs
+++ b/src/sort/rational.rs
@@ -50,6 +50,9 @@ impl Sort for RationalSort {
         add_primitives!(eg, "ceil" = |a: R| -> R { a.ceil() });
         add_primitives!(eg, "round" = |a: R| -> R { a.round() });
         add_primitives!(eg, "rational" = |a: i64, b: i64| -> R { R::new(a, b) });
+        add_primitives!(eg, "numer" = |a: R| -> i64 { *a.numer() });
+        add_primitives!(eg, "denom" = |a: R| -> i64 { *a.denom() });
+
         add_primitives!(eg, "to-f64" = |a: R| -> f64 { a.to_f64().unwrap() });
 
         add_primitives!(eg, "pow" = |a: R, b: R| -> Option<R> {
@@ -101,10 +104,10 @@ impl Sort for RationalSort {
             }
         });
 
-        add_primitives!(eg, "<" = |a: R, b: R| -> Opt { if a < b {Some(())} else {None} }); 
-        add_primitives!(eg, ">" = |a: R, b: R| -> Opt { if a > b {Some(())} else {None} }); 
-        add_primitives!(eg, "<=" = |a: R, b: R| -> Opt { if a <= b {Some(())} else {None} }); 
-        add_primitives!(eg, ">=" = |a: R, b: R| -> Opt { if a >= b {Some(())} else {None} }); 
+        add_primitives!(eg, "<" = |a: R, b: R| -> Opt { if a < b {Some(())} else {None} });
+        add_primitives!(eg, ">" = |a: R, b: R| -> Opt { if a > b {Some(())} else {None} });
+        add_primitives!(eg, "<=" = |a: R, b: R| -> Opt { if a <= b {Some(())} else {None} });
+        add_primitives!(eg, ">=" = |a: R, b: R| -> Opt { if a >= b {Some(())} else {None} });
    }
     fn make_expr(&self, _egraph: &EGraph, value: Value) -> (Cost, Expr) {
         assert!(value.tag == self.name());
diff --git a/tests/rational.egg b/tests/rational.egg
new file mode 100644
index 000000000..af7cae73a
--- /dev/null
+++ b/tests/rational.egg
@@ -0,0 +1,14 @@
+; Test that can run rule matching on rational
+
+(datatype Pretty
+    (pretty-str String)
+    (pretty-rational Rational))
+
+; This will fail with `Unbound variable x in primitive computation` currently:
+; (rewrite (pretty-rational (rational x y)) (pretty-str (+ (to-string x) "/" (to-string y))))
+
+(rewrite (pretty-rational r) (pretty-str (+ (to-string (numer r)) "/" (to-string (denom r)))))
+
+(let z (pretty-rational (rational 1 2)))
+(run 1)
+(check (= z (pretty-str "1/2")))