#include "fusion-phases.h"
module Data.Array.Vector.UArr (
UA, UArr, MUArr, UPrim(..),
lengthU, indexU, sliceU, unitsU, zipU, unzipU, fstU, sndU,
newU, newDynU, newDynResU,
lengthMU, newMU, readMU, writeMU, copyMU, unsafeFreezeMU, unsafeFreezeAllMU,
memcpyMU, memcpyOffMU, memmoveOffMU,
unsafeZipMU, unsafeUnzipMU,
) where
import Control.Monad (liftM, liftM2)
import Data.Complex
import GHC.Real
import Data.Array.Vector.Prim.BUArr (
BUArr, MBUArr, UAE,
lengthBU, indexBU, sliceBU,
lengthMBU, newMBU, readMBU, writeMBU, copyMBU, unsafeFreezeMBU,
memcpyMBU, memcpyOffMBU, memmoveOffMBU)
import System.IO
import GHC.ST
import Data.Word
import Data.Int
import Data.Array.Vector.Prim.Debug
import Data.Array.Vector.Prim.Hyperstrict
infixl 9 `indexU`, `readMU`
class UA e where
data UArr e
data MUArr e :: * -> *
lengthU :: UArr e -> Int
indexU :: UArr e -> Int -> e
sliceU :: UArr e -> Int -> Int -> UArr e
lengthMU :: MUArr e s -> Int
newMU :: Int -> ST s (MUArr e s)
readMU :: MUArr e s -> Int -> ST s e
writeMU :: MUArr e s -> Int -> e -> ST s ()
copyMU :: MUArr e s -> Int -> UArr e -> ST s ()
unsafeFreezeMU :: MUArr e s -> Int -> ST s (UArr e)
memcpyMU :: MUArr e s -> MUArr e s -> Int -> ST s ()
memcpyOffMU :: MUArr e s -> MUArr e s -> Int -> Int -> Int -> ST s ()
memmoveOffMU :: MUArr e s -> MUArr e s -> Int -> Int -> Int -> ST s ()
class UAE e => UPrim e where
mkUAPrim :: BUArr e -> UArr e
unUAPrim :: UArr e -> BUArr e
mkMUAPrim :: MBUArr s e -> MUArr e s
unMUAPrim :: MUArr e s -> MBUArr s e
unsafeFreezeAllMU :: UA e => MUArr e s -> ST s (UArr e)
unsafeFreezeAllMU marr = unsafeFreezeMU marr (lengthMU marr)
newU :: UA e => Int -> (forall s. MUArr e s -> ST s ()) -> UArr e
newU n init = newDynU n (\ma -> init ma >> return n)
newDynU :: UA e => Int -> (forall s. MUArr e s -> ST s Int) -> UArr e
newDynU n init =
runST (do
ma <- newMU n
n' <- init ma
unsafeFreezeMU ma n'
)
newDynResU :: UA e
=> Int -> (forall s. MUArr e s -> ST s (Int :*: r)) -> UArr e :*: r
newDynResU n init =
runST (do
ma <- newMU n
n' :*: r <- init ma
arr <- unsafeFreezeMU ma n'
return (arr :*: r)
)
unitsU :: Int -> UArr ()
unitsU = UAUnit
zipU :: (UA a, UA b) => UArr a -> UArr b -> UArr (a :*: b)
zipU = UAProd
unzipU :: (UA a, UA b) => UArr (a :*: b) -> (UArr a :*: UArr b)
unzipU (UAProd l r) = (l :*: r)
fstU :: (UA a, UA b) => UArr (a :*: b) -> UArr a
fstU (UAProd l r) = l
sndU :: (UA a, UA b) => UArr (a :*: b) -> UArr b
sndU (UAProd l r) = r
unsafeZipMU :: (UA a, UA b) => MUArr a s -> MUArr b s -> MUArr (a :*: b) s
unsafeZipMU = MUAProd
unsafeUnzipMU :: (UA a, UA b) => MUArr (a :*: b) s -> MUArr a s :*: MUArr b s
unsafeUnzipMU (MUAProd mua mub) = mua :*: mub
instance UA () where
newtype UArr () = UAUnit Int
newtype MUArr () s = MUAUnit Int
lengthU (UAUnit n) = n
indexU (UAUnit _) _ = ()
sliceU (UAUnit len) i n = if i == len then UAUnit 0 else UAUnit (min n (len i))
lengthMU (MUAUnit n) = n
newMU n = return $ MUAUnit n
readMU (MUAUnit _) _ = return ()
writeMU (MUAUnit _) _ _ = return ()
copyMU (MUAUnit _) _ (UAUnit _) = return ()
unsafeFreezeMU (MUAUnit _) n = return $ UAUnit n
memcpyMU (MUAUnit _) (MUAUnit _) _ = return ()
memcpyOffMU (MUAUnit _) (MUAUnit _) _ _ _ = return ()
memmoveOffMU (MUAUnit _) (MUAUnit _) _ _ _ = return ()
instance (UA a, UA b) => UA (a :*: b) where
data UArr (a :*: b) = UAProd !(UArr a) !(UArr b)
data MUArr (a :*: b) s = MUAProd !(MUArr a s) !(MUArr b s)
lengthU (UAProd l r) = checkEq "lengthU" "lengths of zipped arrays differ" (lengthU l) (lengthU r)
(lengthU l)
indexU (UAProd l r) i = indexU l i :*: indexU r i
sliceU (UAProd l r) i n = UAProd (sliceU l i n) (sliceU r i n)
lengthMU (MUAProd l r) = lengthMU l
newMU n =
do
a <- newMU n
b <- newMU n
return $ MUAProd a b
readMU (MUAProd a b) i = liftM2 (:*:) (a `readMU` i) (b `readMU` i)
writeMU (MUAProd a b) i (x :*: y) =
do
writeMU a i x
writeMU b i y
copyMU (MUAProd ma mb) i (UAProd a b) =
do
copyMU ma i a
copyMU mb i b
unsafeFreezeMU (MUAProd a b) n =
do
a' <- unsafeFreezeMU a n
b' <- unsafeFreezeMU b n
return $ UAProd a' b'
memcpyMU (MUAProd ma mb) (MUAProd ma' mb') l =
do memcpyMU ma ma' l
memcpyMU mb mb' l
memcpyOffMU (MUAProd ma mb) (MUAProd ma' mb') s d l =
do memcpyOffMU ma ma' s d l
memcpyOffMU mb mb' s d l
memmoveOffMU (MUAProd ma mb) (MUAProd ma' mb') s d l =
do memmoveOffMU ma ma' s d l
memmoveOffMU mb mb' s d l
primLengthU :: UPrim e => UArr e -> Int
primLengthU = lengthBU . unUAPrim
primIndexU :: UPrim e => UArr e -> Int -> e
primIndexU = indexBU . unUAPrim
primSliceU :: UPrim e => UArr e -> Int -> Int -> UArr e
primSliceU arr i = mkUAPrim . sliceBU (unUAPrim arr) i
primLengthMU :: UPrim e => MUArr e s -> Int
primLengthMU = lengthMBU . unMUAPrim
primNewMU :: UPrim e => Int -> ST s (MUArr e s)
primNewMU = liftM mkMUAPrim . newMBU
primReadMU :: UPrim e => MUArr e s -> Int -> ST s e
primReadMU = readMBU . unMUAPrim
primWriteMU :: UPrim e => MUArr e s -> Int -> e -> ST s ()
primWriteMU = writeMBU . unMUAPrim
primCopyMU :: UPrim e => MUArr e s -> Int -> UArr e -> ST s ()
primCopyMU ma i = copyMBU (unMUAPrim ma) i . unUAPrim
primUnsafeFreezeMU :: UPrim e => MUArr e s -> Int -> ST s (UArr e)
primUnsafeFreezeMU ma = liftM mkUAPrim . unsafeFreezeMBU (unMUAPrim ma)
primMemcpyMU :: UPrim e => MUArr e s -> MUArr e s -> Int -> ST s ()
primMemcpyMU src dst l = memcpyMBU (unMUAPrim src) (unMUAPrim dst) l
primMemcpyOffMU :: UPrim e => MUArr e s -> MUArr e s -> Int -> Int -> Int -> ST s ()
primMemcpyOffMU src dst s d l = memcpyOffMBU (unMUAPrim src) (unMUAPrim dst) s d l
primMemmoveOffMU :: UPrim e => MUArr e s -> MUArr e s -> Int -> Int -> Int -> ST s ()
primMemmoveOffMU src dst s d l = memmoveOffMBU (unMUAPrim src) (unMUAPrim dst) s d l
instance UPrim Bool where
mkUAPrim = UABool
unUAPrim (UABool arr) = arr
mkMUAPrim = MUABool
unMUAPrim (MUABool arr) = arr
instance UA Bool where
newtype UArr Bool = UABool (BUArr Bool)
newtype MUArr Bool s = MUABool (MBUArr s Bool)
lengthU = primLengthU
indexU = primIndexU
sliceU = primSliceU
lengthMU = primLengthMU
newMU = primNewMU
readMU = primReadMU
writeMU = primWriteMU
copyMU = primCopyMU
unsafeFreezeMU = primUnsafeFreezeMU
memcpyMU = primMemcpyMU
memcpyOffMU = primMemcpyOffMU
memmoveOffMU = primMemmoveOffMU
instance UPrim Char where
mkUAPrim = UAChar
unUAPrim (UAChar arr) = arr
mkMUAPrim = MUAChar
unMUAPrim (MUAChar arr) = arr
instance UA Char where
newtype UArr Char = UAChar (BUArr Char)
newtype MUArr Char s = MUAChar (MBUArr s Char)
lengthU = primLengthU
indexU = primIndexU
sliceU = primSliceU
lengthMU = primLengthMU
newMU = primNewMU
readMU = primReadMU
writeMU = primWriteMU
copyMU = primCopyMU
unsafeFreezeMU = primUnsafeFreezeMU
memcpyMU = primMemcpyMU
memcpyOffMU = primMemcpyOffMU
memmoveOffMU = primMemmoveOffMU
instance UPrim Int where
mkUAPrim = UAInt
unUAPrim (UAInt arr) = arr
mkMUAPrim = MUAInt
unMUAPrim (MUAInt arr) = arr
instance UA Int where
newtype UArr Int = UAInt (BUArr Int)
newtype MUArr Int s = MUAInt (MBUArr s Int)
lengthU = primLengthU
indexU = primIndexU
sliceU = primSliceU
lengthMU = primLengthMU
newMU = primNewMU
readMU = primReadMU
writeMU = primWriteMU
copyMU = primCopyMU
unsafeFreezeMU = primUnsafeFreezeMU
memcpyMU = primMemcpyMU
memcpyOffMU = primMemcpyOffMU
memmoveOffMU = primMemmoveOffMU
instance UPrim Word where
mkUAPrim = UAWord
unUAPrim (UAWord arr) = arr
mkMUAPrim = MUAWord
unMUAPrim (MUAWord arr) = arr
instance UA Word where
newtype UArr Word = UAWord (BUArr Word)
newtype MUArr Word s = MUAWord (MBUArr s Word)
lengthU = primLengthU
indexU = primIndexU
sliceU = primSliceU
lengthMU = primLengthMU
newMU = primNewMU
readMU = primReadMU
writeMU = primWriteMU
copyMU = primCopyMU
unsafeFreezeMU = primUnsafeFreezeMU
memcpyMU = primMemcpyMU
memcpyOffMU = primMemcpyOffMU
memmoveOffMU = primMemmoveOffMU
instance UPrim Float where
mkUAPrim = UAFloat
unUAPrim (UAFloat arr) = arr
mkMUAPrim = MUAFloat
unMUAPrim (MUAFloat arr) = arr
instance UA Float where
newtype UArr Float = UAFloat (BUArr Float)
newtype MUArr Float s = MUAFloat (MBUArr s Float)
lengthU = primLengthU
indexU = primIndexU
sliceU = primSliceU
lengthMU = primLengthMU
newMU = primNewMU
readMU = primReadMU
writeMU = primWriteMU
copyMU = primCopyMU
unsafeFreezeMU = primUnsafeFreezeMU
memcpyMU = primMemcpyMU
memcpyOffMU = primMemcpyOffMU
memmoveOffMU = primMemmoveOffMU
instance UPrim Double where
mkUAPrim = UADouble
unUAPrim (UADouble arr) = arr
mkMUAPrim = MUADouble
unMUAPrim (MUADouble arr) = arr
instance UA Double where
newtype UArr Double = UADouble (BUArr Double)
newtype MUArr Double s = MUADouble (MBUArr s Double)
lengthU = primLengthU
indexU = primIndexU
sliceU = primSliceU
lengthMU = primLengthMU
newMU = primNewMU
readMU = primReadMU
writeMU = primWriteMU
copyMU = primCopyMU
unsafeFreezeMU = primUnsafeFreezeMU
memcpyMU = primMemcpyMU
memcpyOffMU = primMemcpyOffMU
memmoveOffMU = primMemmoveOffMU
instance UPrim Word8 where
mkUAPrim = UAWord8
unUAPrim (UAWord8 arr) = arr
mkMUAPrim = MUAWord8
unMUAPrim (MUAWord8 arr) = arr
instance UA Word8 where
newtype UArr Word8 = UAWord8 (BUArr Word8)
newtype MUArr Word8 s = MUAWord8 (MBUArr s Word8)
lengthU = primLengthU
indexU = primIndexU
sliceU = primSliceU
lengthMU = primLengthMU
newMU = primNewMU
readMU = primReadMU
writeMU = primWriteMU
copyMU = primCopyMU
unsafeFreezeMU = primUnsafeFreezeMU
memcpyMU = primMemcpyMU
memcpyOffMU = primMemcpyOffMU
memmoveOffMU = primMemmoveOffMU
instance UPrim Word16 where
mkUAPrim = UAWord16
unUAPrim (UAWord16 arr) = arr
mkMUAPrim = MUAWord16
unMUAPrim (MUAWord16 arr) = arr
instance UA Word16 where
newtype UArr Word16 = UAWord16 (BUArr Word16)
newtype MUArr Word16 s = MUAWord16 (MBUArr s Word16)
lengthU = primLengthU
indexU = primIndexU
sliceU = primSliceU
lengthMU = primLengthMU
newMU = primNewMU
readMU = primReadMU
writeMU = primWriteMU
copyMU = primCopyMU
unsafeFreezeMU = primUnsafeFreezeMU
memcpyMU = primMemcpyMU
memcpyOffMU = primMemcpyOffMU
memmoveOffMU = primMemmoveOffMU
instance UPrim Word32 where
mkUAPrim = UAWord32
unUAPrim (UAWord32 arr) = arr
mkMUAPrim = MUAWord32
unMUAPrim (MUAWord32 arr) = arr
instance UA Word32 where
newtype UArr Word32 = UAWord32 (BUArr Word32)
newtype MUArr Word32 s = MUAWord32 (MBUArr s Word32)
lengthU = primLengthU
indexU = primIndexU
sliceU = primSliceU
lengthMU = primLengthMU
newMU = primNewMU
readMU = primReadMU
writeMU = primWriteMU
copyMU = primCopyMU
unsafeFreezeMU = primUnsafeFreezeMU
memcpyMU = primMemcpyMU
memcpyOffMU = primMemcpyOffMU
memmoveOffMU = primMemmoveOffMU
instance UPrim Word64 where
mkUAPrim = UAWord64
unUAPrim (UAWord64 arr) = arr
mkMUAPrim = MUAWord64
unMUAPrim (MUAWord64 arr) = arr
instance UA Word64 where
newtype UArr Word64 = UAWord64 (BUArr Word64)
newtype MUArr Word64 s = MUAWord64 (MBUArr s Word64)
lengthU = primLengthU
indexU = primIndexU
sliceU = primSliceU
lengthMU = primLengthMU
newMU = primNewMU
readMU = primReadMU
writeMU = primWriteMU
copyMU = primCopyMU
unsafeFreezeMU = primUnsafeFreezeMU
memcpyMU = primMemcpyMU
memcpyOffMU = primMemcpyOffMU
memmoveOffMU = primMemmoveOffMU
instance UPrim Int8 where
mkUAPrim = UAInt8
unUAPrim (UAInt8 arr) = arr
mkMUAPrim = MUAInt8
unMUAPrim (MUAInt8 arr) = arr
instance UA Int8 where
newtype UArr Int8 = UAInt8 (BUArr Int8)
newtype MUArr Int8 s = MUAInt8 (MBUArr s Int8)
lengthU = primLengthU
indexU = primIndexU
sliceU = primSliceU
lengthMU = primLengthMU
newMU = primNewMU
readMU = primReadMU
writeMU = primWriteMU
copyMU = primCopyMU
unsafeFreezeMU = primUnsafeFreezeMU
memcpyMU = primMemcpyMU
memcpyOffMU = primMemcpyOffMU
memmoveOffMU = primMemmoveOffMU
instance UPrim Int16 where
mkUAPrim = UAInt16
unUAPrim (UAInt16 arr) = arr
mkMUAPrim = MUAInt16
unMUAPrim (MUAInt16 arr) = arr
instance UA Int16 where
newtype UArr Int16 = UAInt16 (BUArr Int16)
newtype MUArr Int16 s = MUAInt16 (MBUArr s Int16)
lengthU = primLengthU
indexU = primIndexU
sliceU = primSliceU
lengthMU = primLengthMU
newMU = primNewMU
readMU = primReadMU
writeMU = primWriteMU
copyMU = primCopyMU
unsafeFreezeMU = primUnsafeFreezeMU
memcpyMU = primMemcpyMU
memcpyOffMU = primMemcpyOffMU
memmoveOffMU = primMemmoveOffMU
instance UPrim Int32 where
mkUAPrim = UAInt32
unUAPrim (UAInt32 arr) = arr
mkMUAPrim = MUAInt32
unMUAPrim (MUAInt32 arr) = arr
instance UA Int32 where
newtype UArr Int32 = UAInt32 (BUArr Int32)
newtype MUArr Int32 s = MUAInt32 (MBUArr s Int32)
lengthU = primLengthU
indexU = primIndexU
sliceU = primSliceU
lengthMU = primLengthMU
newMU = primNewMU
readMU = primReadMU
writeMU = primWriteMU
copyMU = primCopyMU
unsafeFreezeMU = primUnsafeFreezeMU
memcpyMU = primMemcpyMU
memcpyOffMU = primMemcpyOffMU
memmoveOffMU = primMemmoveOffMU
instance UPrim Int64 where
mkUAPrim = UAInt64
unUAPrim (UAInt64 arr) = arr
mkMUAPrim = MUAInt64
unMUAPrim (MUAInt64 arr) = arr
instance UA Int64 where
newtype UArr Int64 = UAInt64 (BUArr Int64)
newtype MUArr Int64 s = MUAInt64 (MBUArr s Int64)
lengthU = primLengthU
indexU = primIndexU
sliceU = primSliceU
lengthMU = primLengthMU
newMU = primNewMU
readMU = primReadMU
writeMU = primWriteMU
copyMU = primCopyMU
unsafeFreezeMU = primUnsafeFreezeMU
memcpyMU = primMemcpyMU
memcpyOffMU = primMemcpyOffMU
memmoveOffMU = primMemmoveOffMU
instance (RealFloat a, UA a) => UA (Complex a) where
newtype UArr (Complex a) = UAComplex (UArr (a :*: a))
newtype MUArr (Complex a) s = MUAComplex (MUArr (a :*: a) s)
lengthU (UAComplex arr) = lengthU arr
indexU (UAComplex arr) i = case indexU arr i of (a :*: b) -> a :+ b
sliceU (UAComplex arr) i n = UAComplex (sliceU arr i n)
lengthMU (MUAComplex arr) = lengthMU arr
newMU n = return . MUAComplex =<< newMU n
readMU (MUAComplex arr) n = do (a :*: b) <- readMU arr n; return (a :+ b)
writeMU (MUAComplex arr) i (x :+ y) = writeMU arr i (x :*: y)
copyMU (MUAComplex mua) n (UAComplex ua) = copyMU mua n ua
unsafeFreezeMU (MUAComplex arr) n = do arr' <- unsafeFreezeMU arr n; return (UAComplex arr')
memcpyMU (MUAComplex src) (MUAComplex dst) l = memcpyMU src dst l
memcpyOffMU (MUAComplex src) (MUAComplex dst) s d l = memcpyOffMU src dst s d l
memmoveOffMU (MUAComplex src) (MUAComplex dst) s d l = memmoveOffMU src dst s d l
instance (Integral a, UA a) => UA (Ratio a) where
newtype UArr (Ratio a) = UARatio (UArr (a :*: a))
newtype MUArr (Ratio a) s = MUARatio (MUArr (a :*: a) s)
lengthU (UARatio arr) = lengthU arr
indexU (UARatio arr) i = case indexU arr i of (a :*: b) -> a % b
sliceU (UARatio arr) i n = UARatio (sliceU arr i n)
lengthMU (MUARatio arr) = lengthMU arr
newMU n = return . MUARatio =<< newMU n
readMU (MUARatio arr) n = do (a :*: b) <- readMU arr n; return (a % b)
writeMU (MUARatio arr) i (n :% d) = writeMU arr i (n :*: d)
copyMU (MUARatio mua) n (UARatio ua) = copyMU mua n ua
unsafeFreezeMU (MUARatio arr) n = do arr' <- unsafeFreezeMU arr n; return (UARatio arr')
memcpyMU (MUARatio src) (MUARatio dst) l = memcpyMU src dst l
memcpyOffMU (MUARatio src) (MUARatio dst) s d l = memcpyOffMU src dst s d l
memmoveOffMU (MUARatio src) (MUARatio dst) s d l = memmoveOffMU src dst s d l