1 files changed, 99 insertions, 0 deletions

diff --git a/lib/Data/Packed/ST.hs b/lib/Data/Packed/ST.hs
new file mode 100644
index 0000000..3d94014
--- /dev/null
+++ b/lib/Data/Packed/ST.hs
@@ -0,0 +1,99 @@
+{-# OPTIONS -XTypeOperators -XRank2Types  -XFlexibleContexts #-}
+
+-----------------------------------------------------------------------------
+-- |
+-- Module      :  Data.Packed.ST
+-- Copyright   :  (c) Alberto Ruiz 2008
+-- License     :  GPL-style
+--
+-- Maintainer  :  Alberto Ruiz <aruiz@um.es>
+-- Stability   :  provisional
+-- Portability :  portable
+--
+-- In-place manipulation inside the ST monad.
+-- See examples/inplace.hs in the distribution.
+--
+-----------------------------------------------------------------------------
+
+module Data.Packed.ST (
+    STVector, thawVector, freezeVector, runSTVector,
+    readVector, writeVector, modifyVector, liftSTVector,
+    STMatrix, thawMatrix, freezeMatrix, runSTMatrix,
+    readMatrix, writeMatrix, modifyMatrix, liftSTMatrix
+) where
+
+import Data.Packed.Internal
+import Data.Array.Storable
+import Control.Monad.ST
+import Data.Array.ST
+import Foreign
+
+
+ioReadV :: Storable t => Vector t -> Int -> IO t
+ioReadV v k = withForeignPtr (fptr v) $ \s -> peekElemOff s k
+
+ioWriteV :: Storable t => Vector t -> Int -> t -> IO ()
+ioWriteV v k x = withForeignPtr (fptr v) $ \s -> pokeElemOff s k x
+
+newtype STVector s t = Mut (Vector t)
+
+thawVector :: Storable t => Vector t -> ST s (STVector s t)
+thawVector = unsafeIOToST . fmap Mut . cloneVector
+
+unsafeFreezeVector (Mut x) = unsafeIOToST . return $ x
+
+runSTVector :: Storable t => (forall s . ST s (STVector s t)) -> Vector t
+runSTVector st = runST (st >>= unsafeFreezeVector)
+
+readVector :: Storable t => STVector s t -> Int -> ST s t
+readVector   (Mut x) = unsafeIOToST . ioReadV x
+
+writeVector :: Storable t => STVector s t -> Int -> t -> ST s ()
+writeVector  (Mut x) k = unsafeIOToST . ioWriteV x k
+
+modifyVector :: (Storable t) => STVector s t -> Int -> (t -> t) -> ST s ()
+modifyVector x k f = readVector x k >>= return . f >>= writeVector x k
+
+liftSTVector :: (Storable t) => (Vector t -> a) -> STVector s1 t -> ST s2 a
+liftSTVector f (Mut x) = unsafeIOToST . fmap f . cloneVector $ x
+
+freezeVector :: (Storable t) => STVector s1 t -> ST s2 (Vector t)
+freezeVector v = liftSTVector id v
+
+-------------------------------------------------------------------------
+
+ioReadM :: Storable t => Matrix t -> Int -> Int -> IO t
+ioReadM (MC nr nc cv) r c = ioReadV cv (r*nc+c)
+ioReadM (MF nr nc fv) r c = ioReadV fv (c*nr+r)
+
+ioWriteM :: Storable t => Matrix t -> Int -> Int -> t -> IO ()
+ioWriteM (MC nr nc cv) r c val = ioWriteV cv (r*nc+c) val
+ioWriteM (MF nr nc fv) r c val = ioWriteV fv (c*nr+r) val
+
+newtype STMatrix s t = STMatrix (Matrix t)
+
+thawMatrix :: Storable t => Matrix t -> ST s (STMatrix s t)
+thawMatrix = unsafeIOToST . fmap STMatrix . cloneMatrix
+
+unsafeFreezeMatrix (STMatrix x) = unsafeIOToST . return $ x
+
+runSTMatrix :: Storable t => (forall s . ST s (STMatrix s t)) -> Matrix t
+runSTMatrix st = runST (st >>= unsafeFreezeMatrix)
+
+readMatrix :: Storable t => STMatrix s t -> Int -> Int -> ST s t
+readMatrix   (STMatrix x) r = unsafeIOToST . ioReadM x r
+
+writeMatrix :: Storable t => STMatrix s t -> Int -> Int -> t -> ST s ()
+writeMatrix  (STMatrix x) r c = unsafeIOToST . ioWriteM x r c
+
+modifyMatrix :: (Storable t) => STMatrix s t -> Int -> Int -> (t -> t) -> ST s ()
+modifyMatrix x r c f = readMatrix x r c >>= return . f >>= writeMatrix x r c
+
+liftSTMatrix :: (Storable t) => (Matrix t -> a) -> STMatrix s1 t -> ST s2 a
+liftSTMatrix f (STMatrix x) = unsafeIOToST . fmap f . cloneMatrix $ x
+
+freezeMatrix :: (Storable t) => STMatrix s1 t -> ST s2 (Matrix t)
+freezeMatrix m = liftSTMatrix id m
+
+cloneMatrix (MC r c d) = cloneVector d >>= return . MC r c
+cloneMatrix (MF r c d) = cloneVector d >>= return . MF r c