{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE UndecidableInstances #-}

-- |
-- Module      : Jikka.Common.Alpha
-- Description : provides a monad to run alpha-conversion. / alpha 変換用のモナドを提供します。
-- Copyright   : (c) Kimiyuki Onaka, 2020
-- License     : Apache License 2.0
-- Maintainer  : kimiyuki95@gmail.com
-- Stability   : experimental
-- Portability : portable
--
-- `Jikka.Common.Alpha` provides a monad to run alpha-conversion. This monad has only a feature to make unique numbers.
module Jikka.Common.Alpha where

import Control.Arrow (first)
import Control.Monad.Except
import Control.Monad.Identity (Identity (..))
import Control.Monad.Reader
import Control.Monad.Signatures
import Control.Monad.State.Strict
import Control.Monad.Writer.Strict
import Data.Unique
import Language.Haskell.TH (Q)

class Monad m => MonadAlpha m where
  nextCounter :: m Int

newtype AlphaT m a = AlphaT {runAlphaT :: Int -> m (a, Int)}

instance Monad m => MonadAlpha (AlphaT m) where
  nextCounter = AlphaT (\i -> return (i, i + 1))

instance Functor m => Functor (AlphaT m) where
  fmap f (AlphaT x) = AlphaT (\i -> fmap (first f) (x i))

instance Monad m => Applicative (AlphaT m) where
  pure x = AlphaT (\i -> pure (x, i))
  AlphaT f <*> AlphaT x = AlphaT $ \i -> do
    (f, i) <- f i
    (x, i) <- x i
    return (f x, i)

instance Monad m => Monad (AlphaT m) where
  AlphaT x >>= f = AlphaT $ \i -> do
    (x, i) <- x i
    runAlphaT (f x) i

instance MonadFix m => MonadFix (AlphaT m) where
  mfix f = AlphaT (\i -> mfix (\x -> runAlphaT (f (fst x)) i))

liftCatch :: Catch e m (a, Int) -> Catch e (AlphaT m) a
liftCatch catchE m h = AlphaT (\i -> runAlphaT m i `catchE` \e -> runAlphaT (h e) i)

instance MonadTrans AlphaT where
  lift m = AlphaT $ \i -> do
    a <- m
    return (a, i)

instance MonadError e m => MonadError e (AlphaT m) where
  throwError = lift . throwError
  catchError = liftCatch catchError

instance MonadIO m => MonadIO (AlphaT m) where
  liftIO = lift . liftIO

evalAlphaT :: Functor m => AlphaT m a -> Int -> m a
evalAlphaT f i = fst <$> runAlphaT f i

instance MonadAlpha m => MonadAlpha (ExceptT e m) where
  nextCounter = lift nextCounter

instance MonadAlpha m => MonadAlpha (ReaderT r m) where
  nextCounter = lift nextCounter

instance MonadAlpha m => MonadAlpha (StateT s m) where
  nextCounter = lift nextCounter

instance (MonadAlpha m, Monoid w) => MonadAlpha (WriterT w m) where
  nextCounter = lift nextCounter

evalAlpha :: AlphaT Identity a -> Int -> a
evalAlpha f i = runIdentity (evalAlphaT f i)

resetAlphaT :: Monad m => Int -> AlphaT m ()
resetAlphaT i = AlphaT $ \_ -> return ((), i)

instance MonadAlpha IO where
  nextCounter = hashUnique <$> newUnique

instance MonadAlpha Q where
  nextCounter = liftIO nextCounter