Move splitting primitives to serialize library

.gitignore

@@ -19,3 +19,4 @@ _scratch.egg
 
 # racket
 scripts/compiled
+tests/*.json

Cargo.toml

@@ -37,7 +37,7 @@ smallvec = "1.11"
 
 generic_symbolic_expressions = "5.0.4"
 
-egraph-serialize = { version = "0.1.0", features = ["serde", "graphviz"] }
+egraph-serialize = { git = "https://github.com/saulshanabrook/egraph-serialize", rev = "1c205fcc6d3426800b828e9264dbadbd4a5ef6e9", features = ["serde", "graphviz"] }
 serde_json = { optional = true, version = "1.0.100", features = [
   "preserve_order",
 ] }

Makefile

@@ -46,8 +46,13 @@ ${DIST_WASM}: ${RUST_SRC}
 
 graphs: $(patsubst %.egg,%.svg,$(filter-out $(wildcard tests/repro-*.egg),$(wildcard tests/*.egg)))
 
+json: $(patsubst %.egg,%.json,$(filter-out $(wildcard tests/repro-*.egg),$(wildcard tests/*.egg)))
+
 %.svg: %.egg
 	cargo run --release -- --to-dot --to-svg  $^
 
+%.json: %.egg
+	cargo run --release -- --to-json $^
+
 rm-graphs:
 	rm -f tests/*.dot tests/*.svg

src/lib.rs

@@ -33,6 +33,7 @@ use hashbrown::hash_map::Entry;
 use index::ColumnIndex;
 use instant::{Duration, Instant};
 pub use serialize::SerializeConfig;
+pub use serialize::SerializedNode;
 use sort::*;
 pub use termdag::{Term, TermDag, TermId};
 use thiserror::Error;
@@ -1552,24 +1553,6 @@ impl EGraph {
         &self.overall_run_report
     }
 
-    /// Serializes the egraph for export to graphviz.
-    ///
-    /// This will limit the total number of nodes so that visualization does not blow up.
-    pub fn serialize_for_graphviz(
-        &self,
-        split_primitive_outputs: bool,
-        max_functions: usize,
-        max_calls_per_function: usize,
-    ) -> egraph_serialize::EGraph {
-        let config = SerializeConfig {
-            split_primitive_outputs,
-            max_functions: Some(max_functions),
-            max_calls_per_function: Some(max_calls_per_function),
-            ..Default::default()
-        };
-        self.serialize(config)
-    }
-
     pub(crate) fn print_msg(&mut self, msg: String) {
         self.msgs.push(msg);
     }

src/main.rs

@@ -1,5 +1,5 @@
 use clap::Parser;
-use egglog::{EGraph, Error, RunMode, SerializeConfig};
+use egglog::{EGraph, Error, RunMode, SerializeConfig, SerializedNode};
 use std::io::{self, BufRead, BufReader};
 use std::path::PathBuf;
 
@@ -43,6 +43,9 @@ struct Args {
     /// Maximum number of calls per function to render in dot/svg output
     #[clap(long, default_value = "40")]
     max_calls_per_function: usize,
+    /// Number of times to inline leaves
+    #[clap(long, default_value = "0")]
+    serialize_n_inline_leaves: usize,
 }
 
 // test if the current command should be evaluated
@@ -209,31 +212,25 @@ fn main() {
                 std::process::exit(1)
             }
         }
-        // if we are splitting primitive outputs, add `-split` to the end of the file name
-        let serialize_filename = if args.serialize_split_primitive_outputs {
-            input.with_file_name(format!(
-                "{}-split",
-                input.file_stem().unwrap().to_str().unwrap()
-            ))
-        } else {
-            input.clone()
-        };
-        if args.to_json {
-            let json_path = serialize_filename.with_extension("json");
-            let config = SerializeConfig {
-                split_primitive_outputs: args.serialize_split_primitive_outputs,
-                ..SerializeConfig::default()
-            };
-            let serialized = egraph.serialize(config);
-            serialized.to_json_file(json_path).unwrap();
-        }
 
-        if args.to_dot || args.to_svg {
-            let serialized = egraph.serialize_for_graphviz(
-                args.serialize_split_primitive_outputs,
-                args.max_functions,
-                args.max_calls_per_function,
-            );
+        if args.to_json || args.to_dot || args.to_svg {
+            let mut serialized = egraph.serialize(SerializeConfig::default());
+            if args.serialize_split_primitive_outputs {
+                serialized.split_e_classes(|id, _| egraph.from_node_id(id).is_primitive())
+            }
+            for _ in 0..args.serialize_n_inline_leaves {
+                serialized.inline_leaves();
+            }
+
+            // if we are splitting primitive outputs, add `-split` to the end of the file name
+            let serialize_filename = if args.serialize_split_primitive_outputs {
+                input.with_file_name(format!(
+                    "{}-split",
+                    input.file_stem().unwrap().to_str().unwrap()
+                ))
+            } else {
+                input.clone()
+            };
             if args.to_dot {
                 let dot_path = serialize_filename.with_extension("dot");
                 serialized.to_dot_file(dot_path).unwrap()
@@ -242,6 +239,10 @@ fn main() {
                 let svg_path = serialize_filename.with_extension("svg");
                 serialized.to_svg_file(svg_path).unwrap()
             }
+            if args.to_json {
+                let json_path = serialize_filename.with_extension("json");
+                serialized.to_json_file(json_path).unwrap();
+            }
         }
         // no need to drop the egraph if we are going to exit
         if idx == args.inputs.len() - 1 {
@@ -259,12 +260,12 @@ mod tests {
         #[rustfmt::skip]
         let test_cases = vec![
             vec![
-                "(extract", 
-                "\"1", 
-                ")", 
-                "(", 
-                ")))", 
-                "\"", 
+                "(extract",
+                "\"1",
+                ")",
+                "(",
+                ")))",
+                "\"",
                 ";; )",
                 ")"
             ],