{-# LANGUAGE Rank2Types #-}
{- |
Build a lazy storable vector by incrementally adding an element.
This is analogous to Data.Binary.Builder for Data.ByteString.Lazy.

Attention:
This implementation is still almost 3 times slower
than constructing a lazy storable vector using unfoldr
in our Chorus speed test.
-}
module Data.StorableVector.Lazy.Builder (
   Builder,
   toLazyStorableVector,
   put,
   flush,
   ) where

import qualified Data.StorableVector as SV
import qualified Data.StorableVector.Lazy as SVL
import qualified Data.StorableVector.ST.Strict as STV
-- import qualified Data.StorableVector.ST.Lazy as STVL

import Data.StorableVector.Lazy (ChunkSize, )
import Control.Monad (liftM2, )
import Control.Monad.ST.Strict (ST, runST, )
import Data.Monoid (Monoid(mempty, mappend), )
import Data.Semigroup (Semigroup((<>)), )

import Foreign.Storable (Storable, )

import qualified System.Unsafe as Unsafe


{-
Given an initial buffer and a function that generates the rest of the vector,
a 'Builder' generates the whole vector.
The idea is inspired by Data.Binary.Builder.

We use the strict ST monad by default
and only rare 'Unsafe.interleaveST',
since this is more efficient than using lazy ST everywhere.

Before that approach I tried to achieve this with a lazy State monad.
I found this more comprehensible but it was very slow
and had a space leak, when the last chunk shall be handled correctly.
-}
newtype Builder a =
   Builder {forall a.
Builder a
-> forall s.
   ChunkSize
   -> (Buffer s a -> ST s [Vector a]) -> Buffer s a -> ST s [Vector a]
run :: forall s.
      ChunkSize ->
      (Buffer s a -> ST s [SV.Vector a]) ->
      (Buffer s a -> ST s [SV.Vector a])
   }

type Buffer s a = (STV.Vector s a, Int)


-- instance Monoid (Builder a) where
{-
Storable constraint not needed in the current implementation,
but who knows what will be in future ...
-}
instance Storable a => Semigroup (Builder a) where
   {-# INLINE (<>) #-}
   Builder a
x <> :: Builder a -> Builder a -> Builder a
<> Builder a
y = forall a.
(forall s.
 ChunkSize
 -> (Buffer s a -> ST s [Vector a])
 -> Buffer s a
 -> ST s [Vector a])
-> Builder a
Builder (\ChunkSize
cs -> forall a.
Builder a
-> forall s.
   ChunkSize
   -> (Buffer s a -> ST s [Vector a]) -> Buffer s a -> ST s [Vector a]
run Builder a
x ChunkSize
cs forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a.
Builder a
-> forall s.
   ChunkSize
   -> (Buffer s a -> ST s [Vector a]) -> Buffer s a -> ST s [Vector a]
run Builder a
y ChunkSize
cs)

instance Storable a => Monoid (Builder a) where
   {-# INLINE mempty #-}
   {-# INLINE mappend #-}
   mempty :: Builder a
mempty = forall a.
(forall s.
 ChunkSize
 -> (Buffer s a -> ST s [Vector a])
 -> Buffer s a
 -> ST s [Vector a])
-> Builder a
Builder (\ChunkSize
_ -> forall a. a -> a
id)
   mappend :: Builder a -> Builder a -> Builder a
mappend = forall a. Semigroup a => a -> a -> a
(<>)


{- |
> toLazyStorableVector (ChunkSize 7) $ foldMap put ['a'..'z']
-}
{-# INLINE toLazyStorableVector #-}
toLazyStorableVector :: Storable a =>
   ChunkSize -> Builder a -> SVL.Vector a
toLazyStorableVector :: forall a. Storable a => ChunkSize -> Builder a -> Vector a
toLazyStorableVector ChunkSize
cs Builder a
bld =
   forall a. Storable a => [Vector a] -> Vector a
SVL.fromChunks forall a b. (a -> b) -> a -> b
$
   forall a. (forall s. ST s a) -> a
runST (forall a.
Builder a
-> forall s.
   ChunkSize
   -> (Buffer s a -> ST s [Vector a]) -> Buffer s a -> ST s [Vector a]
run Builder a
bld ChunkSize
cs (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall a. a -> [a] -> [a]
:[]) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a s. Storable a => Buffer s a -> ST s (Vector a)
fixVector) forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall a s. Storable a => ChunkSize -> ST s (Buffer s a)
newChunk ChunkSize
cs)


{-# INLINE put #-}
put :: Storable a => a -> Builder a
put :: forall a. Storable a => a -> Builder a
put a
a =
   forall a.
(forall s.
 ChunkSize
 -> (Buffer s a -> ST s [Vector a])
 -> Buffer s a
 -> ST s [Vector a])
-> Builder a
Builder (\ChunkSize
cs Buffer s a -> ST s [Vector a]
cont (Vector s a
v0,Int
i0) ->
      do forall e s. Storable e => Vector s e -> Int -> e -> ST s ()
STV.unsafeWrite Vector s a
v0 Int
i0 a
a
         let i1 :: Int
i1 = forall a. Enum a => a -> a
succ Int
i0
         if Int
i1 forall a. Ord a => a -> a -> Bool
< forall s e. Vector s e -> Int
STV.length Vector s a
v0
           then
             Buffer s a -> ST s [Vector a]
cont (Vector s a
v0, Int
i1)
           else
             forall (m :: * -> *) a1 a2 r.
Monad m =>
(a1 -> a2 -> r) -> m a1 -> m a2 -> m r
liftM2 (:)
                -- we could call 'flush' here, but this requires an extra 'SV.take'
                (forall e s. Storable e => Vector s e -> ST s (Vector e)
STV.unsafeFreeze Vector s a
v0)
                (forall s a. ST s a -> ST s a
Unsafe.interleaveST forall a b. (a -> b) -> a -> b
$
                 Buffer s a -> ST s [Vector a]
cont forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall a s. Storable a => ChunkSize -> ST s (Buffer s a)
newChunk ChunkSize
cs)
   )

{-
put :: Storable a => a -> Builder a
put a =
   Builder (\cs cont (v0,i0) ->
      if i0 < STV.length v0
        then
          do STV.write v0 i0 a
             cont (v0, succ i0)
        else
          liftM2 (:)
             -- we could call 'flush' here, but this requires an extra 'SV.take'
             (STV.unsafeFreeze v0)
             (Unsafe.interleaveST $
              do (v1,i1) <- newChunk cs
                 STV.write v1 i1 a
                 cont (v1, succ i1))
   )
-}

{-
          lazyToStrictST $
          liftM2 (:)
             -- we could call 'flush' here, but this requires an extra 'SV.take'
             (STVL.unsafeFreeze v0)
             (strictToLazyST $
              do (v1,i1) <- newChunk cs
                 STV.write v1 i1 a
                 cont (v1, succ i1))
-}

{-
Prelude Control.Monad.ST.Lazy> Control.Monad.ST.runST (lazyToStrictST $ Monad.liftM2 (,) (strictToLazyST $ return 'a') (strictToLazyST (undefined::Monad m => m Char)))
*** Exception: Prelude.undefined
-}

{- |
Set a laziness break.
-}
{-# INLINE flush #-}
flush :: Storable a => Builder a
flush :: forall a. Storable a => Builder a
flush =
   forall a.
(forall s.
 ChunkSize
 -> (Buffer s a -> ST s [Vector a])
 -> Buffer s a
 -> ST s [Vector a])
-> Builder a
Builder (\ChunkSize
cs Buffer s a -> ST s [Vector a]
cont Buffer s a
vi0 ->
      forall (m :: * -> *) a1 a2 r.
Monad m =>
(a1 -> a2 -> r) -> m a1 -> m a2 -> m r
liftM2 (:)
         (forall a s. Storable a => Buffer s a -> ST s (Vector a)
fixVector Buffer s a
vi0)
         (forall s a. ST s a -> ST s a
Unsafe.interleaveST forall a b. (a -> b) -> a -> b
$ Buffer s a -> ST s [Vector a]
cont forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall a s. Storable a => ChunkSize -> ST s (Buffer s a)
newChunk ChunkSize
cs)
{-
      lazyToStrictST $
      liftM2 (:)
         (strictToLazyST $ fixVector vi0)
         (strictToLazyST $ cont =<< newChunk cs)
-}
   )

{-# INLINE newChunk #-}
newChunk :: (Storable a) =>
   ChunkSize -> ST s (Buffer s a)
newChunk :: forall a s. Storable a => ChunkSize -> ST s (Buffer s a)
newChunk (SVL.ChunkSize Int
size) =
   forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall a b c. (a -> b -> c) -> b -> a -> c
flip (,) Int
0) forall a b. (a -> b) -> a -> b
$ forall e s. Storable e => Int -> ST s (Vector s e)
STV.new_ Int
size

{-# INLINE fixVector #-}
fixVector :: (Storable a) =>
   Buffer s a -> ST s (SV.Vector a)
fixVector :: forall a s. Storable a => Buffer s a -> ST s (Vector a)
fixVector ~(Vector s a
v1,Int
i1) =
   forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall a. Storable a => Int -> Vector a -> Vector a
SV.take Int
i1) forall a b. (a -> b) -> a -> b
$ forall e s. Storable e => Vector s e -> ST s (Vector e)
STV.unsafeFreeze Vector s a
v1