From 7af09ac00964fb7a77f7096c6a6acaee9bc467c3 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Sat, 12 Mar 2022 15:34:32 -0800
Subject: [PATCH] refactor: move equation theorem cache to `Meta/Eqns.lean`

---
 src/Lean/Elab/PreDefinition/Eqns.lean         |  8 --------
 .../Elab/PreDefinition/Structural/Eqns.lean   |  9 ++-------
 src/Lean/Elab/PreDefinition/WF/Eqns.lean      |  9 ++-------
 src/Lean/Meta/Eqns.lean                       | 19 +++++++++++++++++--
 4 files changed, 21 insertions(+), 24 deletions(-)

diff --git a/src/Lean/Elab/PreDefinition/Eqns.lean b/src/Lean/Elab/PreDefinition/Eqns.lean
index b8d29aeba4..baf366a16c 100644
--- a/src/Lean/Elab/PreDefinition/Eqns.lean
+++ b/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
diff --git a/src/Lean/Elab/PreDefinition/Structural/Eqns.lean b/src/Lean/Elab/PreDefinition/Structural/Eqns.lean
index 1eddf5e490..e5603b6d22 100644
--- a/src/Lean/Elab/PreDefinition/Structural/Eqns.lean
+++ b/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
 
diff --git a/src/Lean/Elab/PreDefinition/WF/Eqns.lean b/src/Lean/Elab/PreDefinition/WF/Eqns.lean
index eb8c7cda52..4eef8492cf 100644
--- a/src/Lean/Elab/PreDefinition/WF/Eqns.lean
+++ b/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
 
diff --git a/src/Lean/Meta/Eqns.lean b/src/Lean/Meta/Eqns.lean
index 297d69f5df..1e4b52b8d2 100644
--- a/src/Lean/Meta/Eqns.lean
+++ b/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