From f54c192e6fc137c11b7535249f60e5aeeb7b9898 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Henrik=20B=C3=B6ving?= <hargonix@gmail.com>
Date: Wed, 2 Feb 2022 11:22:15 +0100
Subject: [PATCH] feat: equations

Equation implementation for definitions, right now lots of definitions
simply dont generate equational lemmata at all so lots are left
without them.
---
 DocGen4/Output/Definition.lean | 33 +++++++++++++++++++++++
 DocGen4/Output/Module.lean     |  2 ++
 DocGen4/Process.lean           | 49 +++++++++++++++++++++++++---------
 lean-toolchain                 |  2 +-
 4 files changed, 73 insertions(+), 13 deletions(-)
 create mode 100644 DocGen4/Output/Definition.lean

diff --git a/DocGen4/Output/Definition.lean b/DocGen4/Output/Definition.lean
new file mode 100644
index 0000000..f92627c
--- /dev/null
+++ b/DocGen4/Output/Definition.lean
@@ -0,0 +1,33 @@
+import DocGen4.Output.Template
+
+namespace DocGen4
+namespace Output
+
+open scoped DocGen4.Jsx
+open Lean Widget
+
+def equationToHtml (c : CodeWithInfos) : HtmlM Html := do
+  <li «class»="equation">[←infoFormatToHtml c]</li>
+
+def equationsToHtml (i : DefinitionInfo) : HtmlM (Option Html) := do
+  if let some eqs ← i.equations then
+    let equationsHtml ← eqs.mapM equationToHtml
+    return <details>
+      <summary>Equations</summary>
+      <ul «class»="equations">
+        [equationsHtml]
+      </ul>
+    </details>
+  else
+    return none
+
+def definitionToHtml (i : DefinitionInfo) : HtmlM (Array Html) := do
+  let equationsHtml ← equationsToHtml i
+  if let some equationsHtml ← equationsHtml then
+    #[equationsHtml]
+  else
+    #[]
+
+end Output
+end DocGen4
+
diff --git a/DocGen4/Output/Module.lean b/DocGen4/Output/Module.lean
index 914e26e..1a37ae9 100644
--- a/DocGen4/Output/Module.lean
+++ b/DocGen4/Output/Module.lean
@@ -7,6 +7,7 @@ import DocGen4.Output.Template
 import DocGen4.Output.Inductive
 import DocGen4.Output.Structure
 import DocGen4.Output.Class
+import DocGen4.Output.Definition
 
 namespace DocGen4
 namespace Output
@@ -75,6 +76,7 @@ def docInfoToHtml (module : Name) (doc : DocInfo) : HtmlM Html := do
   | DocInfo.inductiveInfo i => inductiveToHtml i
   | DocInfo.structureInfo i => structureToHtml i
   | DocInfo.classInfo i => classToHtml i
+  | DocInfo.definitionInfo i => definitionToHtml i
   | _ => #[]
 
   return <div «class»="decl" id={doc.getName.toString}>
diff --git a/DocGen4/Process.lean b/DocGen4/Process.lean
index bc9eebd..4fd10a8 100644
--- a/DocGen4/Process.lean
+++ b/DocGen4/Process.lean
@@ -47,6 +47,7 @@ structure DefinitionInfo extends Info where
   --value : CodeWithInfos
   unsafeInformation : DefinitionSafety
   hints : ReducibilityHints
+  equations : Option (Array CodeWithInfos)
   deriving Inhabited
 
 abbrev InstanceInfo := DefinitionInfo
@@ -160,11 +161,33 @@ def OpaqueInfo.ofOpaqueVal (v : OpaqueVal) : MetaM OpaqueInfo := do
 def isInstance (declName : Name) : MetaM Bool := do
   (instanceExtension.getState (←getEnv)).instanceNames.contains declName
 
+
+partial def stripArgs (e : Expr) : Expr :=
+  match e.consumeMData with
+  | Expr.lam name type body data =>
+    let name := name.eraseMacroScopes
+    stripArgs (Expr.instantiate1 body (mkFVar ⟨name⟩))
+  | Expr.forallE name type body data =>
+    let name := name.eraseMacroScopes
+    stripArgs (Expr.instantiate1 body (mkFVar ⟨name⟩))
+  | _ => e
+
+def processEq (eq : Name) : MetaM CodeWithInfos := do
+  let type ← (mkConstWithFreshMVarLevels eq >>= inferType)
+  let final := stripArgs type
+  prettyPrintTerm final
+
 def DefinitionInfo.ofDefinitionVal (v : DefinitionVal) : MetaM DefinitionInfo := do
   let info ← Info.ofConstantVal v.toConstantVal
-  -- Elaborating the value yields weird exceptions
-  --let value ← prettyPrintTerm v.value
-  return DefinitionInfo.mk info v.safety v.hints
+  let mut eqs? := none
+  try
+    eqs? ← getEqnsFor? v.name
+  catch err =>
+    IO.println s!"WARNING: Failed to calculate equational lemmata for {v.name}: {←err.toMessageData.toString}"
+
+  let prettyEqs ← eqs?.mapM (λ eqs => eqs.mapM processEq)
+  return DefinitionInfo.mk info v.safety v.hints prettyEqs
+
 
 def getConstructorType (ctor : Name) : MetaM CodeWithInfos := do
   let env ← getEnv
@@ -254,7 +277,6 @@ def ofConstant : (Name × ConstantInfo) → MetaM (Option DocInfo) := λ (name,
   | ConstantInfo.axiomInfo i => some $ axiomInfo (←AxiomInfo.ofAxiomVal i)
   | ConstantInfo.thmInfo i => some $ theoremInfo (←TheoremInfo.ofTheoremVal i)
   | ConstantInfo.opaqueInfo i => some $ opaqueInfo (←OpaqueInfo.ofOpaqueVal i)
-  -- TODO: Find a way to extract equations nicely
   | ConstantInfo.defnInfo i =>
     if ← (isProjFn i.name) then
       none
@@ -342,14 +364,15 @@ def process : MetaM AnalyzerResult := do
     res := res.insert module (Module.mk module moduleDoc #[])
 
   for cinfo in env.constants.toList do
-    let d := ←DocInfo.ofConstant cinfo
-    match d with
-    | some dinfo =>
-      let some modidx ← env.getModuleIdxFor? cinfo.fst | unreachable!
-      let moduleName := env.allImportedModuleNames.get! modidx
-      let module := res.find! moduleName
-      res := res.insert moduleName {module with members := module.members.push dinfo}
-    | none => ()
+    try
+      let analysis := Prod.fst <$> Meta.MetaM.toIO (DocInfo.ofConstant cinfo) { maxHeartbeats := 5000000, options := ⟨[(`pp.tagAppFns, true)]⟩ } { env := env} {} {}
+      if let some dinfo ← analysis then
+        let some modidx ← env.getModuleIdxFor? dinfo.getName | unreachable!
+        let moduleName := env.allImportedModuleNames.get! modidx
+        let module := res.find! moduleName
+        res := res.insert moduleName {module with members := module.members.push dinfo}
+    catch e =>
+      IO.println s!"WARNING: Failed to obtain information for: {cinfo.fst}: {←e.toMessageData.toString}"
 
   -- TODO: This is definitely not the most efficient way to store this data
   let mut adj := Array.mkArray res.size (Array.mkArray res.size false)
@@ -370,4 +393,6 @@ def process : MetaM AnalyzerResult := do
     importAdj := adj
   }
 
+
+
 end DocGen4
diff --git a/lean-toolchain b/lean-toolchain
index 19c9eec..0085298 100644
--- a/lean-toolchain
+++ b/lean-toolchain
@@ -1 +1 @@
-leanprover/lean4:nightly-2022-01-16
+leanprover/lean4:nightly-2022-02-04