module Generics.MultiRec.Transformations.MemoTable where
import Generics.MultiRec hiding ( show, foldM )
import Control.Monad.State hiding ( foldM, mapM )
import qualified Data.Map as Map
import Data.Map ( Map )
import Generics.MultiRec.Transformations.Path
import Generics.MultiRec.Transformations.Children
type family Ixs (phi :: * -> *) :: [*]
data HList (l :: [*]) where
HNil :: HList '[]
HCons :: h -> HList t -> HList (h ': t)
type family HMap (h :: * -> * -> *) (f :: * -> *) (g :: * -> *) (l :: [*]) :: [*]
type instance HMap h f g '[] = '[]
type instance HMap h f g (t ': ts) = h (f t) (g t) ': HMap h f g ts
type MemoTable phi top ixs = HList (MemoTable' One phi top ixs)
data Type = One | Two
type family MemoTable' (x :: Type) (phi :: * -> *) top (ixs :: [*]) :: [*]
type instance MemoTable' x phi top '[] = '[]
type instance MemoTable' x phi top (t ': ts) = MemCell x phi top t
': MemoTable' x phi top ts
type family MemCell (x :: Type) (phi :: * -> *) top (t :: *) :: *
type instance MemCell One phi top t = Map (MemKey t) (MemVal phi top t)
type MemKey t = (Bool, t, t)
type MemVal phi top t = [Insert phi top t]
data Proxy (t :: k) = Proxy
class Lookup phi (ixs :: [*]) ix where
lookupMT :: Proxy '(phi, ixs)
-> MemoTable phi top ixs -> MemKey ix -> Maybe (MemVal phi top ix)
insertMT :: Proxy '(phi, ixs)
-> MemKey ix -> MemVal phi top ix
-> MemoTable phi top ixs -> MemoTable phi top ixs
instance Lookup phi '[] ix where
lookupMT _ HNil _ = Nothing
insertMT _ _ _ HNil = HNil
instance (Ord t) => Lookup phi (t ': ts) t where
lookupMT _ (HCons mt _) k = Map.lookup k mt
insertMT _ k v (HCons mt tl) = HCons (Map.insert k v mt) tl
instance (Lookup phi ts ix) => Lookup phi (t ': ts) ix where
lookupMT _ (HCons _ mts) = lookupMT (Proxy :: Proxy '(phi, ts)) mts
insertMT _ k v (HCons hd mts) = HCons hd (insertMT (Proxy :: Proxy '(phi, ts)) k v mts)
type Memo phi top a = State (MemoTable phi top (Ixs phi)) a
class EmptyMemo (phi :: * -> *) top (ixs :: [*]) where
emptyMemo :: Proxy '(phi,top,ixs) -> MemoTable phi top ixs
instance EmptyMemo phi top '[] where
emptyMemo _ = HNil
instance (EmptyMemo phi top t) => EmptyMemo phi top (h ': t) where
emptyMemo _ = HCons Map.empty (emptyMemo (Proxy :: Proxy '(phi,top,t)))
runMemo :: forall phi top a. (EmptyMemo phi top (Ixs phi))
=> Proxy '(phi,top) -> Memo phi top a -> a
runMemo _ = flip evalState (emptyMemo (Proxy :: Proxy '(phi,top,Ixs phi)))
recMemo :: forall phi top ix.
(Fam phi, Children phi (PF phi) ix, Lookup phi (Ixs phi) ix, Eq ix, GetChildrenTable phi (Ixs phi) ix) =>
(forall ix. (Children phi (PF phi) ix, Lookup phi (Ixs phi) ix, Eq ix, GetChildrenTable phi (Ixs phi) ix) => Bool -> phi ix -> ix -> ix -> Memo phi top [ Insert phi top ix ])
-> Bool -> phi ix -> ix -> ix -> Memo phi top [ Insert phi top ix ]
recMemo f a p b c = do
mp <- get
let k = (a,b,c)
case lookupMT (Proxy :: Proxy '(phi, Ixs phi)) mp k of
Just r -> return r
Nothing -> do
r <- f a p b c
modify $ insertMT (Proxy :: Proxy '(phi, Ixs phi)) k r
return r
type ChildTable phi top ixs = MemoTable' Two phi top ixs
type instance MemCell Two phi top t = [(Path phi t top, t)]
class ChildrenTable (phi :: * -> *) top (ixs :: [*]) where
childrenTable :: Proxy '(phi,top,ixs) -> top -> HList (ChildTable phi top ixs)
instance ChildrenTable phi top '[] where
childrenTable _ _ = HNil
instance (Fam phi, El phi top, Children phi (PF phi) top, ChildrenTable phi top t)
=> ChildrenTable phi top (top ': t) where
childrenTable _ top = HCons ((Empty,top) : allChildren proof proof top)
(childrenTable (Proxy :: Proxy '(phi,top,t)) top)
instance (Fam phi, El phi h, El phi top, Children phi (PF phi) h, ChildrenTable phi top t)
=> ChildrenTable phi top (h ': t) where
childrenTable _ top = HCons (allChildren proof proof top)
(childrenTable (Proxy :: Proxy '(phi,top,t)) top)
class GetChildrenTable phi (ixs :: [*]) ix where
getChTable :: Proxy '(phi, top, ixs)
-> HList (ChildTable phi top ixs) -> MemCell Two phi top ix
instance GetChildrenTable phi '[] ix where
getChTable _ HNil = error "This shouldn't happen"
instance (Ord t) => GetChildrenTable phi (t ': ts) t where
getChTable _ (HCons mt _) = mt
instance (GetChildrenTable phi ts ix) => GetChildrenTable phi (t ': ts) ix where
getChTable _ (HCons _ mts) = getChTable (Proxy :: Proxy '(phi, top, ts)) mts