feat: save information for generating structural equation theorems later

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

@@ -3,16 +3,53 @@ Copyright (c) 2021 Microsoft Corporation. All rights reserved.
 Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Leonardo de Moura
 -/
+import Lean.Meta.Eqns
+import Lean.Meta.Tactic.Split
 import Lean.Elab.PreDefinition.Basic
 import Lean.Elab.PreDefinition.Structural.Basic
 
 namespace Lean.Elab.Structural
+open Meta
 
-def mkEqns (preDef : PreDefinition) : TermElabM Unit := do
-  trace[Elab.definition.structural.eqns] "mkEqns:\n{preDef.value}"
-  return ()
+structure EqnInfo where
+  declName    : Name
+  levelParams : List Name
+  type        : Expr
+  value       : Expr
+  recArgPos   : Nat
+  deriving Inhabited
+
+def mkEqns (info : EqnInfo) : MetaM (Array Name) := do
+  -- TODO
+  trace[Elab.definition.structural.eqns] "mkEqns:\n{info.value}"
+  return #[]
+
+builtin_initialize eqnInfoExt : MapDeclarationExtension EqnInfo ← mkMapDeclarationExtension `structEqInfo
+
+def registerEqnsInfo (preDef : PreDefinition) (recArgPos : Nat) : CoreM Unit := do
+  modifyEnv fun env => eqnInfoExt.insert env preDef.declName { preDef with recArgPos }
+
+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 {})
+
+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
+  else
+    return none
 
 builtin_initialize
+  registerGetEqnsFn getEqnsFor?
   registerTraceClass `Elab.definition.structural.eqns
 
 end Lean.Elab.Structural

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

@@ -56,7 +56,7 @@ private def getFixedPrefix (declName : Name) (xs : Array Expr) (value : Expr) :
       return true
   numFixedRef.get
 
-private def elimRecursion (preDef : PreDefinition) : M PreDefinition :=
+private def elimRecursion (preDef : PreDefinition) : M (Nat × PreDefinition) :=
   withoutModifyingEnv do lambdaTelescope preDef.value fun xs value => do
     addAsAxiom preDef
     let value ← preprocess value preDef.declName
@@ -72,18 +72,19 @@ private def elimRecursion (preDef : PreDefinition) : M PreDefinition :=
       trace[Elab.definition.structural] "result: {valueNew}"
       -- Recursive applications may still occur in expressions that were not visited by replaceRecApps (e.g., in types)
       let valueNew ← ensureNoRecFn preDef.declName valueNew
-      pure { preDef with value := valueNew }
+      let recArgPos := recArgInfo.fixedParams.size + recArgInfo.pos
+      return (recArgPos, { preDef with value := valueNew })
 
 def structuralRecursion (preDefs : Array PreDefinition) : TermElabM Unit :=
   if preDefs.size != 1 then
     throwError "structural recursion does not handle mutually recursive functions"
   else do
-    let (preDefNonRec, state) ← run <| elimRecursion preDefs[0]
+    let ((recArgPos, preDefNonRec), state) ← run <| elimRecursion preDefs[0]
     state.addMatchers.forM liftM
     mapError (addNonRec preDefNonRec) (fun msg => m!"structural recursion failed, produced type incorrect term{indentD msg}")
     addAndCompilePartialRec preDefs
     addSmartUnfoldingDef preDefs[0] state
-    mkEqns preDefs[0]
+    registerEqnsInfo preDefs[0] recArgPos
 
 builtin_initialize
   registerTraceClass `Elab.definition.structural