module TheatreDev.ExtrasFor.TBQueue where

import Control.Concurrent.STM.TBQueue
import TheatreDev.Prelude

flushNonEmptyTBQueue :: TBQueue a -> STM (NonEmpty a)
flushNonEmptyTBQueue :: forall a. TBQueue a -> STM (NonEmpty a)
flushNonEmptyTBQueue TBQueue a
x = do
  a
head <- forall a. TBQueue a -> STM a
readTBQueue TBQueue a
x
  [a]
tail <- forall a. TBQueue a -> STM [a]
correctFlushTBQueue TBQueue a
x
  return (a
head forall a. a -> [a] -> NonEmpty a
:| [a]
tail)

-- | Get a list of all entries in the queue without removing them.
inspectTBQueue :: TBQueue a -> STM [a]
inspectTBQueue :: forall a. TBQueue a -> STM [a]
inspectTBQueue TBQueue a
queue = do
  [a]
list <- forall a. TBQueue a -> STM [a]
correctFlushTBQueue TBQueue a
queue
  forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [a]
list forall a b. (a -> b) -> a -> b
$ forall a. TBQueue a -> a -> STM ()
writeTBQueue TBQueue a
queue
  return [a]
list

-- | Starting from \"stm\" 2.5.2.0 "flushTBQueue" is broken.
-- We're fixing it here.
correctFlushTBQueue :: TBQueue a -> STM [a]
correctFlushTBQueue :: forall a. TBQueue a -> STM [a]
correctFlushTBQueue TBQueue a
queue =
  [a] -> STM [a]
go []
  where
    go :: [a] -> STM [a]
go ![a]
acc = do
      Maybe a
element <- forall a. TBQueue a -> STM (Maybe a)
tryReadTBQueue TBQueue a
queue
      case Maybe a
element of
        Just a
element -> [a] -> STM [a]
go forall a b. (a -> b) -> a -> b
$ a
element forall a. a -> [a] -> [a]
: [a]
acc
        Maybe a
Nothing -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a. [a] -> [a]
reverse [a]
acc