IORefs for io-sim and io-sim-por

io-classes/io-classes.cabal

@@ -60,6 +60,7 @@ library
                        Control.Monad.Class.MonadAsync
                        Control.Monad.Class.MonadEventlog
                        Control.Monad.Class.MonadFork
+                       Control.Monad.Class.MonadIORef
                        Control.Monad.Class.MonadSay
                        Control.Monad.Class.MonadST
                        Control.Monad.Class.MonadSTM

io-classes/src/Control/Monad/Class/MonadIORef.hs

@@ -0,0 +1,79 @@
+{-# LANGUAGE DefaultSignatures   #-}
+{-# LANGUAGE InstanceSigs        #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE TypeApplications    #-}
+{-# LANGUAGE TypeFamilies        #-}
+
+module Control.Monad.Class.MonadIORef where
+
+import Control.Monad.Reader
+import Data.IORef qualified as IO
+import Data.Kind
+
+class Monad m => MonadIORef m where
+  {-# MINIMAL newIORef, readIORef, writeIORef, atomicModifyIORef, atomicModifyIORef', atomicWriteIORef #-}
+
+  type IORef m :: Type -> Type
+
+  -- | See 'IO.newIORef'.
+  newIORef :: a -> m (IORef m a)
+  -- | See 'IO.readIORef'.
+  readIORef :: IORef m a -> m a
+  -- | See 'IO.writeIORef'.
+  writeIORef :: IORef m a -> a -> m ()
+  -- | See 'IO.modifyIORef'.
+  modifyIORef :: IORef m a -> (a -> a) -> m ()
+  -- | See 'IO.modifyRef''.
+  modifyIORef' :: IORef m a -> (a -> a) -> m ()
+  -- | See 'IO.atomicModifyIORef'.
+  atomicModifyIORef :: IORef m a -> (a -> (a, b)) -> m b
+  -- | See 'IO.atomicModifyIORef''.
+  atomicModifyIORef' :: IORef m a -> (a -> (a, b)) -> m b
+  -- | See 'IO.atomicWriteIORef'.
+  atomicWriteIORef :: IORef m a -> a -> m ()
+
+  default modifyIORef :: IORef m a -> (a -> a) -> m ()
+  modifyIORef ref f = readIORef ref >>= writeIORef ref . f
+
+  default modifyIORef' :: IORef m a -> (a -> a) -> m ()
+  modifyIORef' ref f = do
+    x <- readIORef ref
+    let x' = f x
+    x' `seq` writeIORef ref x'
+
+--
+-- IO instance
+--
+
+instance MonadIORef IO where
+  type IORef IO = IO.IORef
+  newIORef = IO.newIORef
+  readIORef = IO.readIORef
+  writeIORef = IO.writeIORef
+  modifyIORef = IO.modifyIORef
+  modifyIORef' = IO.modifyIORef'
+  atomicModifyIORef = IO.atomicModifyIORef
+  atomicModifyIORef' = IO.atomicModifyIORef'
+  atomicWriteIORef = IO.atomicWriteIORef
+
+--
+-- ReaderT instance
+--
+
+instance MonadIORef m => MonadIORef (ReaderT r m) where
+  type IORef (ReaderT r m) = IORef m
+  newIORef = lift . newIORef
+  readIORef = lift . readIORef
+  writeIORef = lift .: writeIORef
+  modifyIORef = lift .: modifyIORef
+  modifyIORef' = lift .: modifyIORef'
+  atomicModifyIORef = lift .: atomicModifyIORef
+  atomicModifyIORef' = lift .: atomicModifyIORef'
+  atomicWriteIORef = lift .: atomicWriteIORef
+
+--
+-- Utilities
+--
+
+(.:) :: (c -> d) -> (a -> b -> c) -> (a -> b -> d)
+(f .: g) x y = f (g x y)

io-sim/io-sim.cabal

@@ -51,6 +51,7 @@ library
                        Control.Monad.IOSim.Internal,
                        Control.Monad.IOSim.InternalTypes,
                        Control.Monad.IOSim.STM,
+                       Control.Monad.IOSim.IORef,
                        Control.Monad.IOSimPOR.Internal,
                        Control.Monad.IOSimPOR.Types,
                        Control.Monad.IOSimPOR.QuickCheckUtils,

io-sim/src/Control/Monad/IOSim/IORef.hs

@@ -0,0 +1,70 @@
+{-# LANGUAGE BangPatterns  #-}
+{-# LANGUAGE MagicHash     #-}
+{-# LANGUAGE TypeFamilies  #-}
+{-# LANGUAGE UnboxedTuples #-}
+
+{-# OPTIONS_GHC -Wno-orphans #-}
+
+module Control.Monad.IOSim.IORef where
+
+import Control.Monad.Class.MonadIORef
+import Control.Monad.Class.MonadST
+import Control.Monad.IOSim.Types
+import GHC.Exts
+import GHC.ST
+import GHC.STRef
+
+newtype IOSimRef s a = IORef (STRef s a)
+
+instance MonadIORef (IOSim s) where
+  type IORef (IOSim s) = IOSimRef s
+  newIORef v = stToIO (newSTRef v) >>= \ var -> return (IORef var)
+  readIORef (IORef v) = stToIO (readSTRef v)
+  writeIORef (IORef var) v = stToIO (writeSTRef var v)
+  modifyIORef ref f = readIORef ref >>= writeIORef ref . f
+  modifyIORef' ref f = do
+    x <- readIORef ref
+    let x' = f x
+    x' `seq` writeIORef ref x'
+  atomicModifyIORef ref f = do
+    (_old, (_new, res)) <- atomicModifyIORef2 ref f
+    pure res
+  atomicModifyIORef' = Control.Monad.IOSim.IORef.atomicModifyIORef'
+  atomicWriteIORef ref a = do
+    _ <- atomicSwapIORef ref a
+    pure ()
+
+atomicModifyIORef2Lazy :: IORef (IOSim s) a -> (a -> (a,b)) -> IOSim s (a, (a, b))
+atomicModifyIORef2Lazy (IORef (STRef r#)) f = stToIO $
+  ST (\s -> case atomicModifyMutVar2# r# f s of
+              (# s', old, res #) -> (# s', (old, res) #))
+
+atomicModifyIORef2 :: IORef (IOSim s) a -> (a -> (a,b)) -> IOSim s (a, (a, b))
+atomicModifyIORef2 ref f = do
+  r@(_old, (_new, _res)) <- atomicModifyIORef2Lazy ref f
+  return r
+
+atomicModifyIORefP :: IORef (IOSim s) a -> (a -> (a,b)) -> IOSim s b
+atomicModifyIORefP ref f = do
+  (_old, (_,r)) <- atomicModifyIORef2 ref f
+  pure r
+
+atomicModifyIORefLazy_ :: IORef (IOSim s) a -> (a -> a) -> IOSim s (a, a)
+atomicModifyIORefLazy_ (IORef (STRef ref)) f = stToIO $ ST $ \s ->
+  case atomicModifyMutVar_# ref f s of
+    (# s', old, new #) -> (# s', (old, new) #)
+
+atomicModifyIORef'_ :: IORef (IOSim s) a -> (a -> a) -> IOSim s (a, a)
+atomicModifyIORef'_ ref f = do
+  (old, !new) <- atomicModifyIORefLazy_ ref f
+  return (old, new)
+
+atomicSwapIORef :: IORef (IOSim s) a -> a -> IOSim s a
+atomicSwapIORef (IORef (STRef ref)) new = stToIO $ ST (atomicSwapMutVar# ref new)
+
+atomicModifyIORef' :: IORef (IOSim s) a -> (a -> (a,b)) -> IOSim s b
+atomicModifyIORef' ref f = do
+  (_old, (_new, !res)) <- atomicModifyIORef2 ref $
+    \old -> case f old of
+       r@(!_new, _res) -> r
+  pure res