STE.hs

{-# LANGUAGE DeriveDataTypeable         #-}
{-# LANGUAGE DerivingStrategies         #-}
{-# LANGUAGE FlexibleInstances          #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE RankNTypes                 #-}
{-# LANGUAGE ScopedTypeVariables        #-}
{-# LANGUAGE TypeFamilies               #-}
-- | Simplified implementation of the monad-ste package,
-- with a few extras.
module Internal.STE
    ( STE
    , throwSTE
    , runSTE
    , liftST
    , steToST
    , steToIO
    ) where

import Control.Exception
import Control.Monad.Catch     (MonadThrow(..))
import Control.Monad.Primitive (PrimMonad(..))
import Control.Monad.ST
import Control.Monad.ST.Unsafe
import Data.Typeable           (Typeable)

newtype InternalErr e = InternalErr e
    deriving stock (Typeable)

instance Show (InternalErr e) where
    show _ = "(InternalErr _)"

instance Typeable e => Exception (InternalErr e)

newtype STE e s a = STE (IO a)
    deriving newtype (Functor, Applicative, Monad)

instance PrimMonad (STE e s) where
    type PrimState (STE e s) = s
    primitive = liftST . primitive

liftST :: ST s a -> STE e s a
liftST st = STE (unsafeSTToIO st)

throwSTE :: Exception e => e -> STE e s a
throwSTE e = STE (throwIO (InternalErr e))

runSTE :: Exception e => (forall s. STE e s a) -> Either e a
runSTE ste = runST (steToST ste)

steToST :: Typeable e => STE e s a -> ST s (Either e a)
steToST (STE io) = unsafeIOToST $ do
    res <- try io
    case res of
        Left (InternalErr e) -> pure $ Left e
        Right v              -> pure $ Right v

steToIO :: forall e a. Exception e => STE e RealWorld a -> IO a
steToIO (STE io) = do
    res <- try io
    case res of
        Left (InternalErr (e :: e)) -> throwIO e
        Right v                     -> pure v

instance MonadThrow (STE SomeException s) where
    throwM = throwSTE . toException