lean4-htt/library/Init/Lean/MonadCache.lean
2019-11-07 10:40:26 -08:00

109 lines
4.8 KiB
Text
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

/-
Copyright (c) 2019 Microsoft Corporation. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Leonardo de Moura
-/
prelude
import Init.Control.Reader
import Init.Control.EState
import Init.Data.HashMap
namespace Lean
/-- Interface for caching results. -/
class MonadCache (α β : Type) (m : Type → Type) :=
(findCached {} : α → m (Option β))
(cache {} : α → β → m Unit)
/-- If entry `a := b` is already in the cache, then return `b`.
Otherwise, execute `b ← f a`, store `a := b` in the cache and return `b`. -/
@[inline] def checkCache {α β : Type} {m : Type → Type} [MonadCache α β m] [Monad m] (a : α) (f : α → m β) : m β :=
do b? ← MonadCache.findCached a;
match b? with
| some b => pure b
| none => do
b ← f a;
MonadCache.cache a b;
pure b
instance readerLift {α β ρ : Type} {m : Type → Type} [MonadCache α β m] : MonadCache α β (ReaderT ρ m) :=
{ findCached := fun a r => MonadCache.findCached a,
cache := fun a b r => MonadCache.cache a b }
instance exceptLift {α β ε : Type} {m : Type → Type} [MonadCache α β m] [Monad m] : MonadCache α β (ExceptT ε m) :=
{ findCached := fun a => ExceptT.lift $ MonadCache.findCached a,
cache := fun a b => ExceptT.lift $ MonadCache.cache a b }
/-- Adapter for implementing `MonadCache` interface using `HashMap`s.
We just have to specify how to extract/modify the `HashMap`. -/
class MonadHashMapCacheAdapter (α β : Type) (m : Type → Type) [HasBeq α] [Hashable α] :=
(getCache {} : m (HashMap α β))
(modifyCache {} : (HashMap α β → HashMap α β) → m Unit)
namespace MonadHashMapCacheAdapter
@[inline] def findCached {α β : Type} {m : Type → Type} [HasBeq α] [Hashable α] [Monad m] [MonadHashMapCacheAdapter α β m] (a : α) : m (Option β) :=
do c ← getCache;
pure (c.find a)
@[inline] def cache {α β : Type} {m : Type → Type} [HasBeq α] [Hashable α] [MonadHashMapCacheAdapter α β m] (a : α) (b : β) : m Unit :=
modifyCache $ fun s => s.insert a b
instance {α β : Type} {m : Type → Type} [HasBeq α] [Hashable α] [Monad m] [MonadHashMapCacheAdapter α β m] : MonadCache α β m :=
{ findCached := MonadHashMapCacheAdapter.findCached,
cache := MonadHashMapCacheAdapter.cache }
end MonadHashMapCacheAdapter
/-- Auxiliary structure for "adding" a `HashMap` to a state object. -/
structure WithHashMapCache (α β σ : Type) [HasBeq α] [Hashable α] :=
(state : σ)
(cache : HashMap α β := {})
namespace WithHashMapCache
@[inline] def getCache {α β σ : Type} [HasBeq α] [Hashable α] : StateM (WithHashMapCache α β σ) (HashMap α β) :=
do s ← get; pure s.cache
@[inline] def modifyCache {α β σ : Type} [HasBeq α] [Hashable α] (f : HashMap α β → HashMap α β) : StateM (WithHashMapCache α β σ) Unit :=
modify $ fun s => { cache := f s.cache, .. s }
instance stateAdapter (α β σ : Type) [HasBeq α] [Hashable α] : MonadHashMapCacheAdapter α β (StateM (WithHashMapCache α β σ)) :=
{ getCache := WithHashMapCache.getCache,
modifyCache := WithHashMapCache.modifyCache }
@[inline] def getCacheE {α β ε σ : Type} [HasBeq α] [Hashable α] : EStateM ε (WithHashMapCache α β σ) (HashMap α β) :=
do s ← get; pure s.cache
@[inline] def modifyCacheE {α β ε σ : Type} [HasBeq α] [Hashable α] (f : HashMap α β → HashMap α β) : EStateM ε (WithHashMapCache α β σ) Unit :=
modify $ fun s => { cache := f s.cache, .. s }
instance estateAdapter (α β ε σ : Type) [HasBeq α] [Hashable α] : MonadHashMapCacheAdapter α β (EStateM ε (WithHashMapCache α β σ)) :=
{ getCache := WithHashMapCache.getCacheE,
modifyCache := WithHashMapCache.modifyCacheE }
@[inline] def fromState {α β σ δ : Type} [HasBeq α] [Hashable α] (x : StateM σ δ) : StateM (WithHashMapCache α β σ) δ :=
adaptState
(fun (s : WithHashMapCache α β σ) => (s.state, s.cache))
(fun (s : σ) (cache : HashMap α β) => { state := s, cache := cache })
x
@[inline] def toState {α β σ δ : Type} [HasBeq α] [Hashable α] (x : StateM (WithHashMapCache α β σ) δ) : StateM σ δ :=
adaptState'
(fun (s : σ) => ({ state := s } : WithHashMapCache α β σ))
(fun (s : WithHashMapCache α β σ) => s.state)
x
@[inline] def fromEState {α β σ ε δ : Type} [HasBeq α] [Hashable α] (x : EStateM ε σ δ) : EStateM ε (WithHashMapCache α β σ) δ :=
adaptState
(fun (s : WithHashMapCache α β σ) => (s.state, s.cache))
(fun (s : σ) (cache : HashMap α β) => { state := s, cache := cache })
x
@[inline] def toEState {α β σ ε δ : Type} [HasBeq α] [Hashable α] (x : EStateM ε (WithHashMapCache α β σ) δ) : EStateM ε σ δ :=
adaptState'
(fun (s : σ) => ({ state := s } : WithHashMapCache α β σ))
(fun (s : WithHashMapCache α β σ) => s.state)
x
end WithHashMapCache
end Lean