MonadSTM: added flushTQueue

input-output-hk · Sep 27, 2022 · 051d04f · 051d04f
1 parent b38f82f
commit 051d04f
Show file tree

Hide file tree

Showing 5 changed files with 32 additions and 2 deletions.
diff --git a/io-classes/src/Control/Concurrent/Class/MonadSTM/TQueue.hs b/io-classes/src/Control/Concurrent/Class/MonadSTM/TQueue.hs
@@ -11,6 +11,7 @@ module Control.Concurrent.Class.MonadSTM.TQueue
   , tryReadTQueue
   , peekTQueue
   , tryPeekTQueue
+  , flushTQueue
   , writeTQueue
   , unGetTQueue
   , isEmptyTQueue

diff --git a/io-classes/src/Control/Monad/Class/MonadSTM/Internal.hs b/io-classes/src/Control/Monad/Class/MonadSTM/Internal.hs
@@ -64,7 +64,7 @@ import qualified Control.Concurrent.STM.TMVar as STM
 import qualified Control.Concurrent.STM.TQueue as STM
 import qualified Control.Concurrent.STM.TSem as STM
 import qualified Control.Concurrent.STM.TVar as STM
-import           Control.Monad (MonadPlus (..), when)
+import           Control.Monad (MonadPlus (..), unless, when)
 import qualified Control.Monad.STM as STM
 
 import           Control.Monad.Cont (ContT (..))
@@ -151,6 +151,7 @@ class ( Monad m
   tryReadTQueue  :: TQueue m a -> STM m (Maybe a)
   peekTQueue     :: TQueue m a -> STM m a
   tryPeekTQueue  :: TQueue m a -> STM m (Maybe a)
+  flushTQueue    :: TQueue m a -> STM m [a]
   writeTQueue    :: TQueue m a -> a -> STM m ()
   isEmptyTQueue  :: TQueue m a -> STM m Bool
   unGetTQueue    :: TQueue m a -> a -> STM m ()
@@ -285,6 +286,10 @@ class ( Monad m
                         => TQueue m a -> STM m (Maybe a)
   tryPeekTQueue = tryPeekTQueueDefault
 
+  default flushTQueue :: TQueue m ~ TQueueDefault m
+                      => TQueue m a -> STM m [a]
+  flushTQueue = flushTQueueDefault
+
   default newTBQueue :: TBQueue m ~ TBQueueDefault m
                      => Natural -> STM m (TBQueue m a)
   newTBQueue = newTBQueueDefault
@@ -683,7 +688,7 @@ instance MonadSTM IO where
   tryReadTQueue  = STM.tryReadTQueue
   peekTQueue     = STM.peekTQueue
   tryPeekTQueue  = STM.tryPeekTQueue
-  flushTBQueue   = STM.flushTBQueue
+  flushTQueue    = STM.flushTQueue
   writeTQueue    = STM.writeTQueue
   isEmptyTQueue  = STM.isEmptyTQueue
   unGetTQueue    = STM.unGetTQueue
@@ -693,6 +698,7 @@ instance MonadSTM IO where
   peekTBQueue    = STM.peekTBQueue
   tryPeekTBQueue = STM.tryPeekTBQueue
   writeTBQueue   = STM.writeTBQueue
+  flushTBQueue   = STM.flushTBQueue
   lengthTBQueue  = STM.lengthTBQueue
   isEmptyTBQueue = STM.isEmptyTBQueue
   isFullTBQueue  = STM.isFullTBQueue
@@ -940,6 +946,15 @@ tryPeekTQueueDefault (TQueue read _write) = do
       (x:_) -> return (Just x)
       _     -> return Nothing
 
+
+flushTQueueDefault :: MonadSTM m => TQueueDefault m a -> STM m [a]
+flushTQueueDefault (TQueue read write) = do
+  xs <- readTVar read
+  ys <- readTVar write
+  unless (null xs) $ writeTVar read []
+  unless (null ys) $ writeTVar write []
+  return (xs ++ reverse ys)
+
 unGetTQueueDefault :: MonadSTM m => TQueueDefault m a -> a -> STM m ()
 unGetTQueueDefault (TQueue read _write) a = modifyTVar read (a:)
 
@@ -1346,6 +1361,7 @@ instance MonadSTM m => MonadSTM (ContT r m) where
     tryReadTQueue  = WrappedSTM . tryReadTQueue
     peekTQueue     = WrappedSTM . peekTQueue
     tryPeekTQueue  = WrappedSTM . tryPeekTQueue
+    flushTQueue    = WrappedSTM .  flushTQueue
     writeTQueue v  = WrappedSTM . writeTQueue v
     isEmptyTQueue  = WrappedSTM . isEmptyTQueue
     unGetTQueue    = WrappedSTM .: unGetTQueue
@@ -1420,6 +1436,7 @@ instance MonadSTM m => MonadSTM (ReaderT r m) where
     tryReadTQueue  = WrappedSTM .  tryReadTQueue
     peekTQueue     = WrappedSTM .  peekTQueue
     tryPeekTQueue  = WrappedSTM .  tryPeekTQueue
+    flushTQueue    = WrappedSTM .  flushTQueue
     writeTQueue v  = WrappedSTM .  writeTQueue v
     isEmptyTQueue  = WrappedSTM .  isEmptyTQueue
     unGetTQueue    = WrappedSTM .: unGetTQueue
@@ -1494,6 +1511,7 @@ instance (Monoid w, MonadSTM m) => MonadSTM (WriterT w m) where
     tryReadTQueue  = WrappedSTM .  tryReadTQueue
     peekTQueue     = WrappedSTM .  peekTQueue
     tryPeekTQueue  = WrappedSTM .  tryPeekTQueue
+    flushTQueue    = WrappedSTM .  flushTQueue
     writeTQueue v  = WrappedSTM .  writeTQueue v
     isEmptyTQueue  = WrappedSTM .  isEmptyTQueue
     unGetTQueue    = WrappedSTM .: unGetTQueue
@@ -1568,6 +1586,7 @@ instance MonadSTM m => MonadSTM (StateT s m) where
     tryReadTQueue  = WrappedSTM . tryReadTQueue
     peekTQueue     = WrappedSTM . peekTQueue
     tryPeekTQueue  = WrappedSTM . tryPeekTQueue
+    flushTQueue    = WrappedSTM .  flushTQueue
     writeTQueue v  = WrappedSTM . writeTQueue v
     isEmptyTQueue  = WrappedSTM . isEmptyTQueue
     unGetTQueue    = WrappedSTM .: unGetTQueue
@@ -1642,6 +1661,7 @@ instance MonadSTM m => MonadSTM (ExceptT e m) where
     tryReadTQueue  = WrappedSTM .  tryReadTQueue
     peekTQueue     = WrappedSTM .  peekTQueue
     tryPeekTQueue  = WrappedSTM .  tryPeekTQueue
+    flushTQueue    = WrappedSTM .  flushTQueue
     writeTQueue v  = WrappedSTM .  writeTQueue v
     isEmptyTQueue  = WrappedSTM .  isEmptyTQueue
     unGetTQueue    = WrappedSTM .: unGetTQueue
@@ -1716,6 +1736,7 @@ instance (Monoid w, MonadSTM m) => MonadSTM (RWST r w s m) where
     tryReadTQueue  = WrappedSTM .  tryReadTQueue
     peekTQueue     = WrappedSTM .  peekTQueue
     tryPeekTQueue  = WrappedSTM .  tryPeekTQueue
+    flushTQueue    = WrappedSTM .  flushTQueue
     writeTQueue v  = WrappedSTM .  writeTQueue v
     isEmptyTQueue  = WrappedSTM .  isEmptyTQueue
     unGetTQueue    = WrappedSTM .: unGetTQueue

diff --git a/io-sim/src/Control/Monad/IOSim/STM.hs b/io-sim/src/Control/Monad/IOSim/STM.hs
@@ -85,6 +85,9 @@ tryPeekTQueueDefault (TQueue queue) = do
       x :_ -> Just x
       []   -> Nothing
 
+flushTQueueDefault :: MonadSTM m => TQueueDefault m a -> STM m [a]
+flushTQueueDefault (TQueue queue) = uncurry (++) <$> readTVar queue
+
 unGetTQueueDefault :: MonadSTM m => TQueueDefault m a -> a -> STM m ()
 unGetTQueueDefault (TQueue queue) a = do
     (xs, ys) <- readTVar queue

diff --git a/io-sim/src/Control/Monad/IOSim/Types.hs b/io-sim/src/Control/Monad/IOSim/Types.hs
@@ -416,6 +416,7 @@ instance MonadSTM (IOSim s) where
   tryReadTQueue     = tryReadTQueueDefault
   peekTQueue        = peekTQueueDefault
   tryPeekTQueue     = tryPeekTQueueDefault
+  flushTQueue       = flushTQueueDefault
   writeTQueue       = writeTQueueDefault
   isEmptyTQueue     = isEmptyTQueueDefault
   unGetTQueue       = unGetTQueueDefault

diff --git a/strict-stm/src/Control/Concurrent/Class/MonadSTM/Strict/TQueue.hs b/strict-stm/src/Control/Concurrent/Class/MonadSTM/Strict/TQueue.hs
@@ -18,6 +18,7 @@ module Control.Concurrent.Class.MonadSTM.Strict.TQueue
   , tryReadTQueue
   , peekTQueue
   , tryPeekTQueue
+  , flushTQueue
   , writeTQueue
   , unGetTQueue
   , isEmptyTQueue
@@ -78,6 +79,9 @@ peekTQueue = Lazy.peekTQueue . toLazyTQueue
 tryPeekTQueue :: MonadSTM m => StrictTQueue m a -> STM m (Maybe a)
 tryPeekTQueue = Lazy.tryPeekTQueue . toLazyTQueue
 
+flushTQueue :: MonadSTM m => StrictTQueue m a -> STM m [a]
+flushTQueue = Lazy.flushTQueue . toLazyTQueue
+
 writeTQueue :: MonadSTM m => StrictTQueue m a -> a -> STM m ()
 writeTQueue (StrictTQueue tqueue) !a = Lazy.writeTQueue tqueue a