{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE TupleSections #-}
module Control.Schedule.Future where
import qualified Data.Map.Strict as M
import qualified Data.Set as S
import Control.Lens (IndexedTraversal', Lens', at,
contains, indices, (%%~),
(%~), (.~), (?~), (^.))
import Control.Lens.TH (makeLensesFor, makePrisms)
import Control.Monad.Trans.State.Strict (runState, state)
import Data.Function ((&))
import Data.Functor.Compose (Compose (..))
import Data.Schedule (Schedule, Task, TickDelta,
after, cancel_)
import GHC.Generics (Generic)
import GHC.Stack (HasCallStack)
import Safe (fromJustNote)
type OSet a = S.Set a
type OMap k v = M.Map k v
data TimedResult tk r =
TimedOut !tk
| GotResult !r
deriving (Show, Read, Generic, Eq, Ord)
data SFuture wo ro =
SFWaiting !(OSet wo)
| SFResult !ro
deriving (Show, Read, Generic, Eq, Ord)
makePrisms ''SFuture
data SExpect wi ri tk = SExpect {
seExpects :: !(OMap wi (Task tk))
, seResults :: !(OMap wi (TimedResult tk ri))
} deriving (Show, Read, Generic, Eq, Ord)
makeLensesFor ((\x -> (x, "_" <> x)) <$> ["seExpects", "seResults"]) ''SExpect
instance Ord wi => Semigroup (SExpect wi ri tk) where
s1 <> s2 =
SExpect (seExpects s1 <> seExpects s2) (seResults s1 <> seResults s2)
instance Ord wi => Monoid (SExpect wi ri tk) where
mempty = SExpect mempty mempty
data SFStatus e = Expecting e | NotExpecting deriving (Show, Read, Generic, Eq, Ord)
type SFStatusFull wo tk = SFStatus (OSet wo, Task tk)
data SFError =
SFEAlreadyFinished
| SFEInvalidPrecondition {
sfePreExpect :: !(SFStatus ())
, sfePreActual :: !(SFStatus ())
}
deriving (Show, Read, Generic, Eq, Ord)
sCheckStatus
:: (HasCallStack, Ord wi, Ord wo)
=> wi
-> wo
-> Lens' s (SFuture wo r)
-> Lens' s (SExpect wi r tk)
-> s
-> SFStatusFull wo tk
sCheckStatus sfi sei lsf lse s =
case (s ^. lsf, s ^. lse . _seExpects . at sfi) of
(SFResult _ , Just _ ) -> error "SFuture result but SExpect expects"
(SFWaiting waiting, Just lt) -> if waiting ^. contains sei
then Expecting (waiting, lt)
else error "SFuture not waiting but SExpect expects"
(SFResult _ , Nothing) -> NotExpecting
(SFWaiting waiting, Nothing) -> if waiting ^. contains sei
then error "SFuture waiting but SExpect not expects"
else NotExpecting
sExpectFuture
:: (Ord wi, Ord wo)
=> TickDelta
-> tk
-> wi
-> wo
-> Lens' s (SFuture wo r)
-> Lens' s (SExpect wi r tk)
-> Lens' s (Schedule tk)
-> s
-> Either SFError s
sExpectFuture d t sfi sei lsf lse lsch s0 = case status of
Expecting _ -> Left $ SFEInvalidPrecondition NotExpecting (Expecting ())
NotExpecting -> case s0 ^. lsf of
SFWaiting sfWaiting -> do
let (lt, s1) = s0 & lsch %%~ after d t
Right
$ s1
& (lse . _seExpects . at sfi ?~ lt)
& (lsf .~ SFWaiting (sfWaiting & contains sei .~ True))
SFResult r -> do
Right $ s0 & lse . _seResults . at sfi ?~ GotResult r
where status = sCheckStatus sfi sei lsf lse s0
sExpectCancel
:: (Ord wi, Ord wo)
=> wi
-> wo
-> Lens' s (SFuture wo r)
-> Lens' s (SExpect wi r tk)
-> Lens' s (Schedule tk)
-> s
-> Either SFError s
sExpectCancel sfi sei lsf lse lsch s0 = case status of
NotExpecting -> Left $ SFEInvalidPrecondition (Expecting ()) NotExpecting
Expecting (sfWaiting, lt) -> do
Right
$ s0
& (lsch %~ (snd . cancel_ lt))
& (lse . _seExpects . at sfi .~ Nothing)
& (lsf .~ SFWaiting (sfWaiting & contains sei .~ False))
where status = sCheckStatus sfi sei lsf lse s0
sExpectTimeout
:: (HasCallStack, Ord wi, Ord wo)
=> tk
-> wi
-> wo
-> Lens' s (SFuture wo r)
-> Lens' s (SExpect wi r tk)
-> Lens' s (Schedule tk)
-> s
-> Either SFError s
sExpectTimeout tk sfi sei lsf lse lsch s0 = case status of
NotExpecting -> Left $ SFEInvalidPrecondition (Expecting ()) NotExpecting
Expecting (_, lt) -> do
let s1 = s0 & lse . _seResults . at sfi %~ \case
Just _ -> error "SExpect expects but also results"
Nothing -> Just (TimedOut tk)
sExpectCancel sfi sei lsf lse lsch s1
where status = sCheckStatus sfi sei lsf lse s0
sFutureResult
:: (Ord wi, Ord wo)
=> r
-> wi
-> Lens' s (SFuture wo r)
-> IndexedTraversal' wo s (SExpect wi r tk)
-> Lens' s (Schedule tk)
-> s
-> Either SFError s
sFutureResult r sfi lsf lsse lsch s0 = do
(waiting, s1) <- getCompose $ s0 & lsf %%~ \case
SFResult _ -> Compose (Left SFEAlreadyFinished)
SFWaiting w -> Compose (Right (w, SFResult r))
let sch0 = s1 ^. lsch
let (s2, sch1) = f waiting s1 sch0
let s3 = s2 & lsch .~ sch1
Right s3
where
f w s = runState $ s & lsse . indices (`S.member` w) %%~ g
g se = do
let SExpect {..} = se
let (lt', seExpects') = seExpects & at sfi %%~ (, Nothing)
lt = fromJustNote "SFuture idx not found in SExpect expects" lt'
state $ cancel_ lt
let seResults' = seResults & at sfi ?~ GotResult r
pure $ se { seExpects = seExpects', seResults = seResults' }