--module option to count the factor per-module

ImportGraph/HoardFactor.lean

@@ -129,6 +129,59 @@ def Lean.Environment.hoardFactor (env : Environment) (decls : Array Name)
 
   return sum.toFloat / tot.toFloat
 
+def Lean.Environment.moduleHoardFactor (env : Environment) (modules : Array Name)
+    (excludeLean : Bool := true) (excludeMeta : Bool := false) (verbose : Bool := false) : CoreM Float := do
+  if verbose then
+    println!"Constructing the import graph."
+  let mut modules := modules
+  let mut graph := env.importGraph
+  if excludeLean || excludeMeta then
+    let filterLean : Name → Bool := fun n => excludeLean && (
+      Name.isPrefixOf `Lean n ∨
+      Name.isPrefixOf `Init n ∨
+      Name.isPrefixOf `Std n)
+    let filterMeta : Name → Bool := fun n => excludeMeta && (
+      Name.isPrefixOf `Batteries.CodeAction n ∨
+      Name.isPrefixOf `Batteries.Tactic n ∨
+      Name.isPrefixOf `Mathlib.Tactic n ∨
+      Name.isPrefixOf `Mathlib.Lean n ∨
+      Name.isPrefixOf `Mathlib.Mathport n ∨
+      Name.isPrefixOf `Mathlib.Util n)
+    graph := graph.filterMap fun n i =>
+      if filterLean n || filterMeta n then
+        -- not included
+        none
+      else
+        -- include node regularly
+        i.filter fun m => !(filterLean m || filterMeta m)
+  if verbose then
+    println!"Found {graph.size} modules."
+  let transitiveImports := graph.transitiveClosure
+  let (localSizes, importSizes) ← env.importSizes transitiveImports verbose
+
+  let requiredModuleMap ← env.transitivelyRequiredModules' (graph.1.toArray.map Sigma.fst).toList verbose
+  if verbose then
+    println!"Declarations preprocessed. Preprocessing modules."
+
+  let mut sum : Nat := 0
+  let mut tot : Nat := 0
+  let mut iteration : Nat := 0
+
+  if verbose then
+    println!"Entering main loop."
+  for (module, requiredModules) in requiredModuleMap do
+    iteration := iteration + 1
+    if verbose && iteration % 100 == 0 then
+      println!"{module}: {sum.toFloat / tot.toFloat}"
+    tot := tot + importSizes.findD module 0
+    sum := sum + importSizes.findD module 0
+    for n in requiredModules do
+      sum := sum - localSizes.find! n
+  if verbose then
+    println!"Processed {iteration} modules."
+
+  return sum.toFloat / tot.toFloat
+
 open IO.FS IO.Process Name in
 /-- Implementation of the hoard factor command line program. -/
 def hoardFactorCLI (args : Cli.Parsed) : IO UInt32 := do
@@ -139,14 +192,17 @@ def hoardFactorCLI (args : Cli.Parsed) : IO UInt32 := do
 
   let factor ← Core.withImportModules targets do
     let env ← getEnv
-    let mut decls : Array Name := #[]
-    if args.hasFlag "all" then
-      -- TODO: improve the standard library so the next line can be `env.constants.keys`
-      decls := (env.constants.map₁.toArray.append env.constants.map₂.toArray).map (·.fst)
-    else for m in targets do
-      let names := env.header.moduleData[(env.header.moduleNames.getIdx? m).getD 0]!.constNames
-      decls := decls.append names
-    env.hoardFactor decls (excludeLean := !args.hasFlag "include-lean") (excludeMeta := args.hasFlag "exclude-meta") (verbose := args.hasFlag "verbose")
+    if args.hasFlag "module" then
+      env.moduleHoardFactor targets (excludeLean := !args.hasFlag "include-lean") (excludeMeta := args.hasFlag "exclude-meta") (verbose := args.hasFlag "verbose")
+    else
+      let mut decls : Array Name := #[]
+      if args.hasFlag "all" then
+        -- TODO: improve the standard library so the next line can be `env.constants.keys`
+        decls := (env.constants.map₁.toArray.append env.constants.map₂.toArray).map (·.fst)
+      else for m in targets do
+        let names := env.header.moduleData[(env.header.moduleNames.getIdx? m).getD 0]!.constNames
+        decls := decls.append names
+      env.hoardFactor decls (excludeLean := !args.hasFlag "include-lean") (excludeMeta := args.hasFlag "exclude-meta") (verbose := args.hasFlag "verbose")
 
   println!"Hoard factor: {factor}"
   return 0
@@ -163,6 +219,7 @@ def hoard_factor : Cmd := `[Cli|
 The default is to compute the factor for the new declarations in the module itself."
     "include-lean";       "Include the core Lean packages: Init, Lean, Std."
     "exclude-meta";       "Exclude tactic-specific modules during the computations."
+    "module";             "Analyze dependencies for modules as a whole instead of per declaration."
     "verbose";            "Output more information during computations."
 
   ARGS: