From 12971b1572720cac3116237f5383751abfb1e12a Mon Sep 17 00:00:00 2001
From: Marcus Rossel <marcusrossel@icloud.com>
Date: Fri, 20 Dec 2024 16:05:11 +0100
Subject: [PATCH] Fix previous breaking change

---
 Lean/Egg/Core/Gen/TcSpecs.lean                | 18 ++++++---
 .../TC Inst Erasure Blocking TC Spec Gen.lean | 38 +++++++++++++------
 Rust/Egg/src/lean_expr.rs                     |  2 +-
 3 files changed, 40 insertions(+), 18 deletions(-)

diff --git a/Lean/Egg/Core/Gen/TcSpecs.lean b/Lean/Egg/Core/Gen/TcSpecs.lean
index a1635eb..ae1f459 100644
--- a/Lean/Egg/Core/Gen/TcSpecs.lean
+++ b/Lean/Egg/Core/Gen/TcSpecs.lean
@@ -4,8 +4,16 @@ open Lean Meta
 
 namespace Egg
 
-private def Rewrite.tcConditionMVars (rw : Rewrite) : MVarIdSet :=
-  rw.conds.foldl (init := ∅) (·.union ·.mvars.tcInsts)
+def Rewrite.Condition.typeClassInstanceMVar? (cond : Condition) : MetaM (Option MVarId) := do
+  if cond.expr.isMVar && (← isClass? cond.type).isSome
+  then return cond.expr.mvarId!
+  else return none
+
+private def Rewrite.tcConditionMVars (rw : Rewrite) : MetaM MVarIdSet :=
+  rw.conds.foldlM (init := ∅) fun cs c => do
+    if let some m ← c.typeClassInstanceMVar?
+    then return cs.insert m
+    else return cs
 
 -- Important: This function expects the given rewrite to be fresh.
 private partial def genSpecialization
@@ -48,12 +56,12 @@ where
     unless !missing.isEmpty do return none
     let (freshRw, subst) ← rw.freshWithSubst (src := .tcSpec rw.src <| .dir dir)
     let freshMissing := missing.map subst.expr.get!
-    let conds := freshRw.tcConditionMVars
+    let conds ← freshRw.tcConditionMVars
     let (spec, _) ← genSpecialization freshRw (freshMissing.union conds) norm
     return if (spec.validDirs cfg).contains dir then spec else none
   genCondSpecOnly : MetaM (Option Rewrite) := do
     let freshRw ← rw.fresh (src := .tcSpec rw.src .cond)
-    let conds := freshRw.tcConditionMVars
+    let conds ← freshRw.tcConditionMVars
     let (spec, changed) ← genSpecialization freshRw conds norm
     return if changed then spec else none
 
@@ -64,7 +72,7 @@ private def genGoalTypeSpecialization
     return none
   rw ← rw.instantiateMVars
   let missing := rw.mvars.lhs.tcInsts.union rw.mvars.rhs.tcInsts
-  let conds := rw.tcConditionMVars
+  let conds ← rw.tcConditionMVars
   let (spec, changed) ← genSpecialization rw (missing.union conds) norm
   return if changed then spec else none
 
diff --git a/Lean/Egg/Tests/TC Inst Erasure Blocking TC Spec Gen.lean b/Lean/Egg/Tests/TC Inst Erasure Blocking TC Spec Gen.lean
index 9638bbb..d2be8b3 100644
--- a/Lean/Egg/Tests/TC Inst Erasure Blocking TC Spec Gen.lean	
+++ b/Lean/Egg/Tests/TC Inst Erasure Blocking TC Spec Gen.lean	
@@ -16,25 +16,30 @@ info: [egg.rewrites] Rewrites
     [egg.rewrites] #0(⇒): h
       [egg.rewrites] -?x = ?x
       [egg.rewrites] LHS MVars
-          [?inst: [.inTcInstTerm, .isTcInst], ?x: [.unconditionallyVisible]]
+          [?x: [.unconditionallyVisible], ?inst: [.inTcInstTerm, .isTcInst]]
       [egg.rewrites] RHS MVars
           [?x: [.unconditionallyVisible]]
   [egg.rewrites] Tagged (0)
   [egg.rewrites] Builtin (0)
   [egg.rewrites] Derived (1)
     [egg.rewrites] #0<←>(⇔)
-      [egg.rewrites] -?m.38 = ?m.38
+      [egg.rewrites] -?m.50 = ?m.50
       [egg.rewrites] LHS MVars
-          [?m.38: [.unconditionallyVisible]]
+          [?m.50: [.unconditionallyVisible]]
       [egg.rewrites] RHS MVars
-          [?m.38: [.unconditionallyVisible]]
+          [?m.50: [.unconditionallyVisible]]
   [egg.rewrites] Definitional
   [egg.rewrites] Hypotheses (0)
-  [egg.rewrites] Pruned (0)
+  [egg.rewrites] Pruned (1)
+    [egg.rewrites] #0<⊢>(⇔)
+      [egg.rewrites] -?m.57 = ?m.57
+      [egg.rewrites] LHS MVars
+          [?m.57: [.unconditionallyVisible]]
+      [egg.rewrites] RHS MVars [?m.57: [.unconditionallyVisible]]
 -/
 #guard_msgs(info) in
 set_option egg.eraseTCInstances false in
-example (h : ∀ [inst : Neg Int] (x : Int), @Neg.neg Int inst x = x) : true = true := by
+example (h : ∀ [inst : Neg Int] (x : Int), @Neg.neg Int inst x = x) : (0 : Int) = (0 : Int) := by
   egg [h]
 
 /--
@@ -43,24 +48,33 @@ info: [egg.rewrites] Rewrites
     [egg.rewrites] #0(⇔): h
       [egg.rewrites] -?x = ?x
       [egg.rewrites] LHS MVars
-          [?x: [.unconditionallyVisible], ?inst: [.inTcInstTerm, .isTcInst]]
+          [?inst: [.inTcInstTerm, .isTcInst], ?x: [.unconditionallyVisible]]
       [egg.rewrites] RHS MVars
           [?x: [.unconditionallyVisible]]
   [egg.rewrites] Tagged (0)
   [egg.rewrites] Builtin (0)
   [egg.rewrites] Derived (1)
     [egg.rewrites] #0<←>(⇔)
-      [egg.rewrites] -?m.86 = ?m.86
+      [egg.rewrites] -?m.123 = ?m.123
       [egg.rewrites] LHS MVars
-          [?m.86: [.unconditionallyVisible]]
+          [?m.123: [.unconditionallyVisible]]
       [egg.rewrites] RHS MVars
-          [?m.86: [.unconditionallyVisible]]
+          [?m.123: [.unconditionallyVisible]]
   [egg.rewrites] Definitional
     [egg.rewrites] Type Class Instances
   [egg.rewrites] Hypotheses (0)
-  [egg.rewrites] Pruned (0)
+  [egg.rewrites] Pruned (1)
+    [egg.rewrites] #0<⊢>(⇔)
+      [egg.rewrites] -?m.130 = ?m.130
+      [egg.rewrites] LHS MVars
+          [?m.130: [.unconditionallyVisible]]
+      [egg.rewrites] RHS MVars [?m.130: [.unconditionallyVisible]]
 -/
 #guard_msgs(info) in
 set_option egg.eraseTCInstances true in
-example (h : ∀ [inst : Neg Int] (x : Int), @Neg.neg Int inst x = x) : true = true := by
+example (h : ∀ [inst : Neg Int] (x : Int), @Neg.neg Int inst x = x) : (0 : Int) = (0 : Int) := by
   egg [h]
+
+-- BUG: Goal type specialization should generate a rewrite here.
+example (h : ∀ (α) [inst : Neg α] (x : α), @Neg.neg α inst x = x) : (0 : Int) = (0 : Int) := by
+  sorry -- egg [h]
diff --git a/Rust/Egg/src/lean_expr.rs b/Rust/Egg/src/lean_expr.rs
index 2e080cf..01e9810 100644
--- a/Rust/Egg/src/lean_expr.rs
+++ b/Rust/Egg/src/lean_expr.rs
@@ -41,7 +41,7 @@ define_language! {
         // Construct for unknown terms (this is used in η-expansion):
         "_" = Unknown,        
     }
-}
+} 
 
 impl LeanExpr {