feat: add #discr_tree_key command and discr_tree_key tactic (#4447)

Adds a command and tactic to print the `Array <| DiscrTree.Key` for
equalities helping the user to debug perceived `simp` failures.

---------

Co-authored-by: Joachim Breitner <mail@joachim-breitner.de>

src/Init/Notation.lean

@@ -703,6 +703,28 @@ syntax (name := checkSimp) "#check_simp " term "~>" term : command
 -/
 syntax (name := checkSimpFailure) "#check_simp " term "!~>" : command
 
+/--
+`#discr_tree_key  t` prints the discrimination tree keys for a term `t` (or, if it is a single identifier, the type of that constant).
+It uses the default configuration for generating keys.
+
+For example,
+```
+#discr_tree_key (∀ {a n : Nat}, bar a (OfNat.ofNat n))
+-- bar _ (@OfNat.ofNat Nat _ _)
+
+#discr_tree_simp_key Nat.add_assoc
+-- @HAdd.hAdd Nat Nat Nat _ (@HAdd.hAdd Nat Nat Nat _ _ _) _
+```
+
+`#discr_tree_simp_key` is similar to `#discr_tree_key`, but treats the underlying type
+as one of a simp lemma, i.e. transforms it into an equality and produces the key of the
+left-hand side.
+-/
+syntax (name := discrTreeKeyCmd) "#discr_tree_key " term : command
+
+@[inherit_doc discrTreeKeyCmd]
+syntax (name := discrTreeSimpKeyCmd) "#discr_tree_simp_key" term : command
+
 /--
 The `seal foo` command ensures that the definition of `foo` is sealed, meaning it is marked as `[irreducible]`.
 This command is particularly useful in contexts where you want to prevent the reduction of `foo` in proofs.

src/Lean/Elab/Tactic.lean

@@ -40,3 +40,4 @@ import Lean.Elab.Tactic.LibrarySearch
 import Lean.Elab.Tactic.ShowTerm
 import Lean.Elab.Tactic.Rfl
 import Lean.Elab.Tactic.Rewrites
+import Lean.Elab.Tactic.DiscrTreeKey

src/Lean/Elab/Tactic/DiscrTreeKey.lean

@@ -0,0 +1,65 @@
+/-
+Copyright (c) 2024 Matthew Robert Ballard. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Tomas Skrivan, Matthew Robert Ballard
+-/
+prelude
+import Init.Tactics
+import Lean.Elab.Command
+import Lean.Meta.Tactic.Simp.SimpTheorems
+
+namespace Lean.Elab.Tactic.DiscrTreeKey
+
+open Lean.Meta DiscrTree
+open Lean.Elab.Tactic
+open Lean.Elab.Command
+
+private def mkKey (e : Expr) (simp : Bool) : MetaM (Array Key) := do
+  let (_, _, type) ← withReducible <| forallMetaTelescopeReducing e
+  let type ← whnfR type
+  if simp then
+    if let some (_, lhs, _) := type.eq? then
+      mkPath lhs simpDtConfig
+    else if let some (lhs, _) := type.iff? then
+      mkPath lhs simpDtConfig
+    else if let some (_, lhs, _) := type.ne? then
+      mkPath lhs simpDtConfig
+    else if let some p := type.not? then
+      match p.eq? with
+      | some (_, lhs, _) =>
+        mkPath lhs simpDtConfig
+      | _ => mkPath p simpDtConfig
+    else
+      mkPath type simpDtConfig
+  else
+    mkPath type {}
+
+private def getType (t : TSyntax `term) : TermElabM Expr := do
+  if let `($id:ident) := t then
+    if let some ldecl := (← getLCtx).findFromUserName? id.getId then
+      return ldecl.type
+    else
+      let info ← getConstInfo (← realizeGlobalConstNoOverloadWithInfo id)
+      return info.type
+  else
+    Term.elabTerm t none
+
+@[builtin_command_elab Lean.Parser.discrTreeKeyCmd]
+def evalDiscrTreeKeyCmd : CommandElab := fun stx => do
+  Command.liftTermElabM <| do
+    match stx with
+    | `(command| #discr_tree_key $t:term) => do
+      let type ← getType t
+      logInfo (← keysAsPattern <| ← mkKey type false)
+    | _                        => Elab.throwUnsupportedSyntax
+
+@[builtin_command_elab Lean.Parser.discrTreeSimpKeyCmd]
+def evalDiscrTreeSimpKeyCmd : CommandElab := fun stx => do
+  Command.liftTermElabM <| do
+    match stx with
+    | `(command| #discr_tree_simp_key $t:term) => do
+      let type ← getType t
+      logInfo (← keysAsPattern <| ← mkKey type true)
+    | _                        => Elab.throwUnsupportedSyntax
+
+end Lean.Elab.Tactic.DiscrTreeKey

tests/lean/run/discrTreeKey.lean

@@ -0,0 +1,79 @@
+import Init.Data.Nat.Basic
+import Init.Data.List.Lemmas
+
+/-!
+  This file provides examples of use of the commands #discr_tree_key and #discr_tree_simp_key
+  and guards against any breakage of the commands.
+-/
+
+universe u
+
+open Nat List
+
+/-!
+  We can produce `simp` keys for theorems of the form `=`, `≠`, and `↔` by supplying the name
+  of the declaration.
+-/
+
+#check zero_le
+#discr_tree_simp_key zero_le
+
+#check succ_eq_add_one
+#discr_tree_simp_key succ_eq_add_one
+
+#check Nat.mul_one
+#discr_tree_simp_key Nat.mul_one
+
+#check Nat.not_le
+#discr_tree_simp_key Nat.not_le
+
+#check Nat.pred_succ
+#discr_tree_simp_key Nat.pred_succ
+
+#check get?_nil
+#discr_tree_simp_key get?_nil
+
+#check or_cons
+#discr_tree_simp_key or_cons
+
+#check not_mem_nil
+#discr_tree_simp_key not_mem_nil
+
+#check mem_cons
+#discr_tree_simp_key mem_cons
+
+#check singleton_append
+#discr_tree_simp_key singleton_append
+
+#check append_nil
+#discr_tree_simp_key append_eq_nil
+
+#check mapM_nil
+#discr_tree_simp_key mapM_nil
+
+/-!
+  We can produce keys for a general declarations by name using the default configuration
+  for generating keys.
+-/
+
+#check Nat.instIdempotentOpGcd
+#discr_tree_key Nat.instIdempotentOpGcd
+
+#check List.instDecidableMemOfLawfulBEq
+#discr_tree_key List.instDecidableMemOfLawfulBEq
+
+#check List.instForIn'InferInstanceMembership
+#discr_tree_key List.instForIn'InferInstanceMembership
+
+/-!
+  We can also specify a term directly.
+-/
+def bar (_ _ : Nat) : Nat := default
+
+#discr_tree_key (∀ {a n : Nat}, bar a (OfNat.ofNat n) = default)
+
+#discr_tree_simp_key (∀ {a n : Nat}, bar a (no_index (OfNat.ofNat n)) = default)
+
+#discr_tree_simp_key (∀ m : Nat, ∃ n : Nat, m ≠ n)
+
+#discr_tree_simp_key (∀ m : Nat, m > 0 → m ≠ 0)