STM queue

Main.hs

{-# Language BlockArguments, ExplicitNamespaces, DerivingStrategies, ViewPatterns
  , NoFieldSelectors, OverloadedRecordDot, DuplicateRecordFields #-}

{-# OPTIONS_GHC -main-is Main.test -Wall #-}

module Main (
  Queue(push, pop, front, back, size, empty), new, onSome, test) where

import GHC.Conc (type STM, newTVar, readTVar, writeTVar, retry)
import Numeric.Natural (type Natural)
import Data.Bitraversable (bitraverse)
import Data.Function ((&))
import Data.Functor ((<&>))
import Data.Maybe (isNothing)
import Control.Monad (join, when)

-- testing
import GHC.Conc (atomically)
import Control.Concurrent (forkFinally, newEmptyMVar)
import Control.Concurrent.MVar (putMVar, takeMVar)
import Control.Monad (void)
import Data.Foldable (for_)

-- | read-only type for a node cons of the mutable list
data Node elem = Node {
  -- | car
  value :: elem, -- TODO: STM elem ???
  -- | cdr
  next :: !(List elem) }
  deriving stock Functor

-- | read-only type for a mutable list
type List elem = STM (Maybe (Node elem))

-- | mutable list with ability to push to the latest cdr
newtype Push elem = Push {
  -- | push an element to the current last node,
  --
  -- returns new last cons 
  push :: elem -> STM (Push elem) }

-- | a pair of a head and a push of the push-list
type PushList elem = (List elem, Push elem)

-- | smart constructor for 'PushList'
newPushList :: STM (PushList elem)
newPushList = newTVar Nothing <&> \ node -> (readTVar node, Push \ value -> do
  (next, push) <- newPushList
  writeTVar node $ Just Node { value, next }
  pure push)

-- | mutable FIFO queue with the last car exposed
data PopPushQueue elem = PopPushQueue {
  -- | head to read
  head :: !(List elem),
  -- | back & last (TVar TNil) to write
  last :: !(STM (Maybe elem)),
  -- | push an element to the end of the queue
  push :: elem -> STM (),
  -- | pop a first element out of the queue
  pop :: STM (Maybe elem) }

-- | smart constructor for 'PopPushQueue'
newPopPushQueue :: STM (PopPushQueue elem)
newPopPushQueue = do
  (head', push') <- bitraverse newTVar newTVar =<< newPushList
  last' <- newTVar Nothing

  pure PopPushQueue {
    head = join $ readTVar head',
    last = readTVar last',
    push = \ value ->
      readTVar push' <&> (.push) <&> ($ value) & join
      >>= writeTVar push' >> writeTVar last' do Just value,
    pop = readTVar head' & join >>= traverse \ node -> do
      writeTVar head' node.next
      -- optimized: writeTVar last' =<< (>>) <$> node.next <*> readTVar last'
      node.next <&> isNothing >>= flip when do
        writeTVar last' Nothing
      pure node.value }

-- | abstract FIFO queue
data Queue elem = Queue {
  push :: elem -> STM (),
  pop :: STM (Maybe elem),
  front :: STM (Maybe elem),
  back :: STM (Maybe elem),
  size :: STM Natural,
  empty :: STM Bool }

-- | smart constructor for 'Queue'
new :: STM (Queue elem)
new = newPopPushQueue <&> \ !queue -> Queue {
  push = queue.push,
  pop = queue.pop,
  front = queue.head <&> fmap (.value),
  back = queue.last,
  size = let
    size' list counter = list >>= maybe do pure counter
      \ node -> size' node.next $ counter + 1
    in size' queue.head 0,
  empty = queue.head <&> isNothing }

-- | helper for methods that may return 'Nothing'
--
-- computation would continue when there be a non-empty result
--
-- E.g. it can be used like @onSome myQueue.pop@
-- to pop a first element when it became available
onSome :: STM (Maybe elem) -> STM elem
{-# INLINE onSome #-}
onSome = (>>= maybe retry pure)

test :: IO ()
test = do
  q <- atomically $ new @Natural
  let n = 100

  Nothing <- atomically q.front
  Nothing <- atomically q.pop
  Nothing <- atomically q.back
  True <- atomically q.empty

  asyncTest <- newEmptyMVar
  void $ forkFinally
    do
      0 <- atomically $ onSome q.pop
      1 <- atomically $ onSome q.front
      1 <- atomically $ onSome q.back
      Just 1 <- atomically $ q.back
      Just 1 <- atomically $ q.pop
      Nothing <- atomically $ q.front
      atomically $ q.push 0
      0 <- atomically $ onSome q.pop
      Nothing <- atomically $ q.front
      Nothing <- atomically $ q.back
      putStrLn "Async tests successfully passed."
    do const $ putMVar asyncTest ()

  atomically $ q.push 0
  atomically $ q.push 1
  takeMVar asyncTest

  for_ [0 .. n * 2 - 1] \ i -> do
    ((== i) -> True) <- atomically q.size
    atomically $ q.push i
    Just 0 <- atomically q.front
    pure ()

  ((== Just (n * 2 - 1)) -> True) <- atomically q.back

  for_ [0 .. n - 1] \ i -> do
    ((== n * 2 - i) -> True) <- atomically q.size
    ((== Just i) -> True) <-atomically q.front
    ((== Just i) -> True) <- atomically q.pop
    pure ()

  False <- atomically q.empty
  ((== Just n) -> True) <- atomically q.front
  ((== Just (n * 2 - 1)) -> True) <- atomically q.back
  ((== n) -> True) <- atomically q.size

  putStrLn "Tests successfully passed."

You can use a playground to run it.