A monad for scheduling IO actions within STM transactions

Atom.hs

{-# LANGUAGE DerivingVia, BlockArguments #-}

module Atom (
  Atom, atom, atom_,
  stm, io, embed,
  throwAtom, catchAtom,
) where

-- base
import Data.Monoid (Ap(..))
import Data.Functor ((<&>))
import Control.Applicative (Alternative)
import Control.Exception (Exception)

-- stm
import Control.Concurrent.STM (STM, atomically, retry, throwSTM, catchSTM)

-- transformers
import Control.Monad.Trans.Class (MonadTrans(..))
import Control.Monad.Trans.State (StateT(..), modify')


newtype Atom w a = Atom{ unAtom :: IO w -> STM (a, IO w) }
  deriving (Functor, Applicative, Alternative, Monad)
    via StateT (IO w) STM
  deriving (Semigroup, Monoid)
    via Ap (Atom w) a

instance MonadFail (Atom w) where
  fail _ = stm retry

atom :: Monoid w => Atom w a -> IO (w, a)
atom (Atom a) = do
  (x, iow) <- atomically (a mempty)
  iow <&> (,x)

atom_ :: Atom () a -> IO a
atom_ = fmap snd . atom

stm :: STM a -> Atom w a
stm = toAtom . lift

io :: Semigroup w => IO w -> Atom w ()
io new = toAtom $ modify' (<> new)

embed :: (Monoid v, Semigroup w) => (IO v -> IO w) -> Atom v a -> Atom w a
embed vw (Atom a) = Atom \iow ->
  (fmap . fmap) ((iow <>) . vw) (a mempty)

throwAtom :: Exception e => e -> Atom w a
throwAtom = stm . throwSTM

catchAtom :: Exception e => Atom w a -> (e -> Atom w a) -> Atom w a
catchAtom (Atom a) h = Atom \iow ->
  catchSTM (a iow) \e ->
    unAtom (h e) iow

toAtom :: StateT (IO w) STM a -> Atom w a
toAtom = Atom . runStateT