moved to i64 for rational expression solver because some nat expr were to big

Author: SaltyPeppermint

examples/blocking.rs

@@ -1,20 +1,21 @@
 use std::time::Duration;
 
-use egg::{AstSize, RecExpr, Runner, SimpleScheduler};
+use egg::{AstSize, Runner, SimpleScheduler};
 
 use eggshell::eqsat::hooks;
-use eggshell::rewrite_system::rise::{self, PrettyPrint, Rise, RiseRuleset};
-use eggshell::rewrite_system::rise::{BLOCKING_GOAL, MM, SPLIT_GUIDE};
+use eggshell::rewrite_system::rise::{self, BLOCKING_GOAL, MM, RiseRuleset, SPLIT_GUIDE};
 use eggshell::sketch;
+use eggshell::utils;
 
 fn main() {
-    let mm: RecExpr<Rise> = MM.parse().unwrap();
-    let split_guide: RecExpr<Rise> = SPLIT_GUIDE.parse().unwrap();
+    let mm = rise::canon_nat(&MM.parse().unwrap());
+    let split_guide = rise::canon_nat(&SPLIT_GUIDE.parse().unwrap());
+    let blocking_goal = rise::canon_nat(&BLOCKING_GOAL.parse().unwrap());
 
     let runner_1 = Runner::default()
         .with_expr(&mm)
         .with_iter_limit(6)
-        .with_time_limit(Duration::from_secs(30))
+        .with_time_limit(Duration::from_secs(60))
         .with_node_limit(1_000_000)
         .with_scheduler(SimpleScheduler)
         .with_hook(hooks::targe_hook(split_guide.clone()))
@@ -24,32 +25,43 @@ fn main() {
     println!("{}", runner_1.report());
 
     let root_mm = runner_1.egraph.find(runner_1.roots[0]);
-    let split_guide_sketch = rise::sketchify(SPLIT_GUIDE, true);
-    let (_, sketch_extracted_split_guide) =
-        sketch::eclass_extract(&split_guide_sketch, AstSize, &runner_1.egraph, root_mm).unwrap();
-
-    println!("\nGuide Ground Truth");
-    split_guide.pp(false);
-    println!("\nSketch Extracted:");
-    sketch_extracted_split_guide.pp(false);
-
-    assert_eq!(
-        None,
-        eggshell::utils::find_diff(&sketch_extracted_split_guide, &split_guide)
+    let split_guide_sketch = rise::sketchify(&split_guide, true);
+    let sketch_extracted_split_guide = rise::canon_nat(
+        &sketch::eclass_extract(&split_guide_sketch, AstSize, &runner_1.egraph, root_mm)
+            .unwrap()
+            .1,
     );
-    assert_eq!(root_mm, runner_1.egraph.lookup_expr(&split_guide).unwrap());
+    assert!(utils::find_diff(&sketch_extracted_split_guide, &split_guide).is_none());
+
+    // println!("\nGuide Ground Truth");
+    // split_guide.pp(false);
+    // println!("\nSketch Extracted:");
+    // sketch_extracted_split_guide.pp(false);
+
+    // assert_eq!(root_mm, runner_1.egraph.lookup_expr(&split_guide).unwrap());
 
-    let blocking_goal: RecExpr<Rise> = BLOCKING_GOAL.parse().unwrap();
     let runner_2 = Runner::default()
         .with_expr(&split_guide)
-        .with_iter_limit(6)
+        .with_iter_limit(8)
+        .with_time_limit(Duration::from_secs(60))
+        .with_node_limit(1_000_000)
         .with_scheduler(SimpleScheduler)
         .with_hook(hooks::targe_hook(blocking_goal.clone()))
+        .with_hook(hooks::printer_hook)
         .run(&rise::rules(RiseRuleset::MM));
 
+    println!("{}", runner_2.report());
+
     let root_guide = runner_2.egraph.find(runner_2.roots[0]);
-    assert_eq!(
-        root_guide,
-        runner_2.egraph.lookup_expr(&blocking_goal).unwrap()
+    let blocking_goal_sketch = rise::sketchify(&blocking_goal, true);
+    let sketch_extracted_blocking_goal = rise::canon_nat(
+        &sketch::eclass_extract(&blocking_goal_sketch, AstSize, &runner_2.egraph, root_guide)
+            .unwrap()
+            .1,
     );
+    assert!(utils::find_diff(&sketch_extracted_blocking_goal, &blocking_goal).is_none());
+    // assert_eq!(
+    //     root_guide,
+    //     runner_2.egraph.lookup_expr(&blocking_goal).unwrap()
+    // );
 }

src/rewrite_system/rise/analysis.rs

@@ -1,4 +1,4 @@
-use egg::{Analysis, DidMerge, EGraph, Id, Language, RecExpr};
+use egg::{Analysis, DidMerge, EGraph, Language, RecExpr};
 use hashbrown::HashSet;
 
 use super::nat::try_simplify;
@@ -47,28 +47,29 @@ impl Analysis<Rise> for RiseAnalysis {
     fn merge(&mut self, to: &mut AnalysisData, from: AnalysisData) -> DidMerge {
         let before_len = to.free.len();
         to.free.extend(from.free);
-        let free_merge = DidMerge(before_len != to.free.len(), true);
+        let free_merge = before_len != to.free.len();
 
         let beta_merge = if !from.beta_extract.is_empty()
             && (to.beta_extract.is_empty() || to.beta_extract.len() > from.beta_extract.len())
         {
             to.beta_extract = from.beta_extract;
-            DidMerge(true, true)
+            true
         } else {
-            DidMerge(false, true) // TODO: more precise second bool
+            false
         };
 
         // let nat_merge =
         //     egg::merge_option(&mut to.simple_nat, from.simple_nat, |to_nat, from_nat| {
         //         if to_nat.len() > from_nat.len() {
         //             *to_nat = from_nat;
-        //             DidMerge(true, true)
+        //             true
         //         } else {
-        //             DidMerge(false, true)
+        //             false
         //         }
         //     });
 
-        free_merge | beta_merge // | nat_merge
+        // TODO: more precise second bool
+        DidMerge(free_merge || beta_merge, true)
     }
 
     fn make(egraph: &mut EGraph<Rise, RiseAnalysis>, enode: &Rise) -> AnalysisData {
@@ -100,7 +101,8 @@ impl Analysis<Rise> for RiseAnalysis {
             (RecExpr::default(), None)
         } else {
             let expr = enode.join_recexprs(|id| egraph[id].data.beta_extract.as_ref());
-            let simple_nat = try_simplify(&expr).ok();
+            let simple_nat = enode.is_nat().then(|| try_simplify(&expr).ok()).flatten();
+
             (expr, simple_nat)
         };
 

src/rewrite_system/rise/func.rs

@@ -99,7 +99,7 @@ impl<A: Applier<Rise, RiseAnalysis>> Applier<Rise, RiseAnalysis> for VectorizeSc
     }
 }
 
-fn extracted_int(expr: &RecExpr<Rise>) -> i32 {
+fn extracted_int(expr: &RecExpr<Rise>) -> i64 {
     if let Rise::Integer(i) = expr[0.into()] {
         return i;
     }
@@ -110,7 +110,7 @@ fn extracted_int(expr: &RecExpr<Rise>) -> i32 {
 #[expect(clippy::too_many_lines)]
 fn vec_expr(
     expr: &RecExpr<Rise>,
-    n: i32,
+    n: i64,
     v_env: HashSet<DBIndex>,
     type_of_id: Id,
 ) -> Option<(RecExpr<Rise>, Id, Id)> {
@@ -249,7 +249,7 @@ fn vec_expr(
     // }
 }
 
-fn vec_ty(expr: &RecExpr<Rise>, n: i32, id: Id) -> Option<RecExpr<Rise>> {
+fn vec_ty(expr: &RecExpr<Rise>, n: i64, id: Id) -> Option<RecExpr<Rise>> {
     match expr[id] {
         Rise::F32 => {
             let mut vec_ty = RecExpr::default();

src/rewrite_system/rise/lang.rs

@@ -77,8 +77,8 @@ egg::define_language! {
     // "sig" = Sigma([Id; 3]),
     // "phi" = Phi([Id; 3]),
 
-    Integer(i32),
-    Float(NotNan<f32>),
+    Integer(i64),
+    Float(NotNan<f64>),
     // Double(f64),
     // Symbol(Symbol),
   }

src/rewrite_system/rise/mod.rs

@@ -17,7 +17,10 @@ pub use analysis::RiseAnalysis;
 pub use lang::Rise;
 pub use pp::PrettyPrint;
 
-use crate::sketch::{Sketch, SketchLang};
+use crate::{
+    rewrite_system::rise::nat::try_simplify,
+    sketch::{Sketch, SketchLang},
+};
 
 #[derive(Copy, Clone, Eq, PartialEq, Debug, Hash)]
 pub enum RiseRuleset {
@@ -45,9 +48,8 @@ fn add_expr(to: &mut RecExpr<Rise>, e: RecExpr<Rise>) -> Id {
     to.root()
 }
 
-#[expect(clippy::missing_panics_doc)]
 #[must_use]
-pub fn sketchify(term: &str, sketchify_nat_expr: bool) -> Sketch<Rise> {
+pub fn sketchify(expr: &RecExpr<Rise>, sketchify_nat_expr: bool) -> Sketch<Rise> {
     fn rec(
         expr: &RecExpr<Rise>,
         id: Id,
@@ -67,12 +69,29 @@ pub fn sketchify(term: &str, sketchify_nat_expr: bool) -> Sketch<Rise> {
             }
         }
     }
-    let expr: RecExpr<Rise> = term.parse().unwrap();
     let mut sketch = RecExpr::default();
-    rec(&expr, expr.root(), sketchify_nat_expr, &mut sketch);
+    rec(expr, expr.root(), sketchify_nat_expr, &mut sketch);
     sketch
 }
 
+#[must_use]
+pub fn canon_nat(expr: &RecExpr<Rise>) -> RecExpr<Rise> {
+    fn rec(expr: &RecExpr<Rise>, id: Id, canon_expr: &mut RecExpr<Rise>) -> Id {
+        let node = &expr[id];
+        if let Ok(canon_nat_expr) = try_simplify(&node.build_recexpr(|i| expr[i].clone())) {
+            add_expr(canon_expr, canon_nat_expr)
+        } else {
+            let new_node = node
+                .clone()
+                .map_children(|c_id| rec(expr, c_id, canon_expr));
+            canon_expr.add(new_node)
+        }
+    }
+    let mut canon_expr = RecExpr::default();
+    rec(expr, expr.root(), &mut canon_expr);
+    canon_expr
+}
+
 // START TERM
 pub const MM: &str = "(typeOf (natLam (typeOf (natLam (typeOf (natLam (typeOf (lam (typeOf (lam (typeOf (app (typeOf (app (typeOf map (fun (fun (arrT %n0 f32) (arrT %n1 f32)) (fun (arrT %n2 (arrT %n0 f32)) (arrT %n2 (arrT %n1 f32))))) (typeOf (lam (typeOf (app (typeOf (app (typeOf map (fun (fun (arrT %n0 f32) f32) (fun (arrT %n1 (arrT %n0 f32)) (arrT %n1 f32)))) (typeOf (lam (typeOf (app (typeOf (app (typeOf (app (typeOf reduce (fun (fun f32 (fun f32 f32)) (fun f32 (fun (arrT %n0 f32) f32)))) (typeOf add (fun f32 (fun f32 f32)))) (fun f32 (fun (arrT %n0 f32) f32))) (typeOf 0.0 f32)) (fun (arrT %n0 f32) f32)) (typeOf (app (typeOf (app (typeOf map (fun (fun (pairT f32 f32) f32) (fun (arrT %n0 (pairT f32 f32)) (arrT %n0 f32)))) (typeOf (lam (typeOf (app (typeOf (app (typeOf mul (fun f32 (fun f32 f32))) (typeOf (app (typeOf fst (fun (pairT f32 f32) f32)) (typeOf %e0 (pairT f32 f32))) f32)) (fun f32 f32)) (typeOf (app (typeOf snd (fun (pairT f32 f32) f32)) (typeOf %e0 (pairT f32 f32))) f32)) f32)) (fun (pairT f32 f32) f32))) (fun (arrT %n0 (pairT f32 f32)) (arrT %n0 f32))) (typeOf (app (typeOf (app (typeOf zip (fun (arrT %n0 f32) (fun (arrT %n0 f32) (arrT %n0 (pairT f32 f32))))) (typeOf %e1 (arrT %n0 f32))) (fun (arrT %n0 f32) (arrT %n0 (pairT f32 f32)))) (typeOf %e0 (arrT %n0 f32))) (arrT %n0 (pairT f32 f32)))) (arrT %n0 f32))) f32)) (fun (arrT %n0 f32) f32))) (fun (arrT %n1 (arrT %n0 f32)) (arrT %n1 f32))) (typeOf (app (typeOf transpose (fun (arrT %n0 (arrT %n1 f32)) (arrT %n1 (arrT %n0 f32)))) (typeOf %e1 (arrT %n0 (arrT %n1 f32)))) (arrT %n1 (arrT %n0 f32)))) (arrT %n1 f32))) (fun (arrT %n0 f32) (arrT %n1 f32)))) (fun (arrT %n2 (arrT %n0 f32)) (arrT %n2 (arrT %n1 f32)))) (typeOf %e1 (arrT %n2 (arrT %n0 f32)))) (arrT %n2 (arrT %n1 f32)))) (fun (arrT %n0 (arrT %n1 f32)) (arrT %n2 (arrT %n1 f32))))) (fun (arrT %n2 (arrT %n0 f32)) (fun (arrT %n0 (arrT %n1 f32)) (arrT %n2 (arrT %n1 f32)))))) (natFun (fun (arrT %n2 (arrT %n0 f32)) (fun (arrT %n0 (arrT %n1 f32)) (arrT %n2 (arrT %n1 f32))))))) (natFun (natFun (fun (arrT %n2 (arrT %n0 f32)) (fun (arrT %n0 (arrT %n1 f32)) (arrT %n2 (arrT %n1 f32)))))))) (natFun (natFun (natFun (fun (arrT %n2 (arrT %n0 f32)) (fun (arrT %n0 (arrT %n1 f32)) (arrT %n2 (arrT %n1 f32))))))))";
 // GUIDES

src/rewrite_system/rise/nat/mod.rs

@@ -17,6 +17,8 @@ use polynomial::Polynomial;
 pub use applier::{ComputeNat, ComputeNatCheck};
 pub use rational::RationalFunction;
 
+type Ratio = num::rational::Ratio<i64>;
+
 // ============================================================================
 // Helper Functions
 // ============================================================================

src/rewrite_system/rise/nat/monomial.rs

@@ -10,19 +10,19 @@ use crate::rewrite_system::rise::DBIndex;
 #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Default)]
 pub struct Monomial {
     // Map from dbindex to exponent (can be negative for rational expressions)
-    variables: BTreeMap<DBIndex, i32>,
+    variables: BTreeMap<DBIndex, i64>,
 }
 
 impl Monomial {
     pub fn new() -> Self {
         Self::default()
     }
 
-    pub fn variables(&self) -> &BTreeMap<DBIndex, i32> {
+    pub fn variables(&self) -> &BTreeMap<DBIndex, i64> {
         &self.variables
     }
 
-    pub fn with_var(mut self, dbindex: DBIndex, exponent: i32) -> Self {
+    pub fn with_var(mut self, dbindex: DBIndex, exponent: i64) -> Self {
         if exponent != 0 {
             self.variables.insert(dbindex, exponent);
         }

src/rewrite_system/rise/nat/polynomial/from.rs

@@ -1,8 +1,7 @@
 use egg::RecExpr;
-use num::rational::Ratio;
 use num_traits::One;
 
-use super::{Monomial, Polynomial, Rise};
+use super::{Monomial, Polynomial, Ratio, Rise};
 use crate::rewrite_system::rise::DBIndex;
 
 // ============================================================================
@@ -64,14 +63,14 @@ impl From<Polynomial> for RecExpr<Rise> {
 // ============================================================================
 
 /// Create a `Polynomial` from an integer constant
-impl From<i32> for Polynomial {
-    fn from(n: i32) -> Self {
+impl From<i64> for Polynomial {
+    fn from(n: i64) -> Self {
         Self::new().add_term(n.into(), Monomial::new())
     }
 }
 /// Create a `Polynomial` from an integer constant
-impl From<Ratio<i32>> for Polynomial {
-    fn from(r: Ratio<i32>) -> Self {
+impl From<Ratio> for Polynomial {
+    fn from(r: Ratio) -> Self {
         Self::new().add_term(r, Monomial::new())
     }
 }