From a17c98e92241f196e4b5d56e69505139fd92ea07 Mon Sep 17 00:00:00 2001
From: jstoobysmith <72603918+jstoobysmith@users.noreply.github.com>
Date: Mon, 15 Jul 2024 09:35:36 -0400
Subject: [PATCH] feat: Add stats generator

---
 .../StandardModel/HiggsBoson/GaugeAction.lean |  1 +
 lakefile.toml                                 |  4 +
 scripts/stats.lean                            | 87 +++++++++++++++++++
 3 files changed, 92 insertions(+)
 create mode 100644 scripts/stats.lean

diff --git a/HepLean/StandardModel/HiggsBoson/GaugeAction.lean b/HepLean/StandardModel/HiggsBoson/GaugeAction.lean
index a2597e1..2d3257d 100644
--- a/HepLean/StandardModel/HiggsBoson/GaugeAction.lean
+++ b/HepLean/StandardModel/HiggsBoson/GaugeAction.lean
@@ -50,6 +50,7 @@ noncomputable def matrixToLin : Matrix (Fin 2) (Fin 2) ℂ →* (HiggsVec →L[
     Matrix.toLin_mul orthonormBasis.toBasis orthonormBasis.toBasis orthonormBasis.toBasis g h
   map_one' := ContinuousLinearMap.coe_inj.mp $ Matrix.toLin_one orthonormBasis.toBasis
 
+/-- `matrixToLin` commutes with the `star` operation. -/
 lemma matrixToLin_star (g : Matrix (Fin 2) (Fin 2) ℂ) :
     matrixToLin (star g) = star (matrixToLin g) :=
   ContinuousLinearMap.coe_inj.mp $ Matrix.toLin_conjTranspose orthonormBasis orthonormBasis g
diff --git a/lakefile.toml b/lakefile.toml
index 49e2387..6dc2426 100644
--- a/lakefile.toml
+++ b/lakefile.toml
@@ -48,4 +48,8 @@ srcDir = "scripts"
 
 [[lean_exe]]
 name = "mathlib_textLint_on_hepLean"
+srcDir = "scripts"
+
+[[lean_exe]]
+name = "stats"
 srcDir = "scripts"
\ No newline at end of file
diff --git a/scripts/stats.lean b/scripts/stats.lean
new file mode 100644
index 0000000..cdfa2b8
--- /dev/null
+++ b/scripts/stats.lean
@@ -0,0 +1,87 @@
+/-
+Copyright (c) 2024 Joseph Tooby-Smith. All rights reserved.
+Released under Apache 2.0 license.
+Authors: Joseph Tooby-Smith
+-/
+import Batteries.Lean.HashSet
+import Lean
+import Mathlib.Lean.Expr.Basic
+import Mathlib.Lean.CoreM
+import HepLean
+/-!
+
+# HepLean Stats
+
+This file concerns with statistics of HepLean.
+
+-/
+
+open Lean System Meta
+
+structure FileStats where
+  numDefs : Nat
+  numLemsWithoutDocString : Nat
+  numLemsWithDocString : Nat
+
+def FileStats.add (a b : FileStats) : FileStats where
+  numDefs := a.numDefs + b.numDefs
+  numLemsWithoutDocString := a.numLemsWithoutDocString + b.numLemsWithoutDocString
+  numLemsWithDocString := a.numLemsWithDocString + b.numLemsWithDocString
+
+def FileStats.zero : FileStats where
+  numDefs := 0
+  numLemsWithoutDocString := 0
+  numLemsWithDocString := 0
+
+instance : ToString FileStats where
+  toString fs := s!"Number of definitions: {fs.numDefs}\n" ++
+    s!"Number of lemmas without doc string: {fs.numLemsWithoutDocString}\n" ++
+    s!"Number of lemmas with doc string: {fs.numLemsWithDocString}"
+
+def getConstantDoc (constName : Array ConstantInfo) : MetaM (FileStats) := do
+  let env ← getEnv
+  let mut numDefs := 0
+  let mut numLemsWithoutDocString := 0
+  let mut numLemsWithDocString := 0
+  for c in constName do
+    if ¬ Lean.Name.isInternalDetail c.name then
+      if c.isDef then
+        numDefs := numDefs + 1
+      if c.isThm then
+        let doc ← Lean.findDocString? env c.name
+        if doc == none then
+          numLemsWithoutDocString := numLemsWithoutDocString + 1
+        else
+          numLemsWithDocString := numLemsWithDocString + 1
+  return {numDefs := numDefs,
+          numLemsWithoutDocString := numLemsWithoutDocString,
+          numLemsWithDocString := numLemsWithDocString}
+
+def getStats (imp : Import) : MetaM FileStats := do
+  if imp.module == `Init then
+    return {numDefs := 0,
+            numLemsWithoutDocString := 0,
+            numLemsWithDocString := 0}
+  let mFile ← findOLean imp.module
+  let (modData, _) ← readModuleData mFile
+  let cons := modData.constants
+  let out ← (getConstantDoc cons)
+  return out
+
+def getStatsArray (a : Array Import) : MetaM (Array FileStats) := do
+  return ← a.mapM getStats
+
+def main (_ : List String) : IO UInt32 := do
+  initSearchPath (← findSysroot)
+  let mods : Name :=  `HepLean
+  let imp :  Import := {module := mods}
+  let mFile ← findOLean imp.module
+  unless (← mFile.pathExists) do
+        throw <| IO.userError s!"object file '{mFile}' of module {imp.module} does not exist"
+  let (hepLeanMod, _) ← readModuleData mFile
+  let noFiles := hepLeanMod.imports.size - 1
+  let fileStats ← CoreM.withImportModules #[`HepLean] (getStatsArray hepLeanMod.imports).run'
+  let totalStats := Array.foldl FileStats.add FileStats.zero fileStats
+  IO.println s!"Total number of files: {noFiles}"
+  IO.println totalStats
+  pure 0