feat: non-commutative semiring normalizer in grind (#10421)

This PR adds a normalizer for non-commutative semirings to `grind`.
Examples:
```lean
open Lean.Grind
variable (R : Type u) [Semiring R]

example (a b c : R) : a * (b + c) = a * c + a * b := by grind
example (a b : R) : (a + 2 * b)^2 = a^2 + 2 * a * b + 2 * b * a + 4 * b^2 := by grind
example (a b : R) : b^2 + (a + 2 * b)^2 = a^2 + 2 * a * b + b * (1+1) * a * 1 + 5 * b^2 := by grind
example (a b : R) : a^3 + a^2*b + a*b*a + b*a^2 + a*b^2 + b*a*b + b^2*a + b^3 = (a+b)^3 := by grind
```

Author: leodemoura

src/Lean/Meta/Tactic/Grind/Arith/CommRing.lean

@@ -13,6 +13,8 @@ public import Lean.Meta.Tactic.Grind.Arith.CommRing.Internalize
 public import Lean.Meta.Tactic.Grind.Arith.CommRing.ToExpr
 public import Lean.Meta.Tactic.Grind.Arith.CommRing.RingM
 public import Lean.Meta.Tactic.Grind.Arith.CommRing.SemiringM
+public import Lean.Meta.Tactic.Grind.Arith.CommRing.NonCommRingM
+public import Lean.Meta.Tactic.Grind.Arith.CommRing.NonCommSemiringM
 public import Lean.Meta.Tactic.Grind.Arith.CommRing.Functions
 public import Lean.Meta.Tactic.Grind.Arith.CommRing.Reify
 public import Lean.Meta.Tactic.Grind.Arith.CommRing.EqCnstr

src/Lean/Meta/Tactic/Grind/Arith/CommRing/EqCnstr.lean

@@ -38,6 +38,13 @@ private def inSameNonCommRing? (a b : Expr) : GoalM (Option Nat) := do
   unless ringId == ringId' do return none -- This can happen when we have heterogeneous equalities
   return ringId
 
+/-- Returns `some semiringId` if `a` and `b` are elements of the same (non-commutative) semiring. -/
+private def inSameNonCommSemiring? (a b : Expr) : GoalM (Option Nat) := do
+  let some semiringId ← getTermNonCommSemiringId? a | return none
+  let some semiringId' ← getTermNonCommSemiringId? b | return none
+  unless semiringId == semiringId' do return none -- This can happen when we have heterogeneous equalities
+  return semiringId
+
 def mkEqCnstr (p : Poly) (h : EqCnstrProof) : RingM EqCnstr := do
   let id := (← getCommRing).nextId
   let sugar := p.degree
@@ -62,7 +69,7 @@ private def toRingExpr? [Monad m] [MonadLiftT GrindM m] [MonadRing m] (e : Expr)
 Returns the semiring expression denoting the given Lean expression.
 Recall that we compute the semiring expressions during internalization.
 -/
-private def toSemiringExpr? (e : Expr) : SemiringM (Option SemiringExpr) := do
+private def toSemiringExpr? [Monad m] [MonadLiftT GrindM m] [MonadSemiring m] (e : Expr) : m (Option SemiringExpr) := do
   let semiring ← getSemiring
   if let some re := semiring.denote.find? { expr := e } then
     return some re
@@ -353,7 +360,7 @@ def processNewEq (a b : Expr) : GoalM Unit := do
     let some sb ← toSemiringExpr? b | return ()
     let lhs ← sa.denoteAsRingExpr
     let rhs ← sb.denoteAsRingExpr
-    RingM.run (← getSemiring).ringId do
+    RingM.run (← getCommSemiring).ringId do
       let some ra ← reify? lhs (skipVar := false) (gen := (← getGeneration a)) | return ()
       let some rb ← reify? rhs (skipVar := false) (gen := (← getGeneration b)) | return ()
       let p ← (ra.sub rb).toPolyM
@@ -415,7 +422,7 @@ private def processNewDiseqCommSemiring (a b : Expr) : SemiringM Unit := do
     let some sb ← toSemiringExpr? b | return ()
     let lhs ← sa.denoteAsRingExpr
     let rhs ← sb.denoteAsRingExpr
-    RingM.run (← getSemiring).ringId do
+    RingM.run (← getCommSemiring).ringId do
       let some ra ← reify? lhs (skipVar := false) (gen := (← getGeneration a)) | return ()
       let some rb ← reify? rhs (skipVar := false) (gen := (← getGeneration b)) | return ()
       let p ← (ra.sub rb).toPolyM
@@ -443,13 +450,21 @@ private def processNewDiseqNonCommRing (a b : Expr) : NonCommRingM Unit := do
     if ra.toPoly_nc == rb.toPoly_nc then
       setNonCommRingDiseqUnsat a b ra rb
 
+private def processNewDiseqNonCommSemiring (a b : Expr) : NonCommSemiringM Unit := do
+  let some sa ← toSemiringExpr? a | return ()
+  let some sb ← toSemiringExpr? b | return ()
+  if sa.toPolyS_nc == sb.toPolyS_nc then
+    setNonCommSemiringDiseqUnsat a b sa sb
+
 def processNewDiseq (a b : Expr) : GoalM Unit := do
   if let some ringId ← inSameRing? a b then RingM.run ringId do
     processNewDiseqCommRing a b
   else if let some semiringId ← inSameSemiring? a b then SemiringM.run semiringId do
     processNewDiseqCommSemiring a b
   else if let some ncRingId ← inSameNonCommRing? a b then NonCommRingM.run ncRingId do
     processNewDiseqNonCommRing a b
+  else if let some ncSemiringId ← inSameNonCommSemiring? a b then NonCommSemiringM.run ncSemiringId do
+    processNewDiseqNonCommSemiring a b
 
 /--
 Returns `true` if the todo queue is not empty or the `recheck` flag is set to `true`

src/Lean/Meta/Tactic/Grind/Arith/CommRing/Internalize.lean

@@ -147,9 +147,15 @@ def internalize (e : Expr) (parent? : Option Expr) : GoalM Unit := do
     modifySemiring fun s => { s with denote := s.denote.insert { expr := e } re }
   else if let some ncRingId ← getNonCommRingId? type then NonCommRingM.run ncRingId do
     let some re ← ncreify? e | return ()
-    trace_goal[grind.ring.internalize] "(non-comm) ring [ncRingId}]: {e}"
+    trace_goal[grind.ring.internalize] "(non-comm) ring [{ncRingId}]: {e}"
     setTermNonCommRingId e
     ringExt.markTerm e
     modifyRing fun s => { s with denote := s.denote.insert { expr := e } re }
+  else if let some ncSemiringId ← getNonCommSemiringId? type then NonCommSemiringM.run ncSemiringId do
+    let some re ← ncsreify? e | return ()
+    trace_goal[grind.ring.internalize] "(non-comm) semiring [{ncSemiringId}]: {e}"
+    setTermNonCommSemiringId e
+    ringExt.markTerm e
+    modifySemiring fun s => { s with denote := s.denote.insert { expr := e } re }
 
 end Lean.Meta.Grind.Arith.CommRing

src/Lean/Meta/Tactic/Grind/Arith/CommRing/NonCommSemiringM.lean

@@ -0,0 +1,54 @@
+/-
+Copyright (c) 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Leonardo de Moura
+-/
+module
+prelude
+public import Lean.Meta.Tactic.Grind.Arith.CommRing.SemiringM
+public section
+namespace Lean.Meta.Grind.Arith.CommRing
+
+structure NonCommSemiringM.Context where
+  semiringId : Nat
+
+abbrev NonCommSemiringM := ReaderT NonCommSemiringM.Context GoalM
+
+abbrev NonCommSemiringM.run (semiringId : Nat) (x : NonCommSemiringM α) : GoalM α :=
+  x { semiringId }
+
+instance : MonadCanon NonCommSemiringM where
+  canonExpr e := do shareCommon (← canon e)
+  synthInstance? e := Grind.synthInstance? e
+
+protected def NonCommSemiringM.getSemiring : NonCommSemiringM Semiring := do
+  let s ← get'
+  let semiringId := (← read).semiringId
+  if h : semiringId < s.ncSemirings.size then
+    return s.ncSemirings[semiringId]
+  else
+    throwError "`grind` internal error, invalid semiringId"
+
+protected def NonCommSemiringM.modifySemiring (f : Semiring → Semiring) : NonCommSemiringM Unit := do
+  let semiringId := (← read).semiringId
+  modify' fun s => { s with ncSemirings := s.ncSemirings.modify semiringId f }
+
+instance : MonadSemiring NonCommSemiringM where
+  getSemiring := NonCommSemiringM.getSemiring
+  modifySemiring := NonCommSemiringM.modifySemiring
+
+def getTermNonCommSemiringId? (e : Expr) : GoalM (Option Nat) := do
+    return (← get').exprToNCSemiringId.find? { expr := e }
+
+def setTermNonCommSemiringId (e : Expr) : NonCommSemiringM Unit := do
+  let semiringId := (← read).semiringId
+  if let some semiringId' ← getTermNonCommSemiringId? e then
+    unless semiringId' == semiringId do
+      reportIssue! "expression in two different semirings{indentExpr e}"
+    return ()
+  modify' fun s => { s with exprToNCSemiringId := s.exprToNCSemiringId.insert { expr := e } semiringId }
+
+instance : MonadSetTermId NonCommSemiringM where
+  setTermId e := setTermNonCommSemiringId e
+
+end Lean.Meta.Grind.Arith.CommRing

src/Lean/Meta/Tactic/Grind/Arith/CommRing/Proof.lean

@@ -8,6 +8,7 @@ prelude
 public import Lean.Meta.Tactic.Grind.Arith.CommRing.RingId
 public import Lean.Meta.Tactic.Grind.Arith.CommRing.SemiringM
 public import Lean.Meta.Tactic.Grind.Arith.CommRing.NonCommRingM
+public import Lean.Meta.Tactic.Grind.Arith.CommRing.NonCommSemiringM
 import Init.Grind.Ring.CommSemiringAdapter
 import Lean.Data.RArray
 import Lean.Meta.Tactic.Grind.Diseq
@@ -31,7 +32,7 @@ private def toContextExpr [Monad m] [MonadLiftT MetaM m] [MonadCanon m] [MonadRi
   else
     RArray.toExpr ring.type id (RArray.leaf (mkApp (← getNatCastFn) (toExpr 0)))
 
-private def toSContextExpr' (vars : Array Expr) : SemiringM Expr := do
+private def toSContextExpr' [Monad m] [MonadLiftT MetaM m] [MonadCanon m] [MonadSemiring m] (vars : Array Expr) : m Expr := do
   let semiring ← getSemiring
   if h : 0 < vars.size then
     RArray.toExpr semiring.type id (RArray.ofFn (vars[·]) h)
@@ -133,7 +134,7 @@ private def getSemiringIdOf : RingM Nat := do
   return semiringId
 
 private def getSemiringOf : RingM CommSemiring := do
-  SemiringM.run (← getSemiringIdOf) do getSemiring
+  SemiringM.run (← getSemiringIdOf) do getCommSemiring
 
 private def mkSemiringPrefix (declName : Name) : ProofM Expr := do
   let sctx ← getSContext
@@ -322,7 +323,7 @@ Given `a` and `b`, such that `a ≠ b` in the core and `sa` and `sb` their reifi
 terms s.t. `sa.toPoly == sb.toPoly`, close the goal.
 -/
 def setSemiringDiseqUnsat (a b : Expr) (sa sb : SemiringExpr) : SemiringM Unit := do
-  let semiring ← getSemiring
+  let semiring ← getCommSemiring
   let hne ← mkDiseqProof a b
   let usedVars     := sa.collectVars >> sb.collectVars <| {}
   let vars'        := usedVars.toArray
@@ -358,4 +359,23 @@ def setNonCommRingDiseqUnsat (a b : Expr) (ra rb : RingExpr) : NonCommRingM Unit
   let h := mkApp3 h (toExpr ra) (toExpr rb) eagerReflBoolTrue
   closeGoal (mkApp hne h)
 
+/--
+Given `a` and `b`, such that `a ≠ b` in the core and `sa` and `sb` their reified semiring
+terms s.t. `sa.toPolyS_nc == sb.toPolyS_nc`, close the goal.
+-/
+def setNonCommSemiringDiseqUnsat (a b : Expr) (sa sb : SemiringExpr) : NonCommSemiringM Unit := do
+  let semiring ← getSemiring
+  let hne ← mkDiseqProof a b
+  let usedVars     := sa.collectVars >> sb.collectVars <| {}
+  let vars'        := usedVars.toArray
+  let varRename    := mkVarRename vars'
+  let vars         := semiring.vars
+  let vars         := vars'.map fun x => vars[x]!
+  let sa           := sa.renameVars varRename
+  let sb           := sb.renameVars varRename
+  let ctx          ← toSContextExpr' vars
+  let h := mkApp3 (mkConst ``Grind.CommRing.eq_normS_nc [semiring.u]) semiring.type semiring.semiringInst ctx
+  let h := mkApp3 h (toExpr sa) (toExpr sb) eagerReflBoolTrue
+  closeGoal (mkApp hne h)
+
 end Lean.Meta.Grind.Arith.CommRing

src/Lean/Meta/Tactic/Grind/Arith/CommRing/Reify.lean

@@ -8,6 +8,7 @@ prelude
 public import Lean.Meta.Tactic.Grind.Arith.CommRing.RingM
 public import Lean.Meta.Tactic.Grind.Arith.CommRing.SemiringM
 public import Lean.Meta.Tactic.Grind.Arith.CommRing.NonCommRingM
+public import Lean.Meta.Tactic.Grind.Arith.CommRing.NonCommSemiringM
 import Lean.Meta.Tactic.Grind.Simp
 import Lean.Meta.Tactic.Grind.Arith.CommRing.Functions
 public section
@@ -31,7 +32,7 @@ def isNatCastInst (inst : Expr) : m Bool :=
   return isSameExpr (← getNatCastFn).appArg! inst
 
 private def reportAppIssue (e : Expr) : GoalM Unit := do
-  reportIssue! "comm ring term with unexpected instance{indentExpr e}"
+  reportIssue! "ring term with unexpected instance{indentExpr e}"
 
 variable [MonadLiftT GoalM m] [MonadSetTermId m]
 
@@ -119,25 +120,30 @@ partial def reifyCore? (e : Expr) (skipVar : Bool) (gen : Nat) : m (Option RingE
     return some (.num k)
   | _ => toTopVar e
 
-partial def reify? (e : Expr) (skipVar := true) (gen : Nat := 0) : RingM (Option RingExpr) := do
+/-- Reify ring expression. -/
+def reify? (e : Expr) (skipVar := true) (gen : Nat := 0) : RingM (Option RingExpr) := do
   reifyCore? e skipVar gen
 
-partial def ncreify? (e : Expr) (skipVar := true) (gen : Nat := 0) : NonCommRingM (Option RingExpr) := do
+/-- Reify non-commutative ring expression. -/
+def ncreify? (e : Expr) (skipVar := true) (gen : Nat := 0) : NonCommRingM (Option RingExpr) := do
   reifyCore? e skipVar gen
 
 private def reportSAppIssue (e : Expr) : GoalM Unit := do
-  reportIssue! "comm semiring term with unexpected instance{indentExpr e}"
+  reportIssue! "semiring term with unexpected instance{indentExpr e}"
+
+section
+variable [MonadLiftT GoalM m] [MonadError m] [Monad m] [MonadCanon m] [MonadSemiring m] [MonadSetTermId m]
 
 /--
 Similar to `reify?` but for `CommSemiring`
 -/
-partial def sreify? (e : Expr) : SemiringM (Option SemiringExpr) := do
-  let toVar (e : Expr) : SemiringM SemiringExpr := do
-    return .var (← mkSVar e)
-  let asVar (e : Expr) : SemiringM SemiringExpr := do
+partial def sreifyCore? (e : Expr) : m (Option SemiringExpr) := do
+  let toVar (e : Expr) : m SemiringExpr := do
+    return .var (← mkSVarCore e)
+  let asVar (e : Expr) : m SemiringExpr := do
     reportSAppIssue e
-    return .var (← mkSVar e)
-  let rec go (e : Expr) : SemiringM SemiringExpr := do
+    return .var (← mkSVarCore e)
+  let rec go (e : Expr) : m SemiringExpr := do
     match_expr e with
     | HAdd.hAdd _ _ _ i a b =>
       if isSameExpr (← getAddFn').appArg! i then return .add (← go a) (← go b) else asVar e
@@ -156,9 +162,9 @@ partial def sreify? (e : Expr) : SemiringM (Option SemiringExpr) := do
       let some k ← getNatValue? n | toVar e
       return .num k
     | _ => toVar e
-  let toTopVar (e : Expr) : SemiringM (Option SemiringExpr) := do
+  let toTopVar (e : Expr) : m (Option SemiringExpr) := do
     return some (← toVar e)
-  let asTopVar (e : Expr) : SemiringM (Option SemiringExpr) := do
+  let asTopVar (e : Expr) : m (Option SemiringExpr) := do
     reportSAppIssue e
     toTopVar e
   match_expr e with
@@ -180,4 +186,14 @@ partial def sreify? (e : Expr) : SemiringM (Option SemiringExpr) := do
     return some (.num k)
   | _ => toTopVar e
 
+end
+
+/-- Reify semiring expression. -/
+def sreify? (e : Expr) : SemiringM (Option SemiringExpr) := do
+  sreifyCore? e
+
+/-- Reify non-commutative semiring expression. -/
+def ncsreify? (e : Expr) : NonCommSemiringM (Option SemiringExpr) := do
+  sreifyCore? e
+
 end  Lean.Meta.Grind.Arith.CommRing

src/Lean/Meta/Tactic/Grind/Arith/CommRing/RingId.lean

@@ -104,4 +104,21 @@ where
     setCommSemiringId ringId id
     return some id
 
+def getNonCommSemiringId? (type : Expr) : GoalM (Option Nat) := do
+  if let some id? := (← get').ncstypeIdOf.find? { expr := type } then
+    return id?
+  else
+    let id? ← go?
+    modify' fun s => { s with ncstypeIdOf := s.ncstypeIdOf.insert { expr := type } id? }
+    return id?
+where
+  go? : GoalM (Option Nat) := do
+    let u ← getDecLevel type
+    let semiring := mkApp (mkConst ``Grind.Semiring [u]) type
+    let some semiringInst ← synthInstance? semiring | return none
+    let id := (← get').ncSemirings.size
+    let semiring : Semiring := { id, type, u, semiringInst }
+    modify' fun s => { s with ncSemirings := s.ncSemirings.push semiring }
+    return some id
+
 end Lean.Meta.Grind.Arith.CommRing