From bd02f16b43df50c76f921ccaa6737aabacfb771e Mon Sep 17 00:00:00 2001
From: Christian Pehle <christian.pehle@gmail.com>
Date: Wed, 8 Sep 2021 22:03:50 +0200
Subject: [PATCH] feat: optimized ```deriving BEq``` for enumeration types

---
 src/Lean/Elab/Deriving/BEq.lean | 16 +++++++++++++++-
 tests/lean/run/654.lean         |  3 +++
 2 files changed, 18 insertions(+), 1 deletion(-)

diff --git a/src/Lean/Elab/Deriving/BEq.lean b/src/Lean/Elab/Deriving/BEq.lean
index 7357f952f7cf..ac3caf9193be 100644
--- a/src/Lean/Elab/Deriving/BEq.lean
+++ b/src/Lean/Elab/Deriving/BEq.lean
@@ -97,10 +97,24 @@ private def mkBEqInstanceCmds (declNames : Array Name) : TermElabM (Array Syntax
   trace[Elab.Deriving.beq] "\n{cmds}"
   return cmds
 
+private def mkBEqEnumFun (ctx : Context) (name : Name) : TermElabM Syntax := do
+  let auxFunName ← ctx.auxFunNames[0]
+  `(private def $(mkIdent auxFunName):ident  (x y : $(mkIdent name)) : Bool := x.toCtorIdx == y.toCtorIdx)
+
+private def mkBEqEnumCmd (name : Name): TermElabM (Array Syntax) := do
+  let ctx ← mkContext "beq" name
+  let cmds := #[← mkBEqEnumFun ctx name] ++ (← mkInstanceCmds ctx `BEq #[name])
+  trace[Elab.Deriving.beq] "\n{cmds}"
+  return cmds
+
 open Command
 
 def mkBEqInstanceHandler (declNames : Array Name) : CommandElabM Bool := do
-  if (← declNames.allM isInductive) && declNames.size > 0 then
+  if declNames.size == 1 && (← isEnumType declNames[0]) then
+    let cmds ← liftTermElabM none <| mkBEqEnumCmd declNames[0]
+    cmds.forM elabCommand
+    return true
+  else if (← declNames.allM isInductive) && declNames.size > 0 then
     let cmds ← liftTermElabM none <| mkBEqInstanceCmds declNames
     cmds.forM elabCommand
     return true
diff --git a/tests/lean/run/654.lean b/tests/lean/run/654.lean
index 78385ab443d6..2dcfa256754e 100644
--- a/tests/lean/run/654.lean
+++ b/tests/lean/run/654.lean
@@ -278,8 +278,11 @@ inductive CXCursorKind where
   | CXCursor_FirstExtraDecl
   | CXCursor_LastExtraDecl
   | CXCursor_OverloadCandidate
+  deriving BEq
 
 open CXCursorKind
 
 example (h : CXCursor_CUDAGlobalAttr = CXCursor_CUDAHostAttr) : False := by
   contradiction
+
+#eval CXCursor_CUDAGlobalAttr == CXCursor_CUDAHostAttr