Skip to content

Commit f51ba22

Browse files
committed
perf: give currRecDepth its own ReaderT layer in CoreM
This PR stops the recursion-depth guard from rebuilding `Core.Context` on every recursive step. Roughly 96% of `Core.Context` reconstructions exist only to bump the depth counter, and each one paid a reference count increment per pointer field. `withReader` can never reuse the `Context` record, so keeping `currRecDepth` there made every `withIncRecDepth` rebuild it. As its own reader layer the depth is just an extra argument, so a recursion step neither allocates nor touches a reference count. The layer sits inside the `Context` reader so that a bare `read` still resolves to `Context`, and `CoreM.run` gains an initial-depth parameter so nested command elaboration keeps inheriting its caller's depth. Note this changes `CoreM`'s arity, which is externally observable in LCNF output.
1 parent df1b1f0 commit f51ba22

2 files changed

Lines changed: 15 additions & 11 deletions

File tree

src/Lean/CoreM.lean

Lines changed: 13 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -220,7 +220,6 @@ structure Context where
220220
/-- Auxiliary datastructure for converting `String.Pos` into Line/Column number. -/
221221
fileMap : FileMap
222222
options : Options := {}
223-
currRecDepth : Nat := 0
224223
maxRecDepth : Nat := 1000
225224
ref : Syntax := Syntax.missing
226225
currNamespace : Name := Name.anonymous
@@ -252,8 +251,12 @@ The main features it provides are:
252251
- environment state
253252
- Lean options context
254253
- the current open namespace
254+
- the current recursion depth, as its own reader layer rather than a `Context` field: `withReader`
255+
cannot reuse the `Context` record, so keeping the depth there made every recursion step rebuild
256+
`Context` and pay one reference count increment per pointer field. The layer sits *inside* the
257+
`Context` reader so that a bare `read` still resolves to `Context`.
255258
-/
256-
abbrev CoreM := ReaderT Context <| StateRefT State (EIO Exception)
259+
abbrev CoreM := ReaderT Context <| ReaderT Nat <| StateRefT State (EIO Exception)
257260

258261
-- Make the compiler generate specialized `pure`/`bind` so we do not have to optimize through the
259262
-- whole monad stack at every use site. May eventually be covered by `deriving`.
@@ -306,8 +309,8 @@ instance : MonadDeclNameGenerator CoreM where
306309
setDeclNGen ngen := modify fun s => { s with auxDeclNGen := ngen }
307310

308311
instance : MonadRecDepth CoreM where
309-
withRecDepth d x := withReader (fun ctx => { ctx with currRecDepth := d }) x
310-
getRecDepth := return (← read).currRecDepth
312+
withRecDepth d x := fun ctx => withTheReader Nat (fun _ => d) (x ctx)
313+
getRecDepth := fun _ => readThe Nat
311314
getMaxRecDepth := return (← read).maxRecDepth
312315

313316
instance : MonadResolveName CoreM where
@@ -427,11 +430,13 @@ that are conditionally inaccessible, depending on the current value of the `tact
427430
def mkFreshUserName (n : Name) : CoreM Name :=
428431
mkFreshNameImp n
429432

430-
@[inline] def CoreM.run (x : CoreM α) (ctx : Context) (s : State) : EIO Exception (α × State) :=
431-
((withConsistentCtx x) ctx).run s
433+
@[inline] def CoreM.run (x : CoreM α) (ctx : Context) (s : State) (recDepth : Nat := 0) :
434+
EIO Exception (α × State) :=
435+
(((withConsistentCtx x) ctx) recDepth).run s
432436

433-
@[inline] def CoreM.run' (x : CoreM α) (ctx : Context) (s : State) : EIO Exception α :=
434-
Prod.fst <$> x.run ctx s
437+
@[inline] def CoreM.run' (x : CoreM α) (ctx : Context) (s : State) (recDepth : Nat := 0) :
438+
EIO Exception α :=
439+
Prod.fst <$> x.run ctx s recDepth
435440

436441
/--
437442
Run a `CoreM` monad in IO.

src/Lean/Elab/Command.lean

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -267,7 +267,6 @@ private def runCore (x : CoreM α) : CommandElabM α := do
267267
let coreCtx : Core.Context := {
268268
fileName := ctx.fileName
269269
fileMap := ctx.fileMap
270-
currRecDepth := ctx.currRecDepth
271270
maxRecDepth := s.maxRecDepth
272271
ref := ctx.ref
273272
currNamespace := scope.currNamespace
@@ -288,7 +287,7 @@ private def runCore (x : CoreM α) : CommandElabM α := do
288287
infoState.lazyAssignment := s.infoState.lazyAssignment
289288
traceState := s.traceState
290289
snapshotTasks := s.snapshotTasks
291-
}
290+
} (recDepth := ctx.currRecDepth)
292291
let (ea, coreS) ← liftM x
293292
modify fun s => { s with
294293
env := coreS.env
@@ -1096,7 +1095,7 @@ private def liftCommandElabMCore (cmd : CommandElabM α) (throwOnError : Bool) :
10961095
cmd.run {
10971096
fileName := ctx.fileName
10981097
fileMap := ctx.fileMap
1099-
currRecDepth := ctx.currRecDepth
1098+
currRecDepth := (← MonadRecDepth.getRecDepth)
11001099
currMacroScope := ctx.currMacroScope
11011100
ref := ctx.ref
11021101
snap? := none

0 commit comments

Comments
 (0)