Speedup new backend

Speedup new backend by using only one thread and a custom memory allocator
Copied from egraphs-good/egglog#637

Author: saulshanabrook

Cargo.lock

@@ -12,25 +12,26 @@ crate-type = ["cdylib"]
 [dependencies]
 pyo3 = { version = "0.24.2", features = ["extension-module"] }
 
-# include this bugfix https://github.com/egraphs-good/egglog/pull/629
-egglog = { git = "https://github.com/saulshanabrook/egg-smol.git", rev = "1e638e6" }
+egglog = { git = "https://github.com/egraphs-good/egglog.git", rev = "5542549" }
 # egglog = { path = "../egg-smol" }
-egglog-bridge = { git = "https://github.com/egraphs-good/egglog-backend.git", rev = "1a1f913" }
-core-relations = { git = "https://github.com/egraphs-good/egglog-backend.git", rev = "1a1f913" }
-egglog-experimental = { git = "https://github.com/egraphs-good/egglog-experimental", rev = "202078f" }
+egglog-bridge = { git = "https://github.com/egraphs-good/egglog-backend.git", rev = "cd51d04" }
+core-relations = { git = "https://github.com/egraphs-good/egglog-backend.git", rev = "cd51d04" }
+egglog-experimental = { git = "https://github.com/egraphs-good/egglog-experimental", rev = "255b67a" }
 egraph-serialize = { version = "0.2.0", features = ["serde", "graphviz"] }
 serde_json = "1.0.140"
 pyo3-log = "0.12.4"
 log = "0.4.27"
 lalrpop-util = { version = "0.22", features = ["lexer"] }
 ordered-float = "3.7.0"
 uuid = { version = "1.17.0", features = ["v4"] }
+rayon = "1.10.0"
+mimalloc = "0.1.46"
 
 # Use patched version of egglog in experimental
 [patch.'https://github.com/egraphs-good/egglog']
 # egglog = { git = "https://github.com/egraphs-good//egglog.git", rev = "d2fa5b733de0796fb187dc5a27e570d5644aa75a" }
 # egglog = { path = "../egg-smol" }
-egglog = { git = "https://github.com/saulshanabrook/egg-smol.git", rev = "1e638e6" }
+# egglog = { git = "https://github.com/egraphs-good//egglog.git", rev = "5542549" }
 
 # enable debug symbols for easier profiling
 [profile.release]

Cargo.toml

@@ -1,6 +1,7 @@
 // Create wrappers around input types so that convert from pyobjects to them
 // and then from them to the egg_smol types
 use crate::utils::*;
+use egglog::extract::DefaultCost;
 use ordered_float::OrderedFloat;
 use pyo3::prelude::*;
 use pyo3::types::{PyDelta, PyDeltaAccess};
@@ -258,7 +259,7 @@ convert_enums!(
         Include(span: Span, path: String)
             i -> egglog::ast::Command::Include(i.span.clone().into(), (&i.path).into()),
             egglog::ast::Command::Include(span, p) => Include { span: span.into(), path: p.to_string() };
-        Constructor(span: Span, name: String, schema: Schema, cost: Option<usize>, unextractable: bool)
+        Constructor(span: Span, name: String, schema: Schema, cost: Option<DefaultCost>, unextractable: bool)
             c -> egglog::ast::Command::Constructor {
                 span: c.span.clone().into(),
                 name: (&c.name).into(),
@@ -325,7 +326,7 @@ convert_enums!(
             egglog::ast::Subdatatypes::NewSort(name, args) => NewSort { name: name.to_string(), args: args.iter().map(|e| e.into()).collect() }
     };
     egglog::ExtractReport: "{:?}" => ExtractReport {
-        Best(termdag: TermDag, cost: usize, term: Term)
+        Best(termdag: TermDag, cost: DefaultCost, term: Term)
             b -> egglog::ExtractReport::Best {
                 termdag: (&b.termdag).into(),
                 cost: b.cost,
@@ -384,7 +385,7 @@ convert_struct!(
         span: Span,
         name: String,
         types: Vec<String>,
-        cost: Option<usize> = None
+        cost: Option<DefaultCost> = None
     )
         v -> egglog::ast::Variant {span: v.span.clone().into(), name: (&v.name).into(), types: v.types.iter().map(|v| v.into()).collect(), cost: v.cost},
         v -> Variant {span: v.span.clone().into(), name: v.name.to_string(), types: v.types.iter().map(|v| v.to_string()).collect(), cost: v.cost};

src/conversions.rs

@@ -5,8 +5,8 @@ use crate::error::{EggResult, WrappedError};
 use crate::py_object_sort::PyObjectSort;
 use crate::serialize::SerializedEGraph;
 
-use egglog::prelude::add_leaf_sort;
-use egglog::{span, SerializeConfig};
+use egglog::prelude::add_base_sort;
+use egglog::{SerializeConfig, span};
 use log::info;
 use pyo3::prelude::*;
 use std::path::PathBuf;
@@ -38,7 +38,7 @@ impl EGraph {
         egraph.fact_directory = fact_directory;
         egraph.seminaive = seminaive;
         if let Some(py_object_sort) = py_object_sort {
-            add_leaf_sort(&mut egraph, py_object_sort, span!()).unwrap();
+            add_base_sort(&mut egraph, py_object_sort, span!()).unwrap();
         }
         Self {
             egraph,

src/egraph.rs

@@ -6,11 +6,19 @@ mod serialize;
 mod termdag;
 mod utils;
 
+#[global_allocator]
+static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc;
+
 use pyo3::prelude::*;
 
 /// Bindings for egglog rust library
 #[pymodule]
 fn bindings(m: &Bound<'_, PyModule>) -> PyResult<()> {
+    rayon::ThreadPoolBuilder::new()
+        .num_threads(1)
+        .build_global()
+        .unwrap();
+
     pyo3_log::init();
 
     m.add_class::<crate::py_object_sort::PyObjectSort>()?;