Add lazier versions of the TQueue and TBQueue reading functions

Control/Concurrent/STM/TBQueue.hs

@@ -34,9 +34,11 @@ module Control.Concurrent.STM.TBQueue (
         newTBQueue,
         newTBQueueIO,
         readTBQueue,
+        lazyReadTBQueue,
         tryReadTBQueue,
         flushTBQueue,
         peekTBQueue,
+        lazyPeekTBQueue,
         tryPeekTBQueue,
         writeTBQueue,
         unGetTBQueue,
@@ -47,6 +49,7 @@ module Control.Concurrent.STM.TBQueue (
 
 import           Control.Monad   (unless)
 import           Data.Typeable   (Typeable)
+import           Data.Tuple      (Solo(..))
 import           GHC.Conc        (STM, TVar, newTVar, newTVarIO, orElse,
                                   readTVar, retry, writeTVar)
 import           Numeric.Natural (Natural)
@@ -117,29 +120,36 @@ writeTBQueue (TBQueue rsize _read wsize write _size) a = do
   listend <- readTVar write
   writeTVar write (a:listend)
 
--- |Read the next value from the 'TBQueue'.
-readTBQueue :: TBQueue a -> STM a
-readTBQueue (TBQueue rsize read _wsize write _size) = do
+-- |Read the next value from the 'TBQueue', retrying if the channel is empty.
+-- Separates reversing the write-end from evaluating the next value.
+lazyReadTBQueue :: TBQueue a -> STM (Solo a)
+lazyReadTBQueue (TBQueue rsize read _wsize write _size)  = do
   xs <- readTVar read
   r <- readTVar rsize
   writeTVar rsize $! r + 1
   case xs of
     (x:xs') -> do
       writeTVar read xs'
-      return x
+      return (Solo x)
     [] -> do
       ys <- readTVar write
       case ys of
         [] -> retry
         _  -> do
           -- NB. lazy: we want the transaction to be
           -- short, otherwise it will conflict
-          let ~(z,zs) = case reverse ys of
+          let (z,zs) = case reverse ys of
                           z':zs' -> (z',zs')
                           _      -> error "readTBQueue: impossible"
           writeTVar write []
           writeTVar read zs
-          return z
+          return (zs `seq` Solo z)
+
+-- |Read the next value from the 'TBQueue'.
+readTBQueue :: TBQueue a -> STM a
+readTBQueue q = do
+  r <- lazyReadTBQueue q
+  return (case r of Solo z -> z)
 
 -- | A version of 'readTBQueue' which does not retry. Instead it
 -- returns @Nothing@ if no value is available.
@@ -163,13 +173,14 @@ flushTBQueue (TBQueue rsize read wsize write size) = do
       writeTVar wsize size
       return (xs ++ reverse ys)
 
--- | Get the next value from the @TBQueue@ without removing it,
--- retrying if the channel is empty.
-peekTBQueue :: TBQueue a -> STM a
-peekTBQueue (TBQueue _ read _ write _) = do
+-- | Get the next value from the @TBQueue@ without removing it, retrying if the
+-- channel is empty. Separates reversing the write-end from evaluating the next
+-- value.
+lazyPeekTBQueue :: TBQueue a -> STM (Solo a)
+lazyPeekTBQueue (TBQueue _ read _ write _) = do
   xs <- readTVar read
   case xs of
-    (x:_) -> return x
+    (x:_) -> return (Solo x)
     [] -> do
       ys <- readTVar write
       case ys of
@@ -179,7 +190,14 @@ peekTBQueue (TBQueue _ read _ write _) = do
                                   -- short, otherwise it will conflict
           writeTVar write []
           writeTVar read (z:zs)
-          return z
+          return (zs `seq` Solo z)
+
+-- | Get the next value from the @TBQueue@ without removing it,
+-- retrying if the channel is empty.
+peekTBQueue :: TBQueue a -> STM a
+peekTBQueue q = do
+  r <- lazyPeekTBQueue q
+  return (case r of Solo z -> z)
 
 -- | A version of 'peekTBQueue' which does not retry. Instead it
 -- returns @Nothing@ if no value is available.

Control/Concurrent/STM/TQueue.hs

@@ -38,9 +38,11 @@ module Control.Concurrent.STM.TQueue (
         newTQueue,
         newTQueueIO,
         readTQueue,
+        lazyReadTQueue,
         tryReadTQueue,
         flushTQueue,
         peekTQueue,
+        lazyPeekTQueue,
         tryPeekTQueue,
         writeTQueue,
         unGetTQueue,
@@ -49,6 +51,7 @@ module Control.Concurrent.STM.TQueue (
 
 import GHC.Conc
 import Control.Monad (unless)
+import Data.Tuple (Solo(..))
 import Data.Typeable (Typeable)
 
 -- | 'TQueue' is an abstract type representing an unbounded FIFO channel.
@@ -84,14 +87,15 @@ writeTQueue (TQueue _read write) a = do
   listend <- readTVar write
   writeTVar write (a:listend)
 
--- |Read the next value from the 'TQueue'.
-readTQueue :: TQueue a -> STM a
-readTQueue (TQueue read write) = do
+-- |Read the next value from the 'TQueue', retrying if the channel is empty.
+-- Separates reversing the write-end from evaluating the next value.
+lazyReadTQueue :: TQueue a -> STM (Solo a)
+lazyReadTQueue (TQueue read write) = do
   xs <- readTVar read
   case xs of
     (x:xs') -> do
       writeTVar read xs'
-      return x
+      return (Solo x)
     [] -> do
       ys <- readTVar write
       case ys of
@@ -101,7 +105,13 @@ readTQueue (TQueue read write) = do
                                   -- short, otherwise it will conflict
           writeTVar write []
           writeTVar read zs
-          return z
+          return (zs `seq` Solo z)
+
+-- |Read the next value from the 'TQueue'.
+readTQueue :: TQueue a -> STM a
+readTQueue q = do
+  r <- lazyReadTQueue q
+  return (case r of Solo z -> z)
 
 -- | A version of 'readTQueue' which does not retry. Instead it
 -- returns @Nothing@ if no value is available.
@@ -120,13 +130,14 @@ flushTQueue (TQueue read write) = do
   unless (null ys) $ writeTVar write []
   return (xs ++ reverse ys)
 
--- | Get the next value from the @TQueue@ without removing it,
--- retrying if the channel is empty.
-peekTQueue :: TQueue a -> STM a
-peekTQueue (TQueue read write) = do
+-- | Get the next value from the @TQueue@ without removing it, retrying if the
+-- channel is empty. Separates reversing the write-end from evaluating the next
+-- value.
+lazyPeekTQueue :: TQueue a -> STM (Solo a)
+lazyPeekTQueue (TQueue read write) = do
   xs <- readTVar read
   case xs of
-    (x:_) -> return x
+    (x:_) -> return (Solo x)
     [] -> do
       ys <- readTVar write
       case ys of
@@ -136,7 +147,14 @@ peekTQueue (TQueue read write) = do
                                   -- short, otherwise it will conflict
           writeTVar write []
           writeTVar read (z:zs)
-          return z
+          return (zs `seq` Solo z)
+
+-- | Get the next value from the @TQueue@ without removing it,
+-- retrying if the channel is empty.
+peekTQueue :: TQueue a -> STM a
+peekTQueue q = do
+  r <- lazyPeekTQueue q
+  return (case r of Solo z -> z)
 
 -- | A version of 'peekTQueue' which does not retry. Instead it
 -- returns @Nothing@ if no value is available.