fix: more stable choice of representative for atoms in ring and `ab…

…el` (#19119)

Algebraic normalization tactics (`ring`, `abel`, etc.) typically require a concept of "atom", with expressions which are t-defeq (for some transparency t) being identified.  However, the particular representative of this equivalence class which turns up in the final normalized expression is basically random.  (It often ends up being the "first", in some tree sense, occurrence of the equivalence class in the expression being normalized.)

This ends up being particularly unpredictable when multiple expressions are being normalized simultaneously, e.g. with `ring_nf` or `abel_nf`: it can occur that in different expressions, a different representative of the equivalence class is chosen.  For example, on current Mathlib,
```lean
example (x : ℤ) (R : ℤ → ℤ → Prop) : True := by
  let a := x
  have h : R (a + x) (x + a) := sorry
  ring_nf at h
```
the statement of `h` after the ring-normalization step is `h : R (a * 2) (x * 2)`.  Here `a` and `x` are reducibly defeq.  When normalizing `a + x`, the representative `a` was chosen for the equivalence class; when normalizing `x + a`, the representative `x` was chosen.

This PR implements a fix.  The `AtomM` monad (which is used for atom-tracking in `ring`, `abel`, etc.) has its `addAtom` function modified to report, not just the index of an atom, but also the stored form of the atom.  Then the code surrounding `addAtom` calls in the `ring`, `abel` and `module` tactics is modified to use the stored form of the atom, rather than the form actually encountered at that point in the expression.



Co-authored-by: Eric Wieser <[email protected]>

Mathlib/Tactic/Abel.lean

@@ -247,11 +247,11 @@ theorem term_atom_pfg {α} [AddCommGroup α] (x x' : α) (h : x = x') : x = term
 /-- Interpret an expression as an atom for `abel`'s normal form. -/
 def evalAtom (e : Expr) : M (NormalExpr × Expr) := do
   let { expr := e', proof?, .. } ← (← readThe AtomM.Context).evalAtom e
-  let i ← AtomM.addAtom e'
+  let (i, a) ← AtomM.addAtom e'
   let p ← match proof? with
   | none => iapp ``term_atom #[e]
-  | some p => iapp ``term_atom_pf #[e, e', p]
-  return (← term' (← intToExpr 1, 1) (i, e') (← zero'), p)
+  | some p => iapp ``term_atom_pf #[e, a, p]
+  return (← term' (← intToExpr 1, 1) (i, a) (← zero'), p)
 
 theorem unfold_sub {α} [SubtractionMonoid α] (a b c : α) (h : a + -b = c) : a - b = c := by
   rw [sub_eq_add_neg, h]

Mathlib/Tactic/ITauto.lean

@@ -486,7 +486,7 @@ partial def reify (e : Q(Prop)) : AtomM IProp :=
   | ~q(@Ne Prop $a $b) => return .not (.eq (← reify a) (← reify b))
   | e =>
     if e.isArrow then return .imp (← reify e.bindingDomain!) (← reify e.bindingBody!)
-    else return .var (← AtomM.addAtom e)
+    else return .var (← AtomM.addAtom e).1
 
 /-- Once we have a proof object, we have to apply it to the goal. -/
 partial def applyProof (g : MVarId) (Γ : NameMap Expr) (p : Proof) : MetaM Unit :=

Mathlib/Tactic/Linarith/Preprocessing.lean

@@ -274,12 +274,12 @@ partial def findSquares (s : RBSet (Nat × Bool) lexOrd.compare) (e : Expr) :
   | (``HPow.hPow, #[_, _, _, _, a, b]) => match b.numeral? with
     | some 2 => do
       let s ← findSquares s a
-      let ai ← AtomM.addAtom a
+      let (ai, _) ← AtomM.addAtom a
       return (s.insert (ai, true))
     | _ => e.foldlM findSquares s
   | (``HMul.hMul, #[_, _, _, _, a, b]) => do
-    let ai ← AtomM.addAtom a
-    let bi ← AtomM.addAtom b
+    let (ai, _) ← AtomM.addAtom a
+    let (bi, _) ← AtomM.addAtom b
     if ai = bi then do
       let s ← findSquares s a
       return (s.insert (ai, false))

Mathlib/Tactic/Module.lean

@@ -464,9 +464,9 @@ partial def parse (iM : Q(AddCommMonoid $M)) (x : Q($M)) :
     pure ⟨0, q(Nat), q(Nat.instSemiring), q(AddCommGroup.toNatModule), [], q(NF.zero_eq_eval $M)⟩
   /- anything else should be treated as an atom -/
   | _ =>
-    let k : ℕ ← AtomM.addAtom x
-    pure ⟨0, q(Nat), q(Nat.instSemiring), q(AddCommGroup.toNatModule), [((q(1), x), k)],
-      q(NF.atom_eq_eval $x)⟩
+    let (k, ⟨x', _⟩) ← AtomM.addAtomQ x
+    pure ⟨0, q(Nat), q(Nat.instSemiring), q(AddCommGroup.toNatModule), [((q(1), x'), k)],
+      q(NF.atom_eq_eval $x')⟩
 
 /-- Given expressions `R` and `M` representing types such that `M`'s is a module over `R`'s, and
 given two terms `l₁`, `l₂` of type `qNF R M`, i.e. lists of `(Q($R) × Q($M)) × ℕ`s (two `Expr`s

Mathlib/Tactic/Polyrith.lean

@@ -133,7 +133,7 @@ partial def parse {u : Level} {α : Q(Type u)} (sα : Q(CommSemiring $α))
     (c : Ring.Cache sα) (e : Q($α)) : AtomM Poly := do
   let els := do
     try pure <| Poly.const (← (← NormNum.derive e).toRat)
-    catch _ => pure <| Poly.var (← addAtom e)
+    catch _ => pure <| Poly.var (← addAtom e).1
   let .const n _ := (← withReducible <| whnf e).getAppFn | els
   match n, c.rα with
   | ``HAdd.hAdd, _ | ``Add.add, _ => match e with

Mathlib/Tactic/Ring/Basic.lean

@@ -512,9 +512,8 @@ mutual
 partial def ExBase.evalNatCast {a : Q(ℕ)} (va : ExBase sℕ a) : AtomM (Result (ExBase sα) q($a)) :=
   match va with
   | .atom _ => do
-    let a' : Q($α) := q($a)
-    let i ← addAtom a'
-    pure ⟨a', ExBase.atom i, (q(Eq.refl $a') : Expr)⟩
+    let (i, ⟨b', _⟩) ← addAtomQ q($a)
+    pure ⟨b', ExBase.atom i, q(Eq.refl $b')⟩
   | .sum va => do
     let ⟨_, vc, p⟩ ← va.evalNatCast
     pure ⟨_, .sum vc, p⟩
@@ -952,11 +951,11 @@ Evaluates an atom, an expression where `ring` can find no additional structure.
 def evalAtom (e : Q($α)) : AtomM (Result (ExSum sα) e) := do
   let r ← (← read).evalAtom e
   have e' : Q($α) := r.expr
-  let i ← addAtom e'
-  let ve' := (ExBase.atom i (e := e')).toProd (ExProd.mkNat sℕ 1).2 |>.toSum
+  let (i, ⟨a', _⟩) ← addAtomQ e'
+  let ve' := (ExBase.atom i (e := a')).toProd (ExProd.mkNat sℕ 1).2 |>.toSum
   pure ⟨_, ve', match r.proof? with
   | none => (q(atom_pf $e) : Expr)
-  | some (p : Q($e = $e')) => (q(atom_pf' $p) : Expr)⟩
+  | some (p : Q($e = $a')) => (q(atom_pf' $p) : Expr)⟩
 
 theorem inv_mul {R} [DivisionRing R] {a₁ a₂ a₃ b₁ b₃ c}
     (_ : (a₁⁻¹ : R) = b₁) (_ : (a₃⁻¹ : R) = b₃)
@@ -977,9 +976,8 @@ variable (dα : Q(DivisionRing $α))
 
 /-- Applies `⁻¹` to a polynomial to get an atom. -/
 def evalInvAtom (a : Q($α)) : AtomM (Result (ExBase sα) q($a⁻¹)) := do
-  let a' : Q($α) := q($a⁻¹)
-  let i ← addAtom a'
-  pure ⟨a', ExBase.atom i, (q(Eq.refl $a') : Expr)⟩
+  let (i, ⟨b', _⟩) ← addAtomQ q($a⁻¹)
+  pure ⟨b', ExBase.atom i, q(Eq.refl $b')⟩
 
 /-- Inverts a polynomial `va` to get a normalized result polynomial.
 

Mathlib/Tactic/TFAE.lean

@@ -298,18 +298,18 @@ elab_rules : tactic
   goal.withContext do
     let (tfaeListQ, tfaeList) ← getTFAEList (← goal.getType)
     closeMainGoal `tfae_finish <|← AtomM.run .reducible do
-      let is ← tfaeList.mapM AtomM.addAtom
+      let is ← tfaeList.mapM (fun e ↦ Prod.fst <$> AtomM.addAtom e)
       let mut hyps := #[]
       for hyp in ← getLocalHyps do
         let ty ← whnfR <|← instantiateMVars <|← inferType hyp
         if let (``Iff, #[p1, p2]) := ty.getAppFnArgs then
-          let q1 ← AtomM.addAtom p1
-          let q2 ← AtomM.addAtom p2
+          let (q1, _) ← AtomM.addAtom p1
+          let (q2, _) ← AtomM.addAtom p2
           hyps := hyps.push (q1, q2, ← mkAppM ``Iff.mp #[hyp])
           hyps := hyps.push (q2, q1, ← mkAppM ``Iff.mpr #[hyp])
         else if ty.isArrow then
-          let q1 ← AtomM.addAtom ty.bindingDomain!
-          let q2 ← AtomM.addAtom ty.bindingBody!
+          let (q1, _) ← AtomM.addAtom ty.bindingDomain!
+          let (q2, _) ← AtomM.addAtom ty.bindingBody!
           hyps := hyps.push (q1, q2, hyp)
       proveTFAE hyps (← get).atoms is tfaeListQ
 

Mathlib/Util/AtomM.lean

@@ -5,6 +5,7 @@ Authors: Mario Carneiro
 -/
 import Mathlib.Init
 import Lean.Meta.Tactic.Simp.Types
+import Qq
 
 /-!
 # A monad for tracking and deduplicating atoms
@@ -44,13 +45,33 @@ def AtomM.run {α : Type} (red : TransparencyMode) (m : AtomM α)
     MetaM α :=
   (m { red, evalAtom }).run' {}
 
-/-- Get the index corresponding to an atomic expression, if it has already been encountered, or
-put it in the list of atoms and return the new index, otherwise. -/
-def AtomM.addAtom (e : Expr) : AtomM Nat := do
+/-- If an atomic expression has already been encountered, get the index and the stored form of the
+atom (which will be defeq at the specified transparency, but not necessarily syntactically equal).
+If the atomic expression has *not* already been encountered, store it in the list of atoms, and
+return the new index (and the stored form of the atom, which will be itself).
+
+In a normalizing tactic, the expression returned by `addAtom` should be considered the normal form.
+-/
+def AtomM.addAtom (e : Expr) : AtomM (Nat × Expr) := do
   let c ← get
   for h : i in [:c.atoms.size] do
     if ← withTransparency (← read).red <| isDefEq e c.atoms[i] then
-      return i
-  modifyGet fun c ↦ (c.atoms.size, { c with atoms := c.atoms.push e })
+      return (i, c.atoms[i])
+  modifyGet fun c ↦ ((c.atoms.size, e), { c with atoms := c.atoms.push e })
+
+open Qq in
+/-- If an atomic expression has already been encountered, get the index and the stored form of the
+atom (which will be defeq at the specified transparency, but not necessarily syntactically equal).
+If the atomic expression has *not* already been encountered, store it in the list of atoms, and
+return the new index (and the stored form of the atom, which will be itself).
+
+In a normalizing tactic, the expression returned by `addAtomQ` should be considered the normal form.
+
+This is a strongly-typed version of `AtomM.addAtom` for code using `Qq`.
+-/
+def AtomM.addAtomQ {u : Level} {α : Q(Type u)} (e : Q($α)) :
+    AtomM (Nat × {e' : Q($α) // $e =Q $e'}) := do
+  let (n, e') ← AtomM.addAtom e
+  return (n, ⟨e', ⟨⟩⟩)
 
 end Mathlib.Tactic

MathlibTest/abel.lean

@@ -141,3 +141,32 @@ example [AddCommGroup α] (x y z : α) (h : False) (w : x - x = y + z) : False :
   abel_nf at *
   guard_hyp w : 0 = y + z
   assumption
+
+section
+abbrev myId (a : ℤ) : ℤ := a
+
+/-
+Test that when `abel_nf` normalizes multiple expressions which contain a particular atom, it uses a
+form for that atom which is consistent between expressions.
+
+We can't use `guard_hyp h :ₛ` here, as while it does tell apart `x` and `myId x`, it also complains
+about differing instance paths.
+-/
+/--
+info: α : Type _
+a b : α
+x : ℤ
+R : ℤ → ℤ → Prop
+hR : Reflexive R
+h : R (2 • myId x) (2 • myId x)
+⊢ True
+-/
+#guard_msgs (info) in
+set_option pp.mvars false in
+example (x : ℤ) (R : ℤ → ℤ → Prop) (hR : Reflexive R) : True := by
+  have h : R (myId x + x) (x + myId x) := hR ..
+  abel_nf at h
+  trace_state
+  trivial
+
+end

MathlibTest/ring.lean

@@ -178,3 +178,28 @@ example {n : ℝ} (_hn : 0 ≤ n) : (n + 1 / 2) ^ 2 * (n + 1 + 1 / 3) ≤ (n + 1
   ring_nf
   trace_state
   exact test_sorry
+
+section
+abbrev myId (a : ℤ) : ℤ := a
+
+/-
+Test that when `ring_nf` normalizes multiple expressions which contain a particular atom, it uses a
+form for that atom which is consistent between expressions.
+
+We can't use `guard_hyp h :ₛ` here, as while it does tell apart `x` and `myId x`, it also complains
+about differing instance paths.
+-/
+/--
+info: x : ℤ
+R : ℤ → ℤ → Prop
+h : R (myId x * 2) (myId x * 2)
+⊢ True
+-/
+#guard_msgs (info) in
+example (x : ℤ) (R : ℤ → ℤ → Prop) : True := by
+  have h : R (myId x + x) (x + myId x) := test_sorry
+  ring_nf at h
+  trace_state
+  trivial
+
+end