From bc802483bfc3b472e4a5c1987eeda5fe7d8ed84c Mon Sep 17 00:00:00 2001
From: Jad Ghalayini <jeg74@cl.cam.ac.uk>
Date: Thu, 12 Oct 2023 22:11:07 +0100
Subject: [PATCH] Make Analysis::make take in the E-graph mutably (#277)

* Make Analysis::make take in the E-graph mutably

* Strengthen `make` docs

* Update changelog

---------

Co-authored-by: Max Willsey <me@mwillsey.com>
---
 CHANGELOG.md    |  1 +
 src/language.rs | 15 ++++++++++-----
 src/rewrite.rs  |  2 +-
 tests/lambda.rs |  2 +-
 tests/math.rs   |  2 +-
 tests/prop.rs   |  2 +-
 6 files changed, 15 insertions(+), 9 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index bbd14e17..17cc08fe 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,6 +1,7 @@
 # Changes
 
 ## [Unreleased] - ReleaseDate
+- Change the API of `make` to have mutable access to the e-graph for some [advanced uses cases](https://github.com/egraphs-good/egg/pull/277).
 
 
 ## [0.9.5] - 2023-06-29
diff --git a/src/language.rs b/src/language.rs
index 4b829df4..ccfd7895 100644
--- a/src/language.rs
+++ b/src/language.rs
@@ -651,7 +651,7 @@ impl Analysis<SimpleMath> for ConstantFolding {
         egg::merge_max(to, from)
     }
 
-    fn make(egraph: &EGraph<SimpleMath, Self>, enode: &SimpleMath) -> Self::Data {
+    fn make(egraph: &mut EGraph<SimpleMath, Self>, enode: &SimpleMath) -> Self::Data {
         let x = |i: &Id| egraph[*i].data;
         match enode {
             SimpleMath::Num(n) => Some(*n),
@@ -691,10 +691,15 @@ pub trait Analysis<L: Language>: Sized {
     /// The per-[`EClass`] data for this analysis.
     type Data: Debug;
 
-    /// Makes a new [`Analysis`] for a given enode
-    /// [`Analysis`].
+    /// Makes a new [`Analysis`] data for a given e-node.
     ///
-    fn make(egraph: &EGraph<L, Self>, enode: &L) -> Self::Data;
+    /// Note the mutable `egraph` parameter: this is needed for some
+    /// advanced use cases, but most use cases will not need to mutate
+    /// the e-graph in any way.
+    /// It is **not** `make`'s responsiblity to insert the e-node;
+    /// the e-node is "being inserted" when this function is called.
+    /// Doing so will create an infinite loop.
+    fn make(egraph: &mut EGraph<L, Self>, enode: &L) -> Self::Data;
 
     /// An optional hook that allows inspection before a [`union`] occurs.
     /// When explanations are enabled, it gives two ids that represent the two particular terms being unioned, not the canonical ids for the two eclasses.
@@ -751,7 +756,7 @@ pub trait Analysis<L: Language>: Sized {
 
 impl<L: Language> Analysis<L> for () {
     type Data = ();
-    fn make(_egraph: &EGraph<L, Self>, _enode: &L) -> Self::Data {}
+    fn make(_egraph: &mut EGraph<L, Self>, _enode: &L) -> Self::Data {}
     fn merge(&mut self, _: &mut Self::Data, _: Self::Data) -> DidMerge {
         DidMerge(false, false)
     }
diff --git a/src/rewrite.rs b/src/rewrite.rs
index 1687caa6..6b34b48b 100644
--- a/src/rewrite.rs
+++ b/src/rewrite.rs
@@ -255,7 +255,7 @@ where
 ///     fn merge(&mut self, to: &mut Self::Data, from: Self::Data) -> DidMerge {
 ///         merge_min(to, from)
 ///     }
-///     fn make(egraph: &EGraph, enode: &Math) -> Self::Data {
+///     fn make(egraph: &mut EGraph, enode: &Math) -> Self::Data {
 ///         let get_size = |i: Id| egraph[i].data;
 ///         AstSize.cost(enode, get_size)
 ///     }
diff --git a/tests/lambda.rs b/tests/lambda.rs
index 41071e96..80ea4fbd 100644
--- a/tests/lambda.rs
+++ b/tests/lambda.rs
@@ -75,7 +75,7 @@ impl Analysis<Lambda> for LambdaAnalysis {
         })
     }
 
-    fn make(egraph: &EGraph, enode: &Lambda) -> Data {
+    fn make(egraph: &mut EGraph, enode: &Lambda) -> Data {
         let f = |i: &Id| egraph[*i].data.free.iter().cloned();
         let mut free = HashSet::default();
         match enode {
diff --git a/tests/math.rs b/tests/math.rs
index 7f56ded9..a0d8c07a 100644
--- a/tests/math.rs
+++ b/tests/math.rs
@@ -50,7 +50,7 @@ pub struct ConstantFold;
 impl Analysis<Math> for ConstantFold {
     type Data = Option<(Constant, PatternAst<Math>)>;
 
-    fn make(egraph: &EGraph, enode: &Math) -> Self::Data {
+    fn make(egraph: &mut EGraph, enode: &Math) -> Self::Data {
         let x = |i: &Id| egraph[*i].data.as_ref().map(|d| d.0);
         Some(match enode {
             Math::Constant(c) => (*c, format!("{}", c).parse().unwrap()),
diff --git a/tests/prop.rs b/tests/prop.rs
index a5723aca..ed1c7469 100644
--- a/tests/prop.rs
+++ b/tests/prop.rs
@@ -25,7 +25,7 @@ impl Analysis<Prop> for ConstantFold {
         })
     }
 
-    fn make(egraph: &EGraph, enode: &Prop) -> Self::Data {
+    fn make(egraph: &mut EGraph, enode: &Prop) -> Self::Data {
         let x = |i: &Id| egraph[*i].data.as_ref().map(|c| c.0);
         let result = match enode {
             Prop::Bool(c) => Some((*c, c.to_string().parse().unwrap())),