-- |
-- Dynamic gas consumption for free monads.
module Control.Monad.Free.Gas
    ( limitGas
    ) where

import Control.Monad (when)
import Control.Monad.Free (Free, foldFree)
import Control.Monad.Trans.Class (lift)
import Control.Monad.Trans.Maybe (MaybeT (..), runMaybeT)
import Control.Monad.Trans.State (StateT, evalStateT)

import qualified Control.Monad.Trans.State as State

-- |
-- Similar to 'foldFree', but limit the number of steps the interpretation of a
-- free monad action may take. The number of steps is determined dynamically,
-- while the program is running. It is the cummulative result of the cost
-- function given to 'limitGas'. The limit is called the gas, by typical
-- /car analogy/.
--
-- If you want to predict, precisely, how much gas will be used by the
-- computation, then you cannot do this if the computation is monadic, only if
-- it is applicative. See the "Control.Applicative.Free.Gas" module.
--
-- Note that while there is still gas available, the interpretation may cause
-- effects, even if it runs out of gas later!
limitGas
    :: forall f m a
     . Monad m
    => (forall x. f x -> Word) -- ^ Given an action, its cost.
    -> (forall x. f x -> m x)  -- ^ Natural transformation.
    -> Free f a                -- ^ Program to run and cut off.
    -> Word                    -- ^ Initial gas.
    -> m (Maybe a)
limitGas cost nx =
    (runMaybeT .) . evalStateT . foldFree go
    where
    go :: f x -> StateT Word (MaybeT m) x
    go action = do
        gas <- State.get
        when (gas < cost action) $
            lift $ MaybeT (pure Nothing)
        State.modify' (subtract (cost action))
        lift . lift $ nx action