fix: use fvarsSubset

src/Lean/Meta/Tactic/Grind.lean

@@ -12,6 +12,7 @@ import Lean.Meta.Tactic.Grind.Util
 import Lean.Meta.Tactic.Grind.Cases
 import Lean.Meta.Tactic.Grind.Injection
 import Lean.Meta.Tactic.Grind.Core
+import Lean.Meta.Tactic.Grind.Canon
 
 namespace Lean
 

src/Lean/Meta/Tactic/Grind/Canon.lean

@@ -0,0 +1,150 @@
+/-
+Copyright (c) 2024 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
+-/
+prelude
+import Lean.Meta.Basic
+import Lean.Meta.FunInfo
+import Lean.Util.FVarSubset
+import Lean.Util.PtrSet
+import Lean.Util.FVarSubset
+
+namespace Lean.Meta.Grind
+namespace Canon
+
+/-!
+A canonicalizer module for the `grind` tactic. The canonicalizer defined in `Meta/Canonicalizer.lean` is
+not suitable for the `grind` tactic. It was designed for tactics such as `omega`, where the goal is
+to detect when two structurally different atoms are definitionally equal.
+
+The `grind` tactic, on the other hand, uses congruence closure but disregards types, type formers, instances, and proofs.
+Proofs are ignored due to proof irrelevance. Types, type formers, and instances are considered supporting
+elements and are not factored into congruence detection. Instead, `grind` only checks whether
+elements are structurally equal, which, in the context of the `grind` tactic, is equivalent to pointer equality.
+
+This module minimizes the number of `isDefEq` checks by comparing two terms `a` and `b` only if they instances,
+types, or type formers and are the `i`-th arguments of two different `f`-applications. This approach is
+sufficient for the congruence closure procedure used by the `grind` tactic.
+
+To further optimize `isDefEq` checks, instances are compared using `TransparencyMode.instances`, which reduces
+the number of constants that need to be unfolded. If diagnostics are enabled, instances are compared using
+the default transparency mode too for sanity checking, and discrepancies are reported.
+Types and type formers are always checked using default transparency.
+-/
+
+structure State where
+  argMap : PHashMap (Expr × Nat) (List Expr) := {}
+  canon  : PHashMap Expr Expr := {}
+  deriving Inhabited
+
+/--
+Helper function for canonicalizing `e` occurring as the `i`th argument of an `f`-application.
+`isInst` is true if `e` is an type class instance.
+
+Recall that we use `TransparencyMode.instances` for checking whether two instances are definitionally equal or not.
+Thus, if diagnostics are enabled, we also check them using `TransparencyMode.default`. If the result is different
+we report to the user.
+-/
+def canonElemCore (f : Expr) (i : Nat) (e : Expr) (isInst : Bool) : StateT State MetaM Expr := do
+  let s ← get
+  if let some c := s.canon.find? e then
+    return c
+  let key := (f, i)
+  let cs := s.argMap.find? key |>.getD []
+  for c in cs do
+    if (← isDefEq e c) then
+      if c.fvarsSubset e then
+        -- It is not in general safe to replace `e` with `c` if `c` has more free variables than `e`.
+        modify fun s => { s with canon := s.canon.insert e c }
+        return c
+    if isInst then
+      if (← isDiagnosticsEnabled <&&> pure (c.fvarsSubset e) <&&> (withDefault <| isDefEq e c)) then
+        -- TODO: consider storing this information in some structure that can be browsed later.
+        trace[grind.issues] "the following `grind` static elements are definitionally equal with `default` transparency, but not with `instances` transparency{indentExpr e}\nand{indentExpr c}"
+  modify fun s => { s with canon := s.canon.insert e e, argMap := s.argMap.insert key (e::cs) }
+  return e
+
+abbrev canonType (f : Expr) (i : Nat) (e : Expr) := withDefault <| canonElemCore f i e false
+abbrev canonInst (f : Expr) (i : Nat) (e : Expr) := withReducibleAndInstances <| canonElemCore f i e true
+
+/--
+Return type for the `shouldCanon` function.
+-/
+private inductive ShouldCanonResult where
+  | /-
+    Nested proofs are ignored by the canonizer.
+    That is, they are not canonized or recursively visited.
+    -/
+    ignore
+  | /- Nested types (and type formers) are canonicalized. -/
+    canonType
+  | /- Nested instances are canonicalized. -/
+    canonInst
+  | /-
+    Term is not a proof, type (former), nor an instance.
+    Thus, it must be recursively visited by the canonizer.
+    -/
+    visit
+  deriving Inhabited
+
+/--
+See comments at `ShouldCanonResult`.
+-/
+def shouldCanon (pinfos : Array ParamInfo) (i : Nat) (arg : Expr) : MetaM ShouldCanonResult := do
+  if h : i < pinfos.size then
+    let pinfo := pinfos[i]
+    if pinfo.isInstImplicit then
+      return .canonInst
+    else if pinfo.isProp then
+      return .ignore
+  if (← isTypeFormer arg) then
+    return .canonType
+  else
+    return .visit
+
+unsafe def canonImpl (e : Expr) : StateT State MetaM Expr := do
+  visit e |>.run' mkPtrMap
+where
+  visit (e : Expr) : StateRefT (PtrMap Expr Expr) (StateT State MetaM) Expr := do
+    match e with
+    | .bvar .. => unreachable!
+    -- Recall that `grind` treats `let`, `forall`, and `lambda` as atomic terms.
+    | .letE .. | .forallE .. | .lam ..
+    | .const .. | .lit .. | .mvar .. | .sort .. | .fvar ..
+    -- Recall that the `grind` preprocessor uses the `foldProjs` preprocessing step.
+    | .proj ..
+    -- Recall that the `grind` preprocessor uses the `eraseIrrelevantMData` preprocessing step.
+    | .mdata ..  => return e
+    -- We only visit applications
+    | .app .. =>
+      -- Check whether it is cached
+      if let some r := (← get).find? e then
+        return r
+      e.withApp fun f args => do
+        let pinfos := (← getFunInfo f).paramInfo
+        let mut modified := false
+        let mut args := args
+        for i in [:args.size] do
+          let arg := args[i]!
+          let arg' ← match (← shouldCanon pinfos i arg) with
+          | .ignore    => pure arg
+          | .canonType => canonType f i arg
+          | .canonInst => canonInst f i arg
+          | .visit     => visit arg
+          unless ptrEq arg arg' do
+            args := args.set! i arg'
+            modified := true
+        let e' := if modified then mkAppN f args else e
+        modify fun s => s.insert e e'
+        return e'
+
+/--
+Canonicalizes nested types, type formers, and instances in `e`.
+-/
+def canon (e : Expr) : StateT State MetaM Expr :=
+  unsafe canonImpl e
+
+end Canon
+
+end Lean.Meta.Grind

src/Lean/Meta/Tactic/Grind/Types.lean

@@ -8,7 +8,7 @@ import Lean.Util.ShareCommon
 import Lean.Meta.Basic
 import Lean.Meta.CongrTheorems
 import Lean.Meta.AbstractNestedProofs
-import Lean.Meta.Canonicalizer
+import Lean.Meta.Tactic.Grind.Canon
 import Lean.Meta.Tactic.Util
 
 namespace Lean.Meta.Grind
@@ -37,7 +37,7 @@ instance : Hashable CongrTheoremCacheKey where
   hash a := mixHash (unsafe ptrAddrUnsafe a.f).toUInt64 (hash a.numArgs)
 
 structure State where
-  canon      : Canonicalizer.State := {}
+  canon      : Canon.State := {}
   /-- `ShareCommon` (aka `Hashconsing`) state. -/
   scState    : ShareCommon.State.{0} ShareCommon.objectFactory := ShareCommon.State.mk _
   /-- Next index for creating auxiliary theorems. -/
@@ -87,25 +87,13 @@ def shareCommon (e : Expr) : GrindM Expr := do
     (e, { canon, scState, nextThmIdx, congrThms, trueExpr, falseExpr })
 
 /--
-Applies the canonicalizer to all subterms of `e`.
+Canonicalizes nested types, type formers, and instances in `e`.
 -/
--- TODO: the current canonicalizer is not a good solution for `grind`.
--- The problem is that two different applications `@f inst_1 a` and `@f inst_2 b`
--- may still have syntaticaally different instances. Thus, if we learn that `a = b`,
--- congruence closure will fail to see that the two applications are congruent.
 def canon (e : Expr) : GrindM Expr := do
   let canonS ← modifyGet fun s => (s.canon, { s with canon := {} })
-  let (e, canonS) ← Canonicalizer.CanonM.run (canonRec e) (s := canonS)
+  let (e, canonS) ← Canon.canon e |>.run canonS
   modify fun s => { s with canon := canonS }
   return e
-where
-  canonRec (e : Expr) : CanonM Expr := do
-    let post (e : Expr) : CanonM TransformStep := do
-      if e.isApp then
-        return .done (← Meta.canon e)
-      else
-        return .done e
-    transform e post
 
 /--
 Creates a congruence theorem for a `f`-applications with `numArgs` arguments.

tests/lean/run/grindCanonInsts.lean

@@ -0,0 +1,87 @@
+import Lean
+
+open Lean Meta Elab Tactic Grind in
+elab "grind_test" : tactic => withMainContext do
+  let declName := (← Term.getDeclName?).getD `_main
+  Meta.Grind.preprocessAndProbe (← getMainGoal) declName do
+    let nodes ← filterENodes fun e => return e.self.isAppOf ``HMul.hMul
+    logInfo (nodes.toList.map (·.self))
+
+class Semigroup (α : Type u) extends Mul α where
+  mul_assoc (a b c : α) : a * b * c = a * (b * c)
+
+export Semigroup (mul_assoc)
+
+class MulComm (α : Type u)  extends Mul α where
+  mul_comm (a b : α) : a * b = b * a
+
+export MulComm (mul_comm)
+
+class CommSemigroup (α : Type u) extends Semigroup α where
+  mul_comm (a b : α) : a * b = b * a
+
+instance [CommSemigroup α] : MulComm α where
+  mul_comm := CommSemigroup.mul_comm
+
+class One (α : Type u) where
+  one : α
+
+instance [One α] : OfNat α (nat_lit 1) where
+  ofNat := One.one
+
+class Monoid (α : Type u) extends Semigroup α, One α where
+  one_mul (a : α) : 1 * a = a
+  mul_one (a : α) : a * 1 = a
+
+export Monoid (one_mul mul_one)
+
+class CommMonoid (α : Type u) extends Monoid α where
+  mul_comm (a b : α) : a * b = b * a
+
+instance [CommMonoid α] : CommSemigroup α where
+  mul_comm := CommMonoid.mul_comm
+
+instance [CommMonoid α] : MulComm α where
+  mul_comm := CommSemigroup.mul_comm
+
+instance : CommMonoid Nat where
+  mul := Nat.mul
+  one := 1
+  mul_assoc := Nat.mul_assoc
+  mul_comm  := Nat.mul_comm
+  one_mul   := Nat.one_mul
+  mul_one   := Nat.mul_one
+
+theorem left_comm [CommMonoid α] (a b c : α) : a * (b * c) = b * (a * c) := by
+  rw [← Semigroup.mul_assoc, CommMonoid.mul_comm a b, Semigroup.mul_assoc]
+
+/--
+info: [b * c, a * (b * c), d * (b * c)]
+---
+warning: declaration uses 'sorry'
+-/
+#guard_msgs in
+example (a b c d : Nat) : b * (a * c) = d * (b * c) → False := by
+  rw [left_comm] -- Introduces a new (non-canonical) instance for `Mul Nat`
+  grind_test -- State should have only 3 `*`-applications
+  sorry
+
+
+set_option pp.notation false in
+set_option pp.explicit true in
+/--
+info: [@HMul.hMul Nat Nat Nat (@instHMul Nat instMulNat) b a, @HMul.hMul Nat Nat Nat (@instHMul Nat instMulNat) b d]
+---
+warning: declaration uses 'sorry'
+-/
+#guard_msgs in
+example (a b c d : Nat) : b * a = d * b → False := by
+  rw [CommMonoid.mul_comm d b] -- Introduces a new (non-canonical) instance for `Mul Nat`
+  -- See target here
+  guard_target =ₛ
+    @HMul.hMul Nat Nat Nat (@instHMul Nat instMulNat) b a
+    =
+    @HMul.hMul Nat Nat Nat (@instHMul Nat (@Semigroup.toMul Nat (@Monoid.toSemigroup Nat (@CommMonoid.toMonoid Nat instCommMonoidNat)))) b d
+    → False
+  grind_test -- State should have only 2 `*`-applications, and they use the same instance
+  sorry

tests/lean/run/grindCanonTypes.lean

@@ -0,0 +1,28 @@
+import Lean
+
+
+def g (s : Type) := s
+def f (a : α) := a
+
+open Lean Meta Elab Tactic Grind in
+elab "grind_test" : tactic => withMainContext do
+  let declName := (← Term.getDeclName?).getD `_main
+  Meta.Grind.preprocessAndProbe (← getMainGoal) declName do
+    let nodes ← filterENodes fun e => return e.self.isAppOf ``f
+    logInfo (nodes.toList.map (·.self))
+
+
+set_option pp.explicit true
+/--
+info: [@f Nat a, @f Nat b]
+---
+warning: declaration uses 'sorry'
+-/
+#guard_msgs in
+example (a b c d : Nat) : @f Nat a = b → @f (g Nat) a = c → @f (g Nat) b = d → a = b → False := by
+  -- State should have only two `f`-applications: `@f Nat a`, `@f Nat b`
+  -- Note that `@f (g Nat) b` has been canonicalized to `@f Nat b`.
+  -- Thus, if `a` and `b` equivalence classes are merged, `grind` can still detect that
+  -- `@f Nat a` and `@f Nat b` are equal too.
+  grind_test
+  sorry