refactor: move equation theorem cache to Meta/Eqns.lean

src/Lean/Elab/PreDefinition/Eqns.lean

@@ -218,14 +218,6 @@ partial def removeUnusedEqnHypotheses (type value : Expr) : MetaM (Expr × Expr)
         return (type, value)
   | _ => return (type, value)
 
-structure EqnsExtState where
-  map : Std.PHashMap Name (Array Name) := {}
-  deriving Inhabited
-
-/- We generate the equations on demand, and do not save them on .olean files. -/
-builtin_initialize eqnsExt : EnvExtension EqnsExtState ←
-  registerEnvExtension (pure {})
-
 /-- Try to close goal using `rfl` with smart unfolding turned off. -/
 def tryURefl (mvarId : MVarId) : MetaM Bool :=
   withOptions (smartUnfolding.set · false) do

src/Lean/Elab/PreDefinition/Structural/Eqns.lean

@@ -83,13 +83,8 @@ def registerEqnsInfo (preDef : PreDefinition) (recArgPos : Nat) : CoreM Unit :=
   modifyEnv fun env => eqnInfoExt.insert env preDef.declName { preDef with recArgPos }
 
 def getEqnsFor? (declName : Name) : MetaM (Option (Array Name)) := do
-  let env ← getEnv
-  if let some eqs := eqnsExt.getState env |>.map.find? declName then
-    return some eqs
-  else if let some info := eqnInfoExt.find? env declName then
-    let eqs ← mkEqns info
-    modifyEnv fun env => eqnsExt.modifyState env fun s => { s with map := s.map.insert declName eqs }
-    return some eqs
+  if let some info := eqnInfoExt.find? (← getEnv) declName then
+    mkEqns info
   else
     return none
 

src/Lean/Elab/PreDefinition/WF/Eqns.lean

@@ -96,13 +96,8 @@ def registerEqnsInfo (preDefs : Array PreDefinition) (declNameNonRec : Name) : C
       eqnInfoExt.insert env preDef.declName { preDef with declNames, declNameNonRec }
 
 def getEqnsFor? (declName : Name) : MetaM (Option (Array Name)) := do
-  let env ← getEnv
-  if let some eqs := eqnsExt.getState env |>.map.find? declName then
-    return some eqs
-  else if let some info := eqnInfoExt.find? env declName then
-    let eqs ← mkEqns declName info
-    modifyEnv fun env => eqnsExt.modifyState env fun s => { s with map := s.map.insert declName eqs }
-    return some eqs
+  if let some info := eqnInfoExt.find? (← getEnv) declName then
+    mkEqns declName info
   else
     return none
 

src/Lean/Meta/Eqns.lean

@@ -4,7 +4,7 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Leonardo de Moura
 -/
 import Lean.Meta.Basic
-import Lean.Meta.InferType
+import Lean.Meta.AppBuilder
 
 namespace Lean.Meta
 
@@ -42,16 +42,31 @@ def registerGetEqnsFn (f : GetEqnsFn) : IO Unit := do
     throw (IO.userError "failed to register equation getter, this kind of extension can only be registered during initialization")
   getEqnsFnsRef.modify (f :: ·)
 
+/-- Return true iff `declName` is a definition and its type is not a proposition. -/
 private def shouldGenerateEqnThms (declName : Name) : MetaM Bool := do
   if let some (.defnInfo info) := (← getEnv).find? declName then
     return !(← isProp info.type)
   else
     return false
 
+structure EqnsExtState where
+  map : Std.PHashMap Name (Array Name) := {}
+  deriving Inhabited
+
+/- We generate the equations on demand, and do not save them on .olean files. -/
+builtin_initialize eqnsExt : EnvExtension EqnsExtState ←
+  registerEnvExtension (pure {})
+
+/--
+  Return equation theorems for the given declaration.
+-/
 def getEqnsFor? (declName : Name) : MetaM (Option (Array Name)) := do
-  if (← shouldGenerateEqnThms declName) then
+  if let some eqs := eqnsExt.getState (← getEnv) |>.map.find? declName then
+    return some eqs
+  else if (← shouldGenerateEqnThms declName) then
     for f in (← getEqnsFnsRef.get) do
       if let some r ← f declName then
+        modifyEnv fun env => eqnsExt.modifyState env fun s => { s with map := s.map.insert declName r }
         return some r
   return none