{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE UndecidableInstances #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}
module Massiv.Persist
(
putIx
, getIx
, mkSzFail
, putArray
, getArray
, putPrimArray
, getPrimArray
, putStorableArray
, getStorableArray
) where
import Control.DeepSeq (NFData)
import Control.Monad
import qualified Control.Monad.Fail as Fail
import qualified Data.ByteString.Internal as BS
import qualified Data.ByteString.Short.Internal as SBS
import Data.Foldable as F
import Data.Massiv.Array as A
import Data.Massiv.Array.Unsafe
import Data.Persist
import Data.Persist.Internal
import qualified Data.Primitive as Primitive
import qualified Data.Primitive.ByteArray as BA
import Data.Proxy
import Data.Word
import Foreign.Ptr
import Foreign.ForeignPtr
import qualified Foreign.Storable as Storable
#include "MachDeps.h"
instance Persist Comp where
put :: Comp -> Put ()
put Comp
comp =
case Comp
comp of
Comp
Seq -> Word8 -> Put ()
forall t. Persist t => t -> Put ()
put (Word8
0 :: Word8)
ParOn [Int]
xs -> Word8 -> Put ()
forall t. Persist t => t -> Put ()
put (Word8
1 :: Word8) Put () -> Put () -> Put ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> [Int] -> Put ()
forall t. Persist t => t -> Put ()
put [Int]
xs
ParN Word16
n -> Word8 -> Put ()
forall t. Persist t => t -> Put ()
put (Word8
2 :: Word8) Put () -> Put () -> Put ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Word16 -> Put ()
forall t. Persist t => t -> Put ()
put Word16
n
get :: Get Comp
get = do
Word8
ty :: Word8 <- Get Word8
forall t. Persist t => Get t
get
case Word8
ty of
Word8
0 -> Comp -> Get Comp
forall (f :: * -> *) a. Applicative f => a -> f a
pure Comp
Seq
Word8
1 -> [Int] -> Comp
ParOn ([Int] -> Comp) -> Get [Int] -> Get Comp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Get [Int]
forall t. Persist t => Get t
get
Word8
2 -> Word16 -> Comp
ParN (Word16 -> Comp) -> Get Word16 -> Get Comp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Get Word16
forall t. Persist t => Get t
get
Word8
n -> String -> Get Comp
forall (m :: * -> *) a. MonadFail m => String -> m a
Fail.fail (String -> Get Comp) -> String -> Get Comp
forall a b. (a -> b) -> a -> b
$ String
"Unexpected Comp tag: " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Word8 -> String
forall a. Show a => a -> String
show Word8
n
putIx ::
forall ix. Index ix
=> ix
-> Put ()
putIx :: ix -> Put ()
putIx = (Put () -> Int -> Put ()) -> Put () -> ix -> Put ()
forall ix a. Index ix => (a -> Int -> a) -> a -> ix -> a
foldlIndex (\ !Put ()
acc Int
i -> Int -> Put ()
forall t. Persist t => t -> Put ()
put Int
i Put () -> Put () -> Put ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Put ()
acc) (() -> Put ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ())
getIx ::
forall ix. Index ix
=> Get ix
getIx :: Get ix
getIx = do
let getDim :: b -> Dim -> Get b
getDim b
ix Dim
dim = do
Int
i <- Get Int
forall t. Persist t => Get t
get
(SomeException -> Get b)
-> (b -> Get b) -> Either SomeException b -> Get b
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (String -> Get b
forall (m :: * -> *) a. MonadFail m => String -> m a
Fail.fail (String -> Get b)
-> (SomeException -> String) -> SomeException -> Get b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SomeException -> String
forall a. Show a => a -> String
show) b -> Get b
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Either SomeException b -> Get b)
-> Either SomeException b -> Get b
forall a b. (a -> b) -> a -> b
$! b -> Dim -> Int -> Either SomeException b
forall ix (m :: * -> *).
(Index ix, MonadThrow m) =>
ix -> Dim -> Int -> m ix
setDimM b
ix Dim
dim Int
i
(ix -> Dim -> Get ix) -> ix -> [Dim] -> Get ix
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
F.foldlM ix -> Dim -> Get ix
forall b. Index b => b -> Dim -> Get b
getDim ix
forall ix. Index ix => ix
zeroIndex [Dim
1 .. Proxy ix -> Dim
forall ix (proxy :: * -> *). Index ix => proxy ix -> Dim
dimensions (Proxy ix
forall k (t :: k). Proxy t
Proxy :: Proxy ix)]
instance Persist Ix2 where
put :: Ix2 -> Put ()
put = Ix2 -> Put ()
forall ix. Index ix => ix -> Put ()
putIx
get :: Get Ix2
get = Get Ix2
forall ix. Index ix => Get ix
getIx
instance Index (IxN n) => Persist (IxN n) where
put :: IxN n -> Put ()
put = IxN n -> Put ()
forall ix. Index ix => ix -> Put ()
putIx
get :: Get (IxN n)
get = Get (IxN n)
forall ix. Index ix => Get ix
getIx
mkSzFail ::
forall ix m. (Index ix, Fail.MonadFail m)
=> ix
-> m (Sz ix)
mkSzFail :: ix -> m (Sz ix)
mkSzFail ix
ix = do
let guardNegativeOverflow :: b -> b -> m b
guardNegativeOverflow b
i !b
acc = do
Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (b
i b -> b -> Bool
forall a. Ord a => a -> a -> Bool
< b
0) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ String -> m ()
forall (m :: * -> *) a. MonadFail m => String -> m a
Fail.fail (String -> m ()) -> String -> m ()
forall a b. (a -> b) -> a -> b
$ String
"Negative size encountered: " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> b -> String
forall a. Show a => a -> String
show b
i
let acc' :: b
acc' = b
i b -> b -> b
forall a. Num a => a -> a -> a
* b
acc
Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (b
acc' b -> b -> Bool
forall a. Eq a => a -> a -> Bool
/= b
0 Bool -> Bool -> Bool
&& b
acc' b -> b -> Bool
forall a. Ord a => a -> a -> Bool
< b
acc) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ String -> m ()
forall (m :: * -> *) a. MonadFail m => String -> m a
Fail.fail (String -> m ()) -> String -> m ()
forall a b. (a -> b) -> a -> b
$ String
"Overflow detected, size is too big: " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> b -> String
forall a. Show a => a -> String
show b
i
b -> m b
forall (f :: * -> *) a. Applicative f => a -> f a
pure b
acc'
ix -> Sz ix
forall ix. Index ix => ix -> Sz ix
Sz ix
ix Sz ix -> m Int -> m (Sz ix)
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ (m Int -> Int -> m Int) -> m Int -> ix -> m Int
forall ix a. Index ix => (a -> Int -> a) -> a -> ix -> a
foldlIndex (\m Int
acc Int
i -> m Int
acc m Int -> (Int -> m Int) -> m Int
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Int -> Int -> m Int
forall (m :: * -> *) b.
(Ord b, MonadFail m, Show b, Num b) =>
b -> b -> m b
guardNegativeOverflow Int
i) (Int -> m Int
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
1) ix
ix
instance Index ix => Persist (Sz ix) where
put :: Sz ix -> Put ()
put = ix -> Put ()
forall ix. Index ix => ix -> Put ()
putIx (ix -> Put ()) -> (Sz ix -> ix) -> Sz ix -> Put ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Sz ix -> ix
forall ix. Sz ix -> ix
unSz
get :: Get (Sz ix)
get = ix -> Get (Sz ix)
forall ix (m :: * -> *). (Index ix, MonadFail m) => ix -> m (Sz ix)
mkSzFail (ix -> Get (Sz ix)) -> Get ix -> Get (Sz ix)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Get ix
forall ix. Index ix => Get ix
getIx
putArrayHeader :: forall r ix e. (Strategy r, Size r, Index ix) => Array r ix e -> Put ()
Array r ix e
arr = do
Comp -> Put ()
forall t. Persist t => t -> Put ()
put (Array r ix e -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r ix e
arr)
Sz ix -> Put ()
forall t. Persist t => t -> Put ()
put (Array r ix e -> Sz ix
forall r ix e. Size r => Array r ix e -> Sz ix
size Array r ix e
arr)
instance (Index ix, Persist e) => Persist (Array B ix e) where
put :: Array B ix e -> Put ()
put = Array B ix e -> Put ()
forall r e ix.
(Manifest r e, Index ix, Persist e) =>
Array r ix e -> Put ()
putArray
get :: Get (Array B ix e)
get = Get (Array B ix e)
forall r e ix.
(Mutable r e, Index ix, Persist e) =>
Get (Array r ix e)
getArray
instance (Index ix, Persist e) => Persist (Array BL ix e) where
put :: Array BL ix e -> Put ()
put = Array BL ix e -> Put ()
forall r e ix.
(Manifest r e, Index ix, Persist e) =>
Array r ix e -> Put ()
putArray
get :: Get (Array BL ix e)
get = Get (Array BL ix e)
forall r e ix.
(Mutable r e, Index ix, Persist e) =>
Get (Array r ix e)
getArray
instance (Index ix, NFData e, Persist e) => Persist (Array BN ix e) where
put :: Array BN ix e -> Put ()
put = Array BN ix e -> Put ()
forall r e ix.
(Manifest r e, Index ix, Persist e) =>
Array r ix e -> Put ()
putArray
get :: Get (Array BN ix e)
get = Get (Array BN ix e)
forall r e ix.
(Mutable r e, Index ix, Persist e) =>
Get (Array r ix e)
getArray
instance (Index ix, Unbox e, Persist e) => Persist (Array U ix e) where
put :: Array U ix e -> Put ()
put = Array U ix e -> Put ()
forall r e ix.
(Manifest r e, Index ix, Persist e) =>
Array r ix e -> Put ()
putArray
get :: Get (Array U ix e)
get = Get (Array U ix e)
forall r e ix.
(Mutable r e, Index ix, Persist e) =>
Get (Array r ix e)
getArray
putArray :: (Manifest r e, Index ix, Persist e) => Array r ix e -> Put ()
putArray :: Array r ix e -> Put ()
putArray Array r ix e
arr = do
Array r ix e -> Put ()
forall r ix e.
(Strategy r, Size r, Index ix) =>
Array r ix e -> Put ()
putArrayHeader Array r ix e
arr
(e -> Put ()) -> Array r ix e -> Put ()
forall r a ix (m :: * -> *) b.
(Source r a, Index ix, Monad m) =>
(a -> m b) -> Array r ix a -> m ()
A.mapM_ e -> Put ()
forall t. Persist t => t -> Put ()
put Array r ix e
arr
getArray :: (Mutable r e, Index ix, Persist e) => Get (Array r ix e)
getArray :: Get (Array r ix e)
getArray = do
Comp
comp <- Get Comp
forall t. Persist t => Get t
get
Sz ix
sz <- Get (Sz ix)
forall t. Persist t => Get t
get
Comp -> Array r ix e -> Array r ix e
forall r ix e. Strategy r => Comp -> Array r ix e -> Array r ix e
setComp Comp
comp (Array r ix e -> Array r ix e)
-> Get (Array r ix e) -> Get (Array r ix e)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Sz ix -> (ix -> Get e) -> Get (Array r ix e)
forall r ix e (f :: * -> *).
(Manifest r e, Index ix, Applicative f) =>
Sz ix -> (ix -> f e) -> f (Array r ix e)
A.makeArrayA Sz ix
sz (Get e -> ix -> Get e
forall a b. a -> b -> a
const Get e
forall t. Persist t => Get t
get)
putPrimArray :: forall ix e . (Index ix, Prim e, Persist e) => Array P ix e -> Put ()
putPrimArray :: Array P ix e -> Put ()
putPrimArray Array P ix e
arr = do
#ifdef WORDS_BIGENDIAN
putArray
#else
Array P ix e -> Put ()
forall r ix e.
(Strategy r, Size r, Index ix) =>
Array r ix e -> Put ()
putArrayHeader Array P ix e
arr
let eltBytes :: Int
eltBytes = e -> Int
forall a. Prim a => a -> Int
Primitive.sizeOf (e
forall a. HasCallStack => a
undefined :: e)
n :: Int
n = Sz ix -> Int
forall ix. Index ix => Sz ix -> Int
totalElem (Array P ix e -> Sz ix
forall r ix e. Size r => Array r ix e -> Sz ix
size Array P ix e
arr) Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
eltBytes
Put ()
_ = e -> Put ()
forall t. Persist t => t -> Put ()
put (e
forall a. HasCallStack => a
undefined :: e)
s :: ShortByteString
s =
case Array P ix e -> ByteArray
forall ix e. Array P ix e -> ByteArray
unwrapByteArray Array P ix e
arr of
BA.ByteArray ByteArray#
ba -> ByteArray# -> ShortByteString
SBS.SBS ByteArray#
ba
Int -> Put ()
grow Int
n
(PutEnv -> Ptr Word8 -> IO (Ptr Word8 :!: ())) -> Put ()
forall a. (PutEnv -> Ptr Word8 -> IO (Ptr Word8 :!: a)) -> Put a
Put ((PutEnv -> Ptr Word8 -> IO (Ptr Word8 :!: ())) -> Put ())
-> (PutEnv -> Ptr Word8 -> IO (Ptr Word8 :!: ())) -> Put ()
forall a b. (a -> b) -> a -> b
$ \PutEnv
_ Ptr Word8
p -> do
ShortByteString -> Int -> Ptr Word8 -> Int -> IO ()
forall a. ShortByteString -> Int -> Ptr a -> Int -> IO ()
SBS.copyToPtr ShortByteString
s (Array P ix e -> Int
forall ix e. Array P ix e -> Int
unwrapByteArrayOffset Array P ix e
arr Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
eltBytes) Ptr Word8
p Int
n
(Ptr Word8 :!: ()) -> IO (Ptr Word8 :!: ())
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((Ptr Word8 :!: ()) -> IO (Ptr Word8 :!: ()))
-> (Ptr Word8 :!: ()) -> IO (Ptr Word8 :!: ())
forall a b. (a -> b) -> a -> b
$! Ptr Word8
p Ptr Word8 -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
n Ptr Word8 -> () -> Ptr Word8 :!: ()
forall a b. a -> b -> a :!: b
:!: ()
#endif
getPrimArray :: forall ix e . (Index ix, Prim e, Persist e) => Get (Array P ix e)
getPrimArray :: Get (Array P ix e)
getPrimArray = do
#ifdef WORDS_BIGENDIAN
getArray
#else
Comp
comp <- Get Comp
forall t. Persist t => Get t
get
Sz ix
sz <- Get (Sz ix)
forall t. Persist t => Get t
get
let n :: Int
n = Sz ix -> Int
forall ix. Index ix => Sz ix -> Int
totalElem Sz ix
sz Int -> Int -> Int
forall a. Num a => a -> a -> a
* e -> Int
forall a. Prim a => a -> Int
Primitive.sizeOf (e
forall a. HasCallStack => a
undefined :: e)
Put ()
_ = e -> Put ()
forall t. Persist t => t -> Put ()
put (e
forall a. HasCallStack => a
undefined :: e)
SBS.SBS ByteArray#
sbs <- ByteString -> ShortByteString
SBS.toShort (ByteString -> ShortByteString)
-> Get ByteString -> Get ShortByteString
forall (m :: * -> *) a b. Monad m => (a -> b) -> m a -> m b
<$!> Int -> Get ByteString
getBytes Int
n
(SomeException -> Get (Array P ix e))
-> (Array P ix e -> Get (Array P ix e))
-> Either SomeException (Array P ix e)
-> Get (Array P ix e)
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (String -> Get (Array P ix e)
forall (m :: * -> *) a. MonadFail m => String -> m a
Fail.fail (String -> Get (Array P ix e))
-> (SomeException -> String) -> SomeException -> Get (Array P ix e)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SomeException -> String
forall a. Show a => a -> String
show) Array P ix e -> Get (Array P ix e)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Either SomeException (Array P ix e) -> Get (Array P ix e))
-> Either SomeException (Array P ix e) -> Get (Array P ix e)
forall a b. (a -> b) -> a -> b
$
Sz ix -> Array P Int e -> Either SomeException (Array P ix e)
forall r ix ix' e (m :: * -> *).
(MonadThrow m, Index ix', Index ix, Size r) =>
Sz ix' -> Array r ix e -> m (Array r ix' e)
resizeM Sz ix
sz (Comp -> ByteArray -> Array P Int e
forall e. Prim e => Comp -> ByteArray -> Array P Int e
fromByteArray Comp
comp (ByteArray# -> ByteArray
BA.ByteArray ByteArray#
sbs))
#endif
instance (Index ix, Prim e, Persist e) => Persist (Array P ix e) where
put :: Array P ix e -> Put ()
put = Array P ix e -> Put ()
forall ix e.
(Index ix, Prim e, Persist e) =>
Array P ix e -> Put ()
putPrimArray
get :: Get (Array P ix e)
get = Get (Array P ix e)
forall ix e. (Index ix, Prim e, Persist e) => Get (Array P ix e)
getPrimArray
putStorableArray :: forall ix e . (Index ix, Storable e, Persist e) => Array S ix e -> Put ()
putStorableArray :: Array S ix e -> Put ()
putStorableArray Array S ix e
arr = do
#ifdef WORDS_BIGENDIAN
puArray
#else
let Put ()
_ = e -> Put ()
forall t. Persist t => t -> Put ()
put (e
forall a. HasCallStack => a
undefined :: e)
Array S ix e -> Put ()
forall r ix e.
(Strategy r, Size r, Index ix) =>
Array r ix e -> Put ()
putArrayHeader Array S ix e
arr
ByteString -> Put ()
putByteString (ByteString -> Put ()) -> ByteString -> Put ()
forall a b. (a -> b) -> a -> b
$!
case Array S ix e -> (ForeignPtr e, Int)
forall ix e. Index ix => Array S ix e -> (ForeignPtr e, Int)
unsafeArrayToForeignPtr Array S ix e
arr of
(ForeignPtr e
fp, Int
len) ->
ForeignPtr Word8 -> Int -> Int -> ByteString
BS.PS (ForeignPtr e -> ForeignPtr Word8
forall a b. ForeignPtr a -> ForeignPtr b
castForeignPtr ForeignPtr e
fp) Int
0 (Int
len Int -> Int -> Int
forall a. Num a => a -> a -> a
* e -> Int
forall a. Storable a => a -> Int
Storable.sizeOf (e
forall a. HasCallStack => a
undefined :: e))
#endif
getStorableArray :: forall ix e . (Index ix, Storable e, Persist e) => Get (Array S ix e)
getStorableArray :: Get (Array S ix e)
getStorableArray = do
#ifdef WORDS_BIGENDIAN
getArray
#else
Comp
comp <- Get Comp
forall t. Persist t => Get t
get
Sz ix
sz <- Get (Sz ix)
forall t. Persist t => Get t
get
let eltCount :: Int
eltCount = Sz ix -> Int
forall ix. Index ix => Sz ix -> Int
totalElem Sz ix
sz
eltBytes :: Int
eltBytes = e -> Int
forall a. Storable a => a -> Int
Storable.sizeOf (e
forall a. HasCallStack => a
undefined :: e)
Put ()
_ = e -> Put ()
forall t. Persist t => t -> Put ()
put (e
forall a. HasCallStack => a
undefined :: e)
BS.PS ForeignPtr Word8
fp Int
off Int
_ <- Int -> Get ByteString
getByteString (Int
eltCount Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
eltBytes)
(SomeException -> Get (Array S ix e))
-> (Array S ix e -> Get (Array S ix e))
-> Either SomeException (Array S ix e)
-> Get (Array S ix e)
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (String -> Get (Array S ix e)
forall (m :: * -> *) a. MonadFail m => String -> m a
Fail.fail (String -> Get (Array S ix e))
-> (SomeException -> String) -> SomeException -> Get (Array S ix e)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SomeException -> String
forall a. Show a => a -> String
show) Array S ix e -> Get (Array S ix e)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Either SomeException (Array S ix e) -> Get (Array S ix e))
-> Either SomeException (Array S ix e) -> Get (Array S ix e)
forall a b. (a -> b) -> a -> b
$
Sz ix -> Array S Int e -> Either SomeException (Array S ix e)
forall r ix ix' e (m :: * -> *).
(MonadThrow m, Index ix', Index ix, Size r) =>
Sz ix' -> Array r ix e -> m (Array r ix' e)
resizeM Sz ix
sz (Array S Int e -> Either SomeException (Array S ix e))
-> Array S Int e -> Either SomeException (Array S ix e)
forall a b. (a -> b) -> a -> b
$
Comp -> ForeignPtr e -> Sz1 -> Array S Int e
forall e. Comp -> ForeignPtr e -> Sz1 -> Vector S e
unsafeArrayFromForeignPtr0 Comp
comp (ForeignPtr Word8
fp ForeignPtr Word8 -> Int -> ForeignPtr e
forall a b. ForeignPtr a -> Int -> ForeignPtr b
`plusForeignPtr` Int
off) (Int -> Sz1
forall ix. Index ix => ix -> Sz ix
Sz Int
eltCount)
#endif
instance (Index ix, Storable e, Persist e) => Persist (Array S ix e) where
put :: Array S ix e -> Put ()
put = Array S ix e -> Put ()
forall ix e.
(Index ix, Storable e, Persist e) =>
Array S ix e -> Put ()
putStorableArray
get :: Get (Array S ix e)
get = Get (Array S ix e)
forall ix e.
(Index ix, Storable e, Persist e) =>
Get (Array S ix e)
getStorableArray