{-# 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.Unsafe
import qualified Data.List

import Control.Monad (when)
import Control.Monad.ST
import 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 {
    forall s a. ListBuilder s a -> STRef s [a]
start :: STRef s [a]
  , forall s a. ListBuilder s a -> STRef s [a]
end :: STRef s [a]
  , forall s a. ListBuilder s a -> STRef s Int
len :: STRef s Int
}

-- | Create a new, empty 'ListBuilder'
newBuilder :: ST s (ListBuilder s a)
newBuilder :: forall s a. ST s (ListBuilder s a)
newBuilder = do
  STRef s [a]
start <- [a] -> ST s (STRef s [a])
forall a s. a -> ST s (STRef s a)
newSTRef []
  STRef s [a]
end   <- [a] -> ST s (STRef s [a])
forall a s. a -> ST s (STRef s a)
newSTRef []
  STRef s Int
len   <- Int -> ST s (STRef s Int)
forall a s. a -> ST s (STRef s a)
newSTRef Int
0
  ListBuilder s a -> ST s (ListBuilder s a)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ListBuilder s a -> ST s (ListBuilder s a))
-> ListBuilder s a -> ST s (ListBuilder s a)
forall a b. (a -> b) -> a -> b
$
    STRef s [a] -> STRef s [a] -> STRef s Int -> ListBuilder s a
forall s a.
STRef s [a] -> STRef s [a] -> STRef s Int -> ListBuilder s a
ListBuilder STRef s [a]
start STRef s [a]
end STRef s Int
len

-- | Append an item to the back of the 'ListBuilder'
--
--   /O(1)/
append :: a -> ListBuilder s a -> ST s ()
append :: forall a s. a -> ListBuilder s a -> ST s ()
append a
a ListBuilder { STRef s [a]
start :: forall s a. ListBuilder s a -> STRef s [a]
start :: STRef s [a]
start, STRef s [a]
end :: forall s a. ListBuilder s a -> STRef s [a]
end :: STRef s [a]
end, STRef s Int
len :: forall s a. ListBuilder s a -> STRef s Int
len :: STRef s Int
len } = do
  let
    !last' :: [a]
last' = [a
a]
  Int
len' <- STRef s Int -> ST s Int
forall s a. STRef s a -> ST s a
readSTRef STRef s Int
len

  if Int
len' Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 then do
    STRef s [a] -> [a] -> ST s ()
forall s a. STRef s a -> a -> ST s ()
writeSTRef STRef s [a]
start [a]
last'
    STRef s [a] -> [a] -> ST s ()
forall s a. STRef s a -> a -> ST s ()
writeSTRef STRef s [a]
end [a]
last'
  else do
    [a]
end' <- STRef s [a] -> ST s [a]
forall s a. STRef s a -> ST s a
readSTRef STRef s [a]
end
    IO () -> ST s ()
forall a s. IO a -> ST s a
unsafeIOToST (IO () -> ST s ()) -> IO () -> ST s ()
forall a b. (a -> b) -> a -> b
$
      Int -> [a] -> [a] -> IO ()
forall a b. Int -> a -> b -> IO ()
unsafeSetField Int
1 [a]
end' [a]
last'
    STRef s [a] -> [a] -> ST s ()
forall s a. STRef s a -> a -> ST s ()
writeSTRef STRef s [a]
end [a]
last'

  STRef s Int -> (Int -> Int) -> ST s ()
forall s a. STRef s a -> (a -> a) -> ST s ()
modifySTRef' STRef s Int
len (Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)

-- | Prepend an item to the front of the 'ListBuilder'
--
--   /O(1)/
prepend :: a -> ListBuilder s a -> ST s ()
prepend :: forall a s. a -> ListBuilder s a -> ST s ()
prepend a
a ListBuilder { STRef s [a]
start :: forall s a. ListBuilder s a -> STRef s [a]
start :: STRef s [a]
start, STRef s [a]
end :: forall s a. ListBuilder s a -> STRef s [a]
end :: STRef s [a]
end, STRef s Int
len :: forall s a. ListBuilder s a -> STRef s Int
len :: STRef s Int
len } = do
  [a]
front <- STRef s [a] -> ST s [a]
forall s a. STRef s a -> ST s a
readSTRef STRef s [a]
start
  Int
len'  <- STRef s Int -> ST s Int
forall s a. STRef s a -> ST s a
readSTRef STRef s Int
len

  let
    !front' :: [a]
front' = a
a a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
front

  Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
len' Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$
    STRef s [a] -> [a] -> ST s ()
forall s a. STRef s a -> a -> ST s ()
writeSTRef STRef s [a]
end [a]
front'

  STRef s [a] -> [a] -> ST s ()
forall s a. STRef s a -> a -> ST s ()
writeSTRef STRef s [a]
start [a]
front'
  STRef s Int -> (Int -> Int) -> ST s ()
forall s a. STRef s a -> (a -> a) -> ST s ()
modifySTRef' STRef s Int
len (Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)


-- | Internal function. Locates the previous cons cell.
--
--   /O(N)/
locate :: Int -> ListBuilder s a -> ST s [a]
locate :: forall s a. Int -> ListBuilder s a -> ST s [a]
locate Int
0 ListBuilder s a
_ = [Char] -> ST s [a]
forall a. HasCallStack => [Char] -> a
error [Char]
"Internal error: locate called on 0"
locate Int
i ListBuilder { STRef s [a]
start :: forall s a. ListBuilder s a -> STRef s [a]
start :: STRef s [a]
start, STRef s [a]
end :: forall s a. ListBuilder s a -> STRef s [a]
end :: STRef s [a]
end, STRef s Int
len :: forall s a. ListBuilder s a -> STRef s Int
len :: STRef s Int
len } = do
  Int
l <- STRef s Int -> ST s Int
forall s a. STRef s a -> ST s a
readSTRef STRef s Int
len

  if Int
l Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
i then
    STRef s [a] -> ST s [a]
forall s a. STRef s a -> ST s a
readSTRef STRef s [a]
end

  else do
    [a]
start' <- STRef s [a] -> ST s [a]
forall s a. STRef s a -> ST s a
readSTRef STRef s [a]
start
    STRef s [a]
cur    <- [a] -> ST s (STRef s [a])
forall a s. a -> ST s (STRef s a)
newSTRef [a]
start'
    let
      go :: t -> ST s [a]
go t
0 = STRef s [a] -> ST s [a]
forall s a. STRef s a -> ST s a
readSTRef STRef s [a]
cur
      go t
j = do
        STRef s [a] -> ([a] -> [a]) -> ST s ()
forall s a. STRef s a -> (a -> a) -> ST s ()
modifySTRef' STRef s [a]
cur (Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
drop Int
1)
        t -> ST s [a]
go (t
j t -> t -> t
forall a. Num a => a -> a -> a
- t
1)
    Int -> ST s [a]
forall {t}. (Eq t, Num t) => t -> ST s [a]
go (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
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 :: forall a s. Int -> a -> ListBuilder s a -> ST s ()
insert Int
ix a
_ ListBuilder s a
_ | Int
ix Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 = [Char] -> ST s ()
forall a. HasCallStack => [Char] -> a
error [Char]
"Index out of bounds"
insert Int
0 a
a ListBuilder s a
bldr = a -> ListBuilder s a -> ST s ()
forall a s. a -> ListBuilder s a -> ST s ()
prepend a
a ListBuilder s a
bldr
insert Int
ix a
a ListBuilder s a
bldr= do
  Int
len' <- STRef s Int -> ST s Int
forall s a. STRef s a -> ST s a
readSTRef (ListBuilder s a -> STRef s Int
forall s a. ListBuilder s a -> STRef s Int
len ListBuilder s a
bldr)
  if Int
ix Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
len' then
    a -> ListBuilder s a -> ST s ()
forall a s. a -> ListBuilder s a -> ST s ()
append a
a ListBuilder s a
bldr
  else if Int
ix Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
len' then
    [Char] -> ST s ()
forall a. HasCallStack => [Char] -> a
error [Char]
"Index out of bounds"
  else do
    [a]
prev  <- Int -> ListBuilder s a -> ST s [a]
forall s a. Int -> ListBuilder s a -> ST s [a]
locate Int
ix ListBuilder s a
bldr
    let !pn :: [a]
pn = Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
drop Int
1 [a]
prev
    let !nx :: [a]
nx = a
aa -> [a] -> [a]
forall a. a -> [a] -> [a]
:[a]
pn
    IO () -> ST s ()
forall a s. IO a -> ST s a
unsafeIOToST (IO () -> ST s ()) -> IO () -> ST s ()
forall a b. (a -> b) -> a -> b
$
      Int -> [a] -> [a] -> IO ()
forall a b. Int -> a -> b -> IO ()
unsafeSetField Int
1 [a]
prev [a]
nx


-- | The current length of the 'ListBuilder'.
--
--   /O(1)/
readLength :: ListBuilder s a -> ST s Int
readLength :: forall s a. ListBuilder s a -> ST s Int
readLength ListBuilder s a
bldr =
  STRef s Int -> ST s Int
forall s a. STRef s a -> ST s a
readSTRef (ListBuilder s a -> STRef s Int
forall s a. ListBuilder s a -> STRef s Int
len ListBuilder s a
bldr)


-- | Empty the 'ListBuilder' of all values.
--
--   /O(1)/
clear :: ListBuilder s a -> ST s ()
clear :: forall s a. ListBuilder s a -> ST s ()
clear ListBuilder { STRef s [a]
start :: forall s a. ListBuilder s a -> STRef s [a]
start :: STRef s [a]
start, STRef s [a]
end :: forall s a. ListBuilder s a -> STRef s [a]
end :: STRef s [a]
end, STRef s Int
len :: forall s a. ListBuilder s a -> STRef s Int
len :: STRef s Int
len } = do
  STRef s [a] -> [a] -> ST s ()
forall s a. STRef s a -> a -> ST s ()
writeSTRef STRef s [a]
start []
  STRef s [a] -> [a] -> ST s ()
forall s a. STRef s a -> a -> ST s ()
writeSTRef STRef s [a]
end []
  STRef s Int -> Int -> ST s ()
forall s a. STRef s a -> a -> ST s ()
writeSTRef STRef s Int
len Int
0


-- | Filter the 'ListBuilder' with the supplied predicate
--
--   /O(N)/
filterInPlace :: (a -> Bool) -> ListBuilder s a -> ST s ()
filterInPlace :: forall a s. (a -> Bool) -> ListBuilder s a -> ST s ()
filterInPlace a -> Bool
func ListBuilder { STRef s [a]
start :: forall s a. ListBuilder s a -> STRef s [a]
start :: STRef s [a]
start, STRef s [a]
end :: forall s a. ListBuilder s a -> STRef s [a]
end :: STRef s [a]
end, STRef s Int
len :: forall s a. ListBuilder s a -> STRef s Int
len :: STRef s Int
len } = do
  STRef s (Maybe [a])
prev   <- Maybe [a] -> ST s (STRef s (Maybe [a]))
forall a s. a -> ST s (STRef s a)
newSTRef Maybe [a]
forall a. Maybe a
Nothing
  [a]
start' <- STRef s [a] -> ST s [a]
forall s a. STRef s a -> ST s a
readSTRef STRef s [a]
start
  STRef s [a]
cur    <- [a] -> ST s (STRef s [a])
forall a s. a -> ST s (STRef s a)
newSTRef [a]
start'
  let
    go :: ST s ()
go = do
      [a]
cur' <- STRef s [a] -> ST s [a]
forall s a. STRef s a -> ST s a
readSTRef STRef s [a]
cur
      case [a]
cur' of
        [] -> () -> ST s ()
forall a. a -> ST s a
forall (m :: * -> *) a. Monad m => a -> m a
return ()

        (a
h:[a]
follow) -> do
          Maybe [a]
prev' <- STRef s (Maybe [a]) -> ST s (Maybe [a])
forall s a. STRef s a -> ST s a
readSTRef STRef s (Maybe [a])
prev
          if Bool -> Bool
not (a -> Bool
func a
h) then do
            case Maybe [a]
prev' of
              Maybe [a]
Nothing ->
                STRef s [a] -> [a] -> ST s ()
forall s a. STRef s a -> a -> ST s ()
writeSTRef STRef s [a]
start [a]
follow
              Just [a]
y ->
                IO () -> ST s ()
forall a s. IO a -> ST s a
unsafeIOToST (IO () -> ST s ()) -> IO () -> ST s ()
forall a b. (a -> b) -> a -> b
$
                  Int -> [a] -> [a] -> IO ()
forall a b. Int -> a -> b -> IO ()
unsafeSetField Int
1 [a]
y [a]
follow

            STRef s Int -> (Int -> Int) -> ST s ()
forall s a. STRef s a -> (a -> a) -> ST s ()
modifySTRef' STRef s Int
len (\Int
x -> Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
          else
            STRef s (Maybe [a]) -> Maybe [a] -> ST s ()
forall s a. STRef s a -> a -> ST s ()
writeSTRef STRef s (Maybe [a])
prev ([a] -> Maybe [a]
forall a. a -> Maybe a
Just [a]
cur')

          STRef s [a] -> [a] -> ST s ()
forall s a. STRef s a -> a -> ST s ()
writeSTRef STRef s [a]
cur [a]
follow
          ST s ()
go

  ST s ()
go

  Maybe [a]
prev' <- STRef s (Maybe [a]) -> ST s (Maybe [a])
forall s a. STRef s a -> ST s a
readSTRef STRef s (Maybe [a])
prev
  case Maybe [a]
prev' of
    Maybe [a]
Nothing -> do
      STRef s [a] -> [a] -> ST s ()
forall s a. STRef s a -> a -> ST s ()
writeSTRef STRef s [a]
end []
    Just [a]
y ->
      STRef s [a] -> [a] -> ST s ()
forall s a. STRef s a -> a -> ST s ()
writeSTRef STRef s [a]
end [a]
y


-- | Return the current last element in the 'ListBuilder'
--
--   /O(1)/
readLast :: ListBuilder s a -> ST s (Maybe a)
readLast :: forall s a. ListBuilder s a -> ST s (Maybe a)
readLast ListBuilder { STRef s [a]
end :: forall s a. ListBuilder s a -> STRef s [a]
end :: STRef s [a]
end } = do
  [a] -> Maybe a
forall a. [a] -> Maybe a
listToMaybe ([a] -> Maybe a) -> ST s [a] -> ST s (Maybe a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> STRef s [a] -> ST s [a]
forall s a. STRef s a -> ST s a
readSTRef STRef s [a]
end



-- | Return the current first element in the 'ListBuilder'
--
--   /O(1)/
readFirst :: ListBuilder s a -> ST s (Maybe a)
readFirst :: forall s a. ListBuilder s a -> ST s (Maybe a)
readFirst ListBuilder { STRef s [a]
start :: forall s a. ListBuilder s a -> STRef s [a]
start :: STRef s [a]
start } = do
  [a] -> Maybe a
forall a. [a] -> Maybe a
listToMaybe ([a] -> Maybe a) -> ST s [a] -> ST s (Maybe a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> STRef s [a] -> ST s [a]
forall s a. STRef s a -> ST s a
readSTRef STRef s [a]
start


-- | Return the current element at a particular index for
--   the 'ListBuilder'
--
--   /O(N)/
readAt :: ListBuilder s a -> Int -> ST s (Maybe a)
readAt :: forall s a. ListBuilder s a -> Int -> ST s (Maybe a)
readAt ListBuilder { STRef s [a]
start :: forall s a. ListBuilder s a -> STRef s [a]
start :: STRef s [a]
start } Int
ix = do
  ([a] -> Int -> Maybe a
forall a. [a] -> Int -> Maybe a
Data.List.!? Int
ix) ([a] -> Maybe a) -> ST s [a] -> ST s (Maybe a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> STRef s [a] -> ST s [a]
forall s a. STRef s a -> ST s a
readSTRef STRef s [a]
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 :: forall s a. ListBuilder s a -> ST s [a]
unsafeFreeze ListBuilder s a
bldr =
  STRef s [a] -> ST s [a]
forall s a. STRef s a -> ST s a
readSTRef (ListBuilder s a -> STRef s [a]
forall s a. ListBuilder s a -> STRef s [a]
start ListBuilder s a
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 :: forall s a. ListBuilder s a -> ST s [a]
freeze ListBuilder s a
bldr = do
  [a]
aliased <-
    STRef s [a] -> ST s [a]
forall s a. STRef s a -> ST s a
readSTRef (ListBuilder s a -> STRef s [a]
forall s a. ListBuilder s a -> STRef s [a]
start ListBuilder s a
bldr)

  [a] -> ST s [a]
forall a. a -> ST s a
forall (m :: * -> *) a. Monad m => a -> m a
return ([a] -> ST s [a]) -> [a] -> ST s [a]
forall a b. (a -> b) -> a -> b
$!
    (a -> [a] -> [a]) -> [a] -> [a] -> [a]
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr' (:) [] [a]
aliased