Ref.hs

module Control.Monad.Ology.Data.Ref where

import Control.Monad.Ology.Data.Param
import Control.Monad.Ology.Data.Prod
import Control.Monad.Ology.General
import Control.Monad.Ology.Specific.StateT
import Control.Monad.ST.Lazy qualified as Lazy
import Control.Monad.ST.Strict qualified as Strict
import Data.IORef
import Data.STRef.Lazy qualified as Lazy
import Data.STRef.Strict qualified as Strict
import Import

-- | A reference of a monad (as in 'StateT').
data Ref m a = MkRef
    { refGet :: m a
    , refPut :: a -> m ()
    }

instance Functor m => Invariant (Ref m) where
    invmap f g (MkRef gt pt) = MkRef (fmap f gt) (pt . g)

instance Applicative m => Productable (Ref m) where
    rUnit = MkRef (pure ()) (\_ -> pure ())
    ra <***> rb = MkRef (liftA2 (,) (refGet ra) (refGet rb)) $ \(a, b) -> refPut ra a *> refPut rb b

refModify :: Monad m => Ref m a -> (a -> a) -> m ()
refModify ref f = do
    a <- refGet ref
    refPut ref $ f a

refModifyM :: Monad m => Ref m a -> (a -> m a) -> m ()
refModifyM ref f = do
    a <- refGet ref
    a' <- f a
    refPut ref a'

-- | Restore the original value of this reference after the operation.
refRestore :: MonadException m => Ref m a -> m --> m
refRestore ref mr = bracketNoMask (refGet ref) (refPut ref) $ \_ -> mr

-- | Put and restore the original value of this reference after the operation.
refPutRestore :: MonadException m => Ref m a -> a -> m --> m
refPutRestore ref a mr =
    refRestore ref $ do
        refPut ref a
        mr

lensMapRef ::
       forall m a b. Monad m
    => Lens' a b
    -> Ref m a
    -> Ref m b
lensMapRef l ref = let
    refGet' = fmap (\a -> getConst $ l Const a) $ refGet ref
    refPut' b = do
        a <- refGet ref
        refPut ref $ runIdentity $ l (\_ -> Identity b) a
    in MkRef refGet' refPut'

liftRef :: (MonadTrans t, Monad m) => Ref m --> Ref (t m)
liftRef (MkRef g m) = MkRef (lift g) $ \a -> lift $ m a

stateRef :: Monad m => Ref (StateT s m) s
stateRef = MkRef get put

-- | Run a state monad over this reference.
refRunState :: Monad m => Ref m s -> StateT s m --> m
refRunState ref sm = do
    olds <- refGet ref
    (a, news) <- runStateT sm olds
    refPut ref news
    return a

ioRef :: IORef a -> Ref IO a
ioRef r = MkRef (readIORef r) (writeIORef r)

strictSTRef :: Strict.STRef s a -> Ref (Strict.ST s) a
strictSTRef r = MkRef (Strict.readSTRef r) (Strict.writeSTRef r)

lazySTRef :: Lazy.STRef s a -> Ref (Lazy.ST s) a
lazySTRef r = MkRef (Lazy.readSTRef r) (Lazy.writeSTRef r)

-- | Use a reference as a parameter.
refParam ::
       forall m a. MonadException m
    => Ref m a
    -> Param m a
refParam ref = let
    paramAsk = refGet ref
    paramWith :: a -> m --> m
    paramWith = refPutRestore ref
    in MkParam {..}

-- | Use a reference as a product.
refProd ::
       forall m a. (MonadException m, Monoid a)
    => Ref m a
    -> Prod m a
refProd ref = let
    prodTell a = refModify ref $ (<>) a
    prodCollect :: forall r. m r -> m (r, a)
    prodCollect mr =
        refPutRestore ref mempty $ do
            r <- mr
            a <- refGet ref
            return (r, a)
    in MkProd {..}