{-# LANGUAGE DeriveDataTypeable    #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE TypeFamilies          #-}
-----------------------------------------------------------------------------
-- |
-- Module      :  Data.Dense.Mutable
-- Copyright   :  (c) Christopher Chalmers
-- License     :  BSD3
--
-- Maintainer  :  Christopher Chalmers
-- Stability   :  provisional
-- Portability :  non-portable
--
-- This module provides generic functions over mutable multidimensional
-- arrays.
-----------------------------------------------------------------------------
module Data.Dense.Mutable
  (
    -- * Mutable array
    MArray (..)
  , UMArray
  , SMArray
  , BMArray
  , PMArray

  -- * Lenses
  , mlayout
  , mvector

    -- * Creation
  , new
  , replicate
  , replicateM
  , clone

    -- * Standard operations
    -- ** Indexing
  , read
  , linearRead
  , unsafeRead
  , unsafeLinearRead

    -- ** Writing
  , write
  , linearWrite
  , unsafeWrite
  , unsafeLinearWrite

    -- ** Modifying
  , modify
  , linearModify
  , unsafeModify
  , unsafeLinearModify

    -- ** Swap
  , swap
  , linearSwap
  , unsafeSwap
  , unsafeLinearSwap

    -- ** Exchange
  , exchange
  , linearExchange
  , unsafeExchange
  , unsafeLinearExchange

    -- * Misc
  , set
  , clear
  , copy

  ) where

import           Control.Monad                 (liftM)
import           Control.Monad.Primitive
import           Control.Lens                  (IndexedLens, indexed, Lens, (<&>))
import           Data.Foldable                 as F
import           Data.Typeable
import qualified Data.Vector                   as B
import           Data.Vector.Generic.Mutable   (MVector)
import qualified Data.Vector.Generic.Mutable   as GM
import qualified Data.Vector.Primitive.Mutable as P
import qualified Data.Vector.Storable.Mutable  as S
import qualified Data.Vector.Unboxed.Mutable   as U
import           Linear.V1

import           Data.Dense.Index

import           Prelude                       hiding (read, replicate)

-- | A mutable array with a shape.
data MArray v l s a = MArray !(Layout l) !(v s a)
  deriving Typeable

-- | Boxed mutable array.
type BMArray = MArray B.MVector

-- | Unboxed mutable array.
type UMArray = MArray U.MVector

-- | Storable mutable array.
type SMArray = MArray S.MVector

-- | Primitive mutable array.
type PMArray = MArray P.MVector

-- | Lens onto the shape of the vector. The total size of the layout
--   _must_ remain the same or an error is thrown.
mlayout :: (Shape f, Shape f') => Lens (MArray v f s a) (MArray v f' s a) (Layout f) (Layout f')
mlayout :: Lens (MArray v f s a) (MArray v f' s a) (Layout f) (Layout f')
mlayout Layout f -> f (Layout f')
f (MArray Layout f
l v s a
v) = Layout f -> f (Layout f')
f Layout f
l f (Layout f')
-> (Layout f' -> MArray v f' s a) -> f (MArray v f' s a)
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \Layout f'
l' ->
  Int -> Int -> String -> MArray v f' s a -> MArray v f' s a
forall a. Int -> Int -> String -> a -> a
sizeMissmatch (Layout f -> Int
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
F.product Layout f
l) (Layout f' -> Int
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
F.product Layout f'
l')
    (String
"mlayout: trying to replace shape " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Layout f -> String
forall (f :: * -> *). Shape f => f Int -> String
showShape Layout f
l String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
", with " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Layout f' -> String
forall (f :: * -> *). Shape f => f Int -> String
showShape Layout f'
l')
    (MArray v f' s a -> MArray v f' s a)
-> MArray v f' s a -> MArray v f' s a
forall a b. (a -> b) -> a -> b
$ Layout f' -> v s a -> MArray v f' s a
forall (v :: * -> * -> *) (l :: * -> *) s a.
Layout l -> v s a -> MArray v l s a
MArray Layout f'
l' v s a
v
{-# INLINE mlayout #-}

instance Shape f => HasLayout f (MArray v f s a) where
  layout :: (Layout f -> f (Layout f)) -> MArray v f s a -> f (MArray v f s a)
layout = (Layout f -> f (Layout f)) -> MArray v f s a -> f (MArray v f s a)
forall (f :: * -> *) (f' :: * -> *) (v :: * -> * -> *) s a.
(Shape f, Shape f') =>
Lens (MArray v f s a) (MArray v f' s a) (Layout f) (Layout f')
mlayout
  {-# INLINE layout #-}

-- | Indexed lens over the underlying vector of an array. The index is
--   the 'extent' of the array. You must __not__ change the length of
--   the vector, otherwise an error will be thrown.
mvector :: (MVector v a, MVector w b) => IndexedLens (Layout f) (MArray v f s a) (MArray w f t b) (v s a) (w t b)
mvector :: IndexedLens
  (Layout f) (MArray v f s a) (MArray w f t b) (v s a) (w t b)
mvector p (v s a) (f (w t b))
f (MArray Layout f
l v s a
v) =
  p (v s a) (f (w t b)) -> Layout f -> v s a -> f (w t b)
forall i (p :: * -> * -> *) a b.
Indexable i p =>
p a b -> i -> a -> b
indexed p (v s a) (f (w t b))
f Layout f
l v s a
v f (w t b) -> (w t b -> MArray w f t b) -> f (MArray w f t b)
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \w t b
w ->
  Int -> Int -> String -> MArray w f t b -> MArray w f t b
forall a. Int -> Int -> String -> a -> a
sizeMissmatch (v s a -> Int
forall (v :: * -> * -> *) a s. MVector v a => v s a -> Int
GM.length v s a
v) (w t b -> Int
forall (v :: * -> * -> *) a s. MVector v a => v s a -> Int
GM.length w t b
w)
     (String
"mvector: trying to replace vector of length " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show (v s a -> Int
forall (v :: * -> * -> *) a s. MVector v a => v s a -> Int
GM.length v s a
v) String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
", with one of length " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show (w t b -> Int
forall (v :: * -> * -> *) a s. MVector v a => v s a -> Int
GM.length w t b
w))
     (MArray w f t b -> MArray w f t b)
-> MArray w f t b -> MArray w f t b
forall a b. (a -> b) -> a -> b
$ Layout f -> w t b -> MArray w f t b
forall (v :: * -> * -> *) (l :: * -> *) s a.
Layout l -> v s a -> MArray v l s a
MArray Layout f
l w t b
w
{-# INLINE mvector #-}

-- | New mutable array with shape @l@.
new :: (PrimMonad m, Shape f, MVector v a) => Layout f -> m (MArray v f (PrimState m) a)
new :: Layout f -> m (MArray v f (PrimState m) a)
new Layout f
l = Layout f -> v (PrimState m) a -> MArray v f (PrimState m) a
forall (v :: * -> * -> *) (l :: * -> *) s a.
Layout l -> v s a -> MArray v l s a
MArray Layout f
l (v (PrimState m) a -> MArray v f (PrimState m) a)
-> m (v (PrimState m) a) -> m (MArray v f (PrimState m) a)
forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
`liftM` Int -> m (v (PrimState m) a)
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
Int -> m (v (PrimState m) a)
GM.new (Layout f -> Int
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
F.product Layout f
l)
{-# INLINE new #-}

-- | New mutable array with shape @l@ filled with element @a@.
replicate :: (PrimMonad m, Shape f, MVector v a) => Layout f -> a -> m (MArray v f (PrimState m) a)
replicate :: Layout f -> a -> m (MArray v f (PrimState m) a)
replicate Layout f
l a
a = Layout f -> v (PrimState m) a -> MArray v f (PrimState m) a
forall (v :: * -> * -> *) (l :: * -> *) s a.
Layout l -> v s a -> MArray v l s a
MArray Layout f
l (v (PrimState m) a -> MArray v f (PrimState m) a)
-> m (v (PrimState m) a) -> m (MArray v f (PrimState m) a)
forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
`liftM` Int -> a -> m (v (PrimState m) a)
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
Int -> a -> m (v (PrimState m) a)
GM.replicate (Layout f -> Int
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
F.product Layout f
l) a
a
{-# INLINE replicate #-}

-- | New mutable array with shape @l@ filled with result of monadic
--   action @a@.
replicateM :: (PrimMonad m, Shape f, MVector v a) => Layout f -> m a -> m (MArray v f (PrimState m) a)
replicateM :: Layout f -> m a -> m (MArray v f (PrimState m) a)
replicateM Layout f
l m a
a = Layout f -> v (PrimState m) a -> MArray v f (PrimState m) a
forall (v :: * -> * -> *) (l :: * -> *) s a.
Layout l -> v s a -> MArray v l s a
MArray Layout f
l (v (PrimState m) a -> MArray v f (PrimState m) a)
-> m (v (PrimState m) a) -> m (MArray v f (PrimState m) a)
forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
`liftM` Int -> m a -> m (v (PrimState m) a)
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
Int -> m a -> m (v (PrimState m) a)
GM.replicateM (Layout f -> Int
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
F.product Layout f
l) m a
a
{-# INLINE replicateM #-}

-- | Clone a mutable array, making a new, separate mutable array.
clone :: (PrimMonad m, MVector v a) => MArray v f (PrimState m) a -> m (MArray v f (PrimState m) a)
clone :: MArray v f (PrimState m) a -> m (MArray v f (PrimState m) a)
clone (MArray Layout f
l v (PrimState m) a
v) = Layout f -> v (PrimState m) a -> MArray v f (PrimState m) a
forall (v :: * -> * -> *) (l :: * -> *) s a.
Layout l -> v s a -> MArray v l s a
MArray Layout f
l (v (PrimState m) a -> MArray v f (PrimState m) a)
-> m (v (PrimState m) a) -> m (MArray v f (PrimState m) a)
forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
`liftM` v (PrimState m) a -> m (v (PrimState m) a)
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> m (v (PrimState m) a)
GM.clone v (PrimState m) a
v
{-# INLINE clone #-}

-- Individual elements -------------------------------------------------

-- | Clear the elements of a mutable array. This is usually a no-op for
--   unboxed arrays.
clear :: (PrimMonad m, MVector v a) => MArray v f (PrimState m) a -> m ()
clear :: MArray v f (PrimState m) a -> m ()
clear (MArray Layout f
_ v (PrimState m) a
v) = v (PrimState m) a -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> m ()
GM.clear v (PrimState m) a
v
{-# INLINE clear #-}

-- | Read a mutable array at element @l@.
read :: (PrimMonad m, Shape f, MVector v a) => MArray v f (PrimState m) a -> f Int -> m a
read :: MArray v f (PrimState m) a -> f Int -> m a
read (MArray f Int
l v (PrimState m) a
v) f Int
s = f Int -> f Int -> m a -> m a
forall (l :: * -> *) a. Shape l => Layout l -> Layout l -> a -> a
boundsCheck f Int
l f Int
s (m a -> m a) -> m a -> m a
forall a b. (a -> b) -> a -> b
$ v (PrimState m) a -> Int -> m a
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
GM.unsafeRead v (PrimState m) a
v (f Int -> f Int -> Int
forall (f :: * -> *). Shape f => Layout f -> Layout f -> Int
shapeToIndex f Int
l f Int
s)
{-# INLINE read #-}

-- | Write a mutable array at element @l@.
write :: (PrimMonad m, Shape f, MVector v a) => MArray v f (PrimState m) a -> f Int -> a -> m ()
write :: MArray v f (PrimState m) a -> f Int -> a -> m ()
write (MArray f Int
l v (PrimState m) a
v) f Int
s a
a = f Int -> f Int -> m () -> m ()
forall (l :: * -> *) a. Shape l => Layout l -> Layout l -> a -> a
boundsCheck f Int
l f Int
s (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ v (PrimState m) a -> Int -> a -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
GM.unsafeWrite v (PrimState m) a
v (f Int -> f Int -> Int
forall (f :: * -> *). Shape f => Layout f -> Layout f -> Int
shapeToIndex f Int
l f Int
s) a
a
{-# INLINE write #-}

-- | Modify a mutable array at element @l@ by applying a function.
modify :: (PrimMonad m, Shape f, MVector v a) => MArray v f (PrimState m) a -> f Int -> (a -> a) -> m ()
modify :: MArray v f (PrimState m) a -> f Int -> (a -> a) -> m ()
modify (MArray f Int
l v (PrimState m) a
v) f Int
s a -> a
f = f Int -> f Int -> m () -> m ()
forall (l :: * -> *) a. Shape l => Layout l -> Layout l -> a -> a
boundsCheck f Int
l f Int
s (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ v (PrimState m) a -> Int -> m a
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
GM.unsafeRead v (PrimState m) a
v Int
i m a -> (a -> m ()) -> m ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= v (PrimState m) a -> Int -> a -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
GM.unsafeWrite v (PrimState m) a
v Int
i (a -> m ()) -> (a -> a) -> a -> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> a
f
  where i :: Int
i = f Int -> f Int -> Int
forall (f :: * -> *). Shape f => Layout f -> Layout f -> Int
shapeToIndex f Int
l f Int
s
{-# INLINE modify #-}

-- | Swap two elements in a mutable array.
swap :: (PrimMonad m, Shape f, MVector v a) => MArray v f (PrimState m) a -> f Int -> f Int -> m ()
swap :: MArray v f (PrimState m) a -> f Int -> f Int -> m ()
swap (MArray f Int
l v (PrimState m) a
v) f Int
i f Int
j = f Int
-> f Int
-> (f Int -> f Int -> m () -> m ())
-> f Int
-> f Int
-> m ()
-> m ()
forall (l :: * -> *) a. Shape l => Layout l -> Layout l -> a -> a
boundsCheck f Int
l f Int
i f Int -> f Int -> m () -> m ()
forall (l :: * -> *) a. Shape l => Layout l -> Layout l -> a -> a
boundsCheck f Int
l f Int
j (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ v (PrimState m) a -> Int -> Int -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> Int -> m ()
GM.unsafeSwap v (PrimState m) a
v (f Int -> f Int -> Int
forall (f :: * -> *). Shape f => Layout f -> Layout f -> Int
shapeToIndex f Int
l f Int
i) (f Int -> f Int -> Int
forall (f :: * -> *). Shape f => Layout f -> Layout f -> Int
shapeToIndex f Int
l f Int
j)
{-# INLINE swap #-}

-- | Replace the element at the give position and return the old
--   element.
exchange :: (PrimMonad m, Shape f, MVector v a) => MArray v f (PrimState m) a -> f Int -> a -> m a
exchange :: MArray v f (PrimState m) a -> f Int -> a -> m a
exchange (MArray f Int
l v (PrimState m) a
v) f Int
i a
a = f Int -> f Int -> m a -> m a
forall (l :: * -> *) a. Shape l => Layout l -> Layout l -> a -> a
boundsCheck f Int
l f Int
i (m a -> m a) -> m a -> m a
forall a b. (a -> b) -> a -> b
$ v (PrimState m) a -> Int -> a -> m a
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m a
GM.unsafeExchange v (PrimState m) a
v (f Int -> f Int -> Int
forall (f :: * -> *). Shape f => Layout f -> Layout f -> Int
shapeToIndex f Int
l f Int
i) a
a
{-# INLINE exchange #-}

-- | Read a mutable array at element @i@ by indexing the internal
--   vector.
linearRead :: (PrimMonad m, MVector v a) => MArray v f (PrimState m) a -> Int -> m a
linearRead :: MArray v f (PrimState m) a -> Int -> m a
linearRead (MArray Layout f
_ v (PrimState m) a
v) = v (PrimState m) a -> Int -> m a
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
GM.read v (PrimState m) a
v
{-# INLINE linearRead #-}

-- | Write a mutable array at element @i@ by indexing the internal
--   vector.
linearWrite :: (PrimMonad m, MVector v a) => MArray v f (PrimState m) a -> Int -> a -> m ()
linearWrite :: MArray v f (PrimState m) a -> Int -> a -> m ()
linearWrite (MArray Layout f
_ v (PrimState m) a
v) = v (PrimState m) a -> Int -> a -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
GM.write v (PrimState m) a
v
{-# INLINE linearWrite #-}

-- | Swap two elements in a mutable array by indexing the internal
--   vector.
linearSwap :: (PrimMonad m, MVector v a) => MArray v f (PrimState m) a -> Int -> Int -> m ()
linearSwap :: MArray v f (PrimState m) a -> Int -> Int -> m ()
linearSwap (MArray Layout f
_ v (PrimState m) a
v) = v (PrimState m) a -> Int -> Int -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> Int -> m ()
GM.swap v (PrimState m) a
v
{-# INLINE linearSwap #-}

-- | Modify a mutable array at element @i@ by applying a function.
linearModify :: (PrimMonad m, MVector v a) => MArray v f (PrimState m) a -> Int -> (a -> a) -> m ()
linearModify :: MArray v f (PrimState m) a -> Int -> (a -> a) -> m ()
linearModify (MArray Layout f
_ v (PrimState m) a
v) Int
i a -> a
f = v (PrimState m) a -> Int -> m a
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
GM.read v (PrimState m) a
v Int
i m a -> (a -> m ()) -> m ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= v (PrimState m) a -> Int -> a -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
GM.unsafeWrite v (PrimState m) a
v Int
i (a -> m ()) -> (a -> a) -> a -> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> a
f
{-# INLINE linearModify #-}

-- | Replace the element at the give position and return the old
--   element.
linearExchange :: (PrimMonad m, MVector v a) => MArray v f (PrimState m) a -> Int -> a -> m a
linearExchange :: MArray v f (PrimState m) a -> Int -> a -> m a
linearExchange (MArray Layout f
_ v (PrimState m) a
v) Int
i a
a = v (PrimState m) a -> Int -> a -> m a
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m a
GM.exchange v (PrimState m) a
v Int
i a
a
{-# INLINE linearExchange #-}

-- Unsafe varients

-- | 'read' without bounds checking.
unsafeRead :: (PrimMonad m, Shape f, MVector v a) => MArray v f (PrimState m) a -> f Int -> m a
unsafeRead :: MArray v f (PrimState m) a -> f Int -> m a
unsafeRead (MArray f Int
l v (PrimState m) a
v) f Int
s = v (PrimState m) a -> Int -> m a
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
GM.unsafeRead v (PrimState m) a
v (f Int -> f Int -> Int
forall (f :: * -> *). Shape f => Layout f -> Layout f -> Int
shapeToIndex f Int
l f Int
s)
{-# INLINE unsafeRead #-}

-- | 'write' without bounds checking.
unsafeWrite :: (PrimMonad m, Shape f, MVector v a) => MArray v f (PrimState m) a -> f Int -> a -> m ()
unsafeWrite :: MArray v f (PrimState m) a -> f Int -> a -> m ()
unsafeWrite (MArray f Int
l v (PrimState m) a
v) f Int
s = v (PrimState m) a -> Int -> a -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
GM.unsafeWrite v (PrimState m) a
v (f Int -> f Int -> Int
forall (f :: * -> *). Shape f => Layout f -> Layout f -> Int
shapeToIndex f Int
l f Int
s)
{-# INLINE unsafeWrite #-}

-- | 'swap' without bounds checking.
unsafeSwap :: (PrimMonad m, Shape f, MVector v a) => MArray v f (PrimState m) a -> f Int -> f Int -> m ()
unsafeSwap :: MArray v f (PrimState m) a -> f Int -> f Int -> m ()
unsafeSwap (MArray f Int
l v (PrimState m) a
v) f Int
s f Int
j = v (PrimState m) a -> Int -> Int -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> Int -> m ()
GM.unsafeSwap v (PrimState m) a
v (f Int -> f Int -> Int
forall (f :: * -> *). Shape f => Layout f -> Layout f -> Int
shapeToIndex f Int
l f Int
s) (f Int -> f Int -> Int
forall (f :: * -> *). Shape f => Layout f -> Layout f -> Int
shapeToIndex f Int
j f Int
s)
{-# INLINE unsafeSwap #-}

-- | 'modify' without bounds checking.
unsafeModify :: (PrimMonad m, Shape f, MVector v a) => MArray v f (PrimState m) a -> f Int -> (a -> a) -> m ()
unsafeModify :: MArray v f (PrimState m) a -> f Int -> (a -> a) -> m ()
unsafeModify (MArray f Int
l v (PrimState m) a
v) f Int
s a -> a
f = v (PrimState m) a -> Int -> m a
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
GM.unsafeRead v (PrimState m) a
v Int
i m a -> (a -> m ()) -> m ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= v (PrimState m) a -> Int -> a -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
GM.unsafeWrite v (PrimState m) a
v Int
i (a -> m ()) -> (a -> a) -> a -> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> a
f
  where i :: Int
i = f Int -> f Int -> Int
forall (f :: * -> *). Shape f => Layout f -> Layout f -> Int
shapeToIndex f Int
l f Int
s
{-# INLINE unsafeModify #-}

-- | Replace the element at the give position and return the old
--   element.
unsafeExchange :: (PrimMonad m, Shape f, MVector v a) => MArray v f (PrimState m) a -> f Int -> a -> m a
unsafeExchange :: MArray v f (PrimState m) a -> f Int -> a -> m a
unsafeExchange (MArray f Int
l v (PrimState m) a
v) f Int
i a
a = v (PrimState m) a -> Int -> a -> m a
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m a
GM.unsafeExchange v (PrimState m) a
v (f Int -> f Int -> Int
forall (f :: * -> *). Shape f => Layout f -> Layout f -> Int
shapeToIndex f Int
l f Int
i) a
a
{-# INLINE unsafeExchange #-}

-- | 'linearRead' without bounds checking.
unsafeLinearRead :: (PrimMonad m, MVector v a) => MArray v f (PrimState m) a -> Int -> m a
unsafeLinearRead :: MArray v f (PrimState m) a -> Int -> m a
unsafeLinearRead (MArray Layout f
_ v (PrimState m) a
v) = v (PrimState m) a -> Int -> m a
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
GM.unsafeRead v (PrimState m) a
v
{-# INLINE unsafeLinearRead #-}

-- | 'linearWrite' without bounds checking.
unsafeLinearWrite :: (PrimMonad m, MVector v a) => MArray v f (PrimState m) a -> Int -> a -> m ()
unsafeLinearWrite :: MArray v f (PrimState m) a -> Int -> a -> m ()
unsafeLinearWrite (MArray Layout f
_ v (PrimState m) a
v) = v (PrimState m) a -> Int -> a -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
GM.unsafeWrite v (PrimState m) a
v
{-# INLINE unsafeLinearWrite #-}

-- | 'linearSwap' without bounds checking.
unsafeLinearSwap :: (PrimMonad m, MVector v a) => MArray v f (PrimState m) a -> Int -> Int -> m ()
unsafeLinearSwap :: MArray v f (PrimState m) a -> Int -> Int -> m ()
unsafeLinearSwap (MArray Layout f
_ v (PrimState m) a
v) = v (PrimState m) a -> Int -> Int -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> Int -> m ()
GM.unsafeSwap v (PrimState m) a
v
{-# INLINE unsafeLinearSwap #-}

-- | 'linearModify' without bounds checking.
unsafeLinearModify :: (PrimMonad m, MVector v a) => MArray v f (PrimState m) a -> Int -> (a -> a) -> m ()
unsafeLinearModify :: MArray v f (PrimState m) a -> Int -> (a -> a) -> m ()
unsafeLinearModify (MArray Layout f
_ v (PrimState m) a
v) Int
i a -> a
f = v (PrimState m) a -> Int -> m a
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
GM.unsafeRead v (PrimState m) a
v Int
i m a -> (a -> m ()) -> m ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= v (PrimState m) a -> Int -> a -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
GM.unsafeWrite v (PrimState m) a
v Int
i (a -> m ()) -> (a -> a) -> a -> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> a
f
{-# INLINE unsafeLinearModify #-}

-- | Replace the element at the give position and return the old
--   element.
unsafeLinearExchange :: (PrimMonad m, MVector v a) => MArray v f (PrimState m) a -> Int -> a -> m a
unsafeLinearExchange :: MArray v f (PrimState m) a -> Int -> a -> m a
unsafeLinearExchange (MArray Layout f
_ v (PrimState m) a
v) Int
i a
a = v (PrimState m) a -> Int -> a -> m a
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m a
GM.unsafeExchange v (PrimState m) a
v Int
i a
a
{-# INLINE unsafeLinearExchange #-}

-- Filling and copying -------------------------------------------------

-- | Set all elements in a mutable array to a constant value.
set :: (PrimMonad m, MVector v a) => MArray v f (PrimState m) a -> a -> m ()
set :: MArray v f (PrimState m) a -> a -> m ()
set (MArray Layout f
_ v (PrimState m) a
v) = v (PrimState m) a -> a -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> a -> m ()
GM.set v (PrimState m) a
v
{-# INLINE set #-}

-- | Copy all elements from one array into another.
copy :: (PrimMonad m, MVector v a) => MArray v f (PrimState m) a -> MArray v f (PrimState m) a -> m ()
copy :: MArray v f (PrimState m) a -> MArray v f (PrimState m) a -> m ()
copy (MArray Layout f
_ v (PrimState m) a
v) (MArray Layout f
_ v (PrimState m) a
u) = v (PrimState m) a -> v (PrimState m) a -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> v (PrimState m) a -> m ()
GM.copy v (PrimState m) a
v v (PrimState m) a
u
{-# INLINE copy #-}

-- V1 instances --------------------------------------------------------

-- Array v V1 a is essentially v a with a wrapper. Instance is provided
-- for convience.

instance (MVector v a, f ~ V1) => MVector (MArray v f) a where
  {-# INLINE basicLength #-}
  {-# INLINE basicUnsafeSlice #-}
  {-# INLINE basicOverlaps #-}
  {-# INLINE basicUnsafeNew #-}
  {-# INLINE basicUnsafeRead #-}
  {-# INLINE basicUnsafeWrite #-}
  {-# INLINE basicInitialize #-}
  basicLength :: MArray v f s a -> Int
basicLength (MArray (V1 n) v s a
_) = Int
n
  basicUnsafeSlice :: Int -> Int -> MArray v f s a -> MArray v f s a
basicUnsafeSlice Int
i Int
n (MArray Layout f
_ v s a
v) = Layout V1 -> v s a -> MArray v V1 s a
forall (v :: * -> * -> *) (l :: * -> *) s a.
Layout l -> v s a -> MArray v l s a
MArray (Int -> Layout V1
forall a. a -> V1 a
V1 Int
n) (v s a -> MArray v V1 s a) -> v s a -> MArray v V1 s a
forall a b. (a -> b) -> a -> b
$ Int -> Int -> v s a -> v s a
forall (v :: * -> * -> *) a s.
MVector v a =>
Int -> Int -> v s a -> v s a
GM.basicUnsafeSlice Int
i Int
n v s a
v
  basicOverlaps :: MArray v f s a -> MArray v f s a -> Bool
basicOverlaps (MArray Layout f
_ v s a
v) (MArray Layout f
_ v s a
w) = v s a -> v s a -> Bool
forall (v :: * -> * -> *) a s.
MVector v a =>
v s a -> v s a -> Bool
GM.basicOverlaps v s a
v v s a
w
  basicUnsafeNew :: Int -> m (MArray v f (PrimState m) a)
basicUnsafeNew Int
n = Layout V1 -> v (PrimState m) a -> MArray v V1 (PrimState m) a
forall (v :: * -> * -> *) (l :: * -> *) s a.
Layout l -> v s a -> MArray v l s a
MArray (Int -> Layout V1
forall a. a -> V1 a
V1 Int
n) (v (PrimState m) a -> MArray v V1 (PrimState m) a)
-> m (v (PrimState m) a) -> m (MArray v V1 (PrimState m) a)
forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
`liftM` Int -> m (v (PrimState m) a)
forall (v :: * -> * -> *) a (m :: * -> *).
(MVector v a, PrimMonad m) =>
Int -> m (v (PrimState m) a)
GM.basicUnsafeNew Int
n
  basicUnsafeRead :: MArray v f (PrimState m) a -> Int -> m a
basicUnsafeRead (MArray Layout f
_ v (PrimState m) a
v) = v (PrimState m) a -> Int -> m a
forall (v :: * -> * -> *) a (m :: * -> *).
(MVector v a, PrimMonad m) =>
v (PrimState m) a -> Int -> m a
GM.basicUnsafeRead v (PrimState m) a
v
  basicUnsafeWrite :: MArray v f (PrimState m) a -> Int -> a -> m ()
basicUnsafeWrite (MArray Layout f
_ v (PrimState m) a
v) = v (PrimState m) a -> Int -> a -> m ()
forall (v :: * -> * -> *) a (m :: * -> *).
(MVector v a, PrimMonad m) =>
v (PrimState m) a -> Int -> a -> m ()
GM.basicUnsafeWrite v (PrimState m) a
v
  basicInitialize :: MArray v f (PrimState m) a -> m ()
basicInitialize (MArray Layout f
_ v (PrimState m) a
v) = v (PrimState m) a -> m ()
forall (v :: * -> * -> *) a (m :: * -> *).
(MVector v a, PrimMonad m) =>
v (PrimState m) a -> m ()
GM.basicInitialize v (PrimState m) a
v