A mutable single-linked list in ST monad

STList.hs

module STList (
    STList,
    STListPtr,
    new,
    prepend,
    append,
    mark,
    pop,
    popUntil,
    foldr,
    reduce,
) where

import Prelude hiding (foldr)
import Data.STRef
import Control.Monad.ST

newtype STListCell s t = STListCell (Maybe (t, STRef s (STListCell s t)))
type STListPtr s t = STRef s (STListCell s t)
type STListObject s t = (
    STListPtr s t,
    STListPtr s t)
newtype STList s t = STList (STRef s (STListObject s t))

new :: ST s (STList s t)
new = do
    cell <- newSTRef (STListCell Nothing)
    obj <- newSTRef (cell, cell)
    return $ STList obj

prepend :: t -> STList s t -> ST s (STListPtr s t)
prepend x (STList obj) = do
    (first, last) <- readSTRef obj
    cell <- newSTRef (STListCell (Just (x, first)))
    writeSTRef obj (cell, last)
    return $ first

append :: t -> STList s t -> ST s ()
append x (STList obj) = do
    (first, last) <- readSTRef obj
    cell <- newSTRef (STListCell Nothing)
    writeSTRef last (STListCell (Just (x, cell)))
    writeSTRef obj (first, cell)

mark :: STList s t -> ST s (STListPtr s t)
mark (STList obj) = do
    (first, last) <- readSTRef obj
    return $ first

pop :: STListPtr s t -> STList s t -> ST s (Maybe t)
pop ptr (STList obj) = do
    (first, last) <- readSTRef obj
    cellv <- readSTRef first
    case cellv of
        STListCell Nothing -> return $ Nothing
        STListCell (Just (x, next)) -> do
            if next == ptr
                then do
                    writeSTRef obj (next, last)
                    return $ Just x
                else do
                    return $ Nothing

popUntil :: STListPtr s t -> STList s t -> ST s [t]
popUntil ptr (STList obj) = do
    (first, last) <- readSTRef obj
    if first == ptr
        then do
            return []
        else do
            cellv <- readSTRef first
            doPop cellv
    where
    doPop (STListCell Nothing) = return []
    doPop (STListCell (Just (x, next)))
        | next == ptr = do
            modifySTRef' obj (\(first, last) -> (next, last))
            return [x]
        | otherwise = do
            cellv <- readSTRef next
            xs <- doPop cellv
            return (x:xs)

foldr :: (t -> b -> ST s b) -> b -> STList s t -> ST s b
foldr f d (STList obj) = do
    (first, last) <- readSTRef obj
    cellv <- readSTRef first
    doFold cellv
    where
    doFold (STListCell Nothing) = return d
    doFold (STListCell (Just (x, next))) = do
        cellv <- readSTRef next
        right <- doFold cellv
        f x right

for :: STList s t -> (t -> ST s ()) -> ST s ()
for (STList obj) f = do
    (first, last) <- readSTRef obj
    cellv <- readSTRef first
    doFold cellv
    where
    doFold (STListCell Nothing) = return ()
    doFold (STListCell (Just (x, next))) = do
        f x
        cellv <- readSTRef next
        doFold cellv

reduce :: STList s t -> ST s [t]
reduce = foldr (\x xs -> return (x:xs)) []