Added support for specifying max heartbeats via an env. variable or a CLI argument

DocGen4/Load.lean

@@ -22,13 +22,13 @@ def loadInit (imports : Array Name) : IO Hierarchy := do
 Load a list of modules from the current Lean search path into an `Environment`
 to process for documentation.
 -/
-def load (task : Process.AnalyzeTask) : IO (Process.AnalyzerResult × Hierarchy) := do
+def load (task : Process.AnalyzeTask) (maxHeartbeats : Nat := 100000000) : IO (Process.AnalyzerResult × Hierarchy) := do
   initSearchPath (← findSysroot)
   let env ← envOfImports task.getLoad
   let config := {
-    -- TODO: parameterize maxHeartbeats
-    maxHeartbeats := 100000000,
+    maxHeartbeats := maxHeartbeats,
     options := ⟨[
+      (`maxHeartbeats, maxHeartbeats),
       (`pp.tagAppFns, true),
       (`pp.funBinderTypes, true),
       (`debug.skipKernelTC, true),
@@ -39,6 +39,6 @@ def load (task : Process.AnalyzeTask) : IO (Process.AnalyzerResult × Hierarchy)
     fileMap := default,
   }
 
-  Prod.fst <$> Meta.MetaM.toIO (Process.process task) config { env := env } {} {}
+  Prod.fst <$> Meta.MetaM.toIO (Process.process task maxHeartbeats) config { env := env } {} {}
 
 end DocGen4

DocGen4/Process/Analyze.lean

@@ -112,7 +112,7 @@ def mkOptions : IO DocGenOptions := do
 Run the doc-gen analysis on all modules that are loaded into the `Environment`
 of this `MetaM` run and mentioned by the `AnalyzeTask`.
 -/
-def process (task : AnalyzeTask) : MetaM (AnalyzerResult × Hierarchy) := do
+def process (task : AnalyzeTask) (maxHeartbeats : Nat := 5000000) : MetaM (AnalyzerResult × Hierarchy) := do
   let env ← getEnv
   let allModules := env.header.moduleNames
   let relevantModules :=
@@ -135,9 +135,10 @@ def process (task : AnalyzeTask) : MetaM (AnalyzerResult × Hierarchy) := do
 
     res ← tryCatchRuntimeEx
       (do
+        let opts := (← getOptions).setNat `maxHeartbeats maxHeartbeats
         let config := {
-          maxHeartbeats := 5000000,
-          options := ← getOptions,
+          maxHeartbeats := maxHeartbeats,
+          options := opts,
           fileName := ← getFileName,
           fileMap := ← getFileMap,
         }

Main.lean

@@ -4,6 +4,17 @@ import Cli
 
 open DocGen4 Lean Cli
 
+def defaultMaxHeartbeats : Nat := 100_000_000
+
+/-- Get maxHeartbeats from CLI flag, falling back to environment variable, then default. -/
+def getMaxHeartbeats (p : Parsed) : IO Nat := do
+  match p.flag? "max-heartbeats" with
+  | some flag => pure <| flag.as! Nat
+  | none => do
+    match ← IO.getEnv "DOCGEN_MAX_HEARTBEATS" with
+    | some s => pure <| s.toNat?.getD defaultMaxHeartbeats
+    | none => pure defaultMaxHeartbeats
+
 def getTopLevelModules (p : Parsed) : IO (List String) :=  do
   let topLevelModules := p.variableArgsAs! String |>.toList
   if topLevelModules.length == 0 then
@@ -21,9 +32,10 @@ def runSingleCmd (p : Parsed) : IO UInt32 := do
   let buildDir := match p.flag? "build" with
     | some dir => dir.as! String
     | none => ".lake/build"
+  let maxHeartbeats ← getMaxHeartbeats p
   let relevantModules := #[p.positionalArg! "module" |>.as! String |> String.toName]
   let sourceUri := p.positionalArg! "sourceUri" |>.as! String
-  let (doc, hierarchy) ← load <| .analyzeConcreteModules relevantModules
+  let (doc, hierarchy) ← load (.analyzeConcreteModules relevantModules) maxHeartbeats
   let baseConfig ← getSimpleBaseContext buildDir hierarchy
   discard <| htmlOutputResults baseConfig doc (some sourceUri)
   return 0
@@ -41,9 +53,10 @@ def runGenCoreCmd (p : Parsed) : IO UInt32 := do
   let buildDir := match p.flag? "build" with
     | some dir => dir.as! String
     | none => ".lake/build"
+  let maxHeartbeats ← getMaxHeartbeats p
   let manifestOutput? := (p.flag? "manifest").map (·.as! String)
   let module := p.positionalArg! "module" |>.as! String |> String.toName
-  let (doc, hierarchy) ← load <| .analyzePrefixModules module
+  let (doc, hierarchy) ← load (.analyzePrefixModules module) maxHeartbeats
   let baseConfig ← getSimpleBaseContext buildDir hierarchy
   let outputs ← htmlOutputResults baseConfig doc none
   if let .some manifestOutput := manifestOutput? then
@@ -80,6 +93,7 @@ def singleCmd := `[Cli|
 
   FLAGS:
     b, build : String; "Build directory."
+    "max-heartbeats" : Nat; "Maximum heartbeats for elaboration (default: 100_000_000). Can also be set via DOCGEN_MAX_HEARTBEATS env var."
 
   ARGS:
     module : String; "The module to generate the HTML for. Does not have to be part of topLevelModules."
@@ -101,6 +115,7 @@ def genCoreCmd := `[Cli|
   FLAGS:
     b, build : String; "Build directory."
     m, manifest : String; "Manifest output, to list all the files generated."
+    "max-heartbeats" : Nat; "Maximum heartbeats for elaboration (default: 100_000_000). Can also be set via DOCGEN_MAX_HEARTBEATS env var."
 
   ARGS:
     module : String; "The module to generate the HTML for."

README.md

@@ -91,6 +91,12 @@ The different options are:
 ## Disabling equations
 Generation of equations for definitions is enabled by default, but can be disabled by setting the `DISABLE_EQUATIONS` environment variable to `1`.
 
+## Max heartbeats
+The maximum number of heartbeats for elaboration can be configured by setting the `DOCGEN_MAX_HEARTBEATS` environment variable. The default is `100_000_000`. For example:
+```
+DOCGEN_MAX_HEARTBEATS=200_000_000 lake build YourLibraryName:docs
+```
+
 ## How does `docs#Nat.add` from the Lean Zulip work?
 If someone sends a message that contains `docs#Nat.add` on the Lean Zulip this will
 automatically link to `Nat.add` from the `mathlib4` documentation. The way that this

lakefile.lean

@@ -23,6 +23,12 @@ require «UnicodeBasic» from git
 require Cli from git
   "https://github.com/leanprover/lean4-cli" @ "main"
 
+/-- Get max heartbeats CLI args from environment variable if set. -/
+def getMaxHeartbeatsArgs : IO (Array String) := do
+  match ← IO.getEnv "DOCGEN_MAX_HEARTBEATS" with
+  | some value => return #["--max-heartbeats", value]
+  | none => return #[]
+
 /--
 Obtain the subdirectory of the Lean package relative to the root of the enclosing git repository.
 -/
@@ -240,7 +246,7 @@ module_facet docs (mod) : DepSet FilePath := do
             let srcUri ← uriJob.await
             proc {
               cmd := exeFile.toString
-              args := #["single", "--build", buildDir.toString, mod.name.toString, srcUri]
+              args := #["single", "--build", buildDir.toString, mod.name.toString, srcUri] ++ (← getMaxHeartbeatsArgs)
               env := ← getAugmentedEnv
             }
           return DepSet.mk #[docFile] docDeps
@@ -259,7 +265,7 @@ def coreTarget (component : Lean.Name) : FetchM (Job <| Array FilePath) := do
           cmd := exeFile.toString
           args := #["genCore", component.toString,
             "--build", buildDir.toString,
-            "--manifest", manifestFile.toString]
+            "--manifest", manifestFile.toString] ++ (← getMaxHeartbeatsArgs)
           env := ← getAugmentedEnv
         }
       addTrace (← computeTrace dataFile)