Something went wrong. Try again.
Mutable List builder in the ST Monad
Something went wrong. Try again.
6.2 kB · 266 lines
Haskell
at master
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267{-# LANGUAGE DoAndIfThenElse, BangPatterns, NamedFieldPuns #-}-- | Mutable List Builder.---- A @ListBuilder s a@ is like a wrapper over @ST s [a]@, but uses unsafe-- mutation to achieve constant time append as well as prepend.---- As the builder is backed with a standard 'Data.List.List', it-- is light-weight and cheap to return to a list.---- Code from this module is derived from Scala's-- [ListBuffer](https://www.scala-lang.org/api/current/scala/collection/mutable/ListBuffer.html)-- module, using the unsafe set field technique described by-- [Twan van Laarhoven](https://www.twanvl.nl/blog/haskell/unsafe-sequence).module Data.ListBuilder ( -- * Mutable list builder ListBuilder
-- * Construction , newBuilder
-- * Mutations , append , prepend , insert , filterInPlace , clear
-- * Accessors , readLength , readFirst , readLast , readAt
-- * Conversions , freeze , unsafeFreeze ) where
import Data.ListBuilder.Unsafeimport qualified Data.List
import Control.Monad (when)import Control.Monad.STimport Control.Monad.ST.Unsafe
import Data.Foldable (foldr')import Data.Maybe (listToMaybe)import Data.STRef (STRef, newSTRef, readSTRef, writeSTRef, modifySTRef')
-- | A List Builder.---- This builder is backed by a standard haskell 'Data.List.List'.-- It offers predictable (and fast) operations, and doesn't-- pause for grow operations as an array based builder might.data ListBuilder s a = ListBuilder { start :: STRef s [a] , end :: STRef s [a] , len :: STRef s Int}
-- | Create a new, empty 'ListBuilder'newBuilder :: ST s (ListBuilder s a)newBuilder = do start <- newSTRef [] end <- newSTRef [] len <- newSTRef 0 pure $ ListBuilder start end len
-- | Append an item to the back of the 'ListBuilder'---- /O(1)/append :: a -> ListBuilder s a -> ST s ()append a ListBuilder { start, end, len } = do let !last' = [a] len' <- readSTRef len
if len' == 0 then do writeSTRef start last' writeSTRef end last' else do end' <- readSTRef end unsafeIOToST $ unsafeSetField 1 end' last' writeSTRef end last'
modifySTRef' len (+1)
-- | Prepend an item to the front of the 'ListBuilder'---- /O(1)/prepend :: a -> ListBuilder s a -> ST s ()prepend a ListBuilder { start, end, len } = do front <- readSTRef start len' <- readSTRef len
let !front' = a : front
when (len' == 0) $ writeSTRef end front'
writeSTRef start front' modifySTRef' len (+1)
-- | Internal function. Locates the previous cons cell.---- /O(N)/locate :: Int -> ListBuilder s a -> ST s [a]locate 0 _ = error "Internal error: locate called on 0"locate i ListBuilder { start, end, len } = do l <- readSTRef len
if l == i then readSTRef end
else do start' <- readSTRef start cur <- newSTRef start' let go 0 = readSTRef cur go j = do modifySTRef' cur (drop 1) go (j - 1) go (i - 1)
-- | Insert into a location in a 'ListBuilder'.---- This function doesn't create a new spine-- across the list builder, and only allocates-- the new cons cell itself.---- /O(N)/insert :: Int -> a -> ListBuilder s a -> ST s ()insert ix _ _ | ix < 0 = error "Index out of bounds"insert 0 a bldr = prepend a bldrinsert ix a bldr= do len' <- readSTRef (len bldr) if ix == len' then append a bldr else if ix > len' then error "Index out of bounds" else do prev <- locate ix bldr let !pn = drop 1 prev let !nx = a:pn unsafeIOToST $ unsafeSetField 1 prev nx
-- | The current length of the 'ListBuilder'.---- /O(1)/readLength :: ListBuilder s a -> ST s IntreadLength bldr = readSTRef (len bldr)
-- | Empty the 'ListBuilder' of all values.---- /O(1)/clear :: ListBuilder s a -> ST s ()clear ListBuilder { start, end, len } = do writeSTRef start [] writeSTRef end [] writeSTRef len 0
-- | Filter the 'ListBuilder' with the supplied predicate---- /O(N)/filterInPlace :: (a -> Bool) -> ListBuilder s a -> ST s ()filterInPlace func ListBuilder { start, end, len } = do prev <- newSTRef Nothing start' <- readSTRef start cur <- newSTRef start' let go = do cur' <- readSTRef cur case cur' of [] -> return ()
(h:follow) -> do prev' <- readSTRef prev if not (func h) then do case prev' of Nothing -> writeSTRef start follow Just y -> unsafeIOToST $ unsafeSetField 1 y follow
modifySTRef' len (\x -> x - 1) else writeSTRef prev (Just cur')
writeSTRef cur follow go
go
prev' <- readSTRef prev case prev' of Nothing -> do writeSTRef end [] Just y -> writeSTRef end y
-- | Return the current last element in the 'ListBuilder'---- /O(1)/readLast :: ListBuilder s a -> ST s (Maybe a)readLast ListBuilder { end } = do listToMaybe <$> readSTRef end
-- | Return the current first element in the 'ListBuilder'---- /O(1)/readFirst :: ListBuilder s a -> ST s (Maybe a)readFirst ListBuilder { start } = do listToMaybe <$> readSTRef start
-- | Return the current element at a particular index for-- the 'ListBuilder'---- /O(N)/readAt :: ListBuilder s a -> Int -> ST s (Maybe a)readAt ListBuilder { start } ix = do (Data.List.!? ix) <$> readSTRef start
-- | Return the 'Data.List.List' backing the 'ListBuilder'.---- This does /not/ stop mutations made to-- the builder from affecting the resultant-- list. So one must not continue to call the-- mutating functions.---- This function is safe in tail position within a-- call to @runST@.---- /O(1)/unsafeFreeze :: ListBuilder s a -> ST s [a]unsafeFreeze bldr = readSTRef (start bldr)
-- | Freeze the result and return the list.---- This function strictly copies the spine-- of the list.---- /O(n)/freeze :: ListBuilder s a -> ST s [a]freeze bldr = do aliased <- readSTRef (start bldr)
return $! foldr' (:) [] aliased