storablevector-0.2: Data/StorableVector/ST/Strict.hs

{-# OPTIONS_GHC -fglasgow-exts #-}
{- |
Module      : Data.StorableVector.ST.Strict
License     : BSD-style
Maintainer  : haskell@henning-thielemann.de
Stability   : experimental
Portability : portable, requires ffi
Tested with : GHC 6.4.1

Interface for access to a mutable StorableVector.
-}
module Data.StorableVector.ST.Strict (
        Vector,
        new,
        new_,
        read,
        write,
        freeze,
        thaw,
        runSTVector,
        mapST,
        mapSTLazy,
        ) where

import qualified Data.StorableVector.Base as V
import qualified Data.StorableVector as VS
import qualified Data.StorableVector.Lazy as VL

import qualified Control.Monad.ST.Strict as ST
import Control.Monad.ST.Strict (ST, unsafeIOToST, runST, )  -- stToIO,

import Foreign.ForeignPtr       (withForeignPtr, unsafeForeignPtrToPtr, )
import Foreign.Storable         (Storable(peek, poke, pokeElemOff))
import Foreign.Marshal.Array    (advancePtr, )
-- import System.IO.Unsafe         (unsafePerformIO)

-- import Prelude (Int, ($), (+), return, const, )
import Prelude hiding (read, )


newtype Vector s e = SV {vector :: V.Vector e}


{-# INLINE new #-}
{-# INLINE new_ #-}
{-# INLINE read #-}
{-# INLINE write #-}
{-# INLINE freeze #-}
{-# INLINE thaw #-}
{-# INLINE runSTVector #-}
{-# INLINE mapST #-}
{-# INLINE mapSTLazy #-}


-- * access to mutable storable vector

new :: (Storable e) =>
   Int -> e -> ST s (Vector s e)
new n x = return (SV (VS.replicate n x))

new_ :: (Storable e) =>
   Int -> ST s (Vector s e)
new_  =  fmap SV . newVec_

{-# INLINE newVec_ #-}
newVec_ :: (Storable e) =>
   Int -> ST s (VS.Vector e)
newVec_ n =
   -- return (V.unsafeCreate n (const (return ())))
   unsafeIOToST $ V.create n (const (return ()))

{- |
> Control.Monad.ST.runST (do arr <- new_ 10; Monad.zipWithM_ (write arr) [9,8..0] ['a'..]; read arr 3)
-}
read :: (Storable e) =>
   Vector s e -> Int -> ST s e
read (SV xs) n = return (VS.index xs n)

{- |
> VS.unpack $ runSTVector (do arr <- new_ 10; Monad.zipWithM_ (write arr) [9,8..0] ['a'..]; return arr)
-}
write :: (Storable e) =>
   Vector s e -> Int -> e -> ST s ()
write (SV (V.SV v s l)) n x =
   if 0<=n && n<l
     then unsafeIOToST (withForeignPtr v $ \p -> pokeElemOff p (s+n) x)
     else error "StorableVector.ST.Strict.write: index out of range"

freeze :: (Storable e) =>
   Vector s e -> ST s (VS.Vector e)
freeze (SV xs) = return (VS.copy xs)

thaw :: (Storable e) =>
   VS.Vector e -> ST s (Vector s e)
thaw xs = return (SV (VS.copy xs))


runSTVector :: (Storable e) =>
   (forall s. ST s (Vector s e)) -> VS.Vector e
runSTVector m =
   runST (fmap vector m)
--   vector (unsafePerformIO (stToIO m))



-- * operations on immutable storable vector within ST monad

{- |
> :module + Data.STRef
> VS.unpack $ Control.Monad.ST.runST (do ref <- newSTRef 'a'; mapST (\ _n -> do c <- readSTRef ref; modifySTRef ref succ; return c) (VS.pack [1,2,3,4::Data.Int.Int16]))
-}
mapST :: (Storable a, Storable b) =>
   (a -> ST s b) -> VS.Vector a -> ST s (VS.Vector b)
mapST f (V.SV px sx n) =
   let {-# INLINE go #-}
       go l q p =
          if l>0
            then
               do unsafeIOToST . poke p =<< f =<< unsafeIOToST (peek q)
                  go (pred l) (advancePtr q 1) (advancePtr p 1)
            else return ()
   in  do ys@(V.SV py sy _) <- newVec_ n
          go n
              (advancePtr (unsafeForeignPtrToPtr px) sx)
              (advancePtr (unsafeForeignPtrToPtr py) sy)
          return ys

{-
mapST f xs@(V.SV v s l) =
   let go l q p =
          if l>0
            then
               do poke p =<< stToIO . f =<< peek q
                  go (pred l) (advancePtr q 1) (advancePtr p 1)
            else return ()
       n = VS.length xs
   in  return $ V.unsafeCreate n $ \p ->
          withForeignPtr v $ \q -> go n (advancePtr q s) p
-}


{- |
> *Data.StorableVector.ST.Strict Data.STRef> VL.unpack $ Control.Monad.ST.runST (do ref <- newSTRef 'a'; mapSTLazy (\ _n -> do c <- readSTRef ref; modifySTRef ref succ; return c) (VL.pack VL.defaultChunkSize [1,2,3,4::Data.Int.Int16]))
> "abcd"

The following should not work on infinite streams,
since we are in 'ST' with strict '>>='.
But it works. Why?

> *Data.StorableVector.ST.Strict Data.STRef> VL.unpack $ Control.Monad.ST.runST (do ref <- newSTRef 'a'; mapSTLazy (\ _n -> do c <- readSTRef ref; modifySTRef ref succ; return c) (VL.pack VL.defaultChunkSize [0::Data.Int.Int16 ..]))
> "Interrupted.
-}
mapSTLazy :: (Storable a, Storable b) =>
   (a -> ST s b) -> VL.Vector a -> ST s (VL.Vector b)
mapSTLazy f (VL.SV xs) =
   fmap VL.SV $ mapM (mapST f) xs