From 7ea048ead1961229f3aa18a4b794aced7c55f10b Mon Sep 17 00:00:00 2001
From: Nicolas Rouquette <nfr@jpl.nasa.gov>
Date: Sun, 7 Dec 2025 14:04:41 -0800
Subject: [PATCH 1/7] Added support for specifying max heartbeats via an env.
 variable or a CLI argument

---
 DocGen4/Load.lean |  5 ++---
 Main.lean         | 31 +++++++++++++++++++++++++++++--
 README.md         |  6 ++++++
 3 files changed, 37 insertions(+), 5 deletions(-)

diff --git a/DocGen4/Load.lean b/DocGen4/Load.lean
index 2cf3df9..b43e774 100644
--- a/DocGen4/Load.lean
+++ b/DocGen4/Load.lean
@@ -22,12 +22,11 @@ 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 := ⟨[
       (`pp.tagAppFns, true),
       (`pp.funBinderTypes, true),
diff --git a/Main.lean b/Main.lean
index b42aa82..eb365e4 100644
--- a/Main.lean
+++ b/Main.lean
@@ -4,6 +4,25 @@ import Cli
 
 open DocGen4 Lean Cli
 
+def defaultMaxHeartbeats : Nat := 100_000_000
+
+/-- Parse a natural number, allowing underscores as separators (e.g., "100_000_000").
+
+ See discussion: https://leanprover.zulipchat.com/#narrow/channel/270676-lean4/topic/String.2EtoNat.3F.20and.20String.2EtoInt.3F.20handling.20underscores/with/562357458
+ See this related issue: https://github.com/leanprover/lean4/issues/11538
+ -/
+def parseNatWithUnderscores (s : String) : Option Nat :=
+  (s.replace "_" "").toNat?
+
+/-- 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 <| (parseNatWithUnderscores (flag.as! String)).getD defaultMaxHeartbeats
+  | none => do
+    match ← IO.getEnv "DOCGEN_MAX_HEARTBEATS" with
+    | some s => pure <| (parseNatWithUnderscores s).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 +40,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 +61,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
@@ -74,12 +95,17 @@ def runBibPrepassCmd (p : Parsed) : IO UInt32 := do
     | _ => throw <| IO.userError "there should be exactly one source file"
   return 0
 
+/--
+ See discussion: https://leanprover.zulipchat.com/#narrow/channel/270676-lean4/topic/String.2EtoNat.3F.20and.20String.2EtoInt.3F.20handling.20underscores/with/562357458
+ If https://github.com/leanprover/lean4/issues/11538 is resolved, change "max-heartbeats" to be of type Nat directly.
+ -/
 def singleCmd := `[Cli|
   single VIA runSingleCmd;
   "Only generate the documentation for the module it was given, might contain broken links unless all documentation is generated."
 
   FLAGS:
     b, build : String; "Build directory."
+    "max-heartbeats" : String; "Maximum heartbeats for elaboration (default: 100_000_000). Supports underscores as separators. 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 +127,7 @@ def genCoreCmd := `[Cli|
   FLAGS:
     b, build : String; "Build directory."
     m, manifest : String; "Manifest output, to list all the files generated."
+    "max-heartbeats" : String; "Maximum heartbeats for elaboration (default: 100_000_000). Supports underscores as separators. Can also be set via DOCGEN_MAX_HEARTBEATS env var."
 
   ARGS:
     module : String; "The module to generate the HTML for."
diff --git a/README.md b/README.md
index 8cbd6c8..f2c1f1a 100644
--- a/README.md
+++ b/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`. Underscores can be used as separators for readability. 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

From 3eaa508902695eb3f308426486700786627bf8b7 Mon Sep 17 00:00:00 2001
From: Nicolas Rouquette <nfr@jpl.nasa.gov>
Date: Sun, 7 Dec 2025 15:11:46 -0800
Subject: [PATCH 2/7] Pass DOCGEN_MAX_HEARTBEATS env var to Lake facets

Lake's getAugmentedEnv creates a new environment with only specific
variables (LEAN_PATH, LEAN_SRC_PATH, etc.) rather than inheriting
the shell environment. This means DOCGEN_MAX_HEARTBEATS was not
being passed through to the doc-gen4 subprocess.

This fix reads DOCGEN_MAX_HEARTBEATS from the environment in the
lakefile and explicitly passes it as --max-heartbeats CLI argument
to the 'single' and 'genCore' commands.
---
 lakefile.lean | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/lakefile.lean b/lakefile.lean
index 4853d20..64af931 100644
--- a/lakefile.lean
+++ b/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)

From f7eba2b1af830be4186ff4381f4426ba27ea251e Mon Sep 17 00:00:00 2001
From: Nicolas Rouquette <nfr@jpl.nasa.gov>
Date: Sun, 7 Dec 2025 17:15:29 -0800
Subject: [PATCH 3/7] Add debug output

---
 lakefile.lean | 1 +
 1 file changed, 1 insertion(+)

diff --git a/lakefile.lean b/lakefile.lean
index 64af931..888c95a 100644
--- a/lakefile.lean
+++ b/lakefile.lean
@@ -25,6 +25,7 @@ require Cli from git
 
 /-- Get max heartbeats CLI args from environment variable if set. -/
 def getMaxHeartbeatsArgs : IO (Array String) := do
+  IO.eprintln s!"DEBUG: Checking DOCGEN_MAX_HEARTBEATS"
   match ← IO.getEnv "DOCGEN_MAX_HEARTBEATS" with
   | some value => return #["--max-heartbeats", value]
   | none => return #[]

From 69941abc9f999b75c974b7e568c5ae02520122ce Mon Sep 17 00:00:00 2001
From: Nicolas Rouquette <nfr@jpl.nasa.gov>
Date: Sun, 7 Dec 2025 17:24:31 -0800
Subject: [PATCH 4/7] Better debug output

---
 lakefile.lean      |  5 ++--
 test/test_env.lean | 65 ++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 68 insertions(+), 2 deletions(-)
 create mode 100644 test/test_env.lean

diff --git a/lakefile.lean b/lakefile.lean
index 888c95a..d9f2472 100644
--- a/lakefile.lean
+++ b/lakefile.lean
@@ -25,8 +25,9 @@ require Cli from git
 
 /-- Get max heartbeats CLI args from environment variable if set. -/
 def getMaxHeartbeatsArgs : IO (Array String) := do
-  IO.eprintln s!"DEBUG: Checking DOCGEN_MAX_HEARTBEATS"
-  match ← IO.getEnv "DOCGEN_MAX_HEARTBEATS" with
+  let val ← IO.getEnv "DOCGEN_MAX_HEARTBEATS"
+  IO.eprintln s!"DEBUG: DOCGEN_MAX_HEARTBEATS = {val}"
+  match val with
   | some value => return #["--max-heartbeats", value]
   | none => return #[]
 
diff --git a/test/test_env.lean b/test/test_env.lean
new file mode 100644
index 0000000..76dc2c1
--- /dev/null
+++ b/test/test_env.lean
@@ -0,0 +1,65 @@
+/-
+Test script to verify DOCGEN_MAX_HEARTBEATS environment variable handling.
+
+Run with:
+  DOCGEN_MAX_HEARTBEATS=5_000_000_000 lake env lean test/test_env.lean
+
+Or without the env var:
+  lake env lean test/test_env.lean
+-/
+
+def defaultMaxHeartbeats : Nat := 100_000_000
+
+/-- Parse a natural number, allowing underscores as separators (e.g., "100_000_000"). -/
+def parseNatWithUnderscores (s : String) : Option Nat :=
+  (s.replace "_" "").toNat?
+
+def main : IO Unit := do
+  IO.println "=== DOCGEN_MAX_HEARTBEATS Environment Variable Test ==="
+  IO.println ""
+
+  -- Test 1: Check if environment variable is visible
+  let envValue ← IO.getEnv "DOCGEN_MAX_HEARTBEATS"
+  match envValue with
+  | none =>
+    IO.println "❌ DOCGEN_MAX_HEARTBEATS is NOT set"
+    IO.println s!"   Using default: {defaultMaxHeartbeats}"
+  | some value =>
+    IO.println s!"✅ DOCGEN_MAX_HEARTBEATS is set to: \"{value}\""
+
+    -- Test 2: Parse the value
+    match parseNatWithUnderscores value with
+    | none =>
+      IO.println s!"❌ Failed to parse \"{value}\" as Nat"
+      IO.println s!"   Using default: {defaultMaxHeartbeats}"
+    | some n =>
+      IO.println s!"✅ Parsed successfully as: {n}"
+
+  IO.println ""
+
+  -- Test 3: Show all DOCGEN_* environment variables
+  IO.println "=== All DOCGEN_* environment variables ==="
+  for name in ["DOCGEN_MAX_HEARTBEATS", "DOCGEN_SRC", "DISABLE_EQUATIONS"] do
+    let value ← IO.getEnv name
+    match value with
+    | none => IO.println s!"   {name}: (not set)"
+    | some v => IO.println s!"   {name}: \"{v}\""
+
+  IO.println ""
+  IO.println "=== Test parseNatWithUnderscores ==="
+  let testCases := [
+    ("100000000", some 100000000),
+    ("100_000_000", some 100000000),
+    ("5_000_000_000", some 5000000000),
+    ("1_2_3", some 123),
+    ("abc", none),
+    ("", none),
+    ("123abc", none)
+  ]
+  for (input, expected) in testCases do
+    let result := parseNatWithUnderscores input
+    let status := if result == expected then "✅" else "❌"
+    IO.println s!"   {status} parseNatWithUnderscores \"{input}\" = {result} (expected {expected})"
+
+  IO.println ""
+  IO.println "=== Done ==="

From c413a9c98f2186d7c564f41875c1fd5a6bd44836 Mon Sep 17 00:00:00 2001
From: Nicolas Rouquette <nfr@jpl.nasa.gov>
Date: Sun, 7 Dec 2025 18:03:17 -0800
Subject: [PATCH 5/7] Fix: Pass maxHeartbeats through to Process.process

The maxHeartbeats parameter was being passed to Load.load but the
Process.process function had a hardcoded value of 5000000 that was
overriding it. This fix:

1. Adds maxHeartbeats parameter to Process.process (with same default)
2. Uses the parameter instead of hardcoded value
3. Passes maxHeartbeats from Load.load to Process.process

This allows DOCGEN_MAX_HEARTBEATS env var or --max-heartbeats CLI flag
to actually control the heartbeat limit during documentation generation.
---
 DocGen4/Load.lean            |  2 +-
 DocGen4/Process/Analyze.lean |  4 +--
 lakefile.lean                |  4 +--
 test/test_env.lean           | 65 ------------------------------------
 4 files changed, 4 insertions(+), 71 deletions(-)
 delete mode 100644 test/test_env.lean

diff --git a/DocGen4/Load.lean b/DocGen4/Load.lean
index b43e774..f31b5f3 100644
--- a/DocGen4/Load.lean
+++ b/DocGen4/Load.lean
@@ -38,6 +38,6 @@ def load (task : Process.AnalyzeTask) (maxHeartbeats : Nat := 100000000) : IO (P
     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
diff --git a/DocGen4/Process/Analyze.lean b/DocGen4/Process/Analyze.lean
index c6f1f39..9e52336 100644
--- a/DocGen4/Process/Analyze.lean
+++ b/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 :=
@@ -136,7 +136,7 @@ def process (task : AnalyzeTask) : MetaM (AnalyzerResult × Hierarchy) := do
     res ← tryCatchRuntimeEx
       (do
         let config := {
-          maxHeartbeats := 5000000,
+          maxHeartbeats := maxHeartbeats,
           options := ← getOptions,
           fileName := ← getFileName,
           fileMap := ← getFileMap,
diff --git a/lakefile.lean b/lakefile.lean
index d9f2472..64af931 100644
--- a/lakefile.lean
+++ b/lakefile.lean
@@ -25,9 +25,7 @@ require Cli from git
 
 /-- Get max heartbeats CLI args from environment variable if set. -/
 def getMaxHeartbeatsArgs : IO (Array String) := do
-  let val ← IO.getEnv "DOCGEN_MAX_HEARTBEATS"
-  IO.eprintln s!"DEBUG: DOCGEN_MAX_HEARTBEATS = {val}"
-  match val with
+  match ← IO.getEnv "DOCGEN_MAX_HEARTBEATS" with
   | some value => return #["--max-heartbeats", value]
   | none => return #[]
 
diff --git a/test/test_env.lean b/test/test_env.lean
deleted file mode 100644
index 76dc2c1..0000000
--- a/test/test_env.lean
+++ /dev/null
@@ -1,65 +0,0 @@
-/-
-Test script to verify DOCGEN_MAX_HEARTBEATS environment variable handling.
-
-Run with:
-  DOCGEN_MAX_HEARTBEATS=5_000_000_000 lake env lean test/test_env.lean
-
-Or without the env var:
-  lake env lean test/test_env.lean
--/
-
-def defaultMaxHeartbeats : Nat := 100_000_000
-
-/-- Parse a natural number, allowing underscores as separators (e.g., "100_000_000"). -/
-def parseNatWithUnderscores (s : String) : Option Nat :=
-  (s.replace "_" "").toNat?
-
-def main : IO Unit := do
-  IO.println "=== DOCGEN_MAX_HEARTBEATS Environment Variable Test ==="
-  IO.println ""
-
-  -- Test 1: Check if environment variable is visible
-  let envValue ← IO.getEnv "DOCGEN_MAX_HEARTBEATS"
-  match envValue with
-  | none =>
-    IO.println "❌ DOCGEN_MAX_HEARTBEATS is NOT set"
-    IO.println s!"   Using default: {defaultMaxHeartbeats}"
-  | some value =>
-    IO.println s!"✅ DOCGEN_MAX_HEARTBEATS is set to: \"{value}\""
-
-    -- Test 2: Parse the value
-    match parseNatWithUnderscores value with
-    | none =>
-      IO.println s!"❌ Failed to parse \"{value}\" as Nat"
-      IO.println s!"   Using default: {defaultMaxHeartbeats}"
-    | some n =>
-      IO.println s!"✅ Parsed successfully as: {n}"
-
-  IO.println ""
-
-  -- Test 3: Show all DOCGEN_* environment variables
-  IO.println "=== All DOCGEN_* environment variables ==="
-  for name in ["DOCGEN_MAX_HEARTBEATS", "DOCGEN_SRC", "DISABLE_EQUATIONS"] do
-    let value ← IO.getEnv name
-    match value with
-    | none => IO.println s!"   {name}: (not set)"
-    | some v => IO.println s!"   {name}: \"{v}\""
-
-  IO.println ""
-  IO.println "=== Test parseNatWithUnderscores ==="
-  let testCases := [
-    ("100000000", some 100000000),
-    ("100_000_000", some 100000000),
-    ("5_000_000_000", some 5000000000),
-    ("1_2_3", some 123),
-    ("abc", none),
-    ("", none),
-    ("123abc", none)
-  ]
-  for (input, expected) in testCases do
-    let result := parseNatWithUnderscores input
-    let status := if result == expected then "✅" else "❌"
-    IO.println s!"   {status} parseNatWithUnderscores \"{input}\" = {result} (expected {expected})"
-
-  IO.println ""
-  IO.println "=== Done ==="

From 130caf712c0984819a78bc0f13d3fba260fa9dea Mon Sep 17 00:00:00 2001
From: Nicolas Rouquette <nfr@jpl.nasa.gov>
Date: Sun, 7 Dec 2025 18:23:54 -0800
Subject: [PATCH 6/7] Fix: Set maxHeartbeats in options, not just Core.Context

The maxHeartbeats field in Core.Context doesn't affect Lean's internal
operations like equation theorem generation (getEqnsFor?). Those use
the maxHeartbeats *option* which has a default of 200000.

This fix adds maxHeartbeats to the options in both Load.lean and
Analyze.lean, so that elaboration operations respect the configured
limit.
---
 DocGen4/Load.lean            | 1 +
 DocGen4/Process/Analyze.lean | 3 ++-
 Main.lean                    | 1 +
 3 files changed, 4 insertions(+), 1 deletion(-)

diff --git a/DocGen4/Load.lean b/DocGen4/Load.lean
index f31b5f3..ef03398 100644
--- a/DocGen4/Load.lean
+++ b/DocGen4/Load.lean
@@ -28,6 +28,7 @@ def load (task : Process.AnalyzeTask) (maxHeartbeats : Nat := 100000000) : IO (P
   let config := {
     maxHeartbeats := maxHeartbeats,
     options := ⟨[
+      (`maxHeartbeats, maxHeartbeats),
       (`pp.tagAppFns, true),
       (`pp.funBinderTypes, true),
       (`debug.skipKernelTC, true),
diff --git a/DocGen4/Process/Analyze.lean b/DocGen4/Process/Analyze.lean
index 9e52336..2a7423c 100644
--- a/DocGen4/Process/Analyze.lean
+++ b/DocGen4/Process/Analyze.lean
@@ -135,9 +135,10 @@ def process (task : AnalyzeTask) (maxHeartbeats : Nat := 5000000) : MetaM (Analy
 
     res ← tryCatchRuntimeEx
       (do
+        let opts := (← getOptions).setNat `maxHeartbeats maxHeartbeats
         let config := {
           maxHeartbeats := maxHeartbeats,
-          options := ← getOptions,
+          options := opts,
           fileName := ← getFileName,
           fileMap := ← getFileMap,
         }
diff --git a/Main.lean b/Main.lean
index eb365e4..18f6071 100644
--- a/Main.lean
+++ b/Main.lean
@@ -98,6 +98,7 @@ def runBibPrepassCmd (p : Parsed) : IO UInt32 := do
 /--
  See discussion: https://leanprover.zulipchat.com/#narrow/channel/270676-lean4/topic/String.2EtoNat.3F.20and.20String.2EtoInt.3F.20handling.20underscores/with/562357458
  If https://github.com/leanprover/lean4/issues/11538 is resolved, change "max-heartbeats" to be of type Nat directly.
+ See: https://github.com/leanprover/lean4/pull/11541
  -/
 def singleCmd := `[Cli|
   single VIA runSingleCmd;

From 489f16255398be86e0d8eab58f5965e9e874eeda Mon Sep 17 00:00:00 2001
From: Nicolas Rouquette <nfr@jpl.nasa.gov>
Date: Mon, 8 Dec 2025 08:20:17 -0800
Subject: [PATCH 7/7] In Main.lean:

- Removed the parseNatWithUnderscores function and its doc comments with issue references
- Simplified getMaxHeartbeats to use flag.as! Nat directly for the CLI flag and s.toNat? for the environment variable
- Changed the max-heartbeats flag type from String to Nat in both singleCmd and genCoreCmd
- Removed the "Supports underscores as separators" mention from flag descriptions

In README.md:

- Removed the "Underscores can be used as separators for readability" sentence (the examples still show underscores since they will work in a version of Lean4 that includes this merged PR: https://github.com/leanprover/lean4/pull/11541
---
 Main.lean | 21 ++++-----------------
 README.md |  2 +-
 2 files changed, 5 insertions(+), 18 deletions(-)

diff --git a/Main.lean b/Main.lean
index 18f6071..b677a04 100644
--- a/Main.lean
+++ b/Main.lean
@@ -6,21 +6,13 @@ open DocGen4 Lean Cli
 
 def defaultMaxHeartbeats : Nat := 100_000_000
 
-/-- Parse a natural number, allowing underscores as separators (e.g., "100_000_000").
-
- See discussion: https://leanprover.zulipchat.com/#narrow/channel/270676-lean4/topic/String.2EtoNat.3F.20and.20String.2EtoInt.3F.20handling.20underscores/with/562357458
- See this related issue: https://github.com/leanprover/lean4/issues/11538
- -/
-def parseNatWithUnderscores (s : String) : Option Nat :=
-  (s.replace "_" "").toNat?
-
 /-- 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 <| (parseNatWithUnderscores (flag.as! String)).getD defaultMaxHeartbeats
+  | some flag => pure <| flag.as! Nat
   | none => do
     match ← IO.getEnv "DOCGEN_MAX_HEARTBEATS" with
-    | some s => pure <| (parseNatWithUnderscores s).getD defaultMaxHeartbeats
+    | some s => pure <| s.toNat?.getD defaultMaxHeartbeats
     | none => pure defaultMaxHeartbeats
 
 def getTopLevelModules (p : Parsed) : IO (List String) :=  do
@@ -95,18 +87,13 @@ def runBibPrepassCmd (p : Parsed) : IO UInt32 := do
     | _ => throw <| IO.userError "there should be exactly one source file"
   return 0
 
-/--
- See discussion: https://leanprover.zulipchat.com/#narrow/channel/270676-lean4/topic/String.2EtoNat.3F.20and.20String.2EtoInt.3F.20handling.20underscores/with/562357458
- If https://github.com/leanprover/lean4/issues/11538 is resolved, change "max-heartbeats" to be of type Nat directly.
- See: https://github.com/leanprover/lean4/pull/11541
- -/
 def singleCmd := `[Cli|
   single VIA runSingleCmd;
   "Only generate the documentation for the module it was given, might contain broken links unless all documentation is generated."
 
   FLAGS:
     b, build : String; "Build directory."
-    "max-heartbeats" : String; "Maximum heartbeats for elaboration (default: 100_000_000). Supports underscores as separators. Can also be set via DOCGEN_MAX_HEARTBEATS env var."
+    "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."
@@ -128,7 +115,7 @@ def genCoreCmd := `[Cli|
   FLAGS:
     b, build : String; "Build directory."
     m, manifest : String; "Manifest output, to list all the files generated."
-    "max-heartbeats" : String; "Maximum heartbeats for elaboration (default: 100_000_000). Supports underscores as separators. Can also be set via DOCGEN_MAX_HEARTBEATS env var."
+    "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."
diff --git a/README.md b/README.md
index f2c1f1a..79bcdd9 100644
--- a/README.md
+++ b/README.md
@@ -92,7 +92,7 @@ The different options are:
 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`. Underscores can be used as separators for readability. For example:
+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
 ```