Update egglog dependency

Adds support for structured output added in egraphs-good/egglog#653.
Adds a custom output for running schedules to return all reports and print them all.

Author: saulshanabrook

Cargo.lock

@@ -8,7 +8,7 @@ harness = false
 name = "files"
 
 [dependencies]
-egglog = { git = "https://github.com/egraphs-good/egglog.git", rev = "5542549" }
+egglog = { git = "https://github.com/saulshanabrook/egg-smol.git", rev = "afe962958b57b6101f3f48e980f910fa063f3f6e" }
 
 num = "0.4.3"
 lazy_static = "1.4"

Cargo.toml

@@ -1,10 +1,13 @@
-use std::{collections::HashMap, sync::Mutex};
+use std::{
+    collections::HashMap,
+    sync::{Arc, Mutex},
+};
 
 use egglog::{
     ast::{Expr, Fact, Facts, Literal, ParseError},
     prelude::{query, run_ruleset},
     scheduler::{Scheduler, SchedulerId},
-    RunReport, UserDefinedCommand,
+    CommandOutput, RunReport, UserDefinedCommand,
 };
 use lazy_static::lazy_static;
 
@@ -52,9 +55,12 @@ impl ScheduleState {
         };
 
         if let Expr::Var(_, ruleset) = arg {
-            run_ruleset(egraph, ruleset.as_str())?;
-
-            return Ok(egraph.get_run_report().clone().unwrap());
+            let output = run_ruleset(egraph, ruleset.as_str())?;
+            assert!(output.len() == 1);
+            if let CommandOutput::RunSchedule(report) = &output[0] {
+                return Ok(report.clone());
+            }
+            panic!("Expected a RunSchedule, got {:?}", output[0]);
         }
 
         let Expr::Call(span, head, exprs) = arg else {
@@ -193,13 +199,35 @@ impl ScheduleState {
     }
 }
 
+#[derive(Debug, Clone)]
+struct RunExtendedScheduleOutput {
+    reports: Vec<RunReport>,
+}
+
+impl std::fmt::Display for RunExtendedScheduleOutput {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        writeln!(f, "Ran schedules:")?;
+        for report in &self.reports {
+            writeln!(f, "{}", report)?;
+        }
+        Ok(())
+    }
+}
+
 impl UserDefinedCommand for RunExtendedSchedule {
-    fn update(&self, egraph: &mut egglog::EGraph, args: &[Expr]) -> Result<(), egglog::Error> {
+    fn update(
+        &self,
+        egraph: &mut egglog::EGraph,
+        args: &[Expr],
+    ) -> Result<Option<CommandOutput>, egglog::Error> {
         let mut schedule = ScheduleState::new();
+        let mut reports = Vec::new();
         for arg in args {
-            schedule.run(egraph, arg)?;
+            reports.push(schedule.run(egraph, arg)?);
         }
-        Ok(())
+        Ok(Some(CommandOutput::UserDefined(Arc::new(
+            RunExtendedScheduleOutput { reports },
+        ))))
     }
 }
 

src/scheduling.rs

@@ -4,8 +4,9 @@ use egglog::{
     ast::*,
     extract::{CostModel, DefaultCost, Extractor, TreeAdditiveCostModel},
     util::FreshGen,
-    EGraph, Error, Term, TermDag, TypeError, UserDefinedCommand,
+    CommandOutput, EGraph, Error, Term, TermDag, TypeError, UserDefinedCommand,
 };
+use log::log_enabled;
 
 pub fn add_set_cost(egraph: &mut EGraph) {
     egraph
@@ -208,7 +209,11 @@ impl CostModel<DefaultCost> for DynamicCostModel {
 struct CustomExtract;
 
 impl UserDefinedCommand for CustomExtract {
-    fn update(&self, egraph: &mut EGraph, args: &[Expr]) -> Result<(), Error> {
+    fn update(
+        &self,
+        egraph: &mut EGraph,
+        args: &[Expr],
+    ) -> Result<Option<CommandOutput>, egglog::Error> {
         assert!(args.len() <= 2);
         let (sort, value) = egraph.eval_expr(&args[0])?;
         let n = args.get(1).map(|arg| egraph.eval_expr(arg)).transpose()?;
@@ -236,26 +241,15 @@ impl UserDefinedCommand for CustomExtract {
         );
         if n == 0 {
             if let Some((cost, term)) = extractor.extract_best(egraph, &mut termdag, value) {
-                // dont turn termdag into a string if we have messages disabled for performance reasons
-                if egraph.messages_enabled() {
-                    let extracted = termdag.to_string(&term);
-                    log::info!("extracted with cost {cost}: {extracted}");
-                    egraph.print_msg(extracted);
+                if log_enabled!(log::Level::Info) {
+                    log::info!("extracted with cost {cost}: {}", termdag.to_string(&term));
                 }
-                // TODO: egraph.extract_report is private
-                // A future implementation should make a egglog_experimental::EGraph
-                // that provides a similar set of methods and overrides its own extract_report.
-                //
-                // egraph.extract_report = Some(ExtractReport::Best {
-                //     termdag,
-                //     cost,
-                //     term,
-                // });
+                Ok(Some(CommandOutput::ExtractBest(termdag, cost, term)))
             } else {
-                return Err(Error::ExtractError(
+                Err(Error::ExtractError(
                     "Unable to find any valid extraction (likely due to subsume or delete)"
                         .to_string(),
-                ));
+                ))
             }
         } else {
             if n < 0 {
@@ -266,24 +260,8 @@ impl UserDefinedCommand for CustomExtract {
                 .iter()
                 .map(|e| e.1.clone())
                 .collect();
-            // Same as above, avoid turning termdag into a string if we have messages disabled for performance
-            if egraph.messages_enabled() {
-                log::info!("extracted variants:");
-                let mut msg = String::default();
-                msg += "(\n";
-                assert!(!terms.is_empty());
-                for expr in &terms {
-                    let str = termdag.to_string(expr);
-                    log::info!("   {str}");
-                    msg += &format!("   {str}\n");
-                }
-                msg += ")";
-                egraph.print_msg(msg);
-            }
-            // TODO: Same as above. EGraph::extract_report is private.
-            //
-            // egraph.extract_report = Some(ExtractReport::Variants { termdag, terms });
+            log::info!("extracted variants:");
+            Ok(Some(CommandOutput::ExtractVariants(termdag, terms)))
         }
-        Ok(())
     }
 }

src/set_cost.rs

@@ -22,9 +22,8 @@ impl Run {
             );
         } else {
             let mut egraph = new_experimental_egraph();
-            egraph.run_mode = RunMode::ShowDesugaredEgglog;
             let desugared_str = egraph
-                .parse_and_run_program(self.path.to_str().map(String::from), &program)
+                .resugar_program(self.path.to_str().map(String::from), &program)
                 .unwrap()
                 .join("\n");
 
@@ -39,15 +38,19 @@ impl Run {
     fn test_program(&self, filename: Option<String>, program: &str, message: &str) {
         let mut egraph = new_experimental_egraph();
         match egraph.parse_and_run_program(filename, program) {
-            Ok(msgs) => {
+            Ok(outputs) => {
                 if self.should_fail() {
                     panic!(
                         "Program should have failed! Instead, logged:\n {}",
-                        msgs.join("\n")
+                        outputs
+                            .iter()
+                            .map(|output| output.to_string())
+                            .collect::<Vec<_>>()
+                            .join("\n")
                     );
                 } else {
-                    for msg in msgs {
-                        println!("  {}", msg);
+                    for output in outputs {
+                        print!("  {}", output);
                     }
                     // Test graphviz dot generation
                     let mut serialized = egraph.serialize(SerializeConfig {

tests/files.rs

@@ -35,11 +35,11 @@ fn test_extract() {
         .unwrap();
 
     assert_eq!(result.len(), 5);
-    assert_eq!(result[0], "(Add (Num 1) (Num 1))");
-    assert_eq!(result[1], "(Num 2)");
-    assert_eq!(result[2], "(Add (Num 1) (Num 1))");
-    assert_eq!(result[3], "(Add (Num 1) (Num 1))");
-    assert_eq!(result[4], "(Sub (Num 5) (Num 3))");
+    assert_eq!(result[0].to_string(), "(Add (Num 1) (Num 1))\n");
+    assert_eq!(result[1].to_string(), "(Num 2)\n");
+    assert_eq!(result[2].to_string(), "(Add (Num 1) (Num 1))\n");
+    assert_eq!(result[3].to_string(), "(Add (Num 1) (Num 1))\n");
+    assert_eq!(result[4].to_string(), "(Sub (Num 5) (Num 3))\n");
 }
 
 #[test]
@@ -74,7 +74,7 @@ fn test_extract_set_cost_decls() {
             "(with-dynamic-cost
                 (datatype E (Add E E) (Sub E E :cost 200) (Num i64))
                 (constructor Mul (E E) E :cost 100)
-                (datatype* 
+                (datatype*
                   (E2 (Add2 E2 E2) (Sub2 E2 E2 :cost 200) (List VecE2) (Num2 i64))
                   (sort VecE2 (Vec E2))
                 )