Codecs.hs

{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE DerivingVia #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE PartialTypeSignatures #-}
{-# LANGUAGE QuantifiedConstraints #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE TupleSections #-}

module Codecs where

import Control.Applicative (liftA2)
import Data.Functor (($>))
import Data.Kind (Type)

data CodecFor r w i o = Codec
  { read :: r o
  , write :: i -> w o
  }
  deriving stock (Functor)

type Codec r w a = CodecFor r w a a

data BiMap e a b = BiMap
  { forward :: a -> Either e b
  , backward :: b -> Either e a
  }

flipBiMap :: BiMap e b1 b2 -> BiMap e b2 b1
flipBiMap (BiMap fw bw) = BiMap bw fw

toCodec ::
  forall e a b m.
  Monad m =>
  BiMap e a b ->
  Codec
    (StateT [a] (ExceptT (Maybe e) m))
    (WriterT [a] (ExceptT e m))
    b
toCodec (BiMap fw bw) = Codec readOne writeOne
 where
  takeNext =
    get >>= \case
      [] -> pure Nothing
      (x : xs) -> put xs $> Just x
  readOne =
    takeNext >>= \case
      Nothing -> lift $ throw Nothing
      Just i -> case fw i of
        Left e -> lift $ throw (Just e)
        Right b -> pure b
  writeOne b = case bw b of
    Right a -> tell [a] $> b
    Left e -> lift $ throw e

-- Transformers
class
  (forall m. Monad m => Monad (trans m)) =>
  MonadTrans (trans :: (Type -> Type) -> Type -> Type)
  where
  lift :: Monad m => m a -> trans m a

newtype ReaderT v m a = ReaderT (v -> m a)
  deriving stock (Functor)
instance (Applicative m) => Applicative (ReaderT v m) where
  pure a = ReaderT $ \_ -> pure a
  ReaderT rF <*> ReaderT rA = ReaderT $ \v -> rF v <*> rA v
instance (Monad m) => Monad (ReaderT v m) where
  ReaderT rA >>= mkRB = ReaderT $ \v -> do
    a <- rA v
    let ReaderT rB = mkRB a
    rB v
instance MonadTrans (ReaderT b) where
  lift m = ReaderT (const m)

ask :: Applicative m => ReaderT a m a
ask = ReaderT $ \v -> pure v

newtype StateT s m a = StateT (s -> m (s, a))
  deriving stock (Functor)
instance Monad m => Applicative (StateT s m) where
  pure a = StateT (\(!s) -> pure (s, a))
  StateT stF <*> StateT stA = StateT $ \(!s0) -> do
    (!s1, f) <- stF s0
    (!s2, a) <- stA s1
    pure (s2, f a)
instance Monad m => Monad (StateT s m) where
  StateT stA >>= mkStB = StateT $ \(!s0) -> do
    (s1, a) <- stA s0
    let StateT stB = mkStB a
    stB s1
instance MonadTrans (StateT s) where
  lift m = StateT $ \(!s) -> (s,) <$> m

get :: Applicative m => StateT a m a
get = StateT $ \(!s) -> pure (s, s)

put :: Applicative m => s -> StateT s m ()
put !a = StateT $ \_ -> pure (a, ())

data Writer w a = Writer {-# UNPACK #-} !w {-# UNPACK #-} !a
  deriving stock (Functor)
instance (Monoid w) => Applicative (Writer w) where
  pure a = Writer mempty a
  Writer w1 f <*> Writer w2 a = Writer (w1 <> w2) (f a)
instance (Monoid w) => Monad (Writer w) where
  Writer w1 a >>= mkWB =
    let Writer w2 b = mkWB a
    in Writer (w1 <> w2) b

newtype WriterT w m a = WriterT (m (Writer w a))
  deriving stock (Functor)
instance (Monoid w, Monad m) => Applicative (WriterT w m) where
  pure a = WriterT (pure (Writer mempty a))
  WriterT wF <*> WriterT wA = WriterT $ liftA2 (<*>) wF wA
instance (Monoid w, Monad m) => Monad (WriterT w m) where
  WriterT wA >>= mkWB = WriterT $ do
    Writer w0 a <- wA
    let WriterT wB = mkWB a
    (Writer w1 b) <- wB
    pure (Writer (w0 <> w1) b)
instance Monoid w => MonadTrans (WriterT w) where
  lift m = WriterT $ Writer mempty <$> m

tell :: Applicative m => w -> WriterT w m ()
tell w = WriterT (pure (Writer w ()))

newtype ExceptT e m a = ExceptT (m (Either e a))
  deriving stock (Functor)
instance Applicative m => Applicative (ExceptT e m) where
  pure = ExceptT . pure . pure
  ExceptT mbF <*> ExceptT mbA = ExceptT $ liftA2 (<*>) mbF mbA
instance Monad m => Monad (ExceptT e m) where
  ExceptT mbA >>= mkMbB = ExceptT $ do
    mbA >>= \case
      Left e -> pure (Left e)
      Right a -> case mkMbB a of
        ExceptT mbB -> mbB
instance MonadTrans (ExceptT e) where
  lift m = ExceptT $ pure <$> m

throw :: Applicative m => e -> ExceptT e m a
throw e = ExceptT (pure (Left e))

catchE ::
  (Monad m) =>
  -- | the inner computation
  ExceptT e m a ->
  -- | a handler for exceptions in the inner computation
  (e -> ExceptT e' m a) ->
  ExceptT e' m a
ExceptT action `catchE` handle = ExceptT $ do
  action >>= \case
    Left e -> let ExceptT m = handle e in m
    Right a -> pure $ Right a