-- | Dense primitive arrays where the lower index is zero (or the
-- equivalent of zero for newtypes and enumerations).
--
-- Actual @write@s to data structures use a more safe @write@ instead of
-- the unsafe @unsafeWrite@. Writes also tend to occur much less in DP
-- algorithms (say, N^2 writes for an N^3 time algorithm -- mostly reads
-- are being executed).
--
-- TODO consider if we want to force the lower index to be zero, or allow
-- non-zero lower indices. Will have to be considered together with the
-- @Index.Class@ module!
--
-- TODO while @Unboxed@ is, in princile, @Hashable@, we'd need the
-- corresponding @VU.Vector@ instances ...
--
-- TODO rename to Dense.Vector, since there are other possibilities to store,
-- without basing on vector.

module Data.PrimitiveArray.Dense where

import           Control.Lens (makeLenses)
import           Control.DeepSeq
import           Control.Exception (assert)
import           Control.Monad (liftM, forM_, zipWithM_, when)
import           Control.Monad.Primitive (PrimState)
import           Data.Aeson (ToJSON,FromJSON)
import           Data.Binary (Binary)
import           Data.Data
import           Data.Hashable (Hashable)
import           Data.Serialize (Serialize)
import           Data.Typeable (Typeable)
import           Data.Vector.Binary
import           Data.Vector.Generic.Mutable as GM hiding (length)
import           Data.Vector.Serialize
import           Debug.Trace
import           GHC.Generics (Generic)
import qualified Data.Vector as V
import qualified Data.Vector.Fusion.Stream.Monadic as SM
import qualified Data.Vector.Generic as VG
import qualified Data.Vector.Storable as VS
import qualified Data.Vector.Unboxed as VU

import           Data.PrimitiveArray.Class
import           Data.PrimitiveArray.Index.Class



data Dense v sh e = Dense { Dense v sh e -> LimitType sh
_denseLimit :: !(LimitType sh), Dense v sh e -> v e
_denseV :: !(v e) }
makeLenses ''Dense

type Unboxed sh e = Dense VU.Vector sh e

type Storable sh e = Dense VS.Vector sh e

type Boxed sh e = Dense V.Vector sh e



deriving instance (Eq      (LimitType sh), Eq (v e)     ) => Eq      (Dense v sh e)
deriving instance (Generic (LimitType sh), Generic (v e)) => Generic (Dense v sh e)
deriving instance (Read    (LimitType sh), Read (v e)   ) => Read    (Dense v sh e)
deriving instance (Show    (LimitType sh), Show (v e)   ) => Show    (Dense v sh e)
deriving instance (Functor v)                             => Functor (Dense v sh)

deriving instance Typeable (Dense v sh e)

deriving instance (Data (v e), Data (LimitType sh), Data e, Data sh, Typeable sh, Typeable e, Typeable v) => Data (Dense v sh e)

instance (Binary    (LimitType sh), Binary    (v e), Generic (LimitType sh), Generic (v e)) => Binary    (Dense v sh e)
instance (Serialize (LimitType sh), Serialize (v e), Generic (LimitType sh), Generic (v e)) => Serialize (Dense v sh e)
instance (ToJSON    (LimitType sh), ToJSON    (v e), Generic (LimitType sh), Generic (v e)) => ToJSON    (Dense v sh e)
instance (FromJSON  (LimitType sh), FromJSON  (v e), Generic (LimitType sh), Generic (v e)) => FromJSON  (Dense v sh e)
instance (Hashable  (LimitType sh), Hashable  (v e), Generic (LimitType sh), Generic (v e)) => Hashable  (Dense v sh e)

instance (NFData (LimitType sh), NFData (v e)) ⇒ NFData (Dense v sh e) where
  rnf :: Dense v sh e -> ()
rnf (Dense LimitType sh
h v e
xs) = LimitType sh -> ()
forall a. NFData a => a -> ()
rnf LimitType sh
h () -> () -> ()
`seq` v e -> ()
forall a. NFData a => a -> ()
rnf v e
xs
  {-# Inline rnf #-}



data instance MutArr m (Dense v sh e) = MDense !(LimitType sh) !(VG.Mutable v (PrimState m) e)
  deriving ((forall x.
 MutArr m (Dense v sh e) -> Rep (MutArr m (Dense v sh e)) x)
-> (forall x.
    Rep (MutArr m (Dense v sh e)) x -> MutArr m (Dense v sh e))
-> Generic (MutArr m (Dense v sh e))
forall x.
Rep (MutArr m (Dense v sh e)) x -> MutArr m (Dense v sh e)
forall x.
MutArr m (Dense v sh e) -> Rep (MutArr m (Dense v sh e)) x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
forall (m :: * -> *) (v :: * -> *) sh e x.
Rep (MutArr m (Dense v sh e)) x -> MutArr m (Dense v sh e)
forall (m :: * -> *) (v :: * -> *) sh e x.
MutArr m (Dense v sh e) -> Rep (MutArr m (Dense v sh e)) x
$cto :: forall (m :: * -> *) (v :: * -> *) sh e x.
Rep (MutArr m (Dense v sh e)) x -> MutArr m (Dense v sh e)
$cfrom :: forall (m :: * -> *) (v :: * -> *) sh e x.
MutArr m (Dense v sh e) -> Rep (MutArr m (Dense v sh e)) x
Generic,Typeable)

instance (Show (LimitType sh), Show (VG.Mutable v (PrimState m) e), VG.Mutable v (PrimState m) e ~ mv) ⇒ Show (MutArr m (Dense v sh e)) where
  show :: MutArr m (Dense v sh e) -> String
show (MDense sh mv) = (LimitType sh, Mutable v (PrimState m) e) -> String
forall a. Show a => a -> String
show (LimitType sh
sh,Mutable v (PrimState m) e
mv)

instance (NFData (LimitType sh), NFData (VG.Mutable v (PrimState m) e), VG.Mutable v (PrimState m) e ~ mv) ⇒ NFData (MutArr m (Dense v sh e)) where
  rnf :: MutArr m (Dense v sh e) -> ()
rnf (MDense h xs) = LimitType sh -> ()
forall a. NFData a => a -> ()
rnf LimitType sh
h () -> () -> ()
`seq` Mutable v (PrimState m) e -> ()
forall a. NFData a => a -> ()
rnf Mutable v (PrimState m) e
xs
  {-# Inline rnf #-}

{-
instance
  ( Index sh, MutArr m (Dense v sh e) ~ mv
  , GM.MVector (VG.Mutable v) e
#if ADPFUSION_DEBUGOUTPUT
  , Show sh, Show (LimitType sh), Show e
#endif
  ) ⇒ MPrimArrayOps (Dense v) sh e where
-}

instance
  ( Index sh, VG.Vector v e
#if ADPFUSION_DEBUGOUTPUT
  , Show sh, Show (LimitType sh), Show e
#endif
  ) ⇒ PrimArrayOps (Dense v) sh e where

  -- ** pure operations

  {-# Inline upperBound #-}
  upperBound :: Dense v sh e -> LimitType sh
upperBound (Dense LimitType sh
h v e
_) = LimitType sh
h
  {-# Inline unsafeFreezeM #-}
  unsafeFreezeM :: MutArr m (Dense v sh e) -> m (Dense v sh e)
unsafeFreezeM (MDense h mba) = LimitType sh -> v e -> Dense v sh e
forall k (v :: k -> *) sh (e :: k).
LimitType sh -> v e -> Dense v sh e
Dense LimitType sh
h (v e -> Dense v sh e) -> m (v e) -> m (Dense v sh e)
forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
`liftM` Mutable v (PrimState m) e -> m (v e)
forall (m :: * -> *) (v :: * -> *) a.
(PrimMonad m, Vector v a) =>
Mutable v (PrimState m) a -> m (v a)
VG.unsafeFreeze Mutable v (PrimState m) e
mba
  {-# Inline unsafeThawM #-}
  unsafeThawM :: Dense v sh e -> m (MutArr m (Dense v sh e))
unsafeThawM   (Dense LimitType sh
h v e
ba) = LimitType sh
-> Mutable v (PrimState m) e -> MutArr m (Dense v sh e)
forall (m :: * -> *) (v :: * -> *) sh e.
LimitType sh
-> Mutable v (PrimState m) e -> MutArr m (Dense v sh e)
MDense LimitType sh
h (Mutable v (PrimState m) e -> MutArr m (Dense v sh e))
-> m (Mutable v (PrimState m) e) -> m (MutArr m (Dense v sh e))
forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
`liftM` v e -> m (Mutable v (PrimState m) e)
forall (m :: * -> *) (v :: * -> *) a.
(PrimMonad m, Vector v a) =>
v a -> m (Mutable v (PrimState m) a)
VG.unsafeThaw v e
ba
  {-# Inline unsafeIndex #-}
  unsafeIndex :: Dense v sh e -> sh -> e
unsafeIndex  (Dense LimitType sh
h v e
ba) sh
idx = v e -> Int -> e
forall (v :: * -> *) a. Vector v a => v a -> Int -> a
VG.unsafeIndex v e
ba (LimitType sh -> sh -> Int
forall i. Index i => LimitType i -> i -> Int
linearIndex LimitType sh
h sh
idx)
  {-# Inline safeIndex #-}
  safeIndex :: Dense v sh e -> sh -> Maybe e
safeIndex (Dense LimitType sh
h v e
ba) sh
idx = if LimitType sh -> sh -> Bool
forall i. Index i => LimitType i -> i -> Bool
inBounds LimitType sh
h sh
idx then e -> Maybe e
forall a. a -> Maybe a
Just (e -> Maybe e) -> e -> Maybe e
forall a b. (a -> b) -> a -> b
$ Dense v sh e -> sh -> e
forall (arr :: * -> * -> *) sh elm.
PrimArrayOps arr sh elm =>
arr sh elm -> sh -> elm
unsafeIndex (LimitType sh -> v e -> Dense v sh e
forall k (v :: k -> *) sh (e :: k).
LimitType sh -> v e -> Dense v sh e
Dense LimitType sh
h v e
ba) sh
idx else Maybe e
forall a. Maybe a
Nothing
  {-# Inline transformShape #-}
  transformShape :: (LimitType sh -> LimitType sh') -> Dense v sh e -> Dense v sh' e
transformShape LimitType sh -> LimitType sh'
tr (Dense LimitType sh
h v e
ba) = LimitType sh' -> v e -> Dense v sh' e
forall k (v :: k -> *) sh (e :: k).
LimitType sh -> v e -> Dense v sh e
Dense (LimitType sh -> LimitType sh'
tr LimitType sh
h) v e
ba

  -- ** monadic operations

  {-# Inline upperBoundM #-}
  upperBoundM :: MutArr m (Dense v sh e) -> LimitType sh
upperBoundM (MDense h _) = LimitType sh
h
  {-# Inline fromListM #-}
  fromListM :: LimitType sh -> [e] -> m (MutArr m (Dense v sh e))
fromListM LimitType sh
h [e]
xs = do
    MutArr m (Dense v sh e)
ma ← LimitType sh -> m (MutArr m (Dense v sh e))
forall (arr :: * -> * -> *) sh elm (m :: * -> *).
(PrimArrayOps arr sh elm, PrimMonad m) =>
LimitType sh -> m (MutArr m (arr sh elm))
newM LimitType sh
h
    let (MDense _ mba) = MutArr m (Dense v sh e)
ma
    -- there need to be at least as many elements, as we want to fill. There could be more, in debug
    -- tests, we like to do @[0..]@ and this should not trigger the assert.
    (Int -> e -> m ()) -> Stream m Int -> Stream m e -> m ()
forall (m :: * -> *) a b c.
Monad m =>
(a -> b -> m c) -> Stream m a -> Stream m b -> m ()
SM.zipWithM_ (\Int
k e
x → Bool -> m () -> m ()
forall a. (?callStack::CallStack) => Bool -> a -> a
assert ([e] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (Int -> [e] -> [e]
forall a. Int -> [a] -> [a]
Prelude.take (LimitType sh -> Int
forall i. Index i => LimitType i -> Int
size LimitType sh
h) [e]
xs) Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== LimitType sh -> Int
forall i. Index i => LimitType i -> Int
size LimitType sh
h) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ Mutable v (PrimState m) e -> Int -> e -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
unsafeWrite Mutable v (PrimState m) e
mba Int
k e
x) (Int -> Int -> Stream m Int
forall a (m :: * -> *). (Enum a, Monad m) => a -> a -> Stream m a
SM.enumFromTo Int
0 (LimitType sh -> Int
forall i. Index i => LimitType i -> Int
size LimitType sh
h Int -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)) ([e] -> Stream m e
forall (m :: * -> *) a. Monad m => [a] -> Stream m a
SM.fromList [e]
xs)
    MutArr m (Dense v sh e) -> m (MutArr m (Dense v sh e))
forall (m :: * -> *) a. Monad m => a -> m a
return MutArr m (Dense v sh e)
ma
  {-# Inline newM #-}     -- TODO was NoInline, check if anything breaks!
  newM :: LimitType sh -> m (MutArr m (Dense v sh e))
newM LimitType sh
h = LimitType sh
-> Mutable v (PrimState m) e -> MutArr m (Dense v sh e)
forall (m :: * -> *) (v :: * -> *) sh e.
LimitType sh
-> Mutable v (PrimState m) e -> MutArr m (Dense v sh e)
MDense LimitType sh
h (Mutable v (PrimState m) e -> MutArr m (Dense v sh e))
-> m (Mutable v (PrimState m) e) -> m (MutArr m (Dense v sh e))
forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
`liftM` Int -> m (Mutable v (PrimState m) e)
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
Int -> m (v (PrimState m) a)
new (LimitType sh -> Int
forall i. Index i => LimitType i -> Int
size LimitType sh
h)
  {-# Inline newSM #-}
  newSM :: LimitType sh
-> FillStruc (Dense v sh e) -> m (MutArr m (Dense v sh e))
newSM = String
-> LimitType sh
-> FillStruc (Dense v sh e)
-> m (MutArr m (Dense v sh e))
forall a. (?callStack::CallStack) => String -> a
error String
"not implemented, use newM for dense arrays"
  {-# Inline newWithM #-}
  newWithM :: LimitType sh -> e -> m (MutArr m (Dense v sh e))
newWithM LimitType sh
h e
def = do
    MutArr m (Dense v sh e)
ma ← LimitType sh -> m (MutArr m (Dense v sh e))
forall (arr :: * -> * -> *) sh elm (m :: * -> *).
(PrimArrayOps arr sh elm, PrimMonad m) =>
LimitType sh -> m (MutArr m (arr sh elm))
newM LimitType sh
h
    let (MDense _ mba) = MutArr m (Dense v sh e)
ma
    Mutable v (PrimState m) e -> e -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> a -> m ()
GM.set Mutable v (PrimState m) e
mba e
def
    MutArr m (Dense v sh e) -> m (MutArr m (Dense v sh e))
forall (m :: * -> *) a. Monad m => a -> m a
return MutArr m (Dense v sh e)
ma
  {-# Inline newWithSM #-}
  newWithSM :: LimitType sh
-> FillStruc (Dense v sh e) -> e -> m (MutArr m (Dense v sh e))
newWithSM = String
-> LimitType sh
-> FillStruc (Dense v sh e)
-> e
-> m (MutArr m (Dense v sh e))
forall a. (?callStack::CallStack) => String -> a
error String
"not implemented, use newWithSM for dense arrays"
  {-# Inline readM #-}
  readM :: MutArr m (Dense v sh e) -> sh -> m e
readM  (MDense h mba) sh
idx     = Bool -> m e -> m e
forall a. (?callStack::CallStack) => Bool -> a -> a
assert (LimitType sh -> sh -> Bool
forall i. Index i => LimitType i -> i -> Bool
inBounds LimitType sh
h sh
idx) (m e -> m e) -> m e -> m e
forall a b. (a -> b) -> a -> b
$ Mutable v (PrimState m) e -> Int -> m e
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
unsafeRead  Mutable v (PrimState m) e
mba (LimitType sh -> sh -> Int
forall i. Index i => LimitType i -> i -> Int
linearIndex LimitType sh
h sh
idx)
  {-# Inline safeReadM #-}
  safeReadM :: MutArr m (Dense v sh e) -> sh -> m (Maybe e)
safeReadM MutArr m (Dense v sh e)
dense sh
idx = if MutArr m (Dense v sh e) -> sh -> Bool
forall (m :: * -> *) (arr :: * -> * -> *) sh elm.
(Monad m, PrimArrayOps arr sh elm) =>
MutArr m (arr sh elm) -> sh -> Bool
inBoundsM MutArr m (Dense v sh e)
dense sh
idx then e -> Maybe e
forall a. a -> Maybe a
Just (e -> Maybe e) -> m e -> m (Maybe e)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MutArr m (Dense v sh e) -> sh -> m e
forall (arr :: * -> * -> *) sh elm (m :: * -> *).
(PrimArrayOps arr sh elm, PrimMonad m) =>
MutArr m (arr sh elm) -> sh -> m elm
readM MutArr m (Dense v sh e)
dense sh
idx else m (Maybe e)
forall a. (?callStack::CallStack) => a
undefined
  {-# Inline writeM #-}
  writeM :: MutArr m (Dense v sh e) -> sh -> e -> m ()
writeM (MDense h mba) sh
idx e
elm =
#if ADPFUSION_DEBUGOUTPUT
    (if inBounds h idx then id else traceShow ("writeM", h, idx, elm, size h, linearIndex h idx, inBounds h idx))
#endif
    Bool -> m () -> m ()
forall a. (?callStack::CallStack) => Bool -> a -> a
assert (LimitType sh -> sh -> Bool
forall i. Index i => LimitType i -> i -> Bool
inBounds LimitType sh
h sh
idx) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ Mutable v (PrimState m) e -> Int -> e -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
unsafeWrite Mutable v (PrimState m) e
mba (LimitType sh -> sh -> Int
forall i. Index i => LimitType i -> i -> Int
linearIndex LimitType sh
h sh
idx) e
elm
  {-# Inline safeWriteM #-}
  safeWriteM :: MutArr m (Dense v sh e) -> sh -> e -> m ()
safeWriteM MutArr m (Dense v sh e)
dense sh
idx e
elm = Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (MutArr m (Dense v sh e) -> sh -> Bool
forall (m :: * -> *) (arr :: * -> * -> *) sh elm.
(Monad m, PrimArrayOps arr sh elm) =>
MutArr m (arr sh elm) -> sh -> Bool
inBoundsM MutArr m (Dense v sh e)
dense sh
idx) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ MutArr m (Dense v sh e) -> sh -> e -> m ()
forall (arr :: * -> * -> *) sh elm (m :: * -> *).
(PrimArrayOps arr sh elm, PrimMonad m) =>
MutArr m (arr sh elm) -> sh -> elm -> m ()
writeM MutArr m (Dense v sh e)
dense sh
idx e
elm

instance (Index sh, VG.Vector v e, VG.Vector v e') ⇒ PrimArrayMap (Dense v) sh e e' where
  {-# Inline mapArray #-}
  mapArray :: (e -> e') -> Dense v sh e -> Dense v sh e'
mapArray e -> e'
f (Dense LimitType sh
h v e
xs) = LimitType sh -> v e' -> Dense v sh e'
forall k (v :: k -> *) sh (e :: k).
LimitType sh -> v e -> Dense v sh e
Dense LimitType sh
h ((e -> e') -> v e -> v e'
forall (v :: * -> *) a b.
(Vector v a, Vector v b) =>
(a -> b) -> v a -> v b
VG.map e -> e'
f v e
xs)


{-
 -
 - This stuff tells us how to write efficient generics on large data
 - constructors like the Turner and Vienna ctors.
 -

import qualified Data.Generics.TH as T

data Unboxed sh e = Unboxed !sh !(VU.Vector e)
  deriving (Show,Eq,Ord)

data X e = X (Unboxed DIM1 e) (Unboxed DIM1 e)
  deriving (Show,Eq,Ord)

x :: X Int
x = X z z where z = (Unboxed (Z:.10) (VU.fromList [ 0 .. 10] ))

pot :: X Int -> X Double
pot = $( T.thmapT (T.mkTs ['f]) [t| X Int |] ) where
  f :: Unboxed DIM1 Int -> Unboxed DIM1 Double
  f (Unboxed sh xs) = Unboxed sh (VU.map fromIntegral xs)

-}