{-# LANGUAGE TupleSections #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE TypeSynonymInstances #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE UndecidableInstances #-}

module Control.Kernmantle.Caching
  ( CS.ContentStore
  , CS.ContentHashable
  , PureHashable
  , IOHashable
  , ProvidesCaching (..)
  , ProvidesPosCaching (..)
  , AutoIdent (..)
  , SomeHashable (..)
  , StoreWithId (..)
  , LocalStoreWithId
  , CachingContext
  , CS.withStore
  , caching, caching'
  , localStoreWithId
  ) where

import Control.Category
import Control.Arrow
import Control.Kernmantle.Arrow
import Control.Kernmantle.Rope
import Control.Monad.Catch
import Control.Monad.IO.Class
import Control.Monad.IO.Unlift
import qualified Data.CAS.ContentHashable as CS
import qualified Data.CAS.ContentStore as CS
import qualified Data.CAS.RemoteCache as Remote
import Data.Functor.Identity (Identity (..))
import Data.Store (Store)

import Prelude hiding (id, (.))


instance (Monad m) => CS.ContentHashable m SplitId
instance (Monad m) => CS.ContentHashable m ArrowIdent

-- | A value directly hashable
type PureHashable = CS.ContentHashable Identity
-- | A value hashable via some IO action
type IOHashable = CS.ContentHashable IO

data SomeHashable where
  SomePureHashable :: (PureHashable a) => a -> SomeHashable
  SomeIOHashable   :: (IOHashable a)   => a -> SomeHashable
instance CS.ContentHashable IO SomeHashable where
  contentHashUpdate :: Context SHA256 -> SomeHashable -> IO (Context SHA256)
contentHashUpdate Context SHA256
ctx (SomePureHashable a
a) =
    Context SHA256 -> IO (Context SHA256)
forall (m :: * -> *) a. Monad m => a -> m a
return (Context SHA256 -> IO (Context SHA256))
-> Context SHA256 -> IO (Context SHA256)
forall a b. (a -> b) -> a -> b
$ Identity (Context SHA256) -> Context SHA256
forall a. Identity a -> a
runIdentity (Identity (Context SHA256) -> Context SHA256)
-> Identity (Context SHA256) -> Context SHA256
forall a b. (a -> b) -> a -> b
$ Context SHA256 -> a -> Identity (Context SHA256)
forall (m :: * -> *) a.
ContentHashable m a =>
Context SHA256 -> a -> m (Context SHA256)
CS.contentHashUpdate Context SHA256
ctx a
a
  contentHashUpdate Context SHA256
ctx (SomeIOHashable a
a) = Context SHA256 -> a -> IO (Context SHA256)
forall (m :: * -> *) a.
ContentHashable m a =>
Context SHA256 -> a -> m (Context SHA256)
CS.contentHashUpdate Context SHA256
ctx a
a
type CachingContext = [SomeHashable]

-- | A class to cache part of the pipeline
class ProvidesCaching eff where
  usingStore :: (CS.ContentHashable IO a, Store b)
             => eff a b
             -> eff a b
-- | A class to cache part of the pipeline where the hash can depend on the
-- position of the task in the pipeline
class (ProvidesCaching eff) => ProvidesPosCaching eff where
  usingStore' :: (CS.ContentHashable IO a, Store b)
              => eff a b
              -> eff a b

instance {-# OVERLAPPABLE #-} (Functor f, ProvidesCaching eff)
  => ProvidesCaching (f ~> eff) where
  usingStore :: (~>) f eff a b -> (~>) f eff a b
usingStore (Cayley f (eff a b)
f) = f (eff a b) -> (~>) f eff a b
forall k k1 k2 (f :: k -> *) (p :: k1 -> k2 -> k) (a :: k1)
       (b :: k2).
f (p a b) -> Cayley f p a b
Cayley (f (eff a b) -> (~>) f eff a b) -> f (eff a b) -> (~>) f eff a b
forall a b. (a -> b) -> a -> b
$ (eff a b -> eff a b) -> f (eff a b) -> f (eff a b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap eff a b -> eff a b
forall (eff :: * -> * -> *) a b.
(ProvidesCaching eff, ContentHashable IO a, Store b) =>
eff a b -> eff a b
usingStore f (eff a b)
f
instance {-# OVERLAPPABLE #-} (Functor f, ProvidesPosCaching eff)
  => ProvidesPosCaching (f ~> eff) where
  usingStore' :: (~>) f eff a b -> (~>) f eff a b
usingStore' (Cayley f (eff a b)
f) = f (eff a b) -> (~>) f eff a b
forall k k1 k2 (f :: k -> *) (p :: k1 -> k2 -> k) (a :: k1)
       (b :: k2).
f (p a b) -> Cayley f p a b
Cayley (f (eff a b) -> (~>) f eff a b) -> f (eff a b) -> (~>) f eff a b
forall a b. (a -> b) -> a -> b
$ (eff a b -> eff a b) -> f (eff a b) -> f (eff a b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap eff a b -> eff a b
forall (eff :: * -> * -> *) a b.
(ProvidesPosCaching eff, ContentHashable IO a, Store b) =>
eff a b -> eff a b
usingStore' f (eff a b)
f

-- | Bundles together a store with an identifier for the whole pipeline. If
-- identifier is Nothing, no caching will be performed.
data StoreWithId remoteCacher = StoreWithId CS.ContentStore remoteCacher (Maybe Int)

type LocalStoreWithId = StoreWithId Remote.NoCache

-- | A 'StoreWithId' with no remote caching
localStoreWithId :: CS.ContentStore -> Maybe Int -> LocalStoreWithId
localStoreWithId :: ContentStore -> Maybe Int -> LocalStoreWithId
localStoreWithId ContentStore
store Maybe Int
ident = ContentStore -> NoCache -> Maybe Int -> LocalStoreWithId
forall remoteCacher.
ContentStore
-> remoteCacher -> Maybe Int -> StoreWithId remoteCacher
StoreWithId ContentStore
store NoCache
Remote.NoCache Maybe Int
ident

instance (MonadIO m, MonadUnliftIO m, MonadMask m, Remote.Cacher m remoteCacher, HasKleisli m eff)
  => ProvidesCaching (Reader (StoreWithId remoteCacher) ~> eff) where
  usingStore :: (~>) (Reader (StoreWithId remoteCacher)) eff a b
-> (~>) (Reader (StoreWithId remoteCacher)) eff a b
usingStore =
    (StoreWithId remoteCacher -> eff a b -> eff a b)
-> (~>) (Reader (StoreWithId remoteCacher)) eff a b
-> (~>) (Reader (StoreWithId remoteCacher)) eff a b
forall t (eff :: * -> * -> *) a b (eff' :: * -> * -> *) a' b'.
(t -> eff a b -> eff' a' b')
-> (~>) (Reader t) eff a b -> (~>) (Reader t) eff' a' b'
mapReader_ ((StoreWithId remoteCacher -> eff a b -> eff a b)
 -> (~>) (Reader (StoreWithId remoteCacher)) eff a b
 -> (~>) (Reader (StoreWithId remoteCacher)) eff a b)
-> (StoreWithId remoteCacher -> eff a b -> eff a b)
-> (~>) (Reader (StoreWithId remoteCacher)) eff a b
-> (~>) (Reader (StoreWithId remoteCacher)) eff a b
forall a b. (a -> b) -> a -> b
$ \(StoreWithId ContentStore
store remoteCacher
remoteCacher Maybe Int
pipelineId) ->
    ((a -> m b) -> a -> m b) -> eff a b -> eff a b
forall (m :: * -> *) (eff :: * -> * -> *) a b a' b'.
HasKleisli m eff =>
((a -> m b) -> a' -> m b') -> eff a b -> eff a' b'
mapKleisli (((a -> m b) -> a -> m b) -> eff a b -> eff a b)
-> ((a -> m b) -> a -> m b) -> eff a b -> eff a b
forall a b. (a -> b) -> a -> b
$ \a -> m b
act a
input ->
      Maybe Int
-> CacherM m a b
-> ContentStore
-> remoteCacher
-> (a -> m b)
-> a
-> m b
forall (m :: * -> *) remoteCache i o.
(MonadIO m, MonadUnliftIO m, MonadMask m, Cacher m remoteCache) =>
Maybe Int
-> CacherM m i o
-> ContentStore
-> remoteCache
-> (i -> m o)
-> i
-> m o
CS.cacheKleisliIO
       Maybe Int
pipelineId (Int -> CacherM m a b
forall (m :: * -> *) i o.
(MonadIO m, ContentHashable IO i, Store o) =>
Int -> CacherM m i o
CS.defaultIOCacherWithIdent Int
1) ContentStore
store remoteCacher
remoteCacher
       a -> m b
act a
input

instance (Arrow eff, ProvidesCaching eff)
  => ProvidesCaching (Writer CachingContext ~> eff) where
  usingStore :: (~>) (Writer CachingContext) eff a b
-> (~>) (Writer CachingContext) eff a b
usingStore =
    (CachingContext -> eff a b -> eff a b)
-> (~>) (Writer CachingContext) eff a b
-> (~>) (Writer CachingContext) eff a b
forall w (eff :: * -> * -> *) a b (eff' :: * -> * -> *) a' b'.
(w -> eff a b -> eff' a' b')
-> (~>) (Writer w) eff a b -> (~>) (Writer w) eff' a' b'
mapWriter_ ((CachingContext -> eff a b -> eff a b)
 -> (~>) (Writer CachingContext) eff a b
 -> (~>) (Writer CachingContext) eff a b)
-> (CachingContext -> eff a b -> eff a b)
-> (~>) (Writer CachingContext) eff a b
-> (~>) (Writer CachingContext) eff a b
forall a b. (a -> b) -> a -> b
$ \CachingContext
newContext eff a b
eff ->
      (a -> (a, CachingContext)) -> eff a (a, CachingContext)
forall (a :: * -> * -> *) b c. Arrow a => (b -> c) -> a b c
arr (,CachingContext
newContext) eff a (a, CachingContext) -> eff (a, CachingContext) b -> eff a b
forall k (cat :: k -> k -> *) (a :: k) (b :: k) (c :: k).
Category cat =>
cat a b -> cat b c -> cat a c
>>> eff (a, CachingContext) b -> eff (a, CachingContext) b
forall (eff :: * -> * -> *) a b.
(ProvidesCaching eff, ContentHashable IO a, Store b) =>
eff a b -> eff a b
usingStore (eff a b
eff eff a b -> eff (a, CachingContext) a -> eff (a, CachingContext) b
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. ((a, CachingContext) -> a) -> eff (a, CachingContext) a
forall (a :: * -> * -> *) b c. Arrow a => (b -> c) -> a b c
arr (a, CachingContext) -> a
forall a b. (a, b) -> a
fst)
      -- New context is just added as phantom input to the underlying effect
instance (Arrow eff, ProvidesPosCaching eff)
  => ProvidesPosCaching (Writer CachingContext ~> eff) where
  usingStore' :: (~>) (Writer CachingContext) eff a b
-> (~>) (Writer CachingContext) eff a b
usingStore' =
    (CachingContext -> eff a b -> eff a b)
-> (~>) (Writer CachingContext) eff a b
-> (~>) (Writer CachingContext) eff a b
forall w (eff :: * -> * -> *) a b (eff' :: * -> * -> *) a' b'.
(w -> eff a b -> eff' a' b')
-> (~>) (Writer w) eff a b -> (~>) (Writer w) eff' a' b'
mapWriter_ ((CachingContext -> eff a b -> eff a b)
 -> (~>) (Writer CachingContext) eff a b
 -> (~>) (Writer CachingContext) eff a b)
-> (CachingContext -> eff a b -> eff a b)
-> (~>) (Writer CachingContext) eff a b
-> (~>) (Writer CachingContext) eff a b
forall a b. (a -> b) -> a -> b
$ \CachingContext
newContext eff a b
eff ->
      (a -> (a, CachingContext)) -> eff a (a, CachingContext)
forall (a :: * -> * -> *) b c. Arrow a => (b -> c) -> a b c
arr (,CachingContext
newContext) eff a (a, CachingContext) -> eff (a, CachingContext) b -> eff a b
forall k (cat :: k -> k -> *) (a :: k) (b :: k) (c :: k).
Category cat =>
cat a b -> cat b c -> cat a c
>>> eff (a, CachingContext) b -> eff (a, CachingContext) b
forall (eff :: * -> * -> *) a b.
(ProvidesPosCaching eff, ContentHashable IO a, Store b) =>
eff a b -> eff a b
usingStore' (eff a b
eff eff a b -> eff (a, CachingContext) a -> eff (a, CachingContext) b
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. ((a, CachingContext) -> a) -> eff (a, CachingContext) a
forall (a :: * -> * -> *) b c. Arrow a => (b -> c) -> a b c
arr (a, CachingContext) -> a
forall a b. (a, b) -> a
fst)

instance (ProvidesCaching eff) => ProvidesCaching (AutoIdent eff) where
  usingStore :: AutoIdent eff a b -> AutoIdent eff a b
usingStore (AutoIdent ArrowIdent -> eff a b
f) = (ArrowIdent -> eff a b) -> AutoIdent eff a b
forall (arr :: * -> * -> *) a b.
(ArrowIdent -> arr a b) -> AutoIdent arr a b
AutoIdent ((ArrowIdent -> eff a b) -> AutoIdent eff a b)
-> (ArrowIdent -> eff a b) -> AutoIdent eff a b
forall a b. (a -> b) -> a -> b
$ eff a b -> eff a b
forall (eff :: * -> * -> *) a b.
(ProvidesCaching eff, ContentHashable IO a, Store b) =>
eff a b -> eff a b
usingStore (eff a b -> eff a b)
-> (ArrowIdent -> eff a b) -> ArrowIdent -> eff a b
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. ArrowIdent -> eff a b
f
instance (Arrow eff, ProvidesCaching eff) => ProvidesPosCaching (AutoIdent eff) where
  usingStore' :: AutoIdent eff a b -> AutoIdent eff a b
usingStore' (AutoIdent ArrowIdent -> eff a b
f) = (ArrowIdent -> eff a b) -> AutoIdent eff a b
forall (arr :: * -> * -> *) a b.
(ArrowIdent -> arr a b) -> AutoIdent arr a b
AutoIdent ((ArrowIdent -> eff a b) -> AutoIdent eff a b)
-> (ArrowIdent -> eff a b) -> AutoIdent eff a b
forall a b. (a -> b) -> a -> b
$ \ArrowIdent
aid ->
    (a -> (a, ArrowIdent)) -> eff a (a, ArrowIdent)
forall (a :: * -> * -> *) b c. Arrow a => (b -> c) -> a b c
arr (,ArrowIdent
aid) eff a (a, ArrowIdent) -> eff (a, ArrowIdent) b -> eff a b
forall k (cat :: k -> k -> *) (a :: k) (b :: k) (c :: k).
Category cat =>
cat a b -> cat b c -> cat a c
>>> eff (a, ArrowIdent) b -> eff (a, ArrowIdent) b
forall (eff :: * -> * -> *) a b.
(ProvidesCaching eff, ContentHashable IO a, Store b) =>
eff a b -> eff a b
usingStore (ArrowIdent -> eff a b
f ArrowIdent
aid eff a b -> eff (a, ArrowIdent) a -> eff (a, ArrowIdent) b
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. ((a, ArrowIdent) -> a) -> eff (a, ArrowIdent) a
forall (a :: * -> * -> *) b c. Arrow a => (b -> c) -> a b c
arr (a, ArrowIdent) -> a
forall a b. (a, b) -> a
fst)

instance (ProvidesCaching core) => ProvidesCaching (Rope r m core) where
  usingStore :: Rope r m core a b -> Rope r m core a b
usingStore = (core a b -> core a b) -> Rope r m core a b -> Rope r m core a b
forall (core :: * -> * -> *) a b a' b' (r :: RopeRec)
       (m :: [Strand]).
(core a b -> core a' b')
-> Rope r m core a b -> Rope r m core a' b'
mapRopeCore core a b -> core a b
forall (eff :: * -> * -> *) a b.
(ProvidesCaching eff, ContentHashable IO a, Store b) =>
eff a b -> eff a b
usingStore
instance (ProvidesPosCaching core) => ProvidesPosCaching (Rope r m core) where
  usingStore' :: Rope r m core a b -> Rope r m core a b
usingStore' = (core a b -> core a b) -> Rope r m core a b -> Rope r m core a b
forall (core :: * -> * -> *) a b a' b' (r :: RopeRec)
       (m :: [Strand]).
(core a b -> core a' b')
-> Rope r m core a b -> Rope r m core a' b'
mapRopeCore core a b -> core a b
forall (eff :: * -> * -> *) a b.
(ProvidesPosCaching eff, ContentHashable IO a, Store b) =>
eff a b -> eff a b
usingStore'

-- | Any rope whose core provides caching can run cached tasks. The task is
-- identified by its position in the pipeline
caching' :: (ProvidesPosCaching core, CS.ContentHashable IO a, Show a, Store b)
         => Rope r mantle core a b -> Rope r mantle core a b
caching' :: Rope r mantle core a b -> Rope r mantle core a b
caching' = (core a b -> core a b)
-> Rope r mantle core a b -> Rope r mantle core a b
forall (core :: * -> * -> *) a b a' b' (r :: RopeRec)
       (m :: [Strand]).
(core a b -> core a' b')
-> Rope r m core a b -> Rope r m core a' b'
mapRopeCore core a b -> core a b
forall (eff :: * -> * -> *) a b.
(ProvidesPosCaching eff, ContentHashable IO a, Store b) =>
eff a b -> eff a b
usingStore'

-- | Any rope whose core provides caching can run cached tasks. The task is
-- identified by an explicit identifier
caching :: (Arrow core, ProvidesCaching core
           ,CS.ContentHashable IO ident, CS.ContentHashable IO a
           ,Store b)
        => ident -> Rope r mantle core a b -> Rope r mantle core a b
caching :: ident -> Rope r mantle core a b -> Rope r mantle core a b
caching ident
n Rope r mantle core a b
r = (a -> (a, ident)) -> Rope r mantle core a (a, ident)
forall (a :: * -> * -> *) b c. Arrow a => (b -> c) -> a b c
arr (,ident
n) Rope r mantle core a (a, ident)
-> Rope r mantle core (a, ident) b -> Rope r mantle core a b
forall k (cat :: k -> k -> *) (a :: k) (b :: k) (c :: k).
Category cat =>
cat a b -> cat b c -> cat a c
>>> (core (a, ident) b -> core (a, ident) b)
-> Rope r mantle core (a, ident) b
-> Rope r mantle core (a, ident) b
forall (core :: * -> * -> *) a b a' b' (r :: RopeRec)
       (m :: [Strand]).
(core a b -> core a' b')
-> Rope r m core a b -> Rope r m core a' b'
mapRopeCore core (a, ident) b -> core (a, ident) b
forall (eff :: * -> * -> *) a b.
(ProvidesCaching eff, ContentHashable IO a, Store b) =>
eff a b -> eff a b
usingStore (Rope r mantle core a b
r Rope r mantle core a b
-> Rope r mantle core (a, ident) a
-> Rope r mantle core (a, ident) b
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. ((a, ident) -> a) -> Rope r mantle core (a, ident) a
forall (a :: * -> * -> *) b c. Arrow a => (b -> c) -> a b c
arr (a, ident) -> a
forall a b. (a, b) -> a
fst)