-- This file is part of Hoppy.
--
-- Copyright 2015-2021 Bryan Gardiner <bog@khumba.net>
--
-- This program is free software: you can redistribute it and/or modify
-- it under the terms of the GNU Affero General Public License as published by
-- the Free Software Foundation, either version 3 of the License, or
-- (at your option) any later version.
--
-- This program is distributed in the hope that it will be useful,
-- but WITHOUT ANY WARRANTY; without even the implied warranty of
-- MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
-- GNU Affero General Public License for more details.
--
-- You should have received a copy of the GNU Affero General Public License
-- along with this program.  If not, see <http://www.gnu.org/licenses/>.

{-# LANGUAGE CPP, UndecidableInstances #-}

-- | A monad for consuming streams.
module Foreign.Hoppy.Generator.Common.Consume (
  MonadConsume (..),
  ConsumeT,
  runConsumeT,
  evalConsumeT,
  execConsumeT,
  Consume,
  runConsume,
  evalConsume,
  execConsume,
  ) where

#if !MIN_VERSION_base(4,8,0)
import Control.Applicative ((<*>), Applicative, pure)
#endif
import Control.Monad (ap, liftM)
import Control.Monad.Except (ExceptT)
import Control.Monad.Identity (Identity, runIdentity)
import Control.Monad.State (StateT, get, put, runStateT)
import Control.Monad.Trans (MonadTrans, lift)
import Data.Tuple (swap)

-- | A typeclass for monads that can consume items from a stream.
class Monad m => MonadConsume s m | m -> s where
  -- | Attempts to consume an item from the stream.  Returns an item if the
  -- stream is not empty.
  next :: m (Maybe s)

-- | A monad transformer for 'MonadConsume'.
newtype ConsumeT s m a = ConsumeT { ConsumeT s m a -> StateT [s] m a
getConsumeT :: StateT [s] m a }

instance Monad m => Functor (ConsumeT s m) where
  fmap :: (a -> b) -> ConsumeT s m a -> ConsumeT s m b
fmap = (a -> b) -> ConsumeT s m a -> ConsumeT s m b
forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
liftM

instance Monad m => Applicative (ConsumeT s m) where
  pure :: a -> ConsumeT s m a
pure = a -> ConsumeT s m a
forall (m :: * -> *) a. Monad m => a -> m a
return
  <*> :: ConsumeT s m (a -> b) -> ConsumeT s m a -> ConsumeT s m b
(<*>) = ConsumeT s m (a -> b) -> ConsumeT s m a -> ConsumeT s m b
forall (m :: * -> *) a b. Monad m => m (a -> b) -> m a -> m b
ap

instance Monad m => Monad (ConsumeT s m) where
  return :: a -> ConsumeT s m a
return = StateT [s] m a -> ConsumeT s m a
forall s (m :: * -> *) a. StateT [s] m a -> ConsumeT s m a
ConsumeT (StateT [s] m a -> ConsumeT s m a)
-> (a -> StateT [s] m a) -> a -> ConsumeT s m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> StateT [s] m a
forall (m :: * -> *) a. Monad m => a -> m a
return
  ConsumeT s m a
m >>= :: ConsumeT s m a -> (a -> ConsumeT s m b) -> ConsumeT s m b
>>= a -> ConsumeT s m b
f = StateT [s] m b -> ConsumeT s m b
forall s (m :: * -> *) a. StateT [s] m a -> ConsumeT s m a
ConsumeT (StateT [s] m b -> ConsumeT s m b)
-> StateT [s] m b -> ConsumeT s m b
forall a b. (a -> b) -> a -> b
$ ConsumeT s m b -> StateT [s] m b
forall s (m :: * -> *) a. ConsumeT s m a -> StateT [s] m a
getConsumeT (ConsumeT s m b -> StateT [s] m b)
-> (a -> ConsumeT s m b) -> a -> StateT [s] m b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> ConsumeT s m b
f (a -> StateT [s] m b) -> StateT [s] m a -> StateT [s] m b
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< ConsumeT s m a -> StateT [s] m a
forall s (m :: * -> *) a. ConsumeT s m a -> StateT [s] m a
getConsumeT ConsumeT s m a
m

instance MonadTrans (ConsumeT s) where
  lift :: m a -> ConsumeT s m a
lift = StateT [s] m a -> ConsumeT s m a
forall s (m :: * -> *) a. StateT [s] m a -> ConsumeT s m a
ConsumeT (StateT [s] m a -> ConsumeT s m a)
-> (m a -> StateT [s] m a) -> m a -> ConsumeT s m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. m a -> StateT [s] m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift

instance Monad m => MonadConsume s (ConsumeT s m) where
  next :: ConsumeT s m (Maybe s)
next = do
    [s]
stream <- ConsumeT s m [s]
forall (m :: * -> *) s. Monad m => ConsumeT s m [s]
get'
    case [s]
stream of
      [] -> Maybe s -> ConsumeT s m (Maybe s)
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe s
forall a. Maybe a
Nothing
      s
x:[s]
xs -> [s] -> ConsumeT s m ()
forall (m :: * -> *) s. Monad m => [s] -> ConsumeT s m ()
put' [s]
xs ConsumeT s m () -> ConsumeT s m (Maybe s) -> ConsumeT s m (Maybe s)
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Maybe s -> ConsumeT s m (Maybe s)
forall (m :: * -> *) a. Monad m => a -> m a
return (s -> Maybe s
forall a. a -> Maybe a
Just s
x)

instance MonadConsume s m => MonadConsume s (ExceptT e m) where
  next :: ExceptT e m (Maybe s)
next = m (Maybe s) -> ExceptT e m (Maybe s)
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift m (Maybe s)
forall s (m :: * -> *). MonadConsume s m => m (Maybe s)
next

instance MonadConsume s m => MonadConsume s (StateT d m) where
  next :: StateT d m (Maybe s)
next = m (Maybe s) -> StateT d m (Maybe s)
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift m (Maybe s)
forall s (m :: * -> *). MonadConsume s m => m (Maybe s)
next

-- | Runs the consume action, returning the remainder of the stream, and the
-- action's result.
runConsumeT :: Monad m => [s] -> ConsumeT s m a -> m ([s], a)
runConsumeT :: [s] -> ConsumeT s m a -> m ([s], a)
runConsumeT [s]
stream (ConsumeT StateT [s] m a
m) = (a, [s]) -> ([s], a)
forall a b. (a, b) -> (b, a)
swap ((a, [s]) -> ([s], a)) -> m (a, [s]) -> m ([s], a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> StateT [s] m a -> [s] -> m (a, [s])
forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
runStateT StateT [s] m a
m [s]
stream

-- | Runs the consume action, returning the action's result.
evalConsumeT :: Monad m => [s] -> ConsumeT s m a -> m a
evalConsumeT :: [s] -> ConsumeT s m a -> m a
evalConsumeT [s]
stream = (([s], a) -> a) -> m ([s], a) -> m a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ([s], a) -> a
forall a b. (a, b) -> b
snd (m ([s], a) -> m a)
-> (ConsumeT s m a -> m ([s], a)) -> ConsumeT s m a -> m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [s] -> ConsumeT s m a -> m ([s], a)
forall (m :: * -> *) s a.
Monad m =>
[s] -> ConsumeT s m a -> m ([s], a)
runConsumeT [s]
stream

-- | Runs the consume action, returning the remainder of the stream.
execConsumeT :: Monad m => [s] -> ConsumeT s m a -> m [s]
execConsumeT :: [s] -> ConsumeT s m a -> m [s]
execConsumeT [s]
stream = (([s], a) -> [s]) -> m ([s], a) -> m [s]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ([s], a) -> [s]
forall a b. (a, b) -> a
fst (m ([s], a) -> m [s])
-> (ConsumeT s m a -> m ([s], a)) -> ConsumeT s m a -> m [s]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [s] -> ConsumeT s m a -> m ([s], a)
forall (m :: * -> *) s a.
Monad m =>
[s] -> ConsumeT s m a -> m ([s], a)
runConsumeT [s]
stream

type Consume s = ConsumeT s Identity

runConsume :: [s] -> Consume s a -> ([s], a)
runConsume :: [s] -> Consume s a -> ([s], a)
runConsume [s]
stream Consume s a
m = Identity ([s], a) -> ([s], a)
forall a. Identity a -> a
runIdentity (Identity ([s], a) -> ([s], a)) -> Identity ([s], a) -> ([s], a)
forall a b. (a -> b) -> a -> b
$ [s] -> Consume s a -> Identity ([s], a)
forall (m :: * -> *) s a.
Monad m =>
[s] -> ConsumeT s m a -> m ([s], a)
runConsumeT [s]
stream Consume s a
m

evalConsume :: [s] -> Consume s a -> a
evalConsume :: [s] -> Consume s a -> a
evalConsume [s]
stream = ([s], a) -> a
forall a b. (a, b) -> b
snd (([s], a) -> a) -> (Consume s a -> ([s], a)) -> Consume s a -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [s] -> Consume s a -> ([s], a)
forall s a. [s] -> Consume s a -> ([s], a)
runConsume [s]
stream

execConsume :: [s] -> Consume s a -> [s]
execConsume :: [s] -> Consume s a -> [s]
execConsume [s]
stream = ([s], a) -> [s]
forall a b. (a, b) -> a
fst (([s], a) -> [s])
-> (Consume s a -> ([s], a)) -> Consume s a -> [s]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [s] -> Consume s a -> ([s], a)
forall s a. [s] -> Consume s a -> ([s], a)
runConsume [s]
stream

get' :: Monad m => ConsumeT s m [s]
get' :: ConsumeT s m [s]
get' = StateT [s] m [s] -> ConsumeT s m [s]
forall s (m :: * -> *) a. StateT [s] m a -> ConsumeT s m a
ConsumeT StateT [s] m [s]
forall s (m :: * -> *). MonadState s m => m s
get

put' :: Monad m => [s] -> ConsumeT s m ()
put' :: [s] -> ConsumeT s m ()
put' = StateT [s] m () -> ConsumeT s m ()
forall s (m :: * -> *) a. StateT [s] m a -> ConsumeT s m a
ConsumeT (StateT [s] m () -> ConsumeT s m ())
-> ([s] -> StateT [s] m ()) -> [s] -> ConsumeT s m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [s] -> StateT [s] m ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put