{-# LANGUAGE TypeFamilies, StandaloneDeriving, UndecidableInstances, OverloadedStrings, DerivingStrategies, MagicHash #-}

module Control.Monad.Credit.CounterM (CounterM, runCounterM, CounterT, runCounterT) where

import Prelude hiding (lookup)
import Control.Monad.Except
import Control.Monad.Identity
import Control.Monad.State.Lazy
import Control.Monad.ST.Trans

import Control.Monad.Credit.Base

-- Run computation in a state monad

-- State of the computation
-- credits: The total credits consumed so far
newtype St = St Ticks
  deriving (St -> St -> Bool
(St -> St -> Bool) -> (St -> St -> Bool) -> Eq St
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: St -> St -> Bool
== :: St -> St -> Bool
$c/= :: St -> St -> Bool
/= :: St -> St -> Bool
Eq, Eq St
Eq St =>
(St -> St -> Ordering)
-> (St -> St -> Bool)
-> (St -> St -> Bool)
-> (St -> St -> Bool)
-> (St -> St -> Bool)
-> (St -> St -> St)
-> (St -> St -> St)
-> Ord St
St -> St -> Bool
St -> St -> Ordering
St -> St -> St
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
$ccompare :: St -> St -> Ordering
compare :: St -> St -> Ordering
$c< :: St -> St -> Bool
< :: St -> St -> Bool
$c<= :: St -> St -> Bool
<= :: St -> St -> Bool
$c> :: St -> St -> Bool
> :: St -> St -> Bool
$c>= :: St -> St -> Bool
>= :: St -> St -> Bool
$cmax :: St -> St -> St
max :: St -> St -> St
$cmin :: St -> St -> St
min :: St -> St -> St
Ord, Int -> St -> ShowS
[St] -> ShowS
St -> String
(Int -> St -> ShowS)
-> (St -> String) -> ([St] -> ShowS) -> Show St
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> St -> ShowS
showsPrec :: Int -> St -> ShowS
$cshow :: St -> String
show :: St -> String
$cshowList :: [St] -> ShowS
showList :: [St] -> ShowS
Show)

-- | An instance of the counter monad using ST to memoize thunks.
type CounterM s = CounterT s Identity

-- | A monad transformer on the counter monad.
-- Warning! This monad transformer includes the ST monad transformer and
-- should not be used with monads that can contain multiple answers,
-- like the list monad. Safe monads include the monads State, Reader, Writer,
-- Maybe and combinations of their corresponding monad transformers.
newtype CounterT s m a = CounterT { forall s (m :: * -> *) a.
CounterT s m a -> StateT St (ExceptT String (STT s m)) a
runT :: StateT St (ExceptT String (STT s m)) a }

instance Functor m => Functor (CounterT s m) where
  fmap :: forall a b. (a -> b) -> CounterT s m a -> CounterT s m b
fmap a -> b
f (CounterT StateT St (ExceptT String (STT s m)) a
m) = StateT St (ExceptT String (STT s m)) b -> CounterT s m b
forall s (m :: * -> *) a.
StateT St (ExceptT String (STT s m)) a -> CounterT s m a
CounterT ((a -> b)
-> StateT St (ExceptT String (STT s m)) a
-> StateT St (ExceptT String (STT s m)) b
forall a b.
(a -> b)
-> StateT St (ExceptT String (STT s m)) a
-> StateT St (ExceptT String (STT s m)) b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f StateT St (ExceptT String (STT s m)) a
m)

instance Monad m => Applicative (CounterT s m) where
  pure :: forall a. a -> CounterT s m a
pure = StateT St (ExceptT String (STT s m)) a -> CounterT s m a
forall s (m :: * -> *) a.
StateT St (ExceptT String (STT s m)) a -> CounterT s m a
CounterT (StateT St (ExceptT String (STT s m)) a -> CounterT s m a)
-> (a -> StateT St (ExceptT String (STT s m)) a)
-> a
-> CounterT s m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> StateT St (ExceptT String (STT s m)) a
forall a. a -> StateT St (ExceptT String (STT s m)) a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
  CounterT StateT St (ExceptT String (STT s m)) (a -> b)
f <*> :: forall a b.
CounterT s m (a -> b) -> CounterT s m a -> CounterT s m b
<*> CounterT StateT St (ExceptT String (STT s m)) a
x = StateT St (ExceptT String (STT s m)) b -> CounterT s m b
forall s (m :: * -> *) a.
StateT St (ExceptT String (STT s m)) a -> CounterT s m a
CounterT (StateT St (ExceptT String (STT s m)) (a -> b)
f StateT St (ExceptT String (STT s m)) (a -> b)
-> StateT St (ExceptT String (STT s m)) a
-> StateT St (ExceptT String (STT s m)) b
forall a b.
StateT St (ExceptT String (STT s m)) (a -> b)
-> StateT St (ExceptT String (STT s m)) a
-> StateT St (ExceptT String (STT s m)) b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> StateT St (ExceptT String (STT s m)) a
x)

instance Monad m => Monad (CounterT s m) where
  CounterT StateT St (ExceptT String (STT s m)) a
m >>= :: forall a b.
CounterT s m a -> (a -> CounterT s m b) -> CounterT s m b
>>= a -> CounterT s m b
f = StateT St (ExceptT String (STT s m)) b -> CounterT s m b
forall s (m :: * -> *) a.
StateT St (ExceptT String (STT s m)) a -> CounterT s m a
CounterT (StateT St (ExceptT String (STT s m)) a
m StateT St (ExceptT String (STT s m)) a
-> (a -> StateT St (ExceptT String (STT s m)) b)
-> StateT St (ExceptT String (STT s m)) b
forall a b.
StateT St (ExceptT String (STT s m)) a
-> (a -> StateT St (ExceptT String (STT s m)) b)
-> StateT St (ExceptT String (STT s m)) b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= CounterT s m b -> StateT St (ExceptT String (STT s m)) b
forall s (m :: * -> *) a.
CounterT s m a -> StateT St (ExceptT String (STT s m)) a
runT (CounterT s m b -> StateT St (ExceptT String (STT s m)) b)
-> (a -> CounterT s m b)
-> a
-> StateT St (ExceptT String (STT s m)) b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> CounterT s m b
f)

instance Monad m => MonadError String (CounterT s m) where
  throwError :: forall a. String -> CounterT s m a
throwError String
e = StateT St (ExceptT String (STT s m)) a -> CounterT s m a
forall s (m :: * -> *) a.
StateT St (ExceptT String (STT s m)) a -> CounterT s m a
CounterT (String -> StateT St (ExceptT String (STT s m)) a
forall a. String -> StateT St (ExceptT String (STT s m)) a
forall e (m :: * -> *) a. MonadError e m => e -> m a
throwError String
e)
  catchError :: forall a.
CounterT s m a -> (String -> CounterT s m a) -> CounterT s m a
catchError (CounterT StateT St (ExceptT String (STT s m)) a
m) String -> CounterT s m a
h = StateT St (ExceptT String (STT s m)) a -> CounterT s m a
forall s (m :: * -> *) a.
StateT St (ExceptT String (STT s m)) a -> CounterT s m a
CounterT (StateT St (ExceptT String (STT s m)) a
-> (String -> StateT St (ExceptT String (STT s m)) a)
-> StateT St (ExceptT String (STT s m)) a
forall a.
StateT St (ExceptT String (STT s m)) a
-> (String -> StateT St (ExceptT String (STT s m)) a)
-> StateT St (ExceptT String (STT s m)) a
forall e (m :: * -> *) a.
MonadError e m =>
m a -> (e -> m a) -> m a
catchError StateT St (ExceptT String (STT s m)) a
m (CounterT s m a -> StateT St (ExceptT String (STT s m)) a
forall s (m :: * -> *) a.
CounterT s m a -> StateT St (ExceptT String (STT s m)) a
runT (CounterT s m a -> StateT St (ExceptT String (STT s m)) a)
-> (String -> CounterT s m a)
-> String
-> StateT St (ExceptT String (STT s m)) a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> CounterT s m a
h))

instance Monad m => MonadState St (CounterT s m) where
  get :: CounterT s m St
get = StateT St (ExceptT String (STT s m)) St -> CounterT s m St
forall s (m :: * -> *) a.
StateT St (ExceptT String (STT s m)) a -> CounterT s m a
CounterT StateT St (ExceptT String (STT s m)) St
forall s (m :: * -> *). MonadState s m => m s
get
  put :: St -> CounterT s m ()
put St
s = StateT St (ExceptT String (STT s m)) () -> CounterT s m ()
forall s (m :: * -> *) a.
StateT St (ExceptT String (STT s m)) a -> CounterT s m a
CounterT (St -> StateT St (ExceptT String (STT s m)) ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put St
s)

instance MonadTrans (CounterT s) where
  lift :: forall (m :: * -> *) a. Monad m => m a -> CounterT s m a
lift = StateT St (ExceptT String (STT s m)) a -> CounterT s m a
forall s (m :: * -> *) a.
StateT St (ExceptT String (STT s m)) a -> CounterT s m a
CounterT (StateT St (ExceptT String (STT s m)) a -> CounterT s m a)
-> (m a -> StateT St (ExceptT String (STT s m)) a)
-> m a
-> CounterT s m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ExceptT String (STT s m) a
-> StateT St (ExceptT String (STT s m)) a
forall (m :: * -> *) a. Monad m => m a -> StateT St m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (ExceptT String (STT s m) a
 -> StateT St (ExceptT String (STT s m)) a)
-> (m a -> ExceptT String (STT s m) a)
-> m a
-> StateT St (ExceptT String (STT s m)) a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. STT s m a -> ExceptT String (STT s m) a
forall (m :: * -> *) a. Monad m => m a -> ExceptT String m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (STT s m a -> ExceptT String (STT s m) a)
-> (m a -> STT s m a) -> m a -> ExceptT String (STT s m) a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. m a -> STT s m a
forall (m :: * -> *) a. Monad m => m a -> STT s m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift

liftST :: Monad m => STT s m a -> CounterT s m a
liftST :: forall (m :: * -> *) s a. Monad m => STT s m a -> CounterT s m a
liftST = StateT St (ExceptT String (STT s m)) a -> CounterT s m a
forall s (m :: * -> *) a.
StateT St (ExceptT String (STT s m)) a -> CounterT s m a
CounterT (StateT St (ExceptT String (STT s m)) a -> CounterT s m a)
-> (STT s m a -> StateT St (ExceptT String (STT s m)) a)
-> STT s m a
-> CounterT s m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ExceptT String (STT s m) a
-> StateT St (ExceptT String (STT s m)) a
forall (m :: * -> *) a. Monad m => m a -> StateT St m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (ExceptT String (STT s m) a
 -> StateT St (ExceptT String (STT s m)) a)
-> (STT s m a -> ExceptT String (STT s m) a)
-> STT s m a
-> StateT St (ExceptT String (STT s m)) a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. STT s m a -> ExceptT String (STT s m) a
forall (m :: * -> *) a. Monad m => m a -> ExceptT String m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift

instance Monad m => MonadFail (CounterT s m) where
  fail :: forall a. String -> CounterT s m a
fail String
e = String -> CounterT s m a
forall a. String -> CounterT s m a
forall e (m :: * -> *) a. MonadError e m => e -> m a
throwError String
e

instance Monad m => MonadCount (CounterT s m) where
  tick :: CounterT s m ()
tick = do
    (St Ticks
c) <- CounterT s m St
forall s (m :: * -> *). MonadState s m => m s
get
    St -> CounterT s m ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put (Ticks -> St
St (Ticks
c Ticks -> Ticks -> Ticks
forall a. Num a => a -> a -> a
+ Ticks
1))

instance Monad m => MonadLazy (CounterT s m) where
  {-# SPECIALIZE instance MonadLazy (CounterT s Identity) #-}
  {-# SPECIALIZE instance MonadLazy (CounterT s (State st)) #-}
  data Thunk (CounterT s m) t b = Thunk !(STRef s (Either (t b) b))
  delay :: forall (t :: * -> *) a.
t a -> CounterT s m (Thunk (CounterT s m) t a)
delay t a
a = do
    STRef s (Either (t a) a)
s <- STT s m (STRef s (Either (t a) a))
-> CounterT s m (STRef s (Either (t a) a))
forall (m :: * -> *) s a. Monad m => STT s m a -> CounterT s m a
liftST (STT s m (STRef s (Either (t a) a))
 -> CounterT s m (STRef s (Either (t a) a)))
-> STT s m (STRef s (Either (t a) a))
-> CounterT s m (STRef s (Either (t a) a))
forall a b. (a -> b) -> a -> b
$ Either (t a) a -> STT s m (STRef s (Either (t a) a))
forall (m :: * -> *) a s. Applicative m => a -> STT s m (STRef s a)
newSTRef (t a -> Either (t a) a
forall a b. a -> Either a b
Left t a
a)
    Thunk (CounterT s m) t a -> CounterT s m (Thunk (CounterT s m) t a)
forall a. a -> CounterT s m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (STRef s (Either (t a) a) -> Thunk (CounterT s m) t a
forall s (m :: * -> *) (t :: * -> *) b.
STRef s (Either (t b) b) -> Thunk (CounterT s m) t b
Thunk STRef s (Either (t a) a)
s)
  force :: forall (t :: * -> *) a.
HasStep t (CounterT s m) =>
Thunk (CounterT s m) t a -> CounterT s m a
force (Thunk STRef s (Either (t a) a)
t) = do
    Either (t a) a
t' <- STT s m (Either (t a) a) -> CounterT s m (Either (t a) a)
forall (m :: * -> *) s a. Monad m => STT s m a -> CounterT s m a
liftST (STT s m (Either (t a) a) -> CounterT s m (Either (t a) a))
-> STT s m (Either (t a) a) -> CounterT s m (Either (t a) a)
forall a b. (a -> b) -> a -> b
$ STRef s (Either (t a) a) -> STT s m (Either (t a) a)
forall (m :: * -> *) s a. Applicative m => STRef s a -> STT s m a
readSTRef STRef s (Either (t a) a)
t
    case Either (t a) a
t' of
      Left t a
a -> do
        a
b <- t a -> CounterT s m a
forall a. t a -> CounterT s m a
forall {k} (t :: k -> *) (m :: k -> *) (a :: k).
HasStep t m =>
t a -> m a
step t a
a
        STT s m () -> CounterT s m ()
forall (m :: * -> *) s a. Monad m => STT s m a -> CounterT s m a
liftST (STT s m () -> CounterT s m ()) -> STT s m () -> CounterT s m ()
forall a b. (a -> b) -> a -> b
$ STRef s (Either (t a) a) -> Either (t a) a -> STT s m ()
forall (m :: * -> *) s a.
Applicative m =>
STRef s a -> a -> STT s m ()
writeSTRef STRef s (Either (t a) a)
t (a -> Either (t a) a
forall a b. b -> Either a b
Right a
b)
        a -> CounterT s m a
forall a. a -> CounterT s m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
b
      Right a
b -> a -> CounterT s m a
forall a. a -> CounterT s m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
b
  lazymatch :: forall (t :: * -> *) a b.
Thunk (CounterT s m) t a
-> (a -> CounterT s m b)
-> (t a -> CounterT s m b)
-> CounterT s m b
lazymatch (Thunk STRef s (Either (t a) a)
t) a -> CounterT s m b
f t a -> CounterT s m b
g = do
    Either (t a) a
t' <- STT s m (Either (t a) a) -> CounterT s m (Either (t a) a)
forall (m :: * -> *) s a. Monad m => STT s m a -> CounterT s m a
liftST (STT s m (Either (t a) a) -> CounterT s m (Either (t a) a))
-> STT s m (Either (t a) a) -> CounterT s m (Either (t a) a)
forall a b. (a -> b) -> a -> b
$ STRef s (Either (t a) a) -> STT s m (Either (t a) a)
forall (m :: * -> *) s a. Applicative m => STRef s a -> STT s m a
readSTRef STRef s (Either (t a) a)
t
    case Either (t a) a
t' of
      Right a
b -> a -> CounterT s m b
f a
b
      Left t a
a -> t a -> CounterT s m b
g t a
a

instance Monad m => MonadCredit (CounterT s m) where
  {-# SPECIALIZE instance MonadCredit (CounterT s Identity) #-}
  {-# SPECIALIZE instance MonadCredit (CounterT s (State st)) #-}
  creditWith :: forall (t :: * -> *) a.
Thunk (CounterT s m) t a -> Credit -> CounterT s m ()
creditWith Thunk (CounterT s m) t a
_ Credit
_ = () -> CounterT s m ()
forall a. a -> CounterT s m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
  hasAtLeast :: forall (t :: * -> *) a.
Thunk (CounterT s m) t a -> Credit -> CounterT s m ()
hasAtLeast Thunk (CounterT s m) t a
_ Credit
_ = () -> CounterT s m ()
forall a. a -> CounterT s m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

instance Monad m => MonadInherit (CounterT s m) where
  {-# SPECIALIZE instance MonadInherit (CounterT s Identity) #-}
  {-# SPECIALIZE instance MonadInherit (CounterT s (State st)) #-}
  creditAllTo :: forall (t :: * -> *) a. Thunk (CounterT s m) t a -> CounterT s m ()
creditAllTo Thunk (CounterT s m) t a
_ = () -> CounterT s m ()
forall a. a -> CounterT s m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

runStateT' :: Monad m => StateT St m a -> m (a, Ticks)
runStateT' :: forall (m :: * -> *) a. Monad m => StateT St m a -> m (a, Ticks)
runStateT' StateT St m a
m = do
  (a
a, St Ticks
c) <- StateT St m a -> St -> m (a, St)
forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
runStateT StateT St m a
m (St -> m (a, St)) -> St -> m (a, St)
forall a b. (a -> b) -> a -> b
$ Ticks -> St
St Ticks
0
  (a, Ticks) -> m (a, Ticks)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a
a, Ticks
c)

runCounterT :: Monad m => (forall s. CounterT s m a) -> m (Either String (a, Ticks))
runCounterT :: forall (m :: * -> *) a.
Monad m =>
(forall s. CounterT s m a) -> m (Either String (a, Ticks))
runCounterT forall s. CounterT s m a
m = (forall s. STT s m (Either String (a, Ticks)))
-> m (Either String (a, Ticks))
forall (m :: * -> *) a. Monad m => (forall s. STT s m a) -> m a
runSTT ((forall s. STT s m (Either String (a, Ticks)))
 -> m (Either String (a, Ticks)))
-> (forall s. STT s m (Either String (a, Ticks)))
-> m (Either String (a, Ticks))
forall a b. (a -> b) -> a -> b
$ ExceptT String (STT s m) (a, Ticks)
-> STT s m (Either String (a, Ticks))
forall e (m :: * -> *) a. ExceptT e m a -> m (Either e a)
runExceptT (ExceptT String (STT s m) (a, Ticks)
 -> STT s m (Either String (a, Ticks)))
-> ExceptT String (STT s m) (a, Ticks)
-> STT s m (Either String (a, Ticks))
forall a b. (a -> b) -> a -> b
$ StateT St (ExceptT String (STT s m)) a
-> ExceptT String (STT s m) (a, Ticks)
forall (m :: * -> *) a. Monad m => StateT St m a -> m (a, Ticks)
runStateT' (StateT St (ExceptT String (STT s m)) a
 -> ExceptT String (STT s m) (a, Ticks))
-> StateT St (ExceptT String (STT s m)) a
-> ExceptT String (STT s m) (a, Ticks)
forall a b. (a -> b) -> a -> b
$ CounterT s m a -> StateT St (ExceptT String (STT s m)) a
forall s (m :: * -> *) a.
CounterT s m a -> StateT St (ExceptT String (STT s m)) a
runT CounterT s m a
forall s. CounterT s m a
m

runCounterM :: (forall s. CounterM s a) -> Either String (a, Ticks)
runCounterM :: forall a. (forall s. CounterM s a) -> Either String (a, Ticks)
runCounterM forall s. CounterM s a
m = Identity (Either String (a, Ticks)) -> Either String (a, Ticks)
forall a. Identity a -> a
runIdentity (Identity (Either String (a, Ticks)) -> Either String (a, Ticks))
-> Identity (Either String (a, Ticks)) -> Either String (a, Ticks)
forall a b. (a -> b) -> a -> b
$ (forall s. STT s Identity (Either String (a, Ticks)))
-> Identity (Either String (a, Ticks))
forall (m :: * -> *) a. Monad m => (forall s. STT s m a) -> m a
runSTT ((forall s. STT s Identity (Either String (a, Ticks)))
 -> Identity (Either String (a, Ticks)))
-> (forall s. STT s Identity (Either String (a, Ticks)))
-> Identity (Either String (a, Ticks))
forall a b. (a -> b) -> a -> b
$ ExceptT String (STT s Identity) (a, Ticks)
-> STT s Identity (Either String (a, Ticks))
forall e (m :: * -> *) a. ExceptT e m a -> m (Either e a)
runExceptT (ExceptT String (STT s Identity) (a, Ticks)
 -> STT s Identity (Either String (a, Ticks)))
-> ExceptT String (STT s Identity) (a, Ticks)
-> STT s Identity (Either String (a, Ticks))
forall a b. (a -> b) -> a -> b
$ StateT St (ExceptT String (STT s Identity)) a
-> ExceptT String (STT s Identity) (a, Ticks)
forall (m :: * -> *) a. Monad m => StateT St m a -> m (a, Ticks)
runStateT' (StateT St (ExceptT String (STT s Identity)) a
 -> ExceptT String (STT s Identity) (a, Ticks))
-> StateT St (ExceptT String (STT s Identity)) a
-> ExceptT String (STT s Identity) (a, Ticks)
forall a b. (a -> b) -> a -> b
$ CounterT s Identity a
-> StateT St (ExceptT String (STT s Identity)) a
forall s (m :: * -> *) a.
CounterT s m a -> StateT St (ExceptT String (STT s m)) a
runT CounterT s Identity a
forall s. CounterM s a
m