{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UndecidableInstances #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}
-- |
-- Module      : Massiv.Serialise
-- Copyright   : (c) Alexey Kuleshevich 2021
-- License     : BSD3
-- Maintainer  : Alexey Kuleshevich <alexey@kuleshevi.ch>
-- Stability   : experimental
-- Portability : non-portable
--
-- This package provides instances for `Serialise` class for all mutable `Array`
-- representations in [@massiv@](https://hackage.haskell.org/package/massiv) package. These
-- instances are provided as orphans from a separate package in order to avoid direct
-- dependency on [@serialise@](https://hackage.haskell.org/package/serialise) package in
-- @massiv@.
--
-- Array serialisation is done by falling back onto instances for `VG.Vector` types from
-- [@vector@](https://hackage.haskell.org/package/vector) package.
--
-- Below is a simple example how to use it. Note a blank module import: @import
-- Massiv.Serialise ()@, which is the only thing needed from this module in order to use
-- provided orphan instances.
--
-- >>> import Massiv.Serialise ()
-- >>> import Data.Massiv.Array as A
-- >>> let arr = A.fromList A.Seq [72,97,115,107,101,108,108] :: A.Vector A.P Int
-- >>> serialise arr
-- "\NUL\a\135\CANH\CANa\CANs\CANk\CANe\CANl\CANl"
-- >>> deserialise (serialise arr) :: A.Vector A.P Int
-- Array P Seq (Sz1 7)
--   [ 72, 97, 115, 107, 101, 108, 108 ]
--
module Massiv.Serialise
  ( -- * Helper functions used to define Serialise instances
    encodeIx
  , decodeIx
  , mkSzFail
  , encodeArray
  , decodeArray
  ) where

import Codec.Serialise
import Codec.Serialise.Decoding
import Codec.Serialise.Encoding
import Control.DeepSeq (NFData)
import Control.Monad
import qualified Control.Monad.Fail as Fail
import Data.Foldable as F
import Data.Massiv.Array
import Data.Massiv.Array.Manifest.Vector
import Data.Proxy
import Data.Typeable
import qualified Data.Vector as V
import qualified Data.Vector.Generic as VG
import qualified Data.Vector.Primitive as VP
import qualified Data.Vector.Storable as VS
import qualified Data.Vector.Unboxed as VU

instance Serialise Comp where
  encode :: Comp -> Encoding
encode Comp
comp =
    case Comp
comp of
      Comp
Seq      -> Int -> Encoding
forall a. Serialise a => a -> Encoding
encode (Int
0 :: Int)
      ParOn [Int]
xs -> Int -> Encoding
forall a. Serialise a => a -> Encoding
encode (Int
1 :: Int) Encoding -> Encoding -> Encoding
forall a. Semigroup a => a -> a -> a
<> [Int] -> Encoding
forall a. Serialise a => a -> Encoding
encode [Int]
xs
      ParN Word16
n   -> Int -> Encoding
forall a. Serialise a => a -> Encoding
encode (Int
2 :: Int) Encoding -> Encoding -> Encoding
forall a. Semigroup a => a -> a -> a
<> Word16 -> Encoding
forall a. Serialise a => a -> Encoding
encode Word16
n
  decode :: Decoder s Comp
decode = do
    Int
ty :: Int <- Decoder s Int
forall a s. Serialise a => Decoder s a
decode
    case Int
ty of
      Int
0 -> Comp -> Decoder s Comp
forall (f :: * -> *) a. Applicative f => a -> f a
pure Comp
Seq
      Int
1 -> [Int] -> Comp
ParOn ([Int] -> Comp) -> Decoder s [Int] -> Decoder s Comp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Decoder s [Int]
forall a s. Serialise a => Decoder s a
decode
      Int
2 -> Word16 -> Comp
ParN (Word16 -> Comp) -> Decoder s Word16 -> Decoder s Comp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Decoder s Word16
forall a s. Serialise a => Decoder s a
decode
      Int
n -> String -> Decoder s Comp
forall (m :: * -> *) a. MonadFail m => String -> m a
Fail.fail (String -> Decoder s Comp) -> String -> Decoder s Comp
forall a b. (a -> b) -> a -> b
$ String
"Unexpected Comp tag: " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Int -> String
forall a. Show a => a -> String
show Int
n



-- | Encode index
--
-- @since 0.1.0
encodeIx ::
     forall ix. Index ix
  => ix
  -> Encoding
encodeIx :: ix -> Encoding
encodeIx = (Encoding -> Int -> Encoding) -> Encoding -> ix -> Encoding
forall ix a. Index ix => (a -> Int -> a) -> a -> ix -> a
foldlIndex (\ !Encoding
acc Int
i -> Int -> Encoding
forall a. Serialise a => a -> Encoding
encode Int
i Encoding -> Encoding -> Encoding
forall a. Semigroup a => a -> a -> a
<> Encoding
acc) Encoding
forall a. Monoid a => a
mempty

-- | Decode index
--
-- @since 0.1.0
decodeIx ::
     forall s ix. Index ix
  => Decoder s ix
decodeIx :: Decoder s ix
decodeIx = do
  let decodeDim :: b -> Dim -> Decoder s b
decodeDim b
ix Dim
dim = do
        Int
i <- Decoder s Int
forall a s. Serialise a => Decoder s a
decode
        (SomeException -> Decoder s b)
-> (b -> Decoder s b) -> Either SomeException b -> Decoder s b
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (String -> Decoder s b
forall (m :: * -> *) a. MonadFail m => String -> m a
Fail.fail (String -> Decoder s b)
-> (SomeException -> String) -> SomeException -> Decoder s b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SomeException -> String
forall a. Show a => a -> String
show) b -> Decoder s b
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Either SomeException b -> Decoder s b)
-> Either SomeException b -> Decoder s 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 -> Decoder s ix) -> ix -> [Dim] -> Decoder s ix
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
F.foldlM ix -> Dim -> Decoder s ix
forall b s. Index b => b -> Dim -> Decoder s b
decodeDim 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 Serialise Ix2 where
  encode :: Ix2 -> Encoding
encode = Ix2 -> Encoding
forall ix. Index ix => ix -> Encoding
encodeIx
  decode :: Decoder s Ix2
decode = Decoder s Ix2
forall s ix. Index ix => Decoder s ix
decodeIx

instance Index (IxN n) => Serialise (IxN n) where
  encode :: IxN n -> Encoding
encode = IxN n -> Encoding
forall ix. Index ix => ix -> Encoding
encodeIx
  decode :: Decoder s (IxN n)
decode = Decoder s (IxN n)
forall s ix. Index ix => Decoder s ix
decodeIx



-- | Construct size from index verifying its correctness.
--
-- @since 0.1.0
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, Num b, MonadFail m, Show 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, Serialise ix) => Serialise (Sz ix) where
  encode :: Sz ix -> Encoding
encode = ix -> Encoding
forall ix. Index ix => ix -> Encoding
encodeIx (ix -> Encoding) -> (Sz ix -> ix) -> Sz ix -> Encoding
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Sz ix -> ix
forall ix. Sz ix -> ix
unSz
  decode :: Decoder s (Sz ix)
decode = ix -> Decoder s (Sz ix)
forall ix (m :: * -> *). (Index ix, MonadFail m) => ix -> m (Sz ix)
mkSzFail (ix -> Decoder s (Sz ix)) -> Decoder s ix -> Decoder s (Sz ix)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Decoder s ix
forall s ix. Index ix => Decoder s ix
decodeIx



-- | Encode array by using its corresponding vector instance
--
-- @since 0.1.0
encodeArray ::
     forall v r ix e.
     ( Manifest r ix e
     , Mutable (ARepr v) ix e
     , VG.Vector v e
     , VRepr (ARepr v) ~ v
     , Serialise ix
     , Serialise (v e)
     )
  => Array r ix e
  -> Encoding
encodeArray :: Array r ix e -> Encoding
encodeArray Array r ix e
arr =
  Comp -> Encoding
forall a. Serialise a => a -> Encoding
encode (Array r ix e -> Comp
forall r ix e. Load r ix e => Array r ix e -> Comp
getComp Array r ix e
arr) Encoding -> Encoding -> Encoding
forall a. Semigroup a => a -> a -> a
<> Sz ix -> Encoding
forall a. Serialise a => a -> Encoding
encode (Array r ix e -> Sz ix
forall r ix e. Load r ix e => Array r ix e -> Sz ix
size Array r ix e
arr) Encoding -> Encoding -> Encoding
forall a. Semigroup a => a -> a -> a
<> v e -> Encoding
forall a. Serialise a => a -> Encoding
encode (Array r ix e -> v e
forall r ix e (v :: * -> *).
(Manifest r ix e, Mutable (ARepr v) ix e, Vector v e,
 VRepr (ARepr v) ~ v) =>
Array r ix e -> v e
toVector Array r ix e
arr :: v e)

-- | Decode array by using its corresponding vector instance
--
-- @since 0.1.0
decodeArray ::
     forall v r ix e s.
     ( Typeable v
     , VG.Vector v e
     , Mutable (ARepr v) ix e
     , Mutable r ix e
     , Serialise ix
     , Serialise (v e)
     )
  => Decoder s (Array r ix e)
decodeArray :: Decoder s (Array r ix e)
decodeArray = do
  Comp
comp <- Decoder s Comp
forall a s. Serialise a => Decoder s a
decode
  Sz ix
sz <- Decoder s (Sz ix)
forall a s. Serialise a => Decoder s a
decode
  v e
vector :: v e <- Decoder s (v e)
forall a s. Serialise a => Decoder s a
decode
  -- setComp is to workaround a minor bug for boxed arrays in massiv < 0.6
  (SomeException -> Decoder s (Array r ix e))
-> (Array r ix e -> Decoder s (Array r ix e))
-> Either SomeException (Array r ix e)
-> Decoder s (Array r ix e)
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (String -> Decoder s (Array r ix e)
forall (m :: * -> *) a. MonadFail m => String -> m a
Fail.fail (String -> Decoder s (Array r ix e))
-> (SomeException -> String)
-> SomeException
-> Decoder s (Array r ix e)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SomeException -> String
forall a. Show a => a -> String
show) (Array r ix e -> Decoder s (Array r ix e)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Array r ix e -> Decoder s (Array r ix e))
-> (Array r ix e -> Array r ix e)
-> Array r ix e
-> Decoder s (Array r ix e)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Comp -> Array r ix e -> Array r ix e
forall r ix e.
Construct r ix e =>
Comp -> Array r ix e -> Array r ix e
setComp Comp
comp) (Either SomeException (Array r ix e) -> Decoder s (Array r ix e))
-> Either SomeException (Array r ix e) -> Decoder s (Array r ix e)
forall a b. (a -> b) -> a -> b
$ Comp -> Sz ix -> v e -> Either SomeException (Array r ix e)
forall (m :: * -> *) (v :: * -> *) a ix r.
(MonadThrow m, Typeable v, Vector v a, Mutable (ARepr v) ix a,
 Mutable r ix a) =>
Comp -> Sz ix -> v a -> m (Array r ix a)
fromVectorM Comp
comp Sz ix
sz v e
vector

instance (Index ix, Serialise ix, Serialise e) => Serialise (Array B ix e) where
  encode :: Array B ix e -> Encoding
encode = forall r ix e.
(Manifest r ix e, Mutable (ARepr Vector) ix e, Vector Vector e,
 VRepr (ARepr Vector) ~ Vector, Serialise ix,
 Serialise (Vector e)) =>
Array r ix e -> Encoding
forall (v :: * -> *) r ix e.
(Manifest r ix e, Mutable (ARepr v) ix e, Vector v e,
 VRepr (ARepr v) ~ v, Serialise ix, Serialise (v e)) =>
Array r ix e -> Encoding
encodeArray @V.Vector
  decode :: Decoder s (Array B ix e)
decode = forall r ix e s.
(Typeable Vector, Vector Vector e, Mutable (ARepr Vector) ix e,
 Mutable r ix e, Serialise ix, Serialise (Vector e)) =>
Decoder s (Array r ix e)
forall (v :: * -> *) r ix e s.
(Typeable v, Vector v e, Mutable (ARepr v) ix e, Mutable r ix e,
 Serialise ix, Serialise (v e)) =>
Decoder s (Array r ix e)
decodeArray @V.Vector

instance (Index ix, NFData e, Serialise ix, Serialise e) => Serialise (Array N ix e) where
  encode :: Array N ix e -> Encoding
encode = Array B ix e -> Encoding
forall a. Serialise a => a -> Encoding
encode (Array B ix e -> Encoding)
-> (Array N ix e -> Array B ix e) -> Array N ix e -> Encoding
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array N ix e -> Array B ix e
forall ix e. Array N ix e -> Array B ix e
unwrapNormalForm
  decode :: Decoder s (Array N ix e)
decode = Array B ix e -> Array N ix e
forall ix e. (Index ix, NFData e) => Array B ix e -> Array N ix e
evalNormalForm (Array B ix e -> Array N ix e)
-> Decoder s (Array B ix e) -> Decoder s (Array N ix e)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Decoder s (Array B ix e)
forall a s. Serialise a => Decoder s a
decode

instance (Index ix, Storable e, Serialise ix, Serialise e) => Serialise (Array S ix e) where
  encode :: Array S ix e -> Encoding
encode = forall r ix e.
(Manifest r ix e, Mutable (ARepr Vector) ix e, Vector Vector e,
 VRepr (ARepr Vector) ~ Vector, Serialise ix,
 Serialise (Vector e)) =>
Array r ix e -> Encoding
forall (v :: * -> *) r ix e.
(Manifest r ix e, Mutable (ARepr v) ix e, Vector v e,
 VRepr (ARepr v) ~ v, Serialise ix, Serialise (v e)) =>
Array r ix e -> Encoding
encodeArray @VS.Vector
  decode :: Decoder s (Array S ix e)
decode = forall r ix e s.
(Typeable Vector, Vector Vector e, Mutable (ARepr Vector) ix e,
 Mutable r ix e, Serialise ix, Serialise (Vector e)) =>
Decoder s (Array r ix e)
forall (v :: * -> *) r ix e s.
(Typeable v, Vector v e, Mutable (ARepr v) ix e, Mutable r ix e,
 Serialise ix, Serialise (v e)) =>
Decoder s (Array r ix e)
decodeArray @VS.Vector

instance (Index ix, Unbox e, Serialise ix, Serialise e) => Serialise (Array U ix e) where
  encode :: Array U ix e -> Encoding
encode = forall r ix e.
(Manifest r ix e, Mutable (ARepr Vector) ix e, Vector Vector e,
 VRepr (ARepr Vector) ~ Vector, Serialise ix,
 Serialise (Vector e)) =>
Array r ix e -> Encoding
forall (v :: * -> *) r ix e.
(Manifest r ix e, Mutable (ARepr v) ix e, Vector v e,
 VRepr (ARepr v) ~ v, Serialise ix, Serialise (v e)) =>
Array r ix e -> Encoding
encodeArray @VU.Vector
  decode :: Decoder s (Array U ix e)
decode = forall r ix e s.
(Typeable Vector, Vector Vector e, Mutable (ARepr Vector) ix e,
 Mutable r ix e, Serialise ix, Serialise (Vector e)) =>
Decoder s (Array r ix e)
forall (v :: * -> *) r ix e s.
(Typeable v, Vector v e, Mutable (ARepr v) ix e, Mutable r ix e,
 Serialise ix, Serialise (v e)) =>
Decoder s (Array r ix e)
decodeArray @VU.Vector

instance (Index ix, Prim e, Serialise ix, Serialise e) => Serialise (Array P ix e) where
  encode :: Array P ix e -> Encoding
encode = forall r ix e.
(Manifest r ix e, Mutable (ARepr Vector) ix e, Vector Vector e,
 VRepr (ARepr Vector) ~ Vector, Serialise ix,
 Serialise (Vector e)) =>
Array r ix e -> Encoding
forall (v :: * -> *) r ix e.
(Manifest r ix e, Mutable (ARepr v) ix e, Vector v e,
 VRepr (ARepr v) ~ v, Serialise ix, Serialise (v e)) =>
Array r ix e -> Encoding
encodeArray @VP.Vector
  decode :: Decoder s (Array P ix e)
decode = forall r ix e s.
(Typeable Vector, Vector Vector e, Mutable (ARepr Vector) ix e,
 Mutable r ix e, Serialise ix, Serialise (Vector e)) =>
Decoder s (Array r ix e)
forall (v :: * -> *) r ix e s.
(Typeable v, Vector v e, Mutable (ARepr v) ix e, Mutable r ix e,
 Serialise ix, Serialise (v e)) =>
Decoder s (Array r ix e)
decodeArray @VP.Vector