{-# LANGUAGE Rank2Types #-}
{- |
Build a lazy storable vector by incrementally adding an element.
This is analogous to Data.Binary.Builder for Data.ByteString.Lazy.
-}
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.Lazy as STV
import qualified Data.StorableVector.ST.Private as STVP
import qualified Control.Monad.Trans.RWS as RWS

import Foreign.Storable (Storable, )
import Data.StorableVector.Lazy (ChunkSize(ChunkSize), )
import Control.Monad.ST.Lazy (ST, runST, strictToLazyST, )
import Control.Monad.Trans.RWS (RWST, runRWST, )
import Control.Monad.Trans (lift, )
import Data.Monoid (Monoid(mempty, mappend), Endo(Endo), appEndo, )


newtype Builder a =
   Builder {run :: forall s.
      RWST ChunkSize (Endo [SV.Vector a]) (STV.Vector s a, Int) (ST s) ()}


-- instance Monoid (Builder a) where
{-
Storable constraint not need in the current implementation,
but who knows what will be in future ...
-}
instance Storable a => Monoid (Builder a) where
   mempty = Builder (return ())
   mappend x y = Builder (run x >> run y)


{-
SVL.unpack $ toLazyStorableVector (ChunkSize 7) $ Data.Monoid.mconcat $ map put ['a'..'z']
-}
toLazyStorableVector :: Storable a =>
   ChunkSize -> Builder a -> SVL.Vector a
toLazyStorableVector cs@(SVL.ChunkSize size) bld =
   runST (do
      v0 <- STV.new_ size
      (_,vi1,chunks) <- runRWST (run bld) cs (v0,0)
      lastChunk <- fixVector vi1
      return $ SVL.fromChunks $ appEndo chunks [lastChunk])

put :: Storable a => a -> Builder a
put a =
   Builder (
      do (SVL.ChunkSize size) <- RWS.ask
         (v0,i0) <- RWS.get
         (v1,i1) <-
            if i0<size
              then return (v0,i0)
              else
                -- we could call 'flush' here, but this requires an extra 'SV.take'
                do RWS.tell . Endo . (:) =<<
                      (lift $ strictToLazyST $ STVP.unsafeToVector v0)
                   lift $ fmap (flip (,) 0) $ STV.new_ size
         lift $ STV.write v1 i1 a
         RWS.put (v1, succ i1)
   )

{- |
Set a laziness break.
-}
flush :: Storable a => Builder a
flush =
   Builder (
      do RWS.tell . Endo . (:) =<< lift . fixVector =<< RWS.get
         (SVL.ChunkSize size) <- RWS.ask
         v1 <- lift $ STV.new_ size
         RWS.put (v1, 0)
   )

fixVector :: (Storable a) =>
   (STVP.Vector s a, Int) -> ST s (SV.Vector a)
fixVector ~(v1,i1) =
   fmap (SV.take i1) $
   strictToLazyST $ STVP.unsafeToVector v1