{-# LANGUAGE BangPatterns          #-}
{-# LANGUAGE CPP                   #-}
{-# LANGUAGE DeriveDataTypeable    #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE LambdaCase            #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE MultiWayIf            #-}
{-# LANGUAGE RankNTypes            #-}
{-# LANGUAGE TypeFamilies          #-}
-----------------------------------------------------------------------------
-- |
-- Module      :  Data.Dense.Generic
-- Copyright   :  (c) Christopher Chalmers
-- License     :  BSD3
--
-- Maintainer  :  Christopher Chalmers
-- Stability   :  provisional
-- Portability :  non-portable
--
-- This module provides generic functions over multidimensional arrays.
-----------------------------------------------------------------------------
module Data.Dense.Generic
  (
    -- * Array types
    Array
  , Shape (..)
  , BArray
  , UArray
  , SArray
  , PArray

    -- * Layout of an array
  , HasLayout (..)
  , Layout

    -- ** Extracting size
  , extent
  , size

    -- ** Folds over indexes
  , indexes
  , indexesFrom
  , indexesBetween

    -- * Underlying vector
  , vector

    -- ** Traversals
  , values
  , values'
  , valuesBetween

  -- * Construction

  -- ** Flat arrays
  , flat
  , fromList

  -- ** From lists
  , fromListInto
  , fromListInto_

  -- ** From vectors
  , fromVectorInto
  , fromVectorInto_

  -- ** Initialisation
  , replicate
  , generate
  , linearGenerate

  -- ** Monadic initialisation
  , create
  , createT
  , replicateM
  , generateM
  , linearGenerateM

  -- * Functions on arrays

  -- ** Empty arrays
  , empty
  , null

  -- ** Indexing

  , (!)
  , (!?)
  , unsafeIndex
  , linearIndex
  , unsafeLinearIndex

  -- *** Monadic indexing
  , indexM
  , unsafeIndexM
  , linearIndexM
  , unsafeLinearIndexM

  -- ** Modifying arrays

  -- ** Bulk updates
  , (//)

  -- ** Accumulations
  , accum

  -- ** Mapping
  , map
  , imap

  -- * Zipping
  -- ** Tuples
  , Data.Dense.Generic.zip
  , Data.Dense.Generic.zip3

  -- ** Zip with function
  , zipWith
  , zipWith3
  , izipWith
  , izipWith3

  -- ** Slices

  -- *** Matrix
  , ixRow
  , rows
  , ixColumn
  , columns

  -- *** 3D
  , ixPlane
  , planes
  , flattenPlane

  -- *** Ordinals
  , unsafeOrdinals

  -- * Mutable
  , MArray
  , M.BMArray
  , M.UMArray
  , M.SMArray
  , M.PMArray

  , thaw
  , freeze
  , unsafeThaw
  , unsafeFreeze

  -- * Delayed

  , Delayed

  -- ** Generating delayed

  , delayed
  , seqDelayed
  , delay
  , manifest
  , seqManifest
  , genDelayed
  , indexDelayed
  , affirm
  , seqAffirm

  -- * Focused

  , Focused

  -- ** Generating focused

  , focusOn
  , unfocus
  , unfocused
  , extendFocus

  -- ** Focus location
  , locale
  , shiftFocus

  -- ** Boundary
  , Boundary (..)
  , peekB
  , peeksB
  , peekRelativeB

  -- * Fusion
  -- ** Streams
  , streamGenerate
  , streamGenerateM
  , streamIndexes

  -- ** Bundles
  , bundleGenerate
  , bundleGenerateM
  , bundleIndexes

  ) where


#if __GLASGOW_HASKELL__ <= 708
import           Control.Applicative               (Applicative, pure, (<*>))
import           Data.Foldable                     (Foldable)
#endif

import           Control.Comonad
import           Control.Comonad.Store
import           Control.Lens                      hiding (imap)
import           Control.Monad                     (liftM)
import           Control.Monad.Primitive
import           Control.Monad.ST
import qualified Data.Foldable                     as F
import           Data.Functor.Classes
import qualified Data.List                         as L
import           Data.Maybe                        (fromMaybe)
import qualified Data.Traversable                  as T
import           Data.Typeable
import qualified Data.Vector                       as B
import           Data.Vector.Fusion.Bundle         (MBundle)
import qualified Data.Vector.Fusion.Bundle         as Bundle
import qualified Data.Vector.Fusion.Bundle.Monadic as MBundle
import           Data.Vector.Fusion.Bundle.Size
import           Data.Vector.Fusion.Stream.Monadic (Step (..), Stream (..))
import qualified Data.Vector.Fusion.Stream.Monadic as Stream
import           Data.Vector.Generic               (Vector)
import qualified Data.Vector.Generic               as G
import qualified Data.Vector.Generic.Mutable       as GM
import qualified Data.Vector.Primitive             as P
import qualified Data.Vector.Storable              as S
import qualified Data.Vector.Unboxed               as U
import           Linear                            hiding (vector)

import           Data.Dense.Base
import           Data.Dense.Index
import           Data.Dense.Mutable               (MArray (..))
import qualified Data.Dense.Mutable               as M

import           Prelude                           hiding (map, null, replicate,
                                                    zipWith, zipWith3)

-- Aliases -------------------------------------------------------------

-- | 'Boxed' array.
type BArray = Array B.Vector

-- | 'Data.Vector.Unboxed.Unbox'ed array.
type UArray = Array U.Vector

-- | 'Foreign.Storable.Storeable' array.
type SArray = Array S.Vector

-- | 'Data.Primitive.Types.Prim' array.
type PArray = Array P.Vector

-- Lenses --------------------------------------------------------------

-- | Same as 'values' but restrictive in the vector type.
values' :: (Shape f, Vector v a, Vector v b)
       => IndexedTraversal (f Int) (Array v f a) (Array v f b) a b
values' :: IndexedTraversal (f Int) (Array v f a) (Array v f b) a b
values' = p a (f b) -> Array v f a -> f (Array v f b)
forall (f :: * -> *) (v :: * -> *) a (w :: * -> *) b.
(Shape f, Vector v a, Vector w b) =>
IndexedTraversal (f Int) (Array v f a) (Array w f b) a b
values
{-# INLINE values' #-}

-- | Traverse over the 'values' between two indexes.
valuesBetween :: (Shape f, Vector v a) => f Int -> f Int -> IndexedTraversal' (f Int) (Array v f a) a
valuesBetween :: f Int -> f Int -> IndexedTraversal' (f Int) (Array v f a) a
valuesBetween f Int
a f Int
b = [f Int] -> IndexedTraversal' (f Int) (Array v f a) a
forall (v :: * -> *) a (f :: * -> *).
(Vector v a, Shape f) =>
[f Int] -> IndexedTraversal' (f Int) (Array v f a) a
unsafeOrdinals (Getting (Endo [f Int]) (f Int) (f Int) -> f Int -> [f Int]
forall a s. Getting (Endo [a]) s a -> s -> [a]
toListOf (f Int -> IndexedFold Int (f Int) (f Int)
forall (f :: * -> *).
Shape f =>
f Int -> IndexedFold Int (f Int) (f Int)
shapeIndexesFrom f Int
a) f Int
b)
{-# INLINE valuesBetween #-}

-- | 1D arrays are just vectors. You are free to change the length of
--   the vector when going 'over' this 'Iso' (unlike 'linear').
--
--   Note that 'V1' arrays are an instance of 'Vector' so you can use
--   any of the functions in "Data.Vector.Generic" on them without
--   needing to convert.
flat :: Vector w b => Iso (Array v V1 a) (Array w V1 b) (v a) (w b)
flat :: Iso (Array v V1 a) (Array w V1 b) (v a) (w b)
flat = (Array v V1 a -> v a)
-> (w b -> Array w V1 b)
-> Iso (Array v V1 a) (Array w V1 b) (v a) (w b)
forall s a b t. (s -> a) -> (b -> t) -> Iso s t a b
iso (\(Array Layout V1
_ v a
v) -> v a
v) (\w b
v -> Layout V1 -> w b -> Array w V1 b
forall (v :: * -> *) (f :: * -> *) a.
Layout f -> v a -> Array v f a
Array (Int -> Layout V1
forall a. a -> V1 a
V1 (Int -> Layout V1) -> Int -> Layout V1
forall a b. (a -> b) -> a -> b
$ w b -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
G.length w b
v) w b
v)
{-# INLINE flat #-}

-- Constructing vectors ------------------------------------------------

-- | Contruct a flat array from a list. (This is just 'G.fromList' from
--   'Data.Vector.Generic'.)
fromList :: Vector v a => [a] -> Array v V1 a
fromList :: [a] -> Array v V1 a
fromList = [a] -> Array v V1 a
forall (v :: * -> *) a. Vector v a => [a] -> v a
G.fromList
{-# INLINE fromList #-}

-- | O(n) Convert the first @n@ elements of a list to an Array with the
--   given shape. Returns 'Nothing' if there are not enough elements in
--   the list.
fromListInto :: (Shape f, Vector v a) => Layout f -> [a] -> Maybe (Array v f a)
fromListInto :: Layout f -> [a] -> Maybe (Array v f a)
fromListInto Layout f
l [a]
as
  | v a -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
G.length v a
v Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
n = Array v f a -> Maybe (Array v f a)
forall a. a -> Maybe a
Just (Array v f a -> Maybe (Array v f a))
-> Array v f a -> Maybe (Array v f a)
forall a b. (a -> b) -> a -> b
$ Layout f -> v a -> Array v f a
forall (v :: * -> *) (f :: * -> *) a.
Layout f -> v a -> Array v f a
Array Layout f
l v a
v
  | Bool
otherwise       = Maybe (Array v f a)
forall a. Maybe a
Nothing
  where v :: v a
v = Int -> [a] -> v a
forall (v :: * -> *) a. Vector v a => Int -> [a] -> v a
G.fromListN Int
n [a]
as
        n :: Int
n = Layout f -> Int
forall (f :: * -> *). Shape f => Layout f -> Int
shapeSize Layout f
l
{-# INLINE fromListInto #-}

-- | O(n) Convert the first @n@ elements of a list to an Array with the
--   given shape. Throw an error if the list is not long enough.
fromListInto_ :: (Shape f, Vector v a) => Layout f -> [a] -> Array v f a
fromListInto_ :: Layout f -> [a] -> Array v f a
fromListInto_ Layout f
l [a]
as = Array v f a -> Maybe (Array v f a) -> Array v f a
forall a. a -> Maybe a -> a
fromMaybe Array v f a
err (Maybe (Array v f a) -> Array v f a)
-> Maybe (Array v f a) -> Array v f a
forall a b. (a -> b) -> a -> b
$ Layout f -> [a] -> Maybe (Array v f a)
forall (f :: * -> *) (v :: * -> *) a.
(Shape f, Vector v a) =>
Layout f -> [a] -> Maybe (Array v f a)
fromListInto Layout f
l [a]
as
  where
    err :: Array v f a
err = [Char] -> Array v f a
forall a. HasCallStack => [Char] -> a
error ([Char] -> Array v f a) -> [Char] -> Array v f a
forall a b. (a -> b) -> a -> b
$ [Char]
"fromListInto_: shape " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Layout f -> [Char]
forall (f :: * -> *). Shape f => f Int -> [Char]
showShape Layout f
l [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
" is too large for list"
{-# INLINE fromListInto_ #-}

-- | Create an array from a 'vector' and a 'layout'. Return 'Nothing' if
--   the vector is not the right shape.
fromVectorInto :: (Shape f, Vector v a) => Layout f -> v a -> Maybe (Array v f a)
fromVectorInto :: Layout f -> v a -> Maybe (Array v f a)
fromVectorInto Layout f
l v a
v
  | Layout f -> Int
forall (f :: * -> *). Shape f => Layout f -> Int
shapeSize Layout f
l Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== v a -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
G.length v a
v = Array v f a -> Maybe (Array v f a)
forall a. a -> Maybe a
Just (Array v f a -> Maybe (Array v f a))
-> Array v f a -> Maybe (Array v f a)
forall a b. (a -> b) -> a -> b
$! Layout f -> v a -> Array v f a
forall (v :: * -> *) (f :: * -> *) a.
Layout f -> v a -> Array v f a
Array Layout f
l v a
v
  | Bool
otherwise                 = Maybe (Array v f a)
forall a. Maybe a
Nothing
{-# INLINE fromVectorInto #-}

-- | Create an array from a 'vector' and a 'layout'. Throws an error if
--   the vector is not the right shape.
fromVectorInto_ :: (Shape f, Vector v a) => Layout f -> v a -> Array v f a
fromVectorInto_ :: Layout f -> v a -> Array v f a
fromVectorInto_ Layout f
l v a
as = Array v f a -> Maybe (Array v f a) -> Array v f a
forall a. a -> Maybe a -> a
fromMaybe Array v f a
err (Maybe (Array v f a) -> Array v f a)
-> Maybe (Array v f a) -> Array v f a
forall a b. (a -> b) -> a -> b
$ Layout f -> v a -> Maybe (Array v f a)
forall (f :: * -> *) (v :: * -> *) a.
(Shape f, Vector v a) =>
Layout f -> v a -> Maybe (Array v f a)
fromVectorInto Layout f
l v a
as
  where
    err :: Array v f a
err = [Char] -> Array v f a
forall a. HasCallStack => [Char] -> a
error ([Char] -> Array v f a) -> [Char] -> Array v f a
forall a b. (a -> b) -> a -> b
$ [Char]
"fromVectorInto_: shape " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Layout f -> [Char]
forall (f :: * -> *). Shape f => f Int -> [Char]
showShape Layout f
l [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
" is too large for the vector"
{-# INLINE fromVectorInto_ #-}

-- | The empty 'Array' with a 'zero' shape.
empty :: (Vector v a, Additive f) => Array v f a
empty :: Array v f a
empty = Layout f -> v a -> Array v f a
forall (v :: * -> *) (f :: * -> *) a.
Layout f -> v a -> Array v f a
Array Layout f
forall (f :: * -> *) a. (Additive f, Num a) => f a
zero v a
forall (v :: * -> *) a. Vector v a => v a
G.empty
{-# INLINE empty #-}

-- | Test is if the array is 'empty'.
null :: Foldable f => Array v f a -> Bool
null :: Array v f a -> Bool
null (Array Layout f
l v a
_) = (Int -> Bool) -> Layout f -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
F.all (Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
==Int
0) Layout f
l
{-# INLINE null #-}

-- Indexing ------------------------------------------------------------

-- | Index an element of an array. Throws 'IndexOutOfBounds' if the
--   index is out of bounds.
(!) :: (Shape f, Vector v a) => Array v f a -> f Int -> a
(!) (Array f Int
l v a
v) f Int
i = f Int -> f Int -> a -> a
forall (l :: * -> *) a. Shape l => Layout l -> Layout l -> a -> a
boundsCheck f Int
l f Int
i (a -> a) -> a -> a
forall a b. (a -> b) -> a -> b
$ v a -> Int -> a
forall (v :: * -> *) a. Vector v a => v a -> Int -> a
G.unsafeIndex v a
v (f Int -> f Int -> Int
forall (f :: * -> *). Shape f => Layout f -> Layout f -> Int
shapeToIndex f Int
l f Int
i)
{-# INLINE (!) #-}

-- | Safe index of an element.
(!?) :: (Shape f, Vector v a) => Array v f a -> f Int -> Maybe a
Array f Int
l v a
v !? :: Array v f a -> f Int -> Maybe a
!? f Int
i
  | f Int -> f Int -> Bool
forall (f :: * -> *). Shape f => Layout f -> Layout f -> Bool
shapeInRange f Int
l f Int
i = a -> Maybe a
forall a. a -> Maybe a
Just (a -> Maybe a) -> a -> Maybe a
forall a b. (a -> b) -> a -> b
$! v a -> Int -> a
forall (v :: * -> *) a. Vector v a => v a -> Int -> a
G.unsafeIndex v a
v (f Int -> f Int -> Int
forall (f :: * -> *). Shape f => Layout f -> Layout f -> Int
shapeToIndex f Int
l f Int
i)
  | Bool
otherwise   = Maybe a
forall a. Maybe a
Nothing
{-# INLINE (!?) #-}

-- | Index an element of an array without bounds checking.
unsafeIndex :: (Shape f, Vector v a) => Array v f a -> f Int -> a
unsafeIndex :: Array v f a -> f Int -> a
unsafeIndex (Array f Int
l v a
v) f Int
i = v a -> Int -> a
forall (v :: * -> *) a. Vector v a => v a -> Int -> a
G.unsafeIndex v a
v (f Int -> f Int -> Int
forall (f :: * -> *). Shape f => Layout f -> Layout f -> Int
shapeToIndex f Int
l f Int
i)
{-# INLINE unsafeIndex #-}

-- | Index an element of an array while ignoring its shape.
linearIndex :: Vector v a => Array v f a -> Int -> a
linearIndex :: Array v f a -> Int -> a
linearIndex (Array Layout f
_ v a
v) Int
i = v a
v v a -> Int -> a
forall (v :: * -> *) a. Vector v a => v a -> Int -> a
G.! Int
i
{-# INLINE linearIndex #-}

-- | Index an element of an array while ignoring its shape, without
--   bounds checking.
unsafeLinearIndex :: Vector v a => Array v f a -> Int -> a
unsafeLinearIndex :: Array v f a -> Int -> a
unsafeLinearIndex (Array Layout f
_ v a
v) Int
i = v a -> Int -> a
forall (v :: * -> *) a. Vector v a => v a -> Int -> a
G.unsafeIndex v a
v Int
i
{-# INLINE unsafeLinearIndex #-}

-- Monadic indexing ----------------------------------------------------

-- | /O(1)/ Indexing in a monad.
--
--   The monad allows operations to be strict in the vector when necessary.
--   Suppose vector copying is implemented like this:
--
-- > copy mv v = ... write mv i (v ! i) ...
--
--   For lazy vectors, @v ! i@ would not be evaluated which means that
--   @mv@ would unnecessarily retain a reference to @v@ in each element
--   written.
--
--   With 'indexM', copying can be implemented like this instead:
--
-- > copy mv v = ... do
-- >   x <- indexM v i
-- >   write mv i x
--
--   Here, no references to @v@ are retained because indexing (but /not/
--   the elements) is evaluated eagerly.
--
--   Throws an error if the index is out of range.
indexM :: (Shape f, Vector v a, Monad m) => Array v f a -> f Int -> m a
indexM :: Array v f a -> f Int -> m a
indexM (Array f Int
l v a
v) f Int
i = 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 a -> Int -> m a
forall (v :: * -> *) a (m :: * -> *).
(Vector v a, Monad m) =>
v a -> Int -> m a
G.unsafeIndexM v a
v (f Int -> f Int -> Int
forall (f :: * -> *). Shape f => Layout f -> Layout f -> Int
shapeToIndex f Int
l f Int
i)
{-# INLINE indexM #-}

-- | /O(1)/ Indexing in a monad without bounds checks. See 'indexM' for an
--   explanation of why this is useful.
unsafeIndexM :: (Shape f, Vector v a, Monad m) => Array v f a -> f Int -> m a
unsafeIndexM :: Array v f a -> f Int -> m a
unsafeIndexM (Array f Int
l v a
v) f Int
i = v a -> Int -> m a
forall (v :: * -> *) a (m :: * -> *).
(Vector v a, Monad m) =>
v a -> Int -> m a
G.unsafeIndexM v a
v (f Int -> f Int -> Int
forall (f :: * -> *). Shape f => Layout f -> Layout f -> Int
shapeToIndex f Int
l f Int
i)
{-# INLINE unsafeIndexM #-}

-- | /O(1)/ Indexing in a monad. Throws an error if the index is out of
--   range.
linearIndexM :: (Shape f, Vector v a, Monad m) => Array v f a -> Int -> m a
linearIndexM :: Array v f a -> Int -> m a
linearIndexM (Array Layout f
l v a
v) Int
i = Layout f -> Layout f -> m a -> m a
forall (l :: * -> *) a. Shape l => Layout l -> Layout l -> a -> a
boundsCheck Layout f
l (Layout f -> Int -> Layout f
forall (f :: * -> *). Shape f => Layout f -> Int -> Layout f
shapeFromIndex Layout f
l Int
i) (m a -> m a) -> m a -> m a
forall a b. (a -> b) -> a -> b
$ v a -> Int -> m a
forall (v :: * -> *) a (m :: * -> *).
(Vector v a, Monad m) =>
v a -> Int -> m a
G.unsafeIndexM v a
v Int
i
{-# INLINE linearIndexM #-}

-- | /O(1)/ Indexing in a monad without bounds checks. See 'indexM' for an
--   explanation of why this is useful.
unsafeLinearIndexM :: (Vector v a, Monad m) => Array v f a -> Int -> m a
unsafeLinearIndexM :: Array v f a -> Int -> m a
unsafeLinearIndexM (Array Layout f
_ v a
v) = v a -> Int -> m a
forall (v :: * -> *) a (m :: * -> *).
(Vector v a, Monad m) =>
v a -> Int -> m a
G.unsafeIndexM v a
v
{-# INLINE unsafeLinearIndexM #-}

-- Initialisation ------------------------------------------------------

-- | Execute the monadic action and freeze the resulting array.
create :: Vector v a => (forall s. ST s (MArray (G.Mutable v) f s a)) -> Array v f a
create :: (forall s. ST s (MArray (Mutable v) f s a)) -> Array v f a
create forall s. ST s (MArray (Mutable v) f s a)
m = ST Any (MArray (Mutable v) f Any a)
forall s. ST s (MArray (Mutable v) f s a)
m ST Any (MArray (Mutable v) f Any a) -> Array v f a -> Array v f a
`seq` (forall s. ST s (Array v f a)) -> Array v f a
forall a. (forall s. ST s a) -> a
runST (ST s (MArray (Mutable v) f s a)
forall s. ST s (MArray (Mutable v) f s a)
m ST s (MArray (Mutable v) f s a)
-> (MArray (Mutable v) f s a -> ST s (Array v f a))
-> ST s (Array v f a)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= MArray (Mutable v) f s a -> ST s (Array v f a)
forall (m :: * -> *) (v :: * -> *) a (f :: * -> *).
(PrimMonad m, Vector v a) =>
MArray (Mutable v) f (PrimState m) a -> m (Array v f a)
unsafeFreeze)
{-# INLINE create #-}

-- | Execute the monadic action and freeze the resulting array.
createT
  :: (Vector v a, Traversable t)
  => (forall s . ST s (t (MArray (G.Mutable v) f s a)))
  -> t (Array v f a)
createT :: (forall s. ST s (t (MArray (Mutable v) f s a))) -> t (Array v f a)
createT forall s. ST s (t (MArray (Mutable v) f s a))
m = ST Any (t (MArray (Mutable v) f Any a))
forall s. ST s (t (MArray (Mutable v) f s a))
m ST Any (t (MArray (Mutable v) f Any a))
-> t (Array v f a) -> t (Array v f a)
`seq` (forall s. ST s (t (Array v f a))) -> t (Array v f a)
forall a. (forall s. ST s a) -> a
runST (ST s (t (MArray (Mutable v) f s a))
forall s. ST s (t (MArray (Mutable v) f s a))
m ST s (t (MArray (Mutable v) f s a))
-> (t (MArray (Mutable v) f s a) -> ST s (t (Array v f a)))
-> ST s (t (Array v f a))
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (MArray (Mutable v) f s a -> ST s (Array v f a))
-> t (MArray (Mutable v) f s a) -> ST s (t (Array v f a))
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
T.mapM MArray (Mutable v) f s a -> ST s (Array v f a)
forall (m :: * -> *) (v :: * -> *) a (f :: * -> *).
(PrimMonad m, Vector v a) =>
MArray (Mutable v) f (PrimState m) a -> m (Array v f a)
unsafeFreeze)
{-# INLINE createT #-}

-- | O(n) Array of the given shape with the same value in each position.
replicate :: (Shape f, Vector v a) => f Int -> a -> Array v f a
replicate :: f Int -> a -> Array v f a
replicate f Int
l a
a
  | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0     = f Int -> v a -> Array v f a
forall (v :: * -> *) (f :: * -> *) a.
Layout f -> v a -> Array v f a
Array f Int
l (v a -> Array v f a) -> v a -> Array v f a
forall a b. (a -> b) -> a -> b
$ Int -> a -> v a
forall (v :: * -> *) a. Vector v a => Int -> a -> v a
G.replicate Int
n a
a
  | Bool
otherwise = Array v f a
forall (v :: * -> *) a (f :: * -> *).
(Vector v a, Additive f) =>
Array v f a
empty
  where n :: Int
n = f Int -> Int
forall (f :: * -> *). Shape f => Layout f -> Int
shapeSize f Int
l
{-# INLINE replicate #-}

-- | O(n) Construct an array of the given shape by applying the
--   function to each index.
linearGenerate :: (Shape f, Vector v a) => Layout f -> (Int -> a) -> Array v f a
linearGenerate :: Layout f -> (Int -> a) -> Array v f a
linearGenerate Layout f
l Int -> a
f
  | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0     = Layout f -> v a -> Array v f a
forall (v :: * -> *) (f :: * -> *) a.
Layout f -> v a -> Array v f a
Array Layout f
l (v a -> Array v f a) -> v a -> Array v f a
forall a b. (a -> b) -> a -> b
$ Int -> (Int -> a) -> v a
forall (v :: * -> *) a. Vector v a => Int -> (Int -> a) -> v a
G.generate Int
n Int -> a
f
  | Bool
otherwise = Array v f a
forall (v :: * -> *) a (f :: * -> *).
(Vector v a, Additive f) =>
Array v f a
empty
  where n :: Int
n = Layout f -> Int
forall (f :: * -> *). Shape f => Layout f -> Int
shapeSize Layout f
l
{-# INLINE linearGenerate #-}

-- | O(n) Construct an array of the given shape by applying the
--   function to each index.
generate :: (Shape f, Vector v a) => Layout f -> (f Int -> a) -> Array v f a
generate :: Layout f -> (Layout f -> a) -> Array v f a
generate Layout f
l Layout f -> a
f = Layout f -> v a -> Array v f a
forall (v :: * -> *) (f :: * -> *) a.
Layout f -> v a -> Array v f a
Array Layout f
l (v a -> Array v f a) -> v a -> Array v f a
forall a b. (a -> b) -> a -> b
$ Bundle v a -> v a
forall (v :: * -> *) a. Vector v a => Bundle v a -> v a
G.unstream (Layout f -> (Layout f -> a) -> Bundle v a
forall (m :: * -> *) (f :: * -> *) a (v :: * -> *).
(Monad m, Shape f) =>
Layout f -> (Layout f -> a) -> MBundle m v a
bundleGenerate Layout f
l Layout f -> a
f)
{-# INLINE generate #-}

-- Monadic initialisation ----------------------------------------------

-- | O(n) Construct an array of the given shape by filling each position
--   with the monadic value.
replicateM :: (Monad m, Shape f, Vector v a) => Layout f -> m a -> m (Array v f a)
replicateM :: Layout f -> m a -> m (Array v f a)
replicateM Layout f
l m a
a
  | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0     = Layout f -> v a -> Array v f a
forall (v :: * -> *) (f :: * -> *) a.
Layout f -> v a -> Array v f a
Array Layout f
l (v a -> Array v f a) -> m (v a) -> m (Array v f a)
forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
`liftM` Int -> m a -> m (v a)
forall (m :: * -> *) (v :: * -> *) a.
(Monad m, Vector v a) =>
Int -> m a -> m (v a)
G.replicateM Int
n m a
a
  | Bool
otherwise = Array v f a -> m (Array v f a)
forall (m :: * -> *) a. Monad m => a -> m a
return Array v f a
forall (v :: * -> *) a (f :: * -> *).
(Vector v a, Additive f) =>
Array v f a
empty
  where n :: Int
n = Layout f -> Int
forall (f :: * -> *). Shape f => Layout f -> Int
shapeSize Layout f
l
{-# INLINE replicateM #-}

-- | O(n) Construct an array of the given shape by applying the monadic
--   function to each index.
generateM :: (Monad m, Shape f, Vector v a) => Layout f -> (f Int -> m a) -> m (Array v f a)
generateM :: Layout f -> (Layout f -> m a) -> m (Array v f a)
generateM Layout f
l Layout f -> m a
f = Layout f -> v a -> Array v f a
forall (v :: * -> *) (f :: * -> *) a.
Layout f -> v a -> Array v f a
Array Layout f
l (v a -> Array v f a) -> m (v a) -> m (Array v f a)
forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
`liftM` MBundle m Any a -> m (v a)
forall (m :: * -> *) (v :: * -> *) a (u :: * -> *).
(Monad m, Vector v a) =>
MBundle m u a -> m (v a)
unstreamM (Layout f -> (Layout f -> m a) -> MBundle m Any a
forall (m :: * -> *) (f :: * -> *) a (v :: * -> *).
(Monad m, Shape f) =>
Layout f -> (Layout f -> m a) -> MBundle m v a
bundleGenerateM Layout f
l Layout f -> m a
f)
{-# INLINE generateM #-}

-- | O(n) Construct an array of the given shape by applying the monadic
--   function to each index.
linearGenerateM :: (Monad m, Shape f, Vector v a) => Layout f -> (Int -> m a) -> m (Array v f a)
linearGenerateM :: Layout f -> (Int -> m a) -> m (Array v f a)
linearGenerateM Layout f
l Int -> m a
f
  | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0     = Layout f -> v a -> Array v f a
forall (v :: * -> *) (f :: * -> *) a.
Layout f -> v a -> Array v f a
Array Layout f
l (v a -> Array v f a) -> m (v a) -> m (Array v f a)
forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
`liftM` Int -> (Int -> m a) -> m (v a)
forall (m :: * -> *) (v :: * -> *) a.
(Monad m, Vector v a) =>
Int -> (Int -> m a) -> m (v a)
G.generateM Int
n Int -> m a
f
  | Bool
otherwise = Array v f a -> m (Array v f a)
forall (m :: * -> *) a. Monad m => a -> m a
return Array v f a
forall (v :: * -> *) a (f :: * -> *).
(Vector v a, Additive f) =>
Array v f a
empty
  where n :: Int
n = Layout f -> Int
forall (f :: * -> *). Shape f => Layout f -> Int
shapeSize Layout f
l
{-# INLINE linearGenerateM #-}

-- Modifying -----------------------------------------------------------

-- | /O(n)/ Map a function over an array
map :: (Vector v a, Vector v b) => (a -> b) -> Array v f a -> Array v f b
map :: (a -> b) -> Array v f a -> Array v f b
map a -> b
f (Array Layout f
l v a
a) = Layout f -> v b -> Array v f b
forall (v :: * -> *) (f :: * -> *) a.
Layout f -> v a -> Array v f a
Array Layout f
l ((a -> b) -> v a -> v b
forall (v :: * -> *) a b.
(Vector v a, Vector v b) =>
(a -> b) -> v a -> v b
G.map a -> b
f v a
a)
{-# INLINE map #-}

-- | /O(n)/ Apply a function to every element of a vector and its index
imap :: (Shape f, Vector v a, Vector v b) => (f Int -> a -> b) -> Array v f a -> Array v f b
imap :: (f Int -> a -> b) -> Array v f a -> Array v f b
imap f Int -> a -> b
f (Array f Int
l v a
v) =
  f Int -> v b -> Array v f b
forall (v :: * -> *) (f :: * -> *) a.
Layout f -> v a -> Array v f a
Array f Int
l (v b -> Array v f b) -> v b -> Array v f b
forall a b. (a -> b) -> a -> b
$ (Bundle v b -> v b
forall (v :: * -> *) a. Vector v a => Bundle v a -> v a
G.unstream (Bundle v b -> v b) -> (v a -> Bundle v b) -> v a -> v b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (forall (m :: * -> *). Monad m => Stream m a -> Stream m b)
-> (Size -> Size) -> Bundle v a -> Bundle v b
forall a b (v :: * -> *).
(forall (m :: * -> *). Monad m => Stream m a -> Stream m b)
-> (Size -> Size) -> Bundle v a -> Bundle v b
Bundle.inplace ((f Int -> a -> b) -> Stream m (f Int) -> Stream m a -> Stream m b
forall (m :: * -> *) a b c.
Monad m =>
(a -> b -> c) -> Stream m a -> Stream m b -> Stream m c
Stream.zipWith f Int -> a -> b
f (f Int -> Stream m (f Int)
forall (m :: * -> *) (f :: * -> *).
(Monad m, Shape f) =>
Layout f -> Stream m (Layout f)
streamIndexes f Int
l)) Size -> Size
forall a. a -> a
id (Bundle v a -> Bundle v b)
-> (v a -> Bundle v a) -> v a -> Bundle v b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. v a -> Bundle v a
forall (v :: * -> *) a. Vector v a => v a -> Bundle v a
G.stream) v a
v
{-# INLINE imap #-}

-- Bulk updates --------------------------------------------------------

-- | For each pair (i,a) from the list, replace the array element at
--   position i by a.
(//) :: (G.Vector v a, Shape f) => Array v f a -> [(f Int, a)] -> Array v f a
Array f Int
l v a
v // :: Array v f a -> [(f Int, a)] -> Array v f a
// [(f Int, a)]
xs = f Int -> v a -> Array v f a
forall (v :: * -> *) (f :: * -> *) a.
Layout f -> v a -> Array v f a
Array f Int
l (v a -> Array v f a) -> v a -> Array v f a
forall a b. (a -> b) -> a -> b
$ v a
v v a -> [(Int, a)] -> v a
forall (v :: * -> *) a. Vector v a => v a -> [(Int, a)] -> v a
G.// ASetter [(f Int, a)] [(Int, a)] (f Int) Int
-> (f Int -> Int) -> [(f Int, a)] -> [(Int, a)]
forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
over (((f Int, a) -> Identity (Int, a))
-> [(f Int, a)] -> Identity [(Int, a)]
forall s t a b. Each s t a b => Traversal s t a b
each (((f Int, a) -> Identity (Int, a))
 -> [(f Int, a)] -> Identity [(Int, a)])
-> ((f Int -> Identity Int) -> (f Int, a) -> Identity (Int, a))
-> ASetter [(f Int, a)] [(Int, a)] (f Int) Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (f Int -> Identity Int) -> (f Int, a) -> Identity (Int, a)
forall s t a b. Field1 s t a b => Lens s t a b
_1) (f Int -> f Int -> Int
forall (f :: * -> *). Shape f => Layout f -> Layout f -> Int
shapeToIndex f Int
l) [(f Int, a)]
xs

-- Accumilation --------------------------------------------------------

-- | /O(m+n)/ For each pair @(i,b)@ from the list, replace the array element
--   @a@ at position @i@ by @f a b@.
--
accum :: (Shape f, Vector v a)
      => (a -> b -> a)  -- ^ accumulating function @f@
      -> Array v f a    -- ^ initial array
      -> [(f Int, b)]   -- ^ list of index/value pairs (of length @n@)
      -> Array v f a
accum :: (a -> b -> a) -> Array v f a -> [(f Int, b)] -> Array v f a
accum a -> b -> a
f (Array f Int
l v a
v) [(f Int, b)]
us = f Int -> v a -> Array v f a
forall (v :: * -> *) (f :: * -> *) a.
Layout f -> v a -> Array v f a
Array f Int
l (v a -> Array v f a) -> v a -> Array v f a
forall a b. (a -> b) -> a -> b
$ (a -> b -> a) -> v a -> [(Int, b)] -> v a
forall (v :: * -> *) a b.
Vector v a =>
(a -> b -> a) -> v a -> [(Int, b)] -> v a
G.accum a -> b -> a
f v a
v (ASetter [(f Int, b)] [(Int, b)] (f Int) Int
-> (f Int -> Int) -> [(f Int, b)] -> [(Int, b)]
forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
over (((f Int, b) -> Identity (Int, b))
-> [(f Int, b)] -> Identity [(Int, b)]
forall (f :: * -> *) a b. Functor f => Setter (f a) (f b) a b
mapped (((f Int, b) -> Identity (Int, b))
 -> [(f Int, b)] -> Identity [(Int, b)])
-> ((f Int -> Identity Int) -> (f Int, b) -> Identity (Int, b))
-> ASetter [(f Int, b)] [(Int, b)] (f Int) Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (f Int -> Identity Int) -> (f Int, b) -> Identity (Int, b)
forall s t a b. Field1 s t a b => Lens s t a b
_1) (f Int -> f Int -> Int
forall (f :: * -> *). Shape f => Layout f -> Layout f -> Int
shapeToIndex f Int
l) [(f Int, b)]
us)
{-# INLINE accum #-}

------------------------------------------------------------------------
-- Streams
------------------------------------------------------------------------

-- Copied from Data.Vector.Generic because it isn't exported from there.

unstreamM :: (Monad m, Vector v a) => Bundle.MBundle m u a -> m (v a)
{-# INLINE [1] unstreamM #-}
unstreamM :: MBundle m u a -> m (v a)
unstreamM MBundle m u a
s = do
  [a]
xs <- MBundle m u a -> m [a]
forall (m :: * -> *) (v :: * -> *) a.
Monad m =>
Bundle m v a -> m [a]
MBundle.toList MBundle m u a
s
  v a -> m (v a)
forall (m :: * -> *) a. Monad m => a -> m a
return (v a -> m (v a)) -> v a -> m (v a)
forall a b. (a -> b) -> a -> b
$ Bundle v a -> v a
forall (v :: * -> *) a. Vector v a => Bundle v a -> v a
G.unstream (Bundle v a -> v a) -> Bundle v a -> v a
forall a b. (a -> b) -> a -> b
$ Size -> [a] -> Bundle v a
forall a (v :: * -> *). Size -> [a] -> Bundle v a
Bundle.unsafeFromList (MBundle m u a -> Size
forall (m :: * -> *) (v :: * -> *) a. Bundle m v a -> Size
MBundle.size MBundle m u a
s) [a]
xs

unstreamPrimM :: (PrimMonad m, Vector v a) => Bundle.MBundle m u a -> m (v a)
{-# INLINE [1] unstreamPrimM #-}
unstreamPrimM :: MBundle m u a -> m (v a)
unstreamPrimM MBundle m u a
s = MBundle m u a -> m (Mutable v (PrimState m) a)
forall (m :: * -> *) (v :: * -> * -> *) a (u :: * -> *).
(PrimMonad m, MVector v a) =>
MBundle m u a -> m (v (PrimState m) a)
GM.munstream MBundle m u a
s m (Mutable v (PrimState m) a)
-> (Mutable v (PrimState m) a -> m (v a)) -> m (v a)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Mutable v (PrimState m) a -> m (v a)
forall (m :: * -> *) (v :: * -> *) a.
(PrimMonad m, Vector v a) =>
Mutable v (PrimState m) a -> m (v a)
G.unsafeFreeze

-- FIXME: the next two functions are only necessary for the specialisations
unstreamPrimM_IO :: Vector v a => Bundle.MBundle IO u a -> IO (v a)
{-# INLINE unstreamPrimM_IO #-}
unstreamPrimM_IO :: MBundle IO u a -> IO (v a)
unstreamPrimM_IO = MBundle IO u a -> IO (v a)
forall (m :: * -> *) (v :: * -> *) a (u :: * -> *).
(PrimMonad m, Vector v a) =>
MBundle m u a -> m (v a)
unstreamPrimM

unstreamPrimM_ST :: Vector v a => Bundle.MBundle (ST s) u a -> ST s (v a)
{-# INLINE unstreamPrimM_ST #-}
unstreamPrimM_ST :: MBundle (ST s) u a -> ST s (v a)
unstreamPrimM_ST = MBundle (ST s) u a -> ST s (v a)
forall (m :: * -> *) (v :: * -> *) a (u :: * -> *).
(PrimMonad m, Vector v a) =>
MBundle m u a -> m (v a)
unstreamPrimM

{-# RULES

"unstreamM[IO]" unstreamM = unstreamPrimM_IO
"unstreamM[ST]" unstreamM = unstreamPrimM_ST  #-}

-- | Generate a stream from a 'Layout''s indices.
streamGenerate :: (Monad m, Shape f) => Layout f -> (f Int -> a) -> Stream m a
streamGenerate :: Layout f -> (Layout f -> a) -> Stream m a
streamGenerate Layout f
l Layout f -> a
f = Layout f -> (Layout f -> m a) -> Stream m a
forall (m :: * -> *) (f :: * -> *) a.
(Monad m, Shape f) =>
Layout f -> (Layout f -> m a) -> Stream m a
streamGenerateM Layout f
l (a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (a -> m a) -> (Layout f -> a) -> Layout f -> m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Layout f -> a
f)
{-# INLINE streamGenerate #-}

-- | Generate a stream from a 'Layout''s indices.
streamGenerateM :: (Monad m, Shape f) => Layout f -> (f Int -> m a) -> Stream m a
streamGenerateM :: Layout f -> (Layout f -> m a) -> Stream m a
streamGenerateM Layout f
l Layout f -> m a
f = Layout f
l Layout f -> Stream m a -> Stream m a
`seq` (Maybe (Layout f) -> m (Step (Maybe (Layout f)) a))
-> Maybe (Layout f) -> Stream m a
forall (m :: * -> *) a s. (s -> m (Step s a)) -> s -> Stream m a
Stream Maybe (Layout f) -> m (Step (Maybe (Layout f)) a)
step (if Layout f -> Layout f -> Bool
forall (f :: * -> *) a. (Eq1 f, Eq a) => f a -> f a -> Bool
eq1 Layout f
l Layout f
forall (f :: * -> *) a. (Additive f, Num a) => f a
zero then Maybe (Layout f)
forall a. Maybe a
Nothing else Layout f -> Maybe (Layout f)
forall a. a -> Maybe a
Just Layout f
forall (f :: * -> *) a. (Additive f, Num a) => f a
zero)
  where
    {-# INLINE [0] step #-}
    step :: Maybe (Layout f) -> m (Step (Maybe (Layout f)) a)
step (Just Layout f
i) = do
      a
x <- Layout f -> m a
f Layout f
i
      Step (Maybe (Layout f)) a -> m (Step (Maybe (Layout f)) a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (Maybe (Layout f)) a -> m (Step (Maybe (Layout f)) a))
-> Step (Maybe (Layout f)) a -> m (Step (Maybe (Layout f)) a)
forall a b. (a -> b) -> a -> b
$ a -> Maybe (Layout f) -> Step (Maybe (Layout f)) a
forall a s. a -> s -> Step s a
Yield a
x (Layout f -> Layout f -> Maybe (Layout f)
forall (f :: * -> *).
Shape f =>
Layout f -> Layout f -> Maybe (Layout f)
shapeStep Layout f
l Layout f
i)
    step Maybe (Layout f)
Nothing  = Step (Maybe (Layout f)) a -> m (Step (Maybe (Layout f)) a)
forall (m :: * -> *) a. Monad m => a -> m a
return Step (Maybe (Layout f)) a
forall s a. Step s a
Done
{-# INLINE [1] streamGenerateM #-}

-- | Stream a sub-layout of an 'Array'. The layout should be shapeInRange of
--   the array's layout, this is not checked.
unsafeStreamSub :: (Monad m, Shape f, G.Vector v a) => Layout f -> Array v f a -> Stream m a
unsafeStreamSub :: Layout f -> Array v f a -> Stream m a
unsafeStreamSub Layout f
l2 (Array Layout f
l1 v a
v) = Layout f -> (Layout f -> m a) -> Stream m a
forall (m :: * -> *) (f :: * -> *) a.
(Monad m, Shape f) =>
Layout f -> (Layout f -> m a) -> Stream m a
streamGenerateM Layout f
l2 ((Layout f -> m a) -> Stream m a)
-> (Layout f -> m a) -> Stream m a
forall a b. (a -> b) -> a -> b
$ \Layout f
x -> v a -> Int -> m a
forall (v :: * -> *) a (m :: * -> *).
(Vector v a, Monad m) =>
v a -> Int -> m a
G.basicUnsafeIndexM v a
v (Layout f -> Layout f -> Int
forall (f :: * -> *). Shape f => Layout f -> Layout f -> Int
shapeToIndex Layout f
l1 Layout f
x)
{-# INLINE unsafeStreamSub #-}

-- | Stream a sub-layout of an 'Array'.
streamSub :: (Monad m, Shape f, G.Vector v a) => Layout f -> Array v f a -> Stream m a
streamSub :: Layout f -> Array v f a -> Stream m a
streamSub Layout f
l2 arr :: Array v f a
arr@(Array Layout f
l1 v a
_) = Layout f -> Array v f a -> Stream m a
forall (m :: * -> *) (f :: * -> *) (v :: * -> *) a.
(Monad m, Shape f, Vector v a) =>
Layout f -> Array v f a -> Stream m a
unsafeStreamSub (Layout f -> Layout f -> Layout f
forall (f :: * -> *). Shape f => Layout f -> Layout f -> Layout f
shapeIntersect Layout f
l1 Layout f
l2) Array v f a
arr
{-# INLINE streamSub #-}

-- | Make a stream of the indexes of a 'Layout'.
streamIndexes :: (Monad m, Shape f) => Layout f -> Stream m (f Int)
streamIndexes :: Layout f -> Stream m (Layout f)
streamIndexes Layout f
l = (Maybe (Layout f) -> m (Step (Maybe (Layout f)) (Layout f)))
-> Maybe (Layout f) -> Stream m (Layout f)
forall (m :: * -> *) a s. (s -> m (Step s a)) -> s -> Stream m a
Stream Maybe (Layout f) -> m (Step (Maybe (Layout f)) (Layout f))
step (if Layout f -> Layout f -> Bool
forall (f :: * -> *) a. (Eq1 f, Eq a) => f a -> f a -> Bool
eq1 Layout f
l Layout f
forall (f :: * -> *) a. (Additive f, Num a) => f a
zero then Maybe (Layout f)
forall a. Maybe a
Nothing else Layout f -> Maybe (Layout f)
forall a. a -> Maybe a
Just Layout f
forall (f :: * -> *) a. (Additive f, Num a) => f a
zero)
  where
    {-# INLINE [0] step #-}
    step :: Maybe (Layout f) -> m (Step (Maybe (Layout f)) (Layout f))
step (Just Layout f
i) = Step (Maybe (Layout f)) (Layout f)
-> m (Step (Maybe (Layout f)) (Layout f))
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (Maybe (Layout f)) (Layout f)
 -> m (Step (Maybe (Layout f)) (Layout f)))
-> Step (Maybe (Layout f)) (Layout f)
-> m (Step (Maybe (Layout f)) (Layout f))
forall a b. (a -> b) -> a -> b
$ Layout f -> Maybe (Layout f) -> Step (Maybe (Layout f)) (Layout f)
forall a s. a -> s -> Step s a
Yield Layout f
i (Layout f -> Layout f -> Maybe (Layout f)
forall (f :: * -> *).
Shape f =>
Layout f -> Layout f -> Maybe (Layout f)
shapeStep Layout f
l Layout f
i)
    step Maybe (Layout f)
Nothing  = Step (Maybe (Layout f)) (Layout f)
-> m (Step (Maybe (Layout f)) (Layout f))
forall (m :: * -> *) a. Monad m => a -> m a
return Step (Maybe (Layout f)) (Layout f)
forall s a. Step s a
Done
{-# INLINE [1] streamIndexes #-}

------------------------------------------------------------------------
-- Bundles
------------------------------------------------------------------------

-- | Generate a bundle from 'Layout' indices.
bundleGenerate :: (Monad m, Shape f) => Layout f -> (f Int -> a) -> MBundle m v a
bundleGenerate :: Layout f -> (Layout f -> a) -> MBundle m v a
bundleGenerate Layout f
l Layout f -> a
f = Layout f -> (Layout f -> m a) -> MBundle m v a
forall (m :: * -> *) (f :: * -> *) a (v :: * -> *).
(Monad m, Shape f) =>
Layout f -> (Layout f -> m a) -> MBundle m v a
bundleGenerateM Layout f
l (a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (a -> m a) -> (Layout f -> a) -> Layout f -> m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Layout f -> a
f)
{-# INLINE bundleGenerate #-}

-- | Generate a bundle from 'Layout' indices.
bundleGenerateM :: (Monad m, Shape f) => Layout f -> (f Int -> m a) -> MBundle m v a
bundleGenerateM :: Layout f -> (Layout f -> m a) -> MBundle m v a
bundleGenerateM Layout f
l Layout f -> m a
f = Stream m a -> Size -> MBundle m v a
forall (m :: * -> *) a (v :: * -> *).
Monad m =>
Stream m a -> Size -> Bundle m v a
MBundle.fromStream (Layout f -> (Layout f -> m a) -> Stream m a
forall (m :: * -> *) (f :: * -> *) a.
(Monad m, Shape f) =>
Layout f -> (Layout f -> m a) -> Stream m a
streamGenerateM Layout f
l Layout f -> m a
f) (Int -> Size
Exact (Layout f -> Int
forall (f :: * -> *). Shape f => Layout f -> Int
shapeSize Layout f
l))
{-# INLINE [1] bundleGenerateM #-}

-- | Generate a bundle of indexes for the given 'Layout'.
bundleIndexes :: (Monad m, Shape f) => Layout f -> MBundle m v (f Int)
bundleIndexes :: Layout f -> MBundle m v (Layout f)
bundleIndexes Layout f
l = Stream m (Layout f) -> Size -> MBundle m v (Layout f)
forall (m :: * -> *) a (v :: * -> *).
Monad m =>
Stream m a -> Size -> Bundle m v a
MBundle.fromStream (Layout f -> Stream m (Layout f)
forall (m :: * -> *) (f :: * -> *).
(Monad m, Shape f) =>
Layout f -> Stream m (Layout f)
streamIndexes Layout f
l) (Int -> Size
Exact (Layout f -> Int
forall (f :: * -> *). Shape f => Layout f -> Int
shapeSize Layout f
l))
{-# INLINE [1] bundleIndexes #-}

------------------------------------------------------------------------
-- Zipping
------------------------------------------------------------------------

-- Tuple zip -----------------------------------------------------------

-- | Zip two arrays element wise. If the array's don't have the same
--   shape, the new array with be the intersection of the two shapes.
zip :: (Shape f, Vector v a, Vector v b, Vector v (a,b))
    => Array v f a
    -> Array v f b
    -> Array v f (a,b)
zip :: Array v f a -> Array v f b -> Array v f (a, b)
zip = (a -> b -> (a, b))
-> Array v f a -> Array v f b -> Array v f (a, b)
forall (f :: * -> *) (v :: * -> *) a b c.
(Shape f, Vector v a, Vector v b, Vector v c) =>
(a -> b -> c) -> Array v f a -> Array v f b -> Array v f c
zipWith (,)

-- | Zip three arrays element wise. If the array's don't have the same
--   shape, the new array with be the intersection of the two shapes.
zip3 :: (Shape f, Vector v a, Vector v b, Vector v c, Vector v (a,b,c))
     => Array v f a
     -> Array v f b
     -> Array v f c
     -> Array v f (a,b,c)
zip3 :: Array v f a -> Array v f b -> Array v f c -> Array v f (a, b, c)
zip3 = (a -> b -> c -> (a, b, c))
-> Array v f a -> Array v f b -> Array v f c -> Array v f (a, b, c)
forall (f :: * -> *) (v :: * -> *) a b c d.
(Shape f, Vector v a, Vector v b, Vector v c, Vector v d) =>
(a -> b -> c -> d)
-> Array v f a -> Array v f b -> Array v f c -> Array v f d
zipWith3 (,,)

-- Zip with function ---------------------------------------------------

-- | Zip two arrays using the given function. If the array's don't have
--   the same shape, the new array with be the intersection of the two
--   shapes.
zipWith :: (Shape f, Vector v a, Vector v b, Vector v c)
        => (a -> b -> c)
        -> Array v f a
        -> Array v f b
        -> Array v f c
zipWith :: (a -> b -> c) -> Array v f a -> Array v f b -> Array v f c
zipWith a -> b -> c
f a1 :: Array v f a
a1@(Array Layout f
l1 v a
v1) a2 :: Array v f b
a2@(Array Layout f
l2 v b
v2)
  | Layout f -> Layout f -> Bool
forall (f :: * -> *) a. (Eq1 f, Eq a) => f a -> f a -> Bool
eq1 Layout f
l1 Layout f
l2 = Layout f -> v c -> Array v f c
forall (v :: * -> *) (f :: * -> *) a.
Layout f -> v a -> Array v f a
Array Layout f
l1 (v c -> Array v f c) -> v c -> Array v f c
forall a b. (a -> b) -> a -> b
$ (a -> b -> c) -> v a -> v b -> v c
forall (v :: * -> *) a b c.
(Vector v a, Vector v b, Vector v c) =>
(a -> b -> c) -> v a -> v b -> v c
G.zipWith a -> b -> c
f v a
v1 v b
v2
  | Bool
otherwise = Layout f -> v c -> Array v f c
forall (v :: * -> *) (f :: * -> *) a.
Layout f -> v a -> Array v f a
Array Layout f
l' (v c -> Array v f c) -> v c -> Array v f c
forall a b. (a -> b) -> a -> b
$ Bundle v c -> v c
forall (v :: * -> *) a. Vector v a => Bundle v a -> v a
G.unstream (Bundle v c -> v c) -> Bundle v c -> v c
forall a b. (a -> b) -> a -> b
$
      Stream Id c -> Size -> Bundle v c
forall (m :: * -> *) a (v :: * -> *).
Monad m =>
Stream m a -> Size -> Bundle m v a
MBundle.fromStream ((a -> b -> c) -> Stream Id a -> Stream Id b -> Stream Id c
forall (m :: * -> *) a b c.
Monad m =>
(a -> b -> c) -> Stream m a -> Stream m b -> Stream m c
Stream.zipWith a -> b -> c
f (Layout f -> Array v f a -> Stream Id a
forall (m :: * -> *) (f :: * -> *) (v :: * -> *) a.
(Monad m, Shape f, Vector v a) =>
Layout f -> Array v f a -> Stream m a
streamSub Layout f
l' Array v f a
a1) (Layout f -> Array v f b -> Stream Id b
forall (m :: * -> *) (f :: * -> *) (v :: * -> *) a.
(Monad m, Shape f, Vector v a) =>
Layout f -> Array v f a -> Stream m a
streamSub Layout f
l' Array v f b
a2)) (Int -> Size
Exact (Layout f -> Int
forall (f :: * -> *). Shape f => Layout f -> Int
shapeSize Layout f
l'))
  where l' :: Layout f
l' = Layout f -> Layout f -> Layout f
forall (f :: * -> *). Shape f => Layout f -> Layout f -> Layout f
shapeIntersect Layout f
l1 Layout f
l2
{-# INLINE zipWith #-}

-- | Zip three arrays using the given function. If the array's don't
--   have the same shape, the new array with be the intersection of the
--   two shapes.
zipWith3 :: (Shape f, Vector v a, Vector v b, Vector v c, Vector v d)
         => (a -> b -> c -> d)
         -> Array v f a
         -> Array v f b
         -> Array v f c
         -> Array v f d
zipWith3 :: (a -> b -> c -> d)
-> Array v f a -> Array v f b -> Array v f c -> Array v f d
zipWith3 a -> b -> c -> d
f a1 :: Array v f a
a1@(Array Layout f
l1 v a
v1) a2 :: Array v f b
a2@(Array Layout f
l2 v b
v2) a3 :: Array v f c
a3@(Array Layout f
l3 v c
v3)
  | Layout f -> Layout f -> Bool
forall (f :: * -> *) a. (Eq1 f, Eq a) => f a -> f a -> Bool
eq1 Layout f
l1 Layout f
l2 Bool -> Bool -> Bool
&&
    Layout f -> Layout f -> Bool
forall (f :: * -> *) a. (Eq1 f, Eq a) => f a -> f a -> Bool
eq1 Layout f
l2 Layout f
l3 = Layout f -> v d -> Array v f d
forall (v :: * -> *) (f :: * -> *) a.
Layout f -> v a -> Array v f a
Array Layout f
l1 (v d -> Array v f d) -> v d -> Array v f d
forall a b. (a -> b) -> a -> b
$ (a -> b -> c -> d) -> v a -> v b -> v c -> v d
forall (v :: * -> *) a b c d.
(Vector v a, Vector v b, Vector v c, Vector v d) =>
(a -> b -> c -> d) -> v a -> v b -> v c -> v d
G.zipWith3 a -> b -> c -> d
f v a
v1 v b
v2 v c
v3
  | Bool
otherwise = Layout f -> v d -> Array v f d
forall (v :: * -> *) (f :: * -> *) a.
Layout f -> v a -> Array v f a
Array Layout f
l' (v d -> Array v f d) -> v d -> Array v f d
forall a b. (a -> b) -> a -> b
$ Bundle v d -> v d
forall (v :: * -> *) a. Vector v a => Bundle v a -> v a
G.unstream (Bundle v d -> v d) -> Bundle v d -> v d
forall a b. (a -> b) -> a -> b
$
      Stream Id d -> Size -> Bundle v d
forall (m :: * -> *) a (v :: * -> *).
Monad m =>
Stream m a -> Size -> Bundle m v a
MBundle.fromStream ((a -> b -> c -> d)
-> Stream Id a -> Stream Id b -> Stream Id c -> Stream Id d
forall (m :: * -> *) a b c d.
Monad m =>
(a -> b -> c -> d)
-> Stream m a -> Stream m b -> Stream m c -> Stream m d
Stream.zipWith3 a -> b -> c -> d
f (Layout f -> Array v f a -> Stream Id a
forall (m :: * -> *) (f :: * -> *) (v :: * -> *) a.
(Monad m, Shape f, Vector v a) =>
Layout f -> Array v f a -> Stream m a
streamSub Layout f
l' Array v f a
a1) (Layout f -> Array v f b -> Stream Id b
forall (m :: * -> *) (f :: * -> *) (v :: * -> *) a.
(Monad m, Shape f, Vector v a) =>
Layout f -> Array v f a -> Stream m a
streamSub Layout f
l' Array v f b
a2) (Layout f -> Array v f c -> Stream Id c
forall (m :: * -> *) (f :: * -> *) (v :: * -> *) a.
(Monad m, Shape f, Vector v a) =>
Layout f -> Array v f a -> Stream m a
streamSub Layout f
l' Array v f c
a3)) (Int -> Size
Exact (Layout f -> Int
forall (f :: * -> *). Shape f => Layout f -> Int
shapeSize Layout f
l'))
  where l' :: Layout f
l' = Layout f -> Layout f -> Layout f
forall (f :: * -> *). Shape f => Layout f -> Layout f -> Layout f
shapeIntersect (Layout f -> Layout f -> Layout f
forall (f :: * -> *). Shape f => Layout f -> Layout f -> Layout f
shapeIntersect Layout f
l1 Layout f
l2) Layout f
l3
{-# INLINE zipWith3 #-}

-- Indexed zipping -----------------------------------------------------

-- | Zip two arrays using the given function with access to the index.
--   If the array's don't have the same shape, the new array with be the
--   intersection of the two shapes.
izipWith :: (Shape f, Vector v a, Vector v b, Vector v c)
         => (f Int -> a -> b -> c)
         -> Array v f a
         -> Array v f b
         -> Array v f c
izipWith :: (f Int -> a -> b -> c) -> Array v f a -> Array v f b -> Array v f c
izipWith f Int -> a -> b -> c
f a1 :: Array v f a
a1@(Array f Int
l1 v a
v1) a2 :: Array v f b
a2@(Array f Int
l2 v b
v2)
  | f Int -> f Int -> Bool
forall (f :: * -> *) a. (Eq1 f, Eq a) => f a -> f a -> Bool
eq1 f Int
l1 f Int
l2 = f Int -> v c -> Array v f c
forall (v :: * -> *) (f :: * -> *) a.
Layout f -> v a -> Array v f a
Array f Int
l1 (v c -> Array v f c) -> v c -> Array v f c
forall a b. (a -> b) -> a -> b
$ Bundle v c -> v c
forall (v :: * -> *) a. Vector v a => Bundle v a -> v a
G.unstream (Bundle v c -> v c) -> Bundle v c -> v c
forall a b. (a -> b) -> a -> b
$ (f Int -> a -> b -> c)
-> Bundle v (f Int) -> Bundle v a -> Bundle v b -> Bundle v c
forall a b c d (v :: * -> *).
(a -> b -> c -> d)
-> Bundle v a -> Bundle v b -> Bundle v c -> Bundle v d
Bundle.zipWith3 f Int -> a -> b -> c
f (f Int -> Bundle v (f Int)
forall (m :: * -> *) (f :: * -> *) (v :: * -> *).
(Monad m, Shape f) =>
Layout f -> MBundle m v (Layout f)
bundleIndexes f Int
l1) (v a -> Bundle v a
forall (v :: * -> *) a. Vector v a => v a -> Bundle v a
G.stream v a
v1) (v b -> Bundle v b
forall (v :: * -> *) a. Vector v a => v a -> Bundle v a
G.stream v b
v2)
  | Bool
otherwise = f Int -> v c -> Array v f c
forall (v :: * -> *) (f :: * -> *) a.
Layout f -> v a -> Array v f a
Array f Int
l' (v c -> Array v f c) -> v c -> Array v f c
forall a b. (a -> b) -> a -> b
$ Bundle v c -> v c
forall (v :: * -> *) a. Vector v a => Bundle v a -> v a
G.unstream (Bundle v c -> v c) -> Bundle v c -> v c
forall a b. (a -> b) -> a -> b
$
      Stream Id c -> Size -> Bundle v c
forall (m :: * -> *) a (v :: * -> *).
Monad m =>
Stream m a -> Size -> Bundle m v a
MBundle.fromStream ((f Int -> a -> b -> c)
-> Stream Id (f Int) -> Stream Id a -> Stream Id b -> Stream Id c
forall (m :: * -> *) a b c d.
Monad m =>
(a -> b -> c -> d)
-> Stream m a -> Stream m b -> Stream m c -> Stream m d
Stream.zipWith3 f Int -> a -> b -> c
f (f Int -> Stream Id (f Int)
forall (m :: * -> *) (f :: * -> *).
(Monad m, Shape f) =>
Layout f -> Stream m (Layout f)
streamIndexes f Int
l') (f Int -> Array v f a -> Stream Id a
forall (m :: * -> *) (f :: * -> *) (v :: * -> *) a.
(Monad m, Shape f, Vector v a) =>
Layout f -> Array v f a -> Stream m a
streamSub f Int
l' Array v f a
a1) (f Int -> Array v f b -> Stream Id b
forall (m :: * -> *) (f :: * -> *) (v :: * -> *) a.
(Monad m, Shape f, Vector v a) =>
Layout f -> Array v f a -> Stream m a
streamSub f Int
l' Array v f b
a2)) (Int -> Size
Exact (f Int -> Int
forall (f :: * -> *). Shape f => Layout f -> Int
shapeSize f Int
l'))
  where l' :: f Int
l' = f Int -> f Int -> f Int
forall (f :: * -> *). Shape f => Layout f -> Layout f -> Layout f
shapeIntersect f Int
l1 f Int
l2
{-# INLINE izipWith #-}

-- | Zip two arrays using the given function with access to the index.
--   If the array's don't have the same shape, the new array with be the
--   intersection of the two shapes.
izipWith3 :: (Shape f, Vector v a, Vector v b, Vector v c, Vector v d)
          => (f Int -> a -> b -> c -> d)
          -> Array v f a
          -> Array v f b
          -> Array v f c
          -> Array v f d
izipWith3 :: (f Int -> a -> b -> c -> d)
-> Array v f a -> Array v f b -> Array v f c -> Array v f d
izipWith3 f Int -> a -> b -> c -> d
f a1 :: Array v f a
a1@(Array f Int
l1 v a
v1) a2 :: Array v f b
a2@(Array f Int
l2 v b
v2) a3 :: Array v f c
a3@(Array f Int
l3 v c
v3)
  | f Int -> f Int -> Bool
forall (f :: * -> *) a. (Eq1 f, Eq a) => f a -> f a -> Bool
eq1 f Int
l1 f Int
l2 = f Int -> v d -> Array v f d
forall (v :: * -> *) (f :: * -> *) a.
Layout f -> v a -> Array v f a
Array f Int
l1 (v d -> Array v f d) -> v d -> Array v f d
forall a b. (a -> b) -> a -> b
$ Bundle v d -> v d
forall (v :: * -> *) a. Vector v a => Bundle v a -> v a
G.unstream (Bundle v d -> v d) -> Bundle v d -> v d
forall a b. (a -> b) -> a -> b
$ (f Int -> a -> b -> c -> d)
-> Bundle v (f Int)
-> Bundle v a
-> Bundle v b
-> Bundle v c
-> Bundle v d
forall a b c d e (v :: * -> *).
(a -> b -> c -> d -> e)
-> Bundle v a
-> Bundle v b
-> Bundle v c
-> Bundle v d
-> Bundle v e
Bundle.zipWith4 f Int -> a -> b -> c -> d
f (f Int -> Bundle v (f Int)
forall (m :: * -> *) (f :: * -> *) (v :: * -> *).
(Monad m, Shape f) =>
Layout f -> MBundle m v (Layout f)
bundleIndexes f Int
l1) (v a -> Bundle v a
forall (v :: * -> *) a. Vector v a => v a -> Bundle v a
G.stream v a
v1) (v b -> Bundle v b
forall (v :: * -> *) a. Vector v a => v a -> Bundle v a
G.stream v b
v2) (v c -> Bundle v c
forall (v :: * -> *) a. Vector v a => v a -> Bundle v a
G.stream v c
v3)
  | Bool
otherwise =
      f Int -> v d -> Array v f d
forall (v :: * -> *) (f :: * -> *) a.
Layout f -> v a -> Array v f a
Array f Int
l' (v d -> Array v f d) -> v d -> Array v f d
forall a b. (a -> b) -> a -> b
$ Bundle v d -> v d
forall (v :: * -> *) a. Vector v a => Bundle v a -> v a
G.unstream (Bundle v d -> v d) -> Bundle v d -> v d
forall a b. (a -> b) -> a -> b
$ Stream Id d -> Size -> Bundle v d
forall (m :: * -> *) a (v :: * -> *).
Monad m =>
Stream m a -> Size -> Bundle m v a
MBundle.fromStream
        ((f Int -> a -> b -> c -> d)
-> Stream Id (f Int)
-> Stream Id a
-> Stream Id b
-> Stream Id c
-> Stream Id d
forall (m :: * -> *) a b c d e.
Monad m =>
(a -> b -> c -> d -> e)
-> Stream m a
-> Stream m b
-> Stream m c
-> Stream m d
-> Stream m e
Stream.zipWith4 f Int -> a -> b -> c -> d
f (f Int -> Stream Id (f Int)
forall (m :: * -> *) (f :: * -> *).
(Monad m, Shape f) =>
Layout f -> Stream m (Layout f)
streamIndexes f Int
l') (f Int -> Array v f a -> Stream Id a
forall (m :: * -> *) (f :: * -> *) (v :: * -> *) a.
(Monad m, Shape f, Vector v a) =>
Layout f -> Array v f a -> Stream m a
streamSub f Int
l' Array v f a
a1) (f Int -> Array v f b -> Stream Id b
forall (m :: * -> *) (f :: * -> *) (v :: * -> *) a.
(Monad m, Shape f, Vector v a) =>
Layout f -> Array v f a -> Stream m a
streamSub f Int
l' Array v f b
a2) (f Int -> Array v f c -> Stream Id c
forall (m :: * -> *) (f :: * -> *) (v :: * -> *) a.
(Monad m, Shape f, Vector v a) =>
Layout f -> Array v f a -> Stream m a
streamSub f Int
l' Array v f c
a3)) (Int -> Size
Exact (f Int -> Int
forall (f :: * -> *). Shape f => Layout f -> Int
shapeSize f Int
l'))
  where l' :: f Int
l' = f Int -> f Int -> f Int
forall (f :: * -> *). Shape f => Layout f -> Layout f -> Layout f
shapeIntersect (f Int -> f Int -> f Int
forall (f :: * -> *). Shape f => Layout f -> Layout f -> Layout f
shapeIntersect f Int
l1 f Int
l2) f Int
l3
{-# INLINE izipWith3 #-}

------------------------------------------------------------------------
-- Slices
------------------------------------------------------------------------

-- $setup
-- >>> import Debug.SimpleReflect
-- >>> let m = fromListInto_ (V2 3 4) [a,b,c,d,e,f,g,h,i,j,k,l] :: BArray V2 Expr

-- | Indexed traversal over the rows of a matrix. Each row is an
--   efficient 'Data.Vector.Generic.slice' of the original vector.
--
-- >>> traverseOf_ rows print m
-- [a,b,c,d]
-- [e,f,g,h]
-- [i,j,k,l]
rows :: (Vector v a, Vector w b)
     => IndexedTraversal Int (Array v V2 a) (Array w V2 b) (v a) (w b)
rows :: IndexedTraversal Int (Array v V2 a) (Array w V2 b) (v a) (w b)
rows p (v a) (f (w b))
f (Array l :: Layout V2
l@(V2 Int
x Int
y) v a
v) = Layout V2 -> w b -> Array w V2 b
forall (v :: * -> *) (f :: * -> *) a.
Layout f -> v a -> Array v f a
Array Layout V2
l (w b -> Array w V2 b) -> ([w b] -> w b) -> [w b] -> Array w V2 b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [w b] -> w b
forall (v :: * -> *) a. Vector v a => [v a] -> v a
G.concat ([w b] -> Array w V2 b) -> f [w b] -> f (Array w V2 b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Int -> f [w b]
go Int
0 Int
0 where
  go :: Int -> Int -> f [w b]
go Int
i Int
a | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
x    = [w b] -> f [w b]
forall (f :: * -> *) a. Applicative f => a -> f a
pure []
         | Bool
otherwise = (:) (w b -> [w b] -> [w b]) -> f (w b) -> f ([w b] -> [w b])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> p (v a) (f (w b)) -> Int -> v a -> f (w b)
forall i (p :: * -> * -> *) a b.
Indexable i p =>
p a b -> i -> a -> b
indexed p (v a) (f (w b))
f Int
i (Int -> Int -> v a -> v a
forall (v :: * -> *) a. Vector v a => Int -> Int -> v a -> v a
G.slice Int
a Int
y v a
v) f ([w b] -> [w b]) -> f [w b] -> f [w b]
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Int -> Int -> f [w b]
go (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) (Int
aInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
y)
{-# INLINE rows #-}

-- | Affine traversal over a single row in a matrix.
--
-- >>> traverseOf_ rows print $ m & ixRow 1 . each *~ 2
-- [a,b,c,d]
-- [e * 2,f * 2,g * 2,h * 2]
-- [i,j,k,l]
--
--   The row vector should remain the same size to satisfy traversal
--   laws but give reasonable behaviour if the size differs:
--
-- >>> traverseOf_ rows print $ m & ixRow 1 .~ B.fromList [0,1]
-- [a,b,c,d]
-- [0,1,g,h]
-- [i,j,k,l]
--
-- >>> traverseOf_ rows print $ m & ixRow 1 .~ B.fromList [0..100]
-- [a,b,c,d]
-- [0,1,2,3]
-- [i,j,k,l]
ixRow :: Vector v a => Int -> IndexedTraversal' Int (Array v V2 a) (v a)
ixRow :: Int -> IndexedTraversal' Int (Array v V2 a) (v a)
ixRow Int
i p (v a) (f (v a))
f m :: Array v V2 a
m@(Array (l :: Layout V2
l@(V2 Int
x Int
y)) v a
v)
  | Int
y Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
x = Layout V2 -> v a -> Array v V2 a
forall (v :: * -> *) (f :: * -> *) a.
Layout f -> v a -> Array v f a
Array Layout V2
l (v a -> Array v V2 a) -> (v a -> v a) -> v a -> Array v V2 a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. v a -> [(Int, a)] -> v a
forall (v :: * -> *) a. Vector v a => v a -> [(Int, a)] -> v a
G.unsafeUpd v a
v ([(Int, a)] -> v a) -> (v a -> [(Int, a)]) -> v a -> v a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Int] -> [a] -> [(Int, a)]
forall a b. [a] -> [b] -> [(a, b)]
L.zip [Int
a..] ([a] -> [(Int, a)]) -> (v a -> [a]) -> v a -> [(Int, a)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. v a -> [a]
forall (v :: * -> *) a. Vector v a => v a -> [a]
G.toList (v a -> [a]) -> (v a -> v a) -> v a -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> v a -> v a
forall (v :: * -> *) a. Vector v a => Int -> v a -> v a
G.take Int
y (v a -> Array v V2 a) -> f (v a) -> f (Array v V2 a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> p (v a) (f (v a)) -> Int -> v a -> f (v a)
forall i (p :: * -> * -> *) a b.
Indexable i p =>
p a b -> i -> a -> b
indexed p (v a) (f (v a))
f Int
i (Int -> Int -> v a -> v a
forall (v :: * -> *) a. Vector v a => Int -> Int -> v a -> v a
G.slice Int
a Int
y v a
v)
  | Bool
otherwise       = Array v V2 a -> f (Array v V2 a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Array v V2 a
m
  where a :: Int
a  = Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
y
{-# INLINE ixRow #-}

-- | Indexed traversal over the columns of a matrix. Unlike 'rows', each
--   column is a new separate vector.
--
-- >>> traverseOf_ columns print m
-- [a,e,i]
-- [b,f,j]
-- [c,g,k]
-- [d,h,l]
--
-- >>> traverseOf_ rows print $ m & columns . indices odd . each .~ 0
-- [a,0,c,0]
-- [e,0,g,0]
-- [i,0,k,0]
--
--   The vectors should be the same size to be a valid traversal. If the
--   vectors are different sizes, the number of rows in the new array
--   will be the length of the smallest vector.
columns :: (Vector v a, Vector w b)
        => IndexedTraversal Int (Array v V2 a) (Array w V2 b) (v a) (w b)
columns :: IndexedTraversal Int (Array v V2 a) (Array w V2 b) (v a) (w b)
columns p (v a) (f (w b))
f m :: Array v V2 a
m@(Array l :: Layout V2
l@(V2 Int
_ Int
y) v a
_) = Layout V2 -> [w b] -> Array w V2 b
forall (v :: * -> *) a.
Vector v a =>
Layout V2 -> [v a] -> Array v V2 a
transposeConcat Layout V2
l ([w b] -> Array w V2 b) -> f [w b] -> f (Array w V2 b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> f [w b]
go Int
0 where
  go :: Int -> f [w b]
go Int
j | Int
j Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
y    = [w b] -> f [w b]
forall (f :: * -> *) a. Applicative f => a -> f a
pure []
       | Bool
otherwise = (:) (w b -> [w b] -> [w b]) -> f (w b) -> f ([w b] -> [w b])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> p (v a) (f (w b)) -> Int -> v a -> f (w b)
forall i (p :: * -> * -> *) a b.
Indexable i p =>
p a b -> i -> a -> b
indexed p (v a) (f (w b))
f Int
j (Array v V2 a -> Int -> v a
forall (v :: * -> *) a. Vector v a => Array v V2 a -> Int -> v a
getColumn Array v V2 a
m Int
j) f ([w b] -> [w b]) -> f [w b] -> f [w b]
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Int -> f [w b]
go (Int
jInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)
{-# INLINE columns #-}

-- | Affine traversal over a single column in a matrix.
--
-- >>> traverseOf_ rows print $ m & ixColumn 2 . each +~ 1
-- [a,b,c + 1,d]
-- [e,f,g + 1,h]
-- [i,j,k + 1,l]
ixColumn :: Vector v a => Int -> IndexedTraversal' Int (Array v V2 a) (v a)
ixColumn :: Int -> IndexedTraversal' Int (Array v V2 a) (v a)
ixColumn Int
j p (v a) (f (v a))
f m :: Array v V2 a
m@(Array (l :: Layout V2
l@(V2 Int
_ Int
y)) v a
v)
  | Int
j Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Int
j Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
y = Layout V2 -> v a -> Array v V2 a
forall (v :: * -> *) (f :: * -> *) a.
Layout f -> v a -> Array v f a
Array Layout V2
l (v a -> Array v V2 a) -> (v a -> v a) -> v a -> Array v V2 a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. v a -> [(Int, a)] -> v a
forall (v :: * -> *) a. Vector v a => v a -> [(Int, a)] -> v a
G.unsafeUpd v a
v ([(Int, a)] -> v a) -> (v a -> [(Int, a)]) -> v a -> v a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Int] -> [a] -> [(Int, a)]
forall a b. [a] -> [b] -> [(a, b)]
L.zip [Int]
js ([a] -> [(Int, a)]) -> (v a -> [a]) -> v a -> [(Int, a)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. v a -> [a]
forall (v :: * -> *) a. Vector v a => v a -> [a]
G.toList (v a -> [a]) -> (v a -> v a) -> v a -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> v a -> v a
forall (v :: * -> *) a. Vector v a => Int -> v a -> v a
G.take Int
y (v a -> Array v V2 a) -> f (v a) -> f (Array v V2 a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> p (v a) (f (v a)) -> Int -> v a -> f (v a)
forall i (p :: * -> * -> *) a b.
Indexable i p =>
p a b -> i -> a -> b
indexed p (v a) (f (v a))
f Int
j (Array v V2 a -> Int -> v a
forall (v :: * -> *) a. Vector v a => Array v V2 a -> Int -> v a
getColumn Array v V2 a
m Int
j)
  | Bool
otherwise       = Array v V2 a -> f (Array v V2 a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Array v V2 a
m
  where js :: [Int]
js = [Int
j, Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
y .. ]
{-# INLINE ixColumn #-}

getColumn :: Vector v a => Array v V2 a -> Int -> v a
getColumn :: Array v V2 a -> Int -> v a
getColumn (Array (V2 Int
x Int
y) v a
v) Int
j = Int -> (Int -> a) -> v a
forall (v :: * -> *) a. Vector v a => Int -> (Int -> a) -> v a
G.generate Int
x ((Int -> a) -> v a) -> (Int -> a) -> v a
forall a b. (a -> b) -> a -> b
$ \Int
i -> v a -> Int -> a
forall (v :: * -> *) a. Vector v a => v a -> Int -> a
G.unsafeIndex v a
v (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
y Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
j)
{-# INLINE getColumn #-}

transposeConcat :: Vector v a => V2 Int -> [v a] -> Array v V2 a
transposeConcat :: Layout V2 -> [v a] -> Array v V2 a
transposeConcat (V2 Int
_ Int
y) [v a]
vs = Layout V2 -> v a -> Array v V2 a
forall (v :: * -> *) (f :: * -> *) a.
Layout f -> v a -> Array v f a
Array (Int -> Int -> Layout V2
forall a. a -> a -> V2 a
V2 Int
x' Int
y) (v a -> Array v V2 a) -> v a -> Array v V2 a
forall a b. (a -> b) -> a -> b
$ (forall s. ST s (Mutable v s a)) -> v a
forall (v :: * -> *) a.
Vector v a =>
(forall s. ST s (Mutable v s a)) -> v a
G.create ((forall s. ST s (Mutable v s a)) -> v a)
-> (forall s. ST s (Mutable v s a)) -> v a
forall a b. (a -> b) -> a -> b
$ do
  Mutable v s a
mv <- Int -> ST s (Mutable v (PrimState (ST s)) a)
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
Int -> m (v (PrimState m) a)
GM.new (Int
x'Int -> Int -> Int
forall a. Num a => a -> a -> a
*Int
y)
  [v a] -> (Int -> v a -> ST s ()) -> ST s ()
forall i (t :: * -> *) (m :: * -> *) a b.
(FoldableWithIndex i t, Monad m) =>
t a -> (i -> a -> m b) -> m ()
iforM_ [v a]
vs ((Int -> v a -> ST s ()) -> ST s ())
-> (Int -> v a -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
j v a
v ->
    [Int] -> (Int -> ST s ()) -> ST s ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
F.for_ [Int
0..Int
x'Int -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1] ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
i ->
      Mutable v (PrimState (ST s)) a -> Int -> a -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
GM.write Mutable v s a
Mutable v (PrimState (ST s)) a
mv (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
*Int
y Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
j) (v a
v v a -> Int -> a
forall (v :: * -> *) a. Vector v a => v a -> Int -> a
G.! Int
i)
  Mutable v s a -> ST s (Mutable v s a)
forall (m :: * -> *) a. Monad m => a -> m a
return Mutable v s a
mv
  where x' :: Int
x' = [Int] -> Int
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
minimum ([Int] -> Int) -> [Int] -> Int
forall a b. (a -> b) -> a -> b
$ (v a -> Int) -> [v a] -> [Int]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap v a -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
G.length [v a]
vs
{-# INLINE transposeConcat #-}

-- | Traversal over a single plane of a 3D array given a lens onto that
--   plane (like '_xy', '_yz', '_zx').
ixPlane :: Vector v a
        => ALens' (V3 Int) (V2 Int)
        -> Int
        -> IndexedTraversal' Int (Array v V3 a) (Array v V2 a)
ixPlane :: ALens' (V3 Int) (Layout V2)
-> Int -> IndexedTraversal' Int (Array v V3 a) (Array v V2 a)
ixPlane ALens' (V3 Int) (Layout V2)
l32 Int
i p (Array v V2 a) (f (Array v V2 a))
f a :: Array v V3 a
a@(Array V3 Int
l v a
v)
  | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
k = Array v V3 a -> f (Array v V3 a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Array v V3 a
a
  | Bool
otherwise       = V3 Int -> v a -> Array v V3 a
forall (v :: * -> *) (f :: * -> *) a.
Layout f -> v a -> Array v f a
Array V3 Int
l (v a -> Array v V3 a)
-> (Array v V2 a -> v a) -> Array v V2 a -> Array v V3 a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (v a
v v a -> [(Int, a)] -> v a
forall (v :: * -> *) a. Vector v a => v a -> [(Int, a)] -> v a
G.//) ([(Int, a)] -> v a)
-> (Array v V2 a -> [(Int, a)]) -> Array v V2 a -> v a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Int] -> [a] -> [(Int, a)]
forall a b. [a] -> [b] -> [(a, b)]
L.zip [Int]
is ([a] -> [(Int, a)])
-> (Array v V2 a -> [a]) -> Array v V2 a -> [(Int, a)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Getting (Endo [a]) (Array v V2 a) a -> Array v V2 a -> [a]
forall a s. Getting (Endo [a]) s a -> s -> [a]
toListOf Getting (Endo [a]) (Array v V2 a) a
forall (f :: * -> *) (v :: * -> *) a (w :: * -> *) b.
(Shape f, Vector v a, Vector w b) =>
IndexedTraversal (f Int) (Array v f a) (Array w f b) a b
values
                        (Array v V2 a -> Array v V3 a)
-> f (Array v V2 a) -> f (Array v V3 a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> p (Array v V2 a) (f (Array v V2 a))
-> Int -> Array v V2 a -> f (Array v V2 a)
forall i (p :: * -> * -> *) a b.
Indexable i p =>
p a b -> i -> a -> b
indexed p (Array v V2 a) (f (Array v V2 a))
f Int
i (ALens' (V3 Int) (Layout V2) -> Int -> Array v V3 a -> Array v V2 a
forall (v :: * -> *) a.
Vector v a =>
ALens' (V3 Int) (Layout V2) -> Int -> Array v V3 a -> Array v V2 a
getPlane ALens' (V3 Int) (Layout V2)
l32 Int
i Array v V3 a
a)
  where
    is :: [Int]
is = Getting (Endo [Int]) (V3 Int) Int -> V3 Int -> [Int]
forall a s. Getting (Endo [a]) s a -> s -> [a]
toListOf (ALens' (V3 Int) (Layout V2)
-> Lens (V3 Int) (V3 Int) (Layout V2) (Layout V2)
forall s t a b. ALens s t a b -> Lens s t a b
cloneLens ALens' (V3 Int) (Layout V2)
l32 ((Layout V2 -> Const (Endo [Int]) (Layout V2))
 -> V3 Int -> Const (Endo [Int]) (V3 Int))
-> ((Int -> Const (Endo [Int]) Int)
    -> Layout V2 -> Const (Endo [Int]) (Layout V2))
-> Getting (Endo [Int]) (V3 Int) Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Layout V2 -> Const (Endo [Int]) (Layout V2))
-> Layout V2 -> Const (Endo [Int]) (Layout V2)
forall (f :: * -> *).
Shape f =>
IndexedFold Int (Layout f) (Layout f)
shapeIndexes ((Layout V2 -> Const (Endo [Int]) (Layout V2))
 -> Layout V2 -> Const (Endo [Int]) (Layout V2))
-> ((Int -> Const (Endo [Int]) Int)
    -> Layout V2 -> Const (Endo [Int]) (Layout V2))
-> (Int -> Const (Endo [Int]) Int)
-> Layout V2
-> Const (Endo [Int]) (Layout V2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Layout V2 -> Int)
-> (Int -> Const (Endo [Int]) Int)
-> Layout V2
-> Const (Endo [Int]) (Layout V2)
forall (p :: * -> * -> *) (f :: * -> *) s a.
(Profunctor p, Contravariant f) =>
(s -> a) -> Optic' p f s a
to (\Layout V2
x -> V3 Int -> V3 Int -> Int
forall (f :: * -> *). Shape f => Layout f -> Layout f -> Int
shapeToIndex V3 Int
l (V3 Int -> Int) -> V3 Int -> Int
forall a b. (a -> b) -> a -> b
$ Int -> V3 Int
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
i V3 Int -> (V3 Int -> V3 Int) -> V3 Int
forall a b. a -> (a -> b) -> b
& ALens' (V3 Int) (Layout V2)
l32 ALens' (V3 Int) (Layout V2) -> Layout V2 -> V3 Int -> V3 Int
forall s t a b. ALens s t a b -> b -> s -> t
#~ Layout V2
x)) V3 Int
l
    k :: Int
k  = V3 Int -> Int
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
F.sum (V3 Int -> Int) -> V3 Int -> Int
forall a b. (a -> b) -> a -> b
$ V3 Int
l V3 Int -> (V3 Int -> V3 Int) -> V3 Int
forall a b. a -> (a -> b) -> b
& ALens' (V3 Int) (Layout V2)
l32 ALens' (V3 Int) (Layout V2) -> Layout V2 -> V3 Int -> V3 Int
forall s t a b. ALens s t a b -> b -> s -> t
#~ Layout V2
0

-- | Traversal over all planes of 3D array given a lens onto that plane
--   (like '_xy', '_yz', '_zx').
planes :: (Vector v a, Vector w b)
       => ALens' (V3 Int) (V2 Int)
       -> IndexedTraversal Int (Array v V3 a) (Array w V3 b) (Array v V2 a) (Array w V2 b)
planes :: ALens' (V3 Int) (Layout V2)
-> IndexedTraversal
     Int (Array v V3 a) (Array w V3 b) (Array v V2 a) (Array w V2 b)
planes ALens' (V3 Int) (Layout V2)
l32 p (Array v V2 a) (f (Array w V2 b))
f a :: Array v V3 a
a@(Array V3 Int
l v a
_) = V3 Int
-> ALens' (V3 Int) (Layout V2) -> [Array w V2 b] -> Array w V3 b
forall (v :: * -> *) a.
Vector v a =>
V3 Int
-> ALens' (V3 Int) (Layout V2) -> [Array v V2 a] -> Array v V3 a
concatPlanes V3 Int
l ALens' (V3 Int) (Layout V2)
l32 ([Array w V2 b] -> Array w V3 b)
-> f [Array w V2 b] -> f (Array w V3 b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> f [Array w V2 b]
go Int
0 where
  go :: Int -> f [Array w V2 b]
go Int
i | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
k    = [Array w V2 b] -> f [Array w V2 b]
forall (f :: * -> *) a. Applicative f => a -> f a
pure []
       | Bool
otherwise = (:) (Array w V2 b -> [Array w V2 b] -> [Array w V2 b])
-> f (Array w V2 b) -> f ([Array w V2 b] -> [Array w V2 b])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> p (Array v V2 a) (f (Array w V2 b))
-> Int -> Array v V2 a -> f (Array w V2 b)
forall i (p :: * -> * -> *) a b.
Indexable i p =>
p a b -> i -> a -> b
indexed p (Array v V2 a) (f (Array w V2 b))
f Int
i (ALens' (V3 Int) (Layout V2) -> Int -> Array v V3 a -> Array v V2 a
forall (v :: * -> *) a.
Vector v a =>
ALens' (V3 Int) (Layout V2) -> Int -> Array v V3 a -> Array v V2 a
getPlane ALens' (V3 Int) (Layout V2)
l32 Int
i Array v V3 a
a) f ([Array w V2 b] -> [Array w V2 b])
-> f [Array w V2 b] -> f [Array w V2 b]
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Int -> f [Array w V2 b]
go (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)
  k :: Int
k = V3 Int -> Int
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
F.sum (V3 Int -> Int) -> V3 Int -> Int
forall a b. (a -> b) -> a -> b
$ V3 Int
l V3 Int -> (V3 Int -> V3 Int) -> V3 Int
forall a b. a -> (a -> b) -> b
& ALens' (V3 Int) (Layout V2)
l32 ALens' (V3 Int) (Layout V2) -> Layout V2 -> V3 Int -> V3 Int
forall s t a b. ALens s t a b -> b -> s -> t
#~ Layout V2
0
{-# INLINE planes #-}

concatPlanes :: Vector v a => V3 Int -> ALens' (V3 Int) (V2 Int) -> [Array v V2 a] -> Array v V3 a
concatPlanes :: V3 Int
-> ALens' (V3 Int) (Layout V2) -> [Array v V2 a] -> Array v V3 a
concatPlanes V3 Int
l ALens' (V3 Int) (Layout V2)
l32 [Array v V2 a]
as = (forall s. ST s (MArray (Mutable v) V3 s a)) -> Array v V3 a
forall (v :: * -> *) a (f :: * -> *).
Vector v a =>
(forall s. ST s (MArray (Mutable v) f s a)) -> Array v f a
create ((forall s. ST s (MArray (Mutable v) V3 s a)) -> Array v V3 a)
-> (forall s. ST s (MArray (Mutable v) V3 s a)) -> Array v V3 a
forall a b. (a -> b) -> a -> b
$ do
  MArray (Mutable v) V3 s a
arr <- V3 Int -> ST s (MArray (Mutable v) V3 (PrimState (ST s)) a)
forall (m :: * -> *) (f :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, Shape f, MVector v a) =>
Layout f -> m (MArray v f (PrimState m) a)
M.new V3 Int
l
  [Array v V2 a] -> (Int -> Array v V2 a -> ST s ()) -> ST s ()
forall i (t :: * -> *) (m :: * -> *) a b.
(FoldableWithIndex i t, Monad m) =>
t a -> (i -> a -> m b) -> m ()
iforM_ [Array v V2 a]
as ((Int -> Array v V2 a -> ST s ()) -> ST s ())
-> (Int -> Array v V2 a -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
i Array v V2 a
m ->
    IndexedGetting (Layout V2) (Sequenced () (ST s)) (Array v V2 a) a
-> Array v V2 a -> (Layout V2 -> a -> ST s ()) -> ST s ()
forall (m :: * -> *) i r s a.
Monad m =>
IndexedGetting i (Sequenced r m) s a
-> s -> (i -> a -> m r) -> m ()
iforMOf_ IndexedGetting (Layout V2) (Sequenced () (ST s)) (Array v V2 a) a
forall (f :: * -> *) (v :: * -> *) a (w :: * -> *) b.
(Shape f, Vector v a, Vector w b) =>
IndexedTraversal (f Int) (Array v f a) (Array w f b) a b
values Array v V2 a
m ((Layout V2 -> a -> ST s ()) -> ST s ())
-> (Layout V2 -> a -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Layout V2
x a
a -> do
      let w :: V3 Int
w = Int -> V3 Int
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
i V3 Int -> (V3 Int -> V3 Int) -> V3 Int
forall a b. a -> (a -> b) -> b
& ALens' (V3 Int) (Layout V2)
l32 ALens' (V3 Int) (Layout V2) -> Layout V2 -> V3 Int -> V3 Int
forall s t a b. ALens s t a b -> b -> s -> t
#~ Layout V2
x
      MArray (Mutable v) V3 (PrimState (ST s)) a
-> V3 Int -> a -> ST s ()
forall (m :: * -> *) (f :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, Shape f, MVector v a) =>
MArray v f (PrimState m) a -> f Int -> a -> m ()
M.write MArray (Mutable v) V3 s a
MArray (Mutable v) V3 (PrimState (ST s)) a
arr V3 Int
w a
a
  MArray (Mutable v) V3 s a -> ST s (MArray (Mutable v) V3 s a)
forall (m :: * -> *) a. Monad m => a -> m a
return MArray (Mutable v) V3 s a
arr

getPlane :: Vector v a => ALens' (V3 Int) (V2 Int) -> Int -> Array v V3 a -> Array v V2 a
getPlane :: ALens' (V3 Int) (Layout V2) -> Int -> Array v V3 a -> Array v V2 a
getPlane ALens' (V3 Int) (Layout V2)
l32 Int
i Array v V3 a
a = Layout V2 -> (Layout V2 -> a) -> Array v V2 a
forall (f :: * -> *) (v :: * -> *) a.
(Shape f, Vector v a) =>
Layout f -> (Layout f -> a) -> Array v f a
generate (Array v V3 a
a Array v V3 a
-> ALens (Array v V3 a) (Array v V3 a) (Layout V2) (Layout V2)
-> Layout V2
forall s t a b. s -> ALens s t a b -> a
^# (V3 Int -> Pretext (->) (Layout V2) (Layout V2) (V3 Int))
-> Array v V3 a
-> Pretext (->) (Layout V2) (Layout V2) (Array v V3 a)
forall (f :: * -> *) a. HasLayout f a => Lens' a (Layout f)
layout ((V3 Int -> Pretext (->) (Layout V2) (Layout V2) (V3 Int))
 -> Array v V3 a
 -> Pretext (->) (Layout V2) (Layout V2) (Array v V3 a))
-> ALens' (V3 Int) (Layout V2)
-> ALens (Array v V3 a) (Array v V3 a) (Layout V2) (Layout V2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ALens' (V3 Int) (Layout V2)
l32) ((Layout V2 -> a) -> Array v V2 a)
-> (Layout V2 -> a) -> Array v V2 a
forall a b. (a -> b) -> a -> b
$ \Layout V2
x -> Array v V3 a
a Array v V3 a -> V3 Int -> a
forall (f :: * -> *) (v :: * -> *) a.
(Shape f, Vector v a) =>
Array v f a -> f Int -> a
! (Int -> V3 Int
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
i V3 Int -> (V3 Int -> V3 Int) -> V3 Int
forall a b. a -> (a -> b) -> b
& ALens' (V3 Int) (Layout V2)
l32 ALens' (V3 Int) (Layout V2) -> Layout V2 -> V3 Int -> V3 Int
forall s t a b. ALens s t a b -> b -> s -> t
#~ Layout V2
x)

-- | Flatten a plane by reducing a vector in the third dimension to a
--   single value.
flattenPlane :: (Vector v a, Vector w b)
             => ALens' (V3 Int) (V2 Int)
             -> (v a -> b)
             -> Array v V3 a
             -> Array w V2 b
flattenPlane :: ALens' (V3 Int) (Layout V2)
-> (v a -> b) -> Array v V3 a -> Array w V2 b
flattenPlane ALens' (V3 Int) (Layout V2)
l32 v a -> b
f a :: Array v V3 a
a@(Array V3 Int
l v a
_) = Layout V2 -> (Layout V2 -> b) -> Array w V2 b
forall (f :: * -> *) (v :: * -> *) a.
(Shape f, Vector v a) =>
Layout f -> (Layout f -> a) -> Array v f a
generate Layout V2
l' ((Layout V2 -> b) -> Array w V2 b)
-> (Layout V2 -> b) -> Array w V2 b
forall a b. (a -> b) -> a -> b
$ \Layout V2
x -> v a -> b
f (Layout V2 -> v a
getVector Layout V2
x)
  where
    getVector :: Layout V2 -> v a
getVector Layout V2
x = Int -> (Int -> a) -> v a
forall (v :: * -> *) a. Vector v a => Int -> (Int -> a) -> v a
G.generate Int
n ((Int -> a) -> v a) -> (Int -> a) -> v a
forall a b. (a -> b) -> a -> b
$ \Int
i -> Array v V3 a
a Array v V3 a -> V3 Int -> a
forall (f :: * -> *) (v :: * -> *) a.
(Shape f, Vector v a) =>
Array v f a -> f Int -> a
! (Int -> V3 Int
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
i V3 Int -> (V3 Int -> V3 Int) -> V3 Int
forall a b. a -> (a -> b) -> b
& ALens' (V3 Int) (Layout V2)
l32 ALens' (V3 Int) (Layout V2) -> Layout V2 -> V3 Int -> V3 Int
forall s t a b. ALens s t a b -> b -> s -> t
#~ Layout V2
x)
    n :: Int
n  = V3 Int -> Int
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
F.sum (V3 Int -> Int) -> V3 Int -> Int
forall a b. (a -> b) -> a -> b
$ V3 Int
l V3 Int -> (V3 Int -> V3 Int) -> V3 Int
forall a b. a -> (a -> b) -> b
& ALens' (V3 Int) (Layout V2)
l32 ALens' (V3 Int) (Layout V2) -> Layout V2 -> V3 Int -> V3 Int
forall s t a b. ALens s t a b -> b -> s -> t
#~ Layout V2
0
    l' :: Layout V2
l' = V3 Int
l V3 Int -> ALens' (V3 Int) (Layout V2) -> Layout V2
forall s t a b. s -> ALens s t a b -> a
^# ALens' (V3 Int) (Layout V2)
l32
{-# INLINE flattenPlane #-}

-- Ordinals ------------------------------------------------------------

-- | This 'Traversal' should not have any duplicates in the list of
--   indices.
unsafeOrdinals :: (Vector v a, Shape f) => [f Int] -> IndexedTraversal' (f Int) (Array v f a) a
unsafeOrdinals :: [f Int] -> IndexedTraversal' (f Int) (Array v f a) a
unsafeOrdinals [f Int]
is p a (f a)
f (Array f Int
l v a
v) = f Int -> v a -> Array v f a
forall (v :: * -> *) (f :: * -> *) a.
Layout f -> v a -> Array v f a
Array f Int
l (v a -> Array v f a)
-> ([(Int, a)] -> v a) -> [(Int, a)] -> Array v f a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (v a
v v a -> [(Int, a)] -> v a
forall (v :: * -> *) a. Vector v a => v a -> [(Int, a)] -> v a
G.//) ([(Int, a)] -> Array v f a) -> f [(Int, a)] -> f (Array v f a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (f Int -> f (Int, a)) -> [f Int] -> f [(Int, a)]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse f Int -> f (Int, a)
g [f Int]
is
  where g :: f Int -> f (Int, a)
g f Int
x = let i :: Int
i = f Int -> f Int -> Int
forall (f :: * -> *). Shape f => Layout f -> Layout f -> Int
shapeToIndex f Int
l f Int
x in (,) Int
i (a -> (Int, a)) -> f a -> f (Int, a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> p a (f a) -> f Int -> a -> f a
forall i (p :: * -> * -> *) a b.
Indexable i p =>
p a b -> i -> a -> b
indexed p a (f a)
f f Int
x (v a -> Int -> a
forall (v :: * -> *) a. Vector v a => v a -> Int -> a
G.unsafeIndex v a
v Int
i)
{-# INLINE [0] unsafeOrdinals #-}

setOrdinals :: (Indexable (f Int) p, Vector v a, Shape f) => [f Int] -> p a a -> Array v f a -> Array v f a
setOrdinals :: [f Int] -> p a a -> Array v f a -> Array v f a
setOrdinals [f Int]
is p a a
f (Array f Int
l v a
v) = f Int -> v a -> Array v f a
forall (v :: * -> *) (f :: * -> *) a.
Layout f -> v a -> Array v f a
Array f Int
l (v a -> Array v f a) -> v a -> Array v f a
forall a b. (a -> b) -> a -> b
$ v a -> [(Int, a)] -> v a
forall (v :: * -> *) a. Vector v a => v a -> [(Int, a)] -> v a
G.unsafeUpd v a
v ((f Int -> (Int, a)) -> [f Int] -> [(Int, a)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap f Int -> (Int, a)
g [f Int]
is)
  where g :: f Int -> (Int, a)
g f Int
x = let i :: Int
i = f Int -> f Int -> Int
forall (f :: * -> *). Shape f => Layout f -> Layout f -> Int
shapeToIndex f Int
l f Int
x in (,) Int
i (a -> (Int, a)) -> a -> (Int, a)
forall a b. (a -> b) -> a -> b
$ p a a -> f Int -> a -> a
forall i (p :: * -> * -> *) a b.
Indexable i p =>
p a b -> i -> a -> b
indexed p a a
f f Int
x (v a -> Int -> a
forall (v :: * -> *) a. Vector v a => v a -> Int -> a
G.unsafeIndex v a
v Int
i)
{-# INLINE setOrdinals #-}

{-# RULES
"unsafeOrdinals/setOrdinals" forall (is :: [f Int]).
  unsafeOrdinals is = sets (setOrdinals is)
    :: Vector v a => ASetter' (Array v f a) a;
"unsafeOrdinalts/isetOrdintals" forall (is :: [f Int]).
  unsafeOrdinals is = sets (setOrdinals is)
    :: Vector v a => AnIndexedSetter' (f Int) (Array v f a) a
 #-}

-- Mutable -------------------------------------------------------------

-- | O(n) Yield a mutable copy of the immutable vector.
freeze :: (PrimMonad m, Vector v a)
       => MArray (G.Mutable v) f (PrimState m) a -> m (Array v f a)
freeze :: MArray (Mutable v) f (PrimState m) a -> m (Array v f a)
freeze (MArray Layout f
l Mutable v (PrimState m) a
mv) = Layout f -> v a -> Array v f a
forall (v :: * -> *) (f :: * -> *) a.
Layout f -> v a -> Array v f a
Array Layout f
l (v a -> Array v f a) -> m (v a) -> m (Array v f a)
forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
`liftM` Mutable v (PrimState m) a -> m (v a)
forall (m :: * -> *) (v :: * -> *) a.
(PrimMonad m, Vector v a) =>
Mutable v (PrimState m) a -> m (v a)
G.freeze Mutable v (PrimState m) a
mv
{-# INLINE freeze #-}

-- | O(n) Yield an immutable copy of the mutable array.
thaw :: (PrimMonad m, Vector v a)
     => Array v f a -> m (MArray (G.Mutable v) f (PrimState m) a)
thaw :: Array v f a -> m (MArray (Mutable v) f (PrimState m) a)
thaw (Array Layout f
l v a
v) = Layout f
-> Mutable v (PrimState m) a
-> MArray (Mutable v) f (PrimState m) a
forall (v :: * -> * -> *) (l :: * -> *) s a.
Layout l -> v s a -> MArray v l s a
MArray Layout f
l (Mutable v (PrimState m) a -> MArray (Mutable v) f (PrimState m) a)
-> m (Mutable v (PrimState m) a)
-> m (MArray (Mutable v) f (PrimState m) a)
forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
`liftM` v a -> m (Mutable v (PrimState m) a)
forall (m :: * -> *) (v :: * -> *) a.
(PrimMonad m, Vector v a) =>
v a -> m (Mutable v (PrimState m) a)
G.thaw v a
v
{-# INLINE thaw #-}

------------------------------------------------------------------------
-- Delayed
------------------------------------------------------------------------

-- | Isomorphism between an array and its delayed representation.
--   Conversion to the array is done in parallel.
delayed :: (Vector v a, Vector w b, Shape f, Shape g)
        => Iso (Array v f a) (Array w g b) (Delayed f a) (Delayed g b)
delayed :: Iso (Array v f a) (Array w g b) (Delayed f a) (Delayed g b)
delayed = (Array v f a -> Delayed f a)
-> (Delayed g b -> Array w g b)
-> Iso (Array v f a) (Array w g b) (Delayed f a) (Delayed g b)
forall s a b t. (s -> a) -> (b -> t) -> Iso s t a b
iso Array v f a -> Delayed f a
forall (v :: * -> *) a (f :: * -> *).
(Vector v a, Shape f) =>
Array v f a -> Delayed f a
delay Delayed g b -> Array w g b
forall (v :: * -> *) a (f :: * -> *).
(Vector v a, Shape f) =>
Delayed f a -> Array v f a
manifest
{-# INLINE delayed #-}

-- | Isomorphism between an array and its delayed representation.
--   Conversion to the array is done in parallel.
seqDelayed :: (Vector v a, Vector w b, Shape f, Shape g)
        => Iso (Array v f a) (Array w g b) (Delayed f a) (Delayed g b)
seqDelayed :: Iso (Array v f a) (Array w g b) (Delayed f a) (Delayed g b)
seqDelayed = (Array v f a -> Delayed f a)
-> (Delayed g b -> Array w g b)
-> Iso (Array v f a) (Array w g b) (Delayed f a) (Delayed g b)
forall s a b t. (s -> a) -> (b -> t) -> Iso s t a b
iso Array v f a -> Delayed f a
forall (v :: * -> *) a (f :: * -> *).
(Vector v a, Shape f) =>
Array v f a -> Delayed f a
delay Delayed g b -> Array w g b
forall (v :: * -> *) a (f :: * -> *).
(Vector v a, Shape f) =>
Delayed f a -> Array v f a
seqManifest
{-# INLINE seqDelayed #-}

-- | Sequential manifestation of a delayed array.
seqManifest :: (Vector v a, Shape f) => Delayed f a -> Array v f a
seqManifest :: Delayed f a -> Array v f a
seqManifest (Delayed Layout f
l Layout f -> a
f) = Layout f -> (Layout f -> a) -> Array v f a
forall (f :: * -> *) (v :: * -> *) a.
(Shape f, Vector v a) =>
Layout f -> (Layout f -> a) -> Array v f a
generate Layout f
l Layout f -> a
f
{-# INLINE seqManifest #-}

-- | 'manifest' an array to a 'UArray' and delay again. See
--   "Data.Dense.Boxed" or "Data.Dense.Storable" to 'affirm' for other
--   types of arrays.
affirm :: (Shape f, U.Unbox a) => Delayed f a -> Delayed f a
affirm :: Delayed f a -> Delayed f a
affirm = Array Vector f a -> Delayed f a
forall (v :: * -> *) a (f :: * -> *).
(Vector v a, Shape f) =>
Array v f a -> Delayed f a
delay (Array Vector f a -> Delayed f a)
-> (Delayed f a -> Array Vector f a) -> Delayed f a -> Delayed f a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (forall a (f :: * -> *).
(Unbox a, Shape f) =>
Delayed f a -> Array Vector f a
forall (v :: * -> *) a (f :: * -> *).
(Vector v a, Shape f) =>
Delayed f a -> Array v f a
manifest :: (U.Unbox a, Shape f) => Delayed f a -> UArray f a)
{-# INLINE affirm #-}

-- | 'seqManifest' an array to a 'UArray' and delay again. See
--   "Data.Dense.Boxed" or "Data.Dense.Storable" to 'affirm' for other
--   types of arrays.
seqAffirm :: (Shape f, U.Unbox a) => Delayed f a -> Delayed f a
seqAffirm :: Delayed f a -> Delayed f a
seqAffirm = Array Vector f a -> Delayed f a
forall (v :: * -> *) a (f :: * -> *).
(Vector v a, Shape f) =>
Array v f a -> Delayed f a
delay (Array Vector f a -> Delayed f a)
-> (Delayed f a -> Array Vector f a) -> Delayed f a -> Delayed f a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (forall a (f :: * -> *).
(Unbox a, Shape f) =>
Delayed f a -> Array Vector f a
forall (v :: * -> *) a (f :: * -> *).
(Vector v a, Shape f) =>
Delayed f a -> Array v f a
seqManifest :: (U.Unbox a, Shape f) => Delayed f a -> UArray f a)
{-# INLINE seqAffirm #-}

------------------------------------------------------------------------
-- Focused
------------------------------------------------------------------------

-- | Focus on a particular element of a delayed array.
focusOn :: f Int -> Delayed f a -> Focused f a
focusOn :: f Int -> Delayed f a -> Focused f a
focusOn = f Int -> Delayed f a -> Focused f a
forall (f :: * -> *) a. f Int -> Delayed f a -> Focused f a
Focused -- XXX do range checking
{-# INLINE focusOn #-}

-- | Discard the focus to retrieve the delayed array.
unfocus :: Focused f a -> Delayed f a
unfocus :: Focused f a -> Delayed f a
unfocus (Focused f Int
_ Delayed f a
d) = Delayed f a
d
{-# INLINE unfocus #-}

-- | Indexed lens onto the delayed array, indexed at the focus.
unfocused :: IndexedLens (f Int) (Focused f a) (Focused f b) (Delayed f a) (Delayed f b)
unfocused :: p (Delayed f a) (f (Delayed f b)) -> Focused f a -> f (Focused f b)
unfocused p (Delayed f a) (f (Delayed f b))
f (Focused f Int
x Delayed f a
d) = f Int -> Delayed f b -> Focused f b
forall (f :: * -> *) a. f Int -> Delayed f a -> Focused f a
Focused f Int
x (Delayed f b -> Focused f b) -> f (Delayed f b) -> f (Focused f b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> p (Delayed f a) (f (Delayed f b))
-> f Int -> Delayed f a -> f (Delayed f b)
forall i (p :: * -> * -> *) a b.
Indexable i p =>
p a b -> i -> a -> b
indexed p (Delayed f a) (f (Delayed f b))
f f Int
x Delayed f a
d
{-# INLINE unfocused #-}

-- | Modify a 'Delayed' array by extracting a value from a 'Focused'
--   each point.
extendFocus :: Shape f => (Focused f a -> b) -> Delayed f a -> Delayed f b
extendFocus :: (Focused f a -> b) -> Delayed f a -> Delayed f b
extendFocus Focused f a -> b
f = Focused f b -> Delayed f b
forall (f :: * -> *) a. Focused f a -> Delayed f a
unfocus (Focused f b -> Delayed f b)
-> (Delayed f a -> Focused f b) -> Delayed f a -> Delayed f b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Focused f a -> b) -> Focused f a -> Focused f b
forall (w :: * -> *) a b. Comonad w => (w a -> b) -> w a -> w b
extend Focused f a -> b
f (Focused f a -> Focused f b)
-> (Delayed f a -> Focused f a) -> Delayed f a -> Focused f b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. f Int -> Delayed f a -> Focused f a
forall (f :: * -> *) a. f Int -> Delayed f a -> Focused f a
focusOn f Int
forall (f :: * -> *) a. (Additive f, Num a) => f a
zero
{-# INLINE extendFocus #-}

-- | Lens onto the position of a 'ComonadStore'.
--
-- @
-- 'locale' :: 'Lens'' ('Focused' l a) (l 'Int')
-- @
locale :: ComonadStore s w => Lens' (w a) s
locale :: Lens' (w a) s
locale s -> f s
f w a
w = (s -> w a -> w a
forall s (w :: * -> *) a. ComonadStore s w => s -> w a -> w a
`seek` w a
w) (s -> w a) -> f s -> f (w a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> s -> f s
f (w a -> s
forall s (w :: * -> *) a. ComonadStore s w => w a -> s
pos w a
w)
{-# INLINE locale #-}

-- | Focus on a neighbouring element, relative to the current focus.
shiftFocus :: Applicative f => f Int -> Focused f a -> Focused f a
shiftFocus :: f Int -> Focused f a -> Focused f a
shiftFocus f Int
dx (Focused f Int
x d :: Delayed f a
d@(Delayed f Int
l f Int -> a
_)) = f Int -> Delayed f a -> Focused f a
forall (f :: * -> *) a. f Int -> Delayed f a -> Focused f a
Focused f Int
x' Delayed f a
d
  where
    x' :: f Int
x' = Int -> Int -> Int -> Int
forall a. (Num a, Ord a) => a -> a -> a -> a
f (Int -> Int -> Int -> Int) -> f Int -> f (Int -> Int -> Int)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> f Int
l f (Int -> Int -> Int) -> f Int -> f (Int -> Int)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> f Int
x f (Int -> Int) -> f Int -> f Int
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> f Int
dx
    f :: a -> a -> a -> a
f a
k a
i a
di
      | a
i' a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
0    = a
k a -> a -> a
forall a. Num a => a -> a -> a
+ a
i'
      | a
i' a -> a -> Bool
forall a. Ord a => a -> a -> Bool
>= a
k   = a
i' a -> a -> a
forall a. Num a => a -> a -> a
- a
k
      | Bool
otherwise = a
i'
      where i' :: a
i' = a
i a -> a -> a
forall a. Num a => a -> a -> a
+ a
di
{-# INLINE shiftFocus #-}

-- Boundary conditions -------------------------------------------------

-- | The boundary condition used for indexing relative elements in a
--   'Focused'.
data Boundary
  = Clamp  -- ^ clamp coordinates to the extent of the array
  | Mirror -- ^ mirror coordinates beyond the array extent
  | Wrap   -- ^ wrap coordinates around on each dimension
  deriving (Int -> Boundary -> [Char] -> [Char]
[Boundary] -> [Char] -> [Char]
Boundary -> [Char]
(Int -> Boundary -> [Char] -> [Char])
-> (Boundary -> [Char])
-> ([Boundary] -> [Char] -> [Char])
-> Show Boundary
forall a.
(Int -> a -> [Char] -> [Char])
-> (a -> [Char]) -> ([a] -> [Char] -> [Char]) -> Show a
showList :: [Boundary] -> [Char] -> [Char]
$cshowList :: [Boundary] -> [Char] -> [Char]
show :: Boundary -> [Char]
$cshow :: Boundary -> [Char]
showsPrec :: Int -> Boundary -> [Char] -> [Char]
$cshowsPrec :: Int -> Boundary -> [Char] -> [Char]
Show, ReadPrec [Boundary]
ReadPrec Boundary
Int -> ReadS Boundary
ReadS [Boundary]
(Int -> ReadS Boundary)
-> ReadS [Boundary]
-> ReadPrec Boundary
-> ReadPrec [Boundary]
-> Read Boundary
forall a.
(Int -> ReadS a)
-> ReadS [a] -> ReadPrec a -> ReadPrec [a] -> Read a
readListPrec :: ReadPrec [Boundary]
$creadListPrec :: ReadPrec [Boundary]
readPrec :: ReadPrec Boundary
$creadPrec :: ReadPrec Boundary
readList :: ReadS [Boundary]
$creadList :: ReadS [Boundary]
readsPrec :: Int -> ReadS Boundary
$creadsPrec :: Int -> ReadS Boundary
Read, Typeable)

-- Peeking -------------------------------------------------------------

-- | Index a focused using a 'Boundary' condition.
peekB :: Shape f => Boundary -> f Int -> Focused f a -> a
peekB :: Boundary -> f Int -> Focused f a -> a
peekB = \Boundary
b f Int
x -> Boundary -> (f Int -> f Int) -> Focused f a -> a
forall (f :: * -> *) a.
Shape f =>
Boundary -> (f Int -> f Int) -> Focused f a -> a
peeksB Boundary
b (f Int -> f Int -> f Int
forall a b. a -> b -> a
const f Int
x)
{-# INLINE peekB #-}

-- | Index an element relative to the current focus using a 'Boundary'
--   condition.
peekRelativeB :: Shape f => Boundary -> f Int -> Focused f a -> a
peekRelativeB :: Boundary -> f Int -> Focused f a -> a
peekRelativeB = \Boundary
b f Int
i -> Boundary -> (f Int -> f Int) -> Focused f a -> a
forall (f :: * -> *) a.
Shape f =>
Boundary -> (f Int -> f Int) -> Focused f a -> a
peeksB Boundary
b (f Int -> f Int -> f Int
forall (f :: * -> *) a. (Additive f, Num a) => f a -> f a -> f a
^+^ f Int
i)
{-# INLINE peekRelativeB #-}

-- | Index an element by applying a function the current position, using
--   a boundary condition.
peeksB :: Shape f => Boundary -> (f Int -> f Int) -> Focused f a -> a
peeksB :: Boundary -> (f Int -> f Int) -> Focused f a -> a
peeksB = \case
  Boundary
Clamp  -> (f Int -> f Int) -> Focused f a -> a
forall (f :: * -> *) a.
Shape f =>
(f Int -> f Int) -> Focused f a -> a
clampPeeks
  Boundary
Wrap   -> (f Int -> f Int) -> Focused f a -> a
forall (f :: * -> *) a.
Shape f =>
(f Int -> f Int) -> Focused f a -> a
wrapPeeks
  Boundary
Mirror -> (f Int -> f Int) -> Focused f a -> a
forall (f :: * -> *) a.
Shape f =>
(f Int -> f Int) -> Focused f a -> a
mirrorPeeks
{-# INLINE peeksB #-}

-- After much testing, this seems to be the most reliable method to get
-- stencilSum to inline properly.

-- Wrap

wrapPeeks :: Shape f => (f Int -> f Int) -> Focused f a -> a
wrapPeeks :: (f Int -> f Int) -> Focused f a -> a
wrapPeeks f Int -> f Int
f (Focused f Int
x (Delayed f Int
l f Int -> a
ixF)) = f Int -> a
ixF (f Int -> a) -> f Int -> a
forall a b. (a -> b) -> a -> b
$! f Int -> f Int -> f Int
forall (f :: * -> *). Shape f => Layout f -> Layout f -> Layout f
wrapIndex f Int
l (f Int -> f Int
f f Int
x)
{-# INLINE wrapPeeks #-}

wrapIndex :: Shape f => Layout f -> f Int -> f Int
wrapIndex :: Layout f -> Layout f -> Layout f
wrapIndex !Layout f
l !Layout f
x = (Int -> Int -> Int) -> Layout f -> Layout f -> Layout f
forall (f :: * -> *) a b c.
Additive f =>
(a -> b -> c) -> f a -> f b -> f c
liftI2 Int -> Int -> Int
forall a. (Ord a, Num a) => a -> a -> a
f Layout f
l Layout f
x where
  f :: a -> a -> a
f a
n a
i
    | a
i a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
0     = a
n a -> a -> a
forall a. Num a => a -> a -> a
+ a
i
    | a
i a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
n     = a
i
    | Bool
otherwise = a
i a -> a -> a
forall a. Num a => a -> a -> a
- a
n
{-# INLINE wrapIndex #-}

-- Clamp

clampPeeks :: Shape f => (f Int -> f Int) -> Focused f a -> a
clampPeeks :: (f Int -> f Int) -> Focused f a -> a
clampPeeks f Int -> f Int
f (Focused f Int
x (Delayed f Int
l f Int -> a
ixF)) = f Int -> a
ixF (f Int -> a) -> f Int -> a
forall a b. (a -> b) -> a -> b
$! f Int -> f Int -> f Int
forall (f :: * -> *). Shape f => Layout f -> Layout f -> Layout f
clampIndex f Int
l (f Int -> f Int
f f Int
x)
{-# INLINE clampPeeks #-}

clampIndex :: Shape f => Layout f -> f Int -> f Int
clampIndex :: Layout f -> Layout f -> Layout f
clampIndex !Layout f
l !Layout f
x = (Int -> Int -> Int) -> Layout f -> Layout f -> Layout f
forall (f :: * -> *) a b c.
Additive f =>
(a -> b -> c) -> f a -> f b -> f c
liftI2 Int -> Int -> Int
forall a. (Ord a, Num a) => a -> a -> a
f Layout f
l Layout f
x where
  f :: p -> p -> p
f p
n p
i
    | p
i p -> p -> Bool
forall a. Ord a => a -> a -> Bool
< p
0     = p
0
    | p
i p -> p -> Bool
forall a. Ord a => a -> a -> Bool
>= p
n    = p
n p -> p -> p
forall a. Num a => a -> a -> a
- p
1
    | Bool
otherwise = p
i
{-# INLINE clampIndex #-}

-- Mirror

mirrorPeeks :: Shape f => (f Int -> f Int) -> Focused f a -> a
mirrorPeeks :: (f Int -> f Int) -> Focused f a -> a
mirrorPeeks f Int -> f Int
f (Focused f Int
x (Delayed f Int
l f Int -> a
ixF)) = f Int -> a
ixF (f Int -> a) -> f Int -> a
forall a b. (a -> b) -> a -> b
$! f Int -> f Int -> f Int
forall (f :: * -> *). Shape f => Layout f -> Layout f -> Layout f
mirrorIndex f Int
l (f Int -> f Int
f f Int
x)
{-# INLINE mirrorPeeks #-}

mirrorIndex :: Shape f => Layout f -> f Int -> f Int
mirrorIndex :: Layout f -> Layout f -> Layout f
mirrorIndex !Layout f
l !Layout f
x = (Int -> Int -> Int) -> Layout f -> Layout f -> Layout f
forall (f :: * -> *) a b c.
Additive f =>
(a -> b -> c) -> f a -> f b -> f c
liftI2 Int -> Int -> Int
forall a. (Ord a, Num a) => a -> a -> a
f Layout f
l Layout f
x where
  f :: a -> a -> a
f a
n a
i
    | a
i a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
0     = - a
i
    | a
i a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
n     = a
i
    | Bool
otherwise = a
i a -> a -> a
forall a. Num a => a -> a -> a
- a
n
{-# INLINE mirrorIndex #-}