New monads/MonadAdvSTM

{- November 24th, 2006
 
  Demonstration Code by Chris Kuklewicz <haskell@list.mightyreason.com>
  Usual 3 clause BSD Licence
  Copyright 2006
 
  This is inspired by a post by Simon Peyton-Jones on the haskell-cafe
  mailing list, in which the type and semantics of onCommit and
  retryWith were put forth.
 
  The semantics of printing the contents of the TVar "v" created in
  test via retryWith may or may not be well defined.  With GHC 6.6 I get
 
*AdvSTM> main
"hello world"
 
"test"
"onRetry Start"
("onRetry v",7)
"Flipped choice to True to avoid infinite loop"
"onCommit Start"
("onCommit v",42)
("result","foo","retries",1)
 
"testUnlift"
"onRetry Start"
("onRetry v",7)
"Flipped choice to True to avoid infinite loop"
"onCommit Start"
("onCommit v",42)
("result","foo","retries",2)
 
"bye world"
 
  Aside from that I think the unsafeIOToSTM is not really unsafe here
  since it writes to privately created and maintained variables.
 
  Since the implementation is hidden it could be changed from ReaderT
  to some other scheme.
 
  Once could also use MonadBase from
  http://haskell.org/haskellwiki/New_monads/MonadBase to help with the
  lifting, but this has been commented out below.
 
  TODO: figure out semantics of catchAdv.  At least it compiles...
-}
 
module AdvSTM(MonadAdvSTM(..),AdvSTM,retryWith,countRetries
             ,unlifter,unlift,unlift1,unlift2) where
 
-- import MonadBase
import Control.Exception(Exception)
import Control.Monad(MonadPlus(..),liftM)
import Control.Monad.Reader(MonadReader(..),ReaderT,runReaderT,lift,asks)
import Control.Concurrent.MVar(MVar,newEmptyMVar,newMVar,takeMVar,tryTakeMVar,putMVar)
import Control.Concurrent.STM(STM,orElse,retry,catchSTM,atomically)
import Control.Concurrent.STM.TVar(TVar,newTVarIO,newTVar,readTVar,writeTVar)
import Data.Generics(Data)
import Data.Maybe(maybe)
import Data.Typeable(Typeable)
import GHC.Conc(unsafeIOToSTM)
-- for countRetries example
import Data.IORef(IORef,newIORef,readIORef,writeIORef,modifyIORef)
 
class (Monad m) => MonadAdvSTM m where
  onCommit :: IO a -> m ()
  onRetry :: IO a -> m ()
  orElseAdv :: m a -> m a -> m a
  retryAdv :: m a
  atomicAdv :: m a -> IO a
  catchAdv :: m a -> (Exception -> m a) -> m a
  liftAdv :: STM a -> m a
 
-- Export type but not constructor!
newtype AdvSTM a = AdvSTM (ReaderT Env STM a) deriving (Functor,Monad,MonadPlus,Typeable)
type Env = (CommitVar,RetryVar)
type CommitVar = TVar (IO ()->IO ())
type RetryVar = MVar (IO ()->IO ())
 
{- Since lifting retry and `orElse` gives the semantics Simon wants, use deriving MonadPlus instead
instance MonadPlus AdvSTM where
  mzero = retryAdv
  mplus = orElseAdv
-}
 
-- instance MonadBase STM AdvSTM where liftBase = AdvSTM . lift
retryWith :: (Monad m, MonadAdvSTM m) => IO a -> m b
retryWith io = onRetry io >> retryAdv
 
instance MonadAdvSTM AdvSTM where
  onCommit io = do
    commitVar <- AdvSTM $ asks fst
    old <- liftAdv $ readTVar commitVar
    liftAdv $ writeTVar commitVar (old . (io >>))
  onRetry io = do
    retryVar <- AdvSTM $ asks snd
    liftAdv $ unsafeIOToSTM (do
      may'do <- tryTakeMVar retryVar
      let todo = maybe (io >>) (. (io >>)) may'do
      seq todo (putMVar retryVar todo))
  orElseAdv = mplus
  retryAdv = liftAdv retry -- the same as retryAdv = mzero
  atomicAdv = runAdvSTM
  catchAdv action handler = do
    action' <- unlift action
    handler' <- unlift1 handler
    liftAdv $ catchSTM action' handler'
  liftAdv = AdvSTM . lift
 
-- This replaces "atomically"
runAdvSTM :: AdvSTM a -> IO a
runAdvSTM (AdvSTM action) = do
  commitVar <- newTVarIO id
  retryVar <- newMVar id
  let check'retry = do
        may'todo <- unsafeIOToSTM $ tryTakeMVar retryVar
        maybe retry (return . Right) may'todo
  let wrappedAction = (runReaderT (liftM Left action) (commitVar,retryVar))
                      `orElse` (check'retry)
  let attempt = do
        result <- atomically $ wrappedAction
        case result of
          Left answer -> do
            cFun <- atomically (readTVar commitVar)
            cFun (return ())
            return answer
          Right rFun -> do
            rFun (return ())
            attempt
  attempt
 
-- Using ReaderT we can write "unlift" from AdvSTM into STM:
 
-- Do not export runWith
runWith :: Env -> AdvSTM t -> STM t
runWith env (AdvSTM action) = runReaderT action env
 
unlifter :: AdvSTM (AdvSTM a -> STM a)
unlifter = do
  env <- AdvSTM ask
  return (runWith env)
 
unlift :: AdvSTM a -> AdvSTM (STM a)
unlift f = do
  u <- unlifter
  return (u f)
 
unlift1 :: (t -> AdvSTM a) -> AdvSTM (t -> STM a)
unlift1 f = do
  u <- unlifter
  return (\x -> u (f x))
 
unlift2 :: (t -> t1 -> AdvSTM a) -> AdvSTM (t -> t1 -> STM a)
unlift2 f = do
  u <- unlifter
  return (\x y -> u (f x y))
 
-- Example code using the above, lifting into MonadAdvSTM:
test ::(Monad m, MonadAdvSTM m) => TVar Bool -> m [Char]
test todo = do
  onCommit (print "onCommit Start")
  onRetry (print "onRetry Start")
  v <- liftAdv $ newTVar 7
  liftAdv $ writeTVar v 42
  onCommit (atomically (readTVar v) >>= \x->print ("onCommit v",x))
  onRetry (atomically (readTVar v) >>= \x->print ("onRetry v",x))
  choice <- liftAdv $ readTVar todo
  case choice of
    True -> return "foo"
    False -> retryWith $ do
      atomically (writeTVar todo True)
      print "Flipped choice to True to avoid infinite loop"
 
-- Same example as test, but unlifting from AdvSTM
testUnlift :: TVar Bool -> AdvSTM [Char]
testUnlift todo = do
  onCommit <- unlift1 onCommit
  onRetry <- unlift1 onRetry
  retryWith <- unlift1 retryWith
  liftAdv $ do
    onCommit (print "onCommit Start")
    onRetry (print "onRetry Start")
    v <- newTVar 7
    writeTVar v 42
    onCommit (atomically (readTVar v) >>= \x->print ("onCommit v",x))
    onRetry (atomically (readTVar v) >>= \x->print ("onRetry v",x))
    choice <- readTVar todo
    case choice of
      True -> return "foo"
      False -> retryWith $ do 
        atomically (writeTVar todo True)
        print "Flipped choice to True to avoid infinite loop"
 
-- Example similar to Simon's suggested example:
countRetries :: (MonadAdvSTM m, Enum a) => IORef a -> m a1 -> m a1
countRetries ioref action =
  let incr = do old <- readIORef ioref
                writeIORef ioref $! (succ old)
  in action `orElseAdv` (retryWith incr)
 
-- Load this file in GHCI and execute main to run the test:
main = do
  print "hello world"
  putStrLn ""
  counter <- newIORef 0
  todo <- newTVarIO False
  print "test"
  result <- runAdvSTM (countRetries counter $ test todo)
  retries <- readIORef counter
  print ("result",result,"retries",retries)
  atomically (writeTVar todo False)
  putStrLn ""
  print "testUnlift"
  result <- runAdvSTM (countRetries counter $ testUnlift todo)
  retries <- readIORef counter
  print ("result",result,"retries",retries)
  putStrLn ""
  print "bye world"