.hs

{-# Language RebindableSyntax #-}
{-# Language FlexibleInstances #-}
{-# Language NoMonomorphismRestriction #-}
{-# Language TemplateHaskell #-}

import qualified Prelude as P
import Control.Monad.IO.Class
import Prelude hiding ((>>=), return)
import Control.Lens hiding (imap)

class IxFunctor f where
  imap :: (a -> b) -> f j k a -> f j k b
class IxFunctor m => IxPointed m where
  ireturn :: a -> m i i a
class IxPointed m => IxApplicative m where
  iap :: m i j (a -> b) -> m j k a -> m i k b
class IxApplicative m => IxMonad m where
  ibind :: (a -> m j k b) -> m i j a -> m i k b

class IxMonad m => IxMonadState m where
  iget :: m i i i
  iput :: j -> m i j ()

newtype IxStateT m i j a = IxStateT { runIxStateT :: i -> m (a, j) }

instance Monad m => IxFunctor (IxStateT m) where
  imap f s = IxStateT $ \i -> fmap (first f) (runIxStateT s i)
    where first f (a, b) = (f a, b)

instance Monad m => IxPointed (IxStateT m) where
  ireturn a = IxStateT $ \i -> P.return (a, i)

instance Monad m => IxApplicative (IxStateT m) where
  f `iap` v = IxStateT $ \i ->
    runIxStateT f i P.>>= \(f', j) ->
    runIxStateT v j P.>>= \(v', k) ->
    P.return (f' v', k)

instance Monad m => IxMonad (IxStateT m) where
  f `ibind` v = IxStateT $ \i ->
    runIxStateT v i P.>>= \(v', j)  ->
    runIxStateT (f v') j
(>>>=) :: IxMonad m => m i j a -> (a -> m j k b) -> m i k b
(>>>=) = flip ibind

instance Monad m => IxMonadState (IxStateT m) where
  iget = IxStateT $ \i -> P.return (i, i)
  iput v = IxStateT $ \_ -> P.return ((), v)

instance Monad m => Functor (IxStateT m i i) where
  fmap = imap
instance Monad m => Applicative (IxStateT m i i) where
  pure = ireturn
  (<*>) = iap
instance Monad m => Monad (IxStateT m i i) where
  return = ireturn
  m >>= k = ibind k m 
class IxMonadTrans t where
  ilift :: Monad m => m a -> t m i i a
instance IxMonadTrans IxStateT where
  ilift m = IxStateT $ \i -> (flip (,) i) <$> m
instance MonadIO m => MonadIO (IxStateT m i i) where
  liftIO = ilift . liftIO

data Base64 = Base64 String -- Not bytestring because it isn't really necessary for an example
data PrivateKey = PrivateKey String
data Auth = Auth Base64

data Session' a = Session { _keyFetchToken :: a }
makeLenses ''Session'
type SessionWithToken = Session' Base64
type SessionWithoutToken = Session' ()

fetchKey :: Base64 -> IO PrivateKey
fetchKey = undefined
toAuth :: PrivateKey -> Auth
toAuth = undefined

getKeyFetchToken :: IxMonadState m => m SessionWithToken SessionWithoutToken Base64
getKeyFetchToken = do
  session <- iget
  iput $ session & keyFetchToken .~ ()
  return $ session ^. keyFetchToken
  where
    (>>=) = (>>>=)
    f >> g = f >>>= \_ -> g
    return = ireturn

getAuthorization :: (IxMonadState m, MonadIO (m SessionWithoutToken SessionWithoutToken)) => m SessionWithToken SessionWithoutToken Auth
getAuthorization = do
  token <- getKeyFetchToken
  fetched <- liftIO $ fetchKey token
  return $ toAuth fetched
  where
    (>>=) = (>>>=)
    f >> g = f >>>= \_ -> g
    return = ireturn