From fe783cb778540bebb1974025e718a8038164367f Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leomoura@amazon.com>
Date: Thu, 28 Mar 2024 17:03:55 -0700
Subject: [PATCH] feat: save whether a function has been defined by structural
 or well-founded recursion

---
 .../Elab/PreDefinition/Structural/Main.lean   |  1 +
 src/Lean/Elab/PreDefinition/WF/Main.lean      |  1 +
 src/Lean/Meta/Eqns.lean                       | 19 +++++++++++++++++++
 3 files changed, 21 insertions(+)

diff --git a/src/Lean/Elab/PreDefinition/Structural/Main.lean b/src/Lean/Elab/PreDefinition/Structural/Main.lean
index ab7e3af5a7..52be907bd2 100644
--- a/src/Lean/Elab/PreDefinition/Structural/Main.lean
+++ b/src/Lean/Elab/PreDefinition/Structural/Main.lean
@@ -101,6 +101,7 @@ def structuralRecursion (preDefs : Array PreDefinition) : TermElabM Unit :=
         -/
         registerEqnsInfo preDef recArgPos
     addSmartUnfoldingDef preDef recArgPos
+    markAsRecursive preDef.declName
     applyAttributesOf #[preDefNonRec] AttributeApplicationTime.afterCompilation
 
 builtin_initialize
diff --git a/src/Lean/Elab/PreDefinition/WF/Main.lean b/src/Lean/Elab/PreDefinition/WF/Main.lean
index a1cb8e3bf2..f05f560bb7 100644
--- a/src/Lean/Elab/PreDefinition/WF/Main.lean
+++ b/src/Lean/Elab/PreDefinition/WF/Main.lean
@@ -144,6 +144,7 @@ def wfRecursion (preDefs : Array PreDefinition) : TermElabM Unit := do
   let preDefs ← preDefs.mapM (abstractNestedProofs ·)
   registerEqnsInfo preDefs preDefNonRec.declName fixedPrefixSize argsPacker
   for preDef in preDefs do
+    markAsRecursive preDef.declName
     applyAttributesOf #[preDef] AttributeApplicationTime.afterCompilation
 
 builtin_initialize registerTraceClass `Elab.definition.wf
diff --git a/src/Lean/Meta/Eqns.lean b/src/Lean/Meta/Eqns.lean
index ba9b60d8f8..7f65d68971 100644
--- a/src/Lean/Meta/Eqns.lean
+++ b/src/Lean/Meta/Eqns.lean
@@ -9,6 +9,25 @@ import Lean.Meta.Basic
 import Lean.Meta.AppBuilder
 
 namespace Lean.Meta
+/--
+Environment extension for storing which declarations are recursive.
+This information is populated by the `PreDefinition` module, but the simplifier
+uses when unfolding declarations.
+-/
+builtin_initialize recExt : TagDeclarationExtension ← mkTagDeclarationExtension `recExt
+
+/--
+Marks the given declaration as recursive.
+-/
+def markAsRecursive (declName : Name) : CoreM Unit :=
+  modifyEnv (recExt.tag · declName)
+
+/--
+Returns `true` if `declName` was defined using well-founded recursion, or structural recursion.
+-/
+def isRecursiveDefinition (declName : Name) : CoreM Bool :=
+  return recExt.isTagged (← getEnv) declName
+
 def eqnThmSuffixBase := "eq"
 def eqnThmSuffixBasePrefix := eqnThmSuffixBase ++ "_"
 def eqn1ThmSuffix := eqnThmSuffixBasePrefix ++ "1"