From 8dec57987ab13244c67fe03edf004d9e288b4df1 Mon Sep 17 00:00:00 2001
From: Kim Morrison <scott.morrison@gmail.com>
Date: Tue, 7 Jan 2025 03:29:50 +1100
Subject: [PATCH] feat: `grind` tests for basic category theory (#6543)

This PR adds additional tests for `grind`, demonstrating that we can
automate some manual proofs from Mathlib's basic category theory
library, with less reliance on Mathlib's `@[reassoc]` trick.

In several places I've added bidirectional patterns for equational
lemmas.

I've updated some other files to use the new `@[grind_eq]` attribute
(but left as is all cases where we are inspecting the info messages from
`grind_pattern`).

---------

Co-authored-by: Leonardo de Moura <leomoura@amazon.com>
---
 tests/lean/run/grind_cat.lean     | 125 ++++++++++++++++++++++++++++++
 tests/lean/run/grind_ematch1.lean |   9 +--
 tests/lean/run/grind_ematch2.lean |   7 +-
 tests/lean/run/grind_shelf.lean   |  30 +++++++
 4 files changed, 161 insertions(+), 10 deletions(-)
 create mode 100644 tests/lean/run/grind_cat.lean
 create mode 100644 tests/lean/run/grind_shelf.lean

diff --git a/tests/lean/run/grind_cat.lean b/tests/lean/run/grind_cat.lean
new file mode 100644
index 0000000000..9e501f97f9
--- /dev/null
+++ b/tests/lean/run/grind_cat.lean
@@ -0,0 +1,125 @@
+universe v v₁ v₂ v₃ u u₁ u₂ u₃
+
+namespace CategoryTheory
+
+macro "cat_tac" : tactic => `(tactic| (intros; (try ext); grind))
+
+class Category (obj : Type u) : Type max u (v + 1) where
+  Hom : obj → obj → Type v
+  /-- The identity morphism on an object. -/
+  id : ∀ X : obj, Hom X X
+  /-- Composition of morphisms in a category, written `f ≫ g`. -/
+  comp : ∀ {X Y Z : obj}, (Hom X Y) → (Hom Y Z) → (Hom X Z)
+  /-- Identity morphisms are left identities for composition. -/
+  id_comp : ∀ {X Y : obj} (f : Hom X Y), comp (id X) f = f := by cat_tac
+  /-- Identity morphisms are right identities for composition. -/
+  comp_id : ∀ {X Y : obj} (f : Hom X Y), comp f (id Y) = f := by cat_tac
+  /-- Composition in a category is associative. -/
+  assoc : ∀ {W X Y Z : obj} (f : Hom W X) (g : Hom X Y) (h : Hom Y Z), comp (comp f g) h = comp f (comp g h) := by cat_tac
+
+infixr:10 " ⟶ " => Category.Hom
+scoped notation "𝟙" => Category.id  -- type as \b1
+scoped infixr:80 " ≫ " => Category.comp
+
+attribute [simp] Category.id_comp Category.comp_id Category.assoc
+
+attribute [grind =] Category.id_comp Category.comp_id
+attribute [grind _=_] Category.assoc
+
+structure Functor (C : Type u₁) [Category.{v₁} C] (D : Type u₂) [Category.{v₂} D] : Type max v₁ v₂ u₁ u₂ where
+  /-- The action of a functor on objects. -/
+  obj : C → D
+  /-- The action of a functor on morphisms. -/
+  map : ∀ {X Y : C}, (X ⟶ Y) → ((obj X) ⟶ (obj Y))
+  /-- A functor preserves identity morphisms. -/
+  map_id : ∀ X : C, map (𝟙 X) = 𝟙 (obj X) := by cat_tac
+  /-- A functor preserves composition. -/
+  map_comp : ∀ {X Y Z : C} (f : X ⟶ Y) (g : Y ⟶ Z), map (f ≫ g) = (map f) ≫ (map g) := by cat_tac
+
+attribute [simp] Functor.map_id Functor.map_comp
+
+attribute [grind =] Functor.map_id
+attribute [grind _=_] Functor.map_comp
+
+variable {C : Type u₁} [Category.{v₁} C] {D : Type u₂} [Category.{v₂} D] {E : Type u₃} [Category.{v₃} E]
+variable {F G H : Functor C D}
+
+namespace Functor
+
+def comp (F : Functor C D) (G : Functor D E) : Functor C E where
+  obj X := G.obj (F.obj X)
+  map f := G.map (F.map f)
+  -- Note `map_id` and `map_comp` are handled by `cat_tac`.
+
+variable {X Y : C} {G : Functor D E}
+
+@[simp, grind =] theorem comp_obj : (F.comp G).obj X = G.obj (F.obj X) := rfl
+@[simp, grind =] theorem comp_map (f : X ⟶ Y) : (F.comp G).map f = G.map (F.map f) := rfl
+
+end Functor
+
+@[ext]
+structure NatTrans [Category.{v₁, u₁} C] [Category.{v₂, u₂} D] (F G : Functor C D) : Type max u₁ v₂ where
+  /-- The component of a natural transformation. -/
+  app : ∀ X : C, F.obj X ⟶ G.obj X
+  /-- The naturality square for a given morphism. -/
+  naturality : ∀ ⦃X Y : C⦄ (f : X ⟶ Y), F.map f ≫ app Y = app X ≫ G.map f := by cat_tac
+
+attribute [simp, grind =] NatTrans.naturality
+
+namespace NatTrans
+
+variable {X : C}
+
+protected def id (F : Functor C D) : NatTrans F F where app X := 𝟙 (F.obj X)
+
+@[simp, grind =] theorem id_app : (NatTrans.id F).app X = 𝟙 (F.obj X) := rfl
+
+protected def vcomp (α : NatTrans F G) (β : NatTrans G H) : NatTrans F H where
+  app X := α.app X ≫ β.app X
+  -- `naturality` is now handled by `grind`; in Mathlib this relies on `@[reassoc]` attributes.
+  -- Manual proof:
+  -- rw [← Category.assoc]
+  -- rw [α.naturality f]
+  -- rw [Category.assoc]
+  -- rw [β.naturality f]
+  -- rw [← Category.assoc]
+
+@[simp, grind =] theorem vcomp_app (α : NatTrans F G) (β : NatTrans G H) (X : C) :
+    (α.vcomp β).app X = α.app X ≫ β.app X := rfl
+
+end NatTrans
+
+instance Functor.category : Category.{max u₁ v₂} (Functor C D) where
+  Hom F G := NatTrans F G
+  id F := NatTrans.id F
+  comp α β := NatTrans.vcomp α β
+  -- Here we're okay: all the proofs are handled by `cat_tac`.
+
+@[simp, grind =]
+theorem id_app (F : Functor C D) (X : C) : (𝟙 F : F ⟶ F).app X = 𝟙 (F.obj X) := rfl
+
+@[simp, grind _=_]
+theorem comp_app {F G H : Functor C D} (α : F ⟶ G) (β : G ⟶ H) (X : C) :
+    (α ≫ β).app X = α.app X ≫ β.app X := rfl
+
+theorem app_naturality {F G : Functor C (Functor D E)} (T : F ⟶ G) (X : C) {Y Z : D} (f : Y ⟶ Z) :
+    (F.obj X).map f ≫ (T.app X).app Z = (T.app X).app Y ≫ (G.obj X).map f := by
+  cat_tac
+
+theorem naturality_app {F G : Functor C (Functor D E)} (T : F ⟶ G) (Z : D) {X Y : C} (f : X ⟶ Y) :
+    (F.map f).app Z ≫ (T.app Y).app Z = (T.app X).app Z ≫ (G.map f).app Z := by
+  cat_tac -- this is done manually in Mathlib!
+  -- rw [← comp_app]
+  -- rw [T.naturality f]
+  -- rw [comp_app]
+
+open Category Functor NatTrans
+
+def hcomp {H I : Functor D E} (α : F ⟶ G) (β : H ⟶ I) : F.comp H ⟶ G.comp I where
+  app := fun X : C => β.app (F.obj X) ≫ I.map (α.app X)
+  -- `grind` can now handle `naturality`, while Mathlib does this manually:
+  -- rw [Functor.comp_map, Functor.comp_map, ← assoc, naturality, assoc, ← I.map_comp, naturality,
+  --   map_comp, assoc]
+
+end CategoryTheory
diff --git a/tests/lean/run/grind_ematch1.lean b/tests/lean/run/grind_ematch1.lean
index 7f8a9c053b..b2a79778f8 100644
--- a/tests/lean/run/grind_ematch1.lean
+++ b/tests/lean/run/grind_ematch1.lean
@@ -1,5 +1,6 @@
 set_option trace.grind.ematch.pattern true
-grind_pattern Array.size_set => Array.set a i v h
+
+attribute [grind =] Array.size_set
 
 set_option grind.debug true
 
@@ -21,12 +22,10 @@ example (as bs : Array α) (v : α)
 
 set_option trace.grind.ematch.instance true
 
-grind_pattern Array.get_set_eq  => a.set i v h
-grind_pattern Array.get_set_ne => (a.set i v hi)[j]
+attribute [grind =] Array.get_set_ne
 
 /--
-info: [grind.ematch.instance] Array.get_set_eq: (as.set i v ⋯)[i] = v
-[grind.ematch.instance] Array.size_set: (as.set i v ⋯).size = as.size
+info: [grind.ematch.instance] Array.size_set: (as.set i v ⋯).size = as.size
 [grind.ematch.instance] Array.get_set_ne: ∀ (hj : j < as.size), i ≠ j → (as.set i v ⋯)[j] = as[j]
 -/
 #guard_msgs (info) in
diff --git a/tests/lean/run/grind_ematch2.lean b/tests/lean/run/grind_ematch2.lean
index 62a4ba2161..3341d3266a 100644
--- a/tests/lean/run/grind_ematch2.lean
+++ b/tests/lean/run/grind_ematch2.lean
@@ -1,6 +1,4 @@
-grind_pattern Array.size_set => Array.set a i v h
-grind_pattern Array.get_set_eq  => a.set i v h
-grind_pattern Array.get_set_ne => (a.set i v hi)[j]
+attribute [grind =] Array.size_set Array.get_set_eq Array.get_set_ne
 
 set_option grind.debug true
 set_option trace.grind.ematch.pattern true
@@ -57,8 +55,7 @@ example (as bs cs ds : Array α) (v₁ v₂ v₃ : α)
   grind
 
 opaque f (a b : α) : α := a
-theorem fx : f x (f x x) = x := sorry
-grind_pattern fx => f x (f x x)
+@[grind =] theorem fx : f x (f x x) = x := sorry
 
 /--
 info: [grind.ematch.instance] fx: f a (f a a) = a
diff --git a/tests/lean/run/grind_shelf.lean b/tests/lean/run/grind_shelf.lean
new file mode 100644
index 0000000000..8c5f633176
--- /dev/null
+++ b/tests/lean/run/grind_shelf.lean
@@ -0,0 +1,30 @@
+class One (α : Type u) where
+  one : α
+
+instance (priority := 300) One.toOfNat1 {α} [One α] : OfNat α (nat_lit 1) where
+  ofNat := ‹One α›.1
+
+class Shelf (α : Type u) where
+  act : α → α → α
+  self_distrib : ∀ {x y z : α}, act x (act y z) = act (act x y) (act x z)
+
+class UnitalShelf (α : Type u) extends Shelf α, One α where
+  one_act : ∀ a : α, act 1 a = a
+  act_one : ∀ a : α, act a 1 = a
+
+infixr:65 " ◃ " => Shelf.act
+
+-- Mathlib proof from UnitalShelf.act_act_self_eq
+example {S} [UnitalShelf S] (x y : S) : (x ◃ y) ◃ x = x ◃ y := by
+  have h : (x ◃ y) ◃ x = (x ◃ y) ◃ (x ◃ 1) := by rw [UnitalShelf.act_one]
+  rw [h, ← Shelf.self_distrib, UnitalShelf.act_one]
+
+attribute [grind =] UnitalShelf.one_act UnitalShelf.act_one
+
+-- We actually want the reverse direction of `Shelf.self_distrib`, so don't use the `grind_eq` attribute.
+grind_pattern Shelf.self_distrib => self.act (self.act x y) (self.act x z)
+
+-- Proof using `grind`:
+example {S} [UnitalShelf S] (x y : S) : (x ◃ y) ◃ x = x ◃ y := by
+  have h : (x ◃ y) ◃ x = (x ◃ y) ◃ (x ◃ 1) := by grind
+  grind