{-# LANGUAGE FunctionalDependencies,
             GeneralisedNewtypeDeriving,
             DerivingStrategies,
             UndecidableInstances #-}
module Parsley.Internal.Common.Fresh (
    VFreshT, HFreshT, VFresh, HFresh,
    runFreshT, runFresh,
    evalFreshT, evalFresh,
    execFreshT, execFresh,
    MonadFresh(..), construct, mapVFreshT,
    RunFreshT
  ) where

import Control.Applicative        (liftA2)
import Control.Monad.Fix          (MonadFix(..))
import Control.Monad.Identity     (Identity, runIdentity)
import Control.Monad.Reader.Class (MonadReader(..))
import Control.Monad.State.Class  (MonadState(..))
import Control.Monad.Trans        (MonadTrans(..), MonadIO(..))

-- Fresh operations
class Monad m => MonadFresh x m | m -> x where
  newVar :: m x
  newScope :: m a -> m a

construct :: MonadFresh x m => (x -> a) -> m a
construct f = fmap f newVar

mapVFreshT :: (m (a, x, x) -> n (b, x, x)) -> VFreshT x m a -> VFreshT x n b
mapVFreshT f m = vFreshT (\cur max -> f (unVFreshT m cur max))

class (Monad n, Monad m) => RunFreshT x n m | m -> x, m -> n where
  runFreshT :: m a -> x -> n (a, x)

evalFreshT :: RunFreshT x n m => m a -> x -> n a
evalFreshT m init = fst <$> runFreshT m init

execFreshT :: RunFreshT x n m => m a -> x -> n x
execFreshT m init = snd <$> runFreshT m init

-- Fresh type
type HFresh x = HFreshT x Identity
type VFresh x = VFreshT x Identity
-- TODO Nominals
newtype VFreshT x m a = VFreshT (FreshT x m a) deriving newtype (Functor, Applicative, Monad, MonadFix, MonadTrans, MonadIO, MonadReader r, MonadState s, RunFreshT x m)
newtype HFreshT x m a = HFreshT (FreshT x m a) deriving newtype (Functor, Applicative, Monad, MonadFix, MonadTrans, MonadIO, MonadReader r, MonadState s, RunFreshT x m)
newtype FreshT x m a = FreshT {unFreshT :: x -> x -> m (a, x, x)}

instance Monad n => RunFreshT x n (FreshT x n) where
  runFreshT k init =
    do (x, _, max) <- unFreshT k init init
       return (x, max)

runFresh :: RunFreshT x Identity m => m a -> x -> (a, x)
runFresh mx = runIdentity . runFreshT mx

evalFresh :: RunFreshT x Identity m => m a -> x -> a
evalFresh mx = runIdentity . evalFreshT mx

execFresh :: RunFreshT x Identity m => m a -> x -> x
execFresh mx = runIdentity . execFreshT mx

vFreshT :: (x -> x -> m (a, x, x)) -> VFreshT x m a
vFreshT = VFreshT . FreshT

unVFreshT :: VFreshT x m a -> x -> x -> m (a, x, x)
unVFreshT (VFreshT f) = unFreshT f

hFreshT :: (x -> x -> m (a, x, x)) -> HFreshT x m a
hFreshT = HFreshT . FreshT

unHFreshT :: HFreshT x m a -> x -> x -> m (a, x, x)
unHFreshT (HFreshT f) = unFreshT f

instance Functor f => Functor (FreshT x f) where
  {-# INLINE fmap #-}
  fmap f (FreshT k) = FreshT (\cur max -> fmap (\(x, cur', max') -> (f x, cur', max')) (k cur max))

instance Monad m => Applicative (FreshT x m) where
  {-# INLINE pure #-}
  pure x = FreshT (\cur max -> pure (x, cur, max))
  {-# INLINE liftA2 #-}
  liftA2 f (FreshT mx) (FreshT my) = FreshT (\cur max ->
    do (x, cur', max') <- mx cur max
       (y, cur'', max'') <- my cur' max'
       return (f x y, cur'', max''))

instance Monad m => Monad (FreshT x m) where
  {-# INLINE return #-}
  return = pure
  {-# INLINE (>>=) #-}
  (FreshT mx) >>= f = FreshT (\cur max ->
    do (x, cur', max') <- mx cur max
       unFreshT (f x) cur' max')

instance MonadFix m => MonadFix (FreshT x m) where
  {-# INLINE mfix #-}
  mfix f = FreshT (\cur max -> mfix (\ ~(x, _, _) -> unFreshT (f x) cur max))

instance MonadTrans (FreshT x) where
  {-# INLINE lift #-}
  lift m = FreshT (\cur max ->
    do x <- m
       return (x, cur, max))

instance (Monad m, Ord x, Enum x) => MonadFresh x (VFreshT x m) where
  newVar = vFreshT (\cur m -> return (cur, cur, max m cur))
  newScope scoped = vFreshT (\cur max ->
    do (x, _, max') <- unVFreshT scoped (succ cur) max
       return (x, cur, max'))

instance (Monad m, Ord x, Enum x) => MonadFresh x (HFreshT x m) where
  newVar = hFreshT (\cur m -> return (cur, succ cur, max m cur))
  newScope scoped = hFreshT (\cur max ->
    do (x, _, max') <- unHFreshT scoped cur max
       return (x, cur, max'))

instance MonadIO m => MonadIO (FreshT x m) where liftIO = lift . liftIO
instance MonadReader r m => MonadReader r (FreshT x m) where
  ask = lift ask
  local f m = FreshT (\cur next -> local f (unFreshT m cur next))
instance MonadState s m => MonadState s (FreshT x m) where
  get = lift get
  put = lift . put