{-# LANGUAGE Rank2Types #-}
{- |
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,
        modify,
        maybeRead,
        maybeWrite,
        maybeModify,
        unsafeRead,
        unsafeWrite,
        unsafeModify,
        freeze,
        unsafeFreeze,
        thaw,
        length,
        runSTVector,
        mapST,
        mapSTLazy,
        ) where

import Data.StorableVector.ST.Private
          (Vector(SV), create, unsafeCreate, unsafeToVector, )
import qualified Data.StorableVector.Base as V
import qualified Data.StorableVector as VS
import qualified Data.StorableVector.Lazy as VL

import Control.Monad.ST.Strict (ST, runST, )

import Foreign.Ptr              (Ptr, )
import Foreign.ForeignPtr       (withForeignPtr, )
import Foreign.Storable         (Storable(peek, poke))
import Foreign.Marshal.Array    (advancePtr, copyArray, )
import qualified System.Unsafe as Unsafe

import qualified Data.Traversable as Traversable
import Data.Maybe.HT (toMaybe, )
import Data.Maybe (isJust, )

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


-- * access to mutable storable vector

{-# INLINE new #-}
new :: (Storable e) =>
   Int -> e -> ST s (Vector s e)
new :: forall e s. Storable e => Int -> e -> ST s (Vector s e)
new Int
n e
x =
   forall a s.
Storable a =>
Int -> (Ptr a -> IO ()) -> ST s (Vector s a)
unsafeCreate Int
n forall a b. (a -> b) -> a -> b
$
   let {-# INLINE go #-}
       go :: t -> Ptr e -> IO ()
go t
m Ptr e
p =
         if t
mforall a. Ord a => a -> a -> Bool
>t
0
           then forall a. Storable a => Ptr a -> a -> IO ()
poke Ptr e
p e
x forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> t -> Ptr e -> IO ()
go (forall a. Enum a => a -> a
pred t
m) (forall a. Storable a => Ptr a -> Ptr a
V.incPtr Ptr e
p)
           else forall (m :: * -> *) a. Monad m => a -> m a
return ()
   in  forall {t}. (Ord t, Num t, Enum t) => t -> Ptr e -> IO ()
go Int
n

{-# INLINE new_ #-}
new_ :: (Storable e) =>
   Int -> ST s (Vector s e)
new_ :: forall e s. Storable e => Int -> ST s (Vector s e)
new_ Int
n =
   forall a s.
Storable a =>
Int -> (Ptr a -> IO ()) -> ST s (Vector s a)
unsafeCreate Int
n (forall a b. a -> b -> a
const (forall (m :: * -> *) a. Monad m => a -> m a
return ()))


{- |
> Control.Monad.ST.runST (do arr <- new_ 10; Monad.zipWithM_ (write arr) [9,8..0] ['a'..]; read arr 3)
-}
{-# INLINE read #-}
read :: (Storable e) =>
   Vector s e -> Int -> ST s e
read :: forall e s. Storable e => Vector s e -> Int -> ST s e
read Vector s e
v Int
n =
   forall e s a.
Storable e =>
String -> Vector s e -> Int -> ST s a -> ST s a
access String
"read" Vector s e
v Int
n forall a b. (a -> b) -> a -> b
$ forall e s. Storable e => Vector s e -> Int -> ST s e
unsafeRead Vector s e
v Int
n

{- |
> VS.unpack $ runSTVector (do arr <- new_ 10; Monad.zipWithM_ (write arr) [9,8..0] ['a'..]; return arr)
-}
{-# INLINE write #-}
write :: (Storable e) =>
   Vector s e -> Int -> e -> ST s ()
write :: forall e s. Storable e => Vector s e -> Int -> e -> ST s ()
write Vector s e
v Int
n e
x =
   forall e s a.
Storable e =>
String -> Vector s e -> Int -> ST s a -> ST s a
access String
"write" Vector s e
v Int
n forall a b. (a -> b) -> a -> b
$ forall e s. Storable e => Vector s e -> Int -> e -> ST s ()
unsafeWrite Vector s e
v Int
n e
x

{- |
> VS.unpack $ runSTVector (do arr <- new 10 'a'; Monad.mapM_ (\n -> modify arr (mod n 8) succ) [0..10]; return arr)
-}
{-# INLINE modify #-}
modify :: (Storable e) =>
   Vector s e -> Int -> (e -> e) -> ST s ()
modify :: forall e s. Storable e => Vector s e -> Int -> (e -> e) -> ST s ()
modify Vector s e
v Int
n e -> e
f =
   forall e s a.
Storable e =>
String -> Vector s e -> Int -> ST s a -> ST s a
access String
"modify" Vector s e
v Int
n forall a b. (a -> b) -> a -> b
$ forall e s. Storable e => Vector s e -> Int -> (e -> e) -> ST s ()
unsafeModify Vector s e
v Int
n e -> e
f

{-# INLINE access #-}
access :: (Storable e) =>
   String -> Vector s e -> Int -> ST s a -> ST s a
access :: forall e s a.
Storable e =>
String -> Vector s e -> Int -> ST s a -> ST s a
access String
name (SV ForeignPtr e
_v Int
l) Int
n ST s a
act =
   if Int
0forall a. Ord a => a -> a -> Bool
<=Int
n Bool -> Bool -> Bool
&& Int
nforall a. Ord a => a -> a -> Bool
<Int
l
     then ST s a
act
     else forall a. HasCallStack => String -> a
error (String
"StorableVector.ST." forall a. [a] -> [a] -> [a]
++ String
name forall a. [a] -> [a] -> [a]
++ String
": index out of range")


{- |
Returns @Just e@, when the element @e@ could be read
and 'Nothing' if the index was out of range.
This way you can avoid duplicate index checks
that may be needed when using 'read'.

> Control.Monad.ST.runST (do arr <- new_ 10; Monad.zipWithM_ (write arr) [9,8..0] ['a'..]; read arr 3)

In future 'maybeRead' will replace 'read'.
-}
{-# INLINE maybeRead #-}
maybeRead :: (Storable e) =>
   Vector s e -> Int -> ST s (Maybe e)
maybeRead :: forall e s. Storable e => Vector s e -> Int -> ST s (Maybe e)
maybeRead Vector s e
v Int
n =
   forall e s a.
Storable e =>
Vector s e -> Int -> ST s a -> ST s (Maybe a)
maybeAccess Vector s e
v Int
n forall a b. (a -> b) -> a -> b
$ forall e s. Storable e => Vector s e -> Int -> ST s e
unsafeRead Vector s e
v Int
n

{- |
Returns 'True' if the element could be written
and 'False' if the index was out of range.

> runSTVector (do arr <- new_ 10; foldr (\c go i -> maybeWrite arr i c >>= \cont -> if cont then go (succ i) else return arr) (error "unreachable") ['a'..] 0)

In future 'maybeWrite' will replace 'write'.
-}
{-# INLINE maybeWrite #-}
maybeWrite :: (Storable e) =>
   Vector s e -> Int -> e -> ST s Bool
maybeWrite :: forall e s. Storable e => Vector s e -> Int -> e -> ST s Bool
maybeWrite Vector s e
v Int
n e
x =
   forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a. Maybe a -> Bool
isJust forall a b. (a -> b) -> a -> b
$ forall e s a.
Storable e =>
Vector s e -> Int -> ST s a -> ST s (Maybe a)
maybeAccess Vector s e
v Int
n forall a b. (a -> b) -> a -> b
$ forall e s. Storable e => Vector s e -> Int -> e -> ST s ()
unsafeWrite Vector s e
v Int
n e
x

{- |
Similar to 'maybeWrite'.

In future 'maybeModify' will replace 'modify'.
-}
{-# INLINE maybeModify #-}
maybeModify :: (Storable e) =>
   Vector s e -> Int -> (e -> e) -> ST s Bool
maybeModify :: forall e s.
Storable e =>
Vector s e -> Int -> (e -> e) -> ST s Bool
maybeModify Vector s e
v Int
n e -> e
f =
   forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a. Maybe a -> Bool
isJust forall a b. (a -> b) -> a -> b
$ forall e s a.
Storable e =>
Vector s e -> Int -> ST s a -> ST s (Maybe a)
maybeAccess Vector s e
v Int
n forall a b. (a -> b) -> a -> b
$ forall e s. Storable e => Vector s e -> Int -> (e -> e) -> ST s ()
unsafeModify Vector s e
v Int
n e -> e
f

{-# INLINE maybeAccess #-}
maybeAccess :: (Storable e) =>
   Vector s e -> Int -> ST s a -> ST s (Maybe a)
maybeAccess :: forall e s a.
Storable e =>
Vector s e -> Int -> ST s a -> ST s (Maybe a)
maybeAccess (SV ForeignPtr e
_v Int
l) Int
n ST s a
act =
   forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
Traversable.sequence forall a b. (a -> b) -> a -> b
$ forall a. Bool -> a -> Maybe a
toMaybe (Int
0forall a. Ord a => a -> a -> Bool
<=Int
n Bool -> Bool -> Bool
&& Int
nforall a. Ord a => a -> a -> Bool
<Int
l) ST s a
act
{-
   if 0<=n && n<l
     then fmap Just act
     else return Nothing
-}

{-# INLINE unsafeRead #-}
unsafeRead :: (Storable e) =>
   Vector s e -> Int -> ST s e
unsafeRead :: forall e s. Storable e => Vector s e -> Int -> ST s e
unsafeRead Vector s e
v Int
n =
   forall e s a.
Storable e =>
Vector s e -> Int -> (Ptr e -> IO a) -> ST s a
unsafeAccess Vector s e
v Int
n forall a b. (a -> b) -> a -> b
$ forall a. Storable a => Ptr a -> IO a
peek

{-# INLINE unsafeWrite #-}
unsafeWrite :: (Storable e) =>
   Vector s e -> Int -> e -> ST s ()
unsafeWrite :: forall e s. Storable e => Vector s e -> Int -> e -> ST s ()
unsafeWrite Vector s e
v Int
n e
x =
   forall e s a.
Storable e =>
Vector s e -> Int -> (Ptr e -> IO a) -> ST s a
unsafeAccess Vector s e
v Int
n forall a b. (a -> b) -> a -> b
$ \Ptr e
p -> forall a. Storable a => Ptr a -> a -> IO ()
poke Ptr e
p e
x

{-# INLINE unsafeModify #-}
unsafeModify :: (Storable e) =>
   Vector s e -> Int -> (e -> e) -> ST s ()
unsafeModify :: forall e s. Storable e => Vector s e -> Int -> (e -> e) -> ST s ()
unsafeModify Vector s e
v Int
n e -> e
f =
   forall e s a.
Storable e =>
Vector s e -> Int -> (Ptr e -> IO a) -> ST s a
unsafeAccess Vector s e
v Int
n forall a b. (a -> b) -> a -> b
$ \Ptr e
p -> forall a. Storable a => Ptr a -> a -> IO ()
poke Ptr e
p forall b c a. (b -> c) -> (a -> b) -> a -> c
. e -> e
f forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall a. Storable a => Ptr a -> IO a
peek Ptr e
p

{-# INLINE unsafeAccess #-}
unsafeAccess :: (Storable e) =>
   Vector s e -> Int -> (Ptr e -> IO a) -> ST s a
unsafeAccess :: forall e s a.
Storable e =>
Vector s e -> Int -> (Ptr e -> IO a) -> ST s a
unsafeAccess (SV ForeignPtr e
v Int
_l) Int
n Ptr e -> IO a
act =
   forall a s. IO a -> ST s a
Unsafe.ioToST (forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr e
v forall a b. (a -> b) -> a -> b
$ \Ptr e
p -> Ptr e -> IO a
act (forall a. Storable a => Ptr a -> Int -> Ptr a
advancePtr Ptr e
p Int
n))


{-# INLINE freeze #-}
freeze :: (Storable e) =>
   Vector s e -> ST s (VS.Vector e)
freeze :: forall e s. Storable e => Vector s e -> ST s (Vector e)
freeze (SV ForeignPtr e
x Int
l) =
   forall a s. IO a -> ST s a
Unsafe.ioToST forall a b. (a -> b) -> a -> b
$
   forall a. Storable a => Int -> (Ptr a -> IO ()) -> IO (Vector a)
V.create Int
l forall a b. (a -> b) -> a -> b
$ \Ptr e
p ->
   forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr e
x forall a b. (a -> b) -> a -> b
$ \Ptr e
f ->
   forall a. Storable a => Ptr a -> Ptr a -> Int -> IO ()
copyArray Ptr e
p Ptr e
f (forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
l)

{- |
This is like 'freeze' but it does not copy the vector.
You must make sure that you never write again to the array.
It is best to use 'unsafeFreeze' only at the end of a block,
that is run by 'runST'.
-}
{-# INLINE unsafeFreeze #-}
unsafeFreeze :: (Storable e) =>
   Vector s e -> ST s (VS.Vector e)
unsafeFreeze :: forall e s. Storable e => Vector s e -> ST s (Vector e)
unsafeFreeze = forall s a. Vector s a -> ST s (Vector a)
unsafeToVector


{-# INLINE thaw #-}
thaw :: (Storable e) =>
   VS.Vector e -> ST s (Vector s e)
thaw :: forall e s. Storable e => Vector e -> ST s (Vector s e)
thaw Vector e
v =
   forall a s. IO a -> ST s a
Unsafe.ioToST forall a b. (a -> b) -> a -> b
$
   forall a b.
Storable a =>
Vector a -> (Ptr a -> Int -> IO b) -> IO b
V.withStartPtr Vector e
v forall a b. (a -> b) -> a -> b
$ \Ptr e
f Int
l ->
   forall a s.
Storable a =>
Int -> (Ptr a -> IO ()) -> IO (Vector s a)
create Int
l forall a b. (a -> b) -> a -> b
$ \Ptr e
p ->
   forall a. Storable a => Ptr a -> Ptr a -> Int -> IO ()
copyArray Ptr e
p Ptr e
f (forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
l)


{-# INLINE length #-}
length ::
   Vector s e -> Int
length :: forall s e. Vector s e -> Int
length (SV ForeignPtr e
_v Int
l) = Int
l


{-# INLINE runSTVector #-}
runSTVector :: (Storable e) =>
   (forall s. ST s (Vector s e)) -> VS.Vector e
runSTVector :: forall e. Storable e => (forall s. ST s (Vector s e)) -> Vector e
runSTVector forall s. ST s (Vector s e)
m =
   forall a. (forall s. ST s a) -> a
runST (forall s a. Vector s a -> ST s (Vector a)
unsafeToVector forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall s. ST s (Vector s e)
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]))
-}
{-# INLINE mapST #-}
mapST :: (Storable a, Storable b) =>
   (a -> ST s b) -> VS.Vector a -> ST s (VS.Vector b)
mapST :: forall a b s.
(Storable a, Storable b) =>
(a -> ST s b) -> Vector a -> ST s (Vector b)
mapST a -> ST s b
f (V.SV ForeignPtr a
px Int
sx Int
n) =
   let {-# INLINE go #-}
       go :: t -> Ptr a -> Ptr b -> ST s ()
go t
l Ptr a
q Ptr b
p =
          if t
lforall a. Ord a => a -> a -> Bool
>t
0
            then
               do forall a s. IO a -> ST s a
Unsafe.ioToST forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Storable a => Ptr a -> a -> IO ()
poke Ptr b
p forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< a -> ST s b
f forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall a s. IO a -> ST s a
Unsafe.ioToST (forall a. Storable a => Ptr a -> IO a
peek Ptr a
q)
                  t -> Ptr a -> Ptr b -> ST s ()
go (forall a. Enum a => a -> a
pred t
l) (forall a. Storable a => Ptr a -> Int -> Ptr a
advancePtr Ptr a
q Int
1) (forall a. Storable a => Ptr a -> Int -> Ptr a
advancePtr Ptr b
p Int
1)
            else forall (m :: * -> *) a. Monad m => a -> m a
return ()
   in  do ys :: Vector s b
ys@(SV ForeignPtr b
py Int
_) <- forall e s. Storable e => Int -> ST s (Vector s e)
new_ Int
n
          forall {t}.
(Ord t, Num t, Enum t) =>
t -> Ptr a -> Ptr b -> ST s ()
go Int
n
              (forall a. ForeignPtr a -> Ptr a
Unsafe.foreignPtrToPtr ForeignPtr a
px forall a. Storable a => Ptr a -> Int -> Ptr a
`advancePtr` Int
sx)
              (forall a. ForeignPtr a -> Ptr a
Unsafe.foreignPtrToPtr ForeignPtr b
py)
          forall s a. Vector s a -> ST s (Vector a)
unsafeToVector Vector s b
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.
-}
{-# INLINE mapSTLazy #-}
mapSTLazy :: (Storable a, Storable b) =>
   (a -> ST s b) -> VL.Vector a -> ST s (VL.Vector b)
mapSTLazy :: forall a b s.
(Storable a, Storable b) =>
(a -> ST s b) -> Vector a -> ST s (Vector b)
mapSTLazy a -> ST s b
f (VL.SV [Vector a]
xs) =
   forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a. [Vector a] -> Vector a
VL.SV forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall a b s.
(Storable a, Storable b) =>
(a -> ST s b) -> Vector a -> ST s (Vector b)
mapST a -> ST s b
f) [Vector a]
xs