Add code to dump egraphs in egraph-serialize format

egg-herbie/Cargo.toml

@@ -6,6 +6,7 @@ edition = "2021"
 
 [dependencies]
 egg = { git = "https://github.com/egraphs-good/egg.git", rev = "c11d03db646aca5f55df2632927eec958e1f4d4d" }
+egraph-serialize = "0.1.0"
 
 log = "0.4"
 indexmap = "1"

egg-herbie/main.rkt

@@ -22,6 +22,7 @@
          egraph_is_unsound_detected
          egraph_get_times_applied
          egraph_get_proof
+         egraph_dump_to_file
          (struct-out iteration-data)
          (struct-out FFIRule)
          make-ffi-rule)
@@ -299,3 +300,9 @@
                       _pointer ;; name of the rule
                       ->
                       _uint))
+
+(define-eggmath egraph_dump_to_file
+                (_fun _egraph-pointer
+                      _string ;; file name
+                      ->
+                      _void))

egg-herbie/src/lib.rs

@@ -13,6 +13,7 @@ use std::mem::{self, ManuallyDrop};
 use std::os::raw::c_char;
 use std::time::Duration;
 use std::{slice, sync::atomic::Ordering};
+use std::path::Path;
 
 pub struct Context {
     iteration: usize,
@@ -481,3 +482,36 @@ pub unsafe extern "C" fn egraph_get_size(ptr: *mut Context) -> u32 {
         .map(|iteration| iteration.egraph_nodes as u32)
         .unwrap_or_default()
 }
+
+#[no_mangle]
+pub unsafe extern "C" fn egraph_dump_to_file(ptr : *mut Context, path_str : *mut c_char) {
+    let context = ManuallyDrop::new(Box::from_raw(ptr));
+    let path_cstr = CStr::from_ptr(path_str).to_str().expect("Non-UTF-8 path");
+    let path = Path::new(&path_cstr);
+
+    use egraph_serialize::*;
+    let mut out = EGraph::default();
+    for class in context.runner.egraph.classes() {
+        for (i, node) in class.nodes.iter().enumerate() {
+            out.add_node(
+                format!("{}.{}", class.id, i),
+                Node {
+                    op: node.to_string(),
+                    children: node
+                        .children()
+                        .iter()
+                        .map(|id| NodeId::from(format!("{}.0", id)))
+                        .collect(),
+                    eclass: ClassId::from(format!("{}", class.id)),
+                    cost: Cost::new(1.0).unwrap(),
+                },
+            )
+        }
+    }
+
+    for root in context.runner.roots.clone() {
+        out.root_eclasses.push(ClassId::from(format!("{}", context.runner.egraph.find(root))));
+    }
+
+    out.to_json_file(path).expect("Failed to write file");
+}

src/config.rkt

@@ -19,7 +19,8 @@
                        numerics
                        special
                        bools
-                       branches)]))
+                       branches)]
+        [dump . (egg)]))
 
 (define default-flags
   #hash([precision . ()]
@@ -36,7 +37,8 @@
                        numerics
                        special
                        bools
-                       branches)]))
+                       branches)]
+        [dump . ()]))
 
 (define (check-flag-deprecated! category flag)
   (match* (category flag)

src/core/egg-herbie.rkt

@@ -1251,6 +1251,15 @@
   (define ctx (egg-runner-ctx runner))
   (define-values (root-ids egg-graph)
     (egraph-run-schedule (egg-runner-exprs runner) (egg-runner-schedule runner) ctx))
+  (when (and (flag-set? 'dump 'egg) (set-member? '(single multi) (car cmd)))
+    (define dump-dir "dump-egg")
+    (unless (directory-exists? dump-dir)
+      (make-directory dump-dir))
+    (define name
+      (for/first ([i (in-naturals)]
+                  #:unless (file-exists? (build-path dump-dir (format "~a.json" i))))
+        (build-path dump-dir (format "~a.json" i))))
+    (egraph_dump_to_file (egraph-data-egraph-pointer egg-graph) (path->string name)))
   ; Perform extraction
   (match cmd
     [`(single . ,extractor) ; single expression extraction