diff --git a/src/Init/Lean/Meta/Basic.lean b/src/Init/Lean/Meta/Basic.lean index 286f11c593..e1946f6420 100644 --- a/src/Init/Lean/Meta/Basic.lean +++ b/src/Init/Lean/Meta/Basic.lean @@ -108,6 +108,7 @@ structure Cache := (inferType : PersistentExprStructMap Expr := {}) (funInfo : PersistentHashMap InfoCacheKey FunInfo := {}) (synthInstance : PersistentHashMap Expr (Option Expr) := {}) +(whnfDefault : PersistentExprStructMap Expr := {}) -- cache for closed terms and `TransparencyMode.default` structure Context := (config : Config := {}) diff --git a/src/Init/Lean/Meta/WHNF.lean b/src/Init/Lean/Meta/WHNF.lean index f78ac35fae..127912f1a1 100644 --- a/src/Init/Lean/Meta/WHNF.lean +++ b/src/Init/Lean/Meta/WHNF.lean @@ -85,21 +85,46 @@ else match e with else pure none | _ => pure none + +@[inline] private def useWHNFCache (e : Expr) : MetaM Bool := do +-- We cache only consed terms +if e.hasFVar then pure false +else do + ctx ← read; + pure $ ctx.config.transparency == TransparencyMode.default + +@[inline] private def cached? (useCache : Bool) (e : Expr) : MetaM (Option Expr) := do +if useCache then do + s ← get; + pure $ s.cache.whnfDefault.find? e +else + pure none + +private def cache (useCache : Bool) (e r : Expr) : MetaM Expr := do +when useCache $ + modify $ fun s => { cache := { whnfDefault := s.cache.whnfDefault.insert e r, .. s.cache }, .. s }; +pure r + partial def whnfImpl : Expr → MetaM Expr | e => Lean.WHNF.whnfEasyCases getLocalDecl getExprMVarAssignment? e $ fun e => do - e ← whnfCore e; - v? ← reduceNat? e; - match v? with - | some v => pure v - | none => do - v? ← reduceNative? e; + useCache ← useWHNFCache e; + e? ← cached? useCache e; + match e? with + | some e' => pure e' + | none => do + e' ← whnfCore e; + v? ← reduceNat? e'; match v? with - | some v => pure v - | none => do - e? ← unfoldDefinition? e; - match e? with - | some e => whnfImpl e - | none => pure e + | some v => cache useCache e v + | none => do + v? ← reduceNative? e'; + match v? with + | some v => cache useCache e v + | none => do + e? ← unfoldDefinition? e'; + match e? with + | some e => whnfImpl e + | none => cache useCache e e' @[init] def setWHNFRef : IO Unit := whnfRef.set whnfImpl