{-# LANGUAGE DefaultSignatures      #-}
{-# LANGUAGE DeriveDataTypeable     #-}
{-# LANGUAGE FlexibleContexts       #-}
{-# LANGUAGE FlexibleInstances      #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE LambdaCase             #-}
{-# LANGUAGE MultiParamTypeClasses  #-}
{-# LANGUAGE Rank2Types             #-}
{-# LANGUAGE TypeFamilies           #-}
{-# LANGUAGE UndecidableInstances   #-}
-----------------------------------------------------------------------------
-- |
-- Module      :  Data.Dense.Mutable
-- Copyright   :  (c) Christopher Chalmers
-- License     :  BSD3
--
-- Maintainer  :  Christopher Chalmers
-- Stability   :  provisional
-- Portability :  non-portable
--
-- This module provides a class for types that can be converted to and
-- from linear indexes.
--
-- The default instances are defined in row-major order.
-----------------------------------------------------------------------------
module Data.Dense.Index
  ( -- * Shape class
    Layout
  , Shape (..)
  , indexIso
  , shapeIndexes
  , shapeIndexesFrom
  , shapeIndexesBetween

    -- * HasLayout
  , HasLayout (..)
  , extent
  , size
  , indexes
  , indexesBetween
  , indexesFrom

    -- * Exceptions

    -- (* Bounds checking
  , ArrayException (IndexOutOfBounds)
  , _IndexOutOfBounds
  , boundsCheck

    -- (* Size missmatch
  , SizeMissmatch (..)
  , AsSizeMissmatch (..)
  , sizeMissmatch

    -- * Utilities
  , showShape
  ) where

import           Control.Applicative
import           Control.Exception
import           Control.Exception.Lens
import           Control.Lens
import           Control.Lens.Internal.Getter
import           Data.Foldable                as F
import           Data.Typeable

import           Data.Functor.Classes
import           Data.Traversable
import           Linear

-- | A 'Layout' is the full size of an array. This alias is used to help
--   distinguish between the layout of an array and an index (usually
--   just @l Int@) in a type signature.
type Layout f = f Int

------------------------------------------------------------------------
-- Shape class
------------------------------------------------------------------------

-- | Class for types that can be converted to and from linear indexes.
class (Eq1 f, Additive f, Traversable f) => Shape f where
  -- | Convert a shape to its linear index using the 'Layout'.
  shapeToIndex :: Layout f -> f Int -> Int
  shapeToIndex Layout f
l Layout f
x = ((Int, Int) -> Int -> Int) -> Int -> f (Int, Int) -> Int
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
F.foldr (\(Int
e, Int
a) Int
k -> Int
a Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
eInt -> Int -> Int
forall a. Num a => a -> a -> a
*Int
k) Int
0 ((Int -> Int -> (Int, Int)) -> Layout f -> Layout f -> f (Int, Int)
forall (f :: * -> *) a b c.
Additive f =>
(a -> b -> c) -> f a -> f b -> f c
liftI2 (,) Layout f
l Layout f
x)
  {-# INLINE shapeToIndex #-}

  -- | Convert a linear index to a shape the 'Layout'.
  shapeFromIndex :: Layout f -> Int -> f Int
  shapeFromIndex Layout f
l Int
i = (Int, Layout f) -> Layout f
forall a b. (a, b) -> b
snd ((Int, Layout f) -> Layout f) -> (Int, Layout f) -> Layout f
forall a b. (a -> b) -> a -> b
$ (Int -> Int -> (Int, Int)) -> Int -> Layout f -> (Int, Layout f)
forall (t :: * -> *) a b c.
Traversable t =>
(a -> b -> (a, c)) -> a -> t b -> (a, t c)
mapAccumL Int -> Int -> (Int, Int)
forall a. Integral a => a -> a -> (a, a)
quotRem Int
i Layout f
l
  {-# INLINE shapeFromIndex #-}

  -- | Calculate the intersection of two shapes.
  shapeIntersect :: Layout f -> Layout f -> Layout f
  shapeIntersect = (Int -> Int -> Int) -> Layout f -> Layout f -> Layout f
forall (f :: * -> *) a.
Additive f =>
(a -> a -> a) -> f a -> f a -> f a
liftU2 Int -> Int -> Int
forall a. Ord a => a -> a -> a
min
  {-# INLINE shapeIntersect #-}

  -- | Increment a shape by one. It is assumed that the provided index
  --   is 'inRange'.
  unsafeShapeStep :: Layout f -> f Int -> f Int
  unsafeShapeStep Layout f
l
    = Layout f -> Int -> Layout f
forall (f :: * -> *). Shape f => Layout f -> Int -> Layout f
shapeFromIndex Layout f
l
    (Int -> Layout f) -> (Layout f -> Int) -> Layout f -> Layout f
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)
    (Int -> Int) -> (Layout f -> Int) -> Layout f -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Layout f -> Layout f -> Int
forall (f :: * -> *). Shape f => Layout f -> Layout f -> Int
shapeToIndex Layout f
l
  {-# INLINE unsafeShapeStep #-}

  -- | Increment a shape by one. It is assumed that the provided index
  --   is 'inRange'.
  shapeStep :: Layout f -> f Int -> Maybe (f Int)
  shapeStep Layout f
l = (Int -> Layout f) -> Maybe Int -> Maybe (Layout f)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Layout f -> Int -> Layout f
forall (f :: * -> *). Shape f => Layout f -> Int -> Layout f
shapeFromIndex Layout f
l)
              (Maybe Int -> Maybe (Layout f))
-> (Layout f -> Maybe Int) -> Layout f -> Maybe (Layout f)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int -> Bool) -> Int -> Maybe Int
forall (f :: * -> *) a. Alternative f => (a -> Bool) -> a -> f a
guardPure (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Layout f -> Int
forall (f :: * -> *). Shape f => Layout f -> Int
shapeSize Layout f
l)
              (Int -> Maybe Int) -> (Layout f -> Int) -> Layout f -> Maybe Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)
              (Int -> Int) -> (Layout f -> Int) -> Layout f -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Layout f -> Layout f -> Int
forall (f :: * -> *). Shape f => Layout f -> Layout f -> Int
shapeToIndex Layout f
l
  {-# INLINE shapeStep #-}

  -- | Increment a shape by one between the two bounds
  shapeStepBetween :: f Int -> Layout f -> f Int -> Maybe (f Int)
  shapeStepBetween Layout f
a Layout f
l = (Layout f -> Layout f) -> Maybe (Layout f) -> Maybe (Layout f)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Layout f -> Layout f -> Layout f
forall (f :: * -> *) a. (Additive f, Num a) => f a -> f a -> f a
^+^ Layout f
a) (Maybe (Layout f) -> Maybe (Layout f))
-> (Layout f -> Maybe (Layout f)) -> Layout f -> Maybe (Layout f)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Layout f -> Layout f -> Maybe (Layout f)
forall (f :: * -> *).
Shape f =>
Layout f -> Layout f -> Maybe (Layout f)
shapeStep Layout f
l (Layout f -> Maybe (Layout f))
-> (Layout f -> Layout f) -> Layout f -> Maybe (Layout f)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Layout f -> Layout f -> Layout f
forall (f :: * -> *) a. (Additive f, Num a) => f a -> f a -> f a
^-^ Layout f
a)
  {-# INLINE shapeStepBetween #-}

  -- | @inRange ex i@ checks @i < ex@ for every coordinate of @f@.
  shapeInRange :: Layout f -> f Int -> Bool
  shapeInRange Layout f
l Layout f
i = f Bool -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
F.and (f Bool -> Bool) -> f Bool -> Bool
forall a b. (a -> b) -> a -> b
$ (Int -> Int -> Bool) -> Layout f -> Layout f -> f Bool
forall (f :: * -> *) a b c.
Additive f =>
(a -> b -> c) -> f a -> f b -> f c
liftI2 (\Int
ii Int
li -> Int
ii Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Int
ii Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
li) Layout f
i Layout f
l
  {-# INLINE shapeInRange #-}

  -- | The number of elements in a shape.
  shapeSize :: Layout f -> Int
  shapeSize = Layout f -> Int
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
F.product
  {-# INLINE shapeSize #-}

guardPure :: Alternative f => (a -> Bool) -> a -> f a
guardPure :: (a -> Bool) -> a -> f a
guardPure a -> Bool
p a
a = if a -> Bool
p a
a then a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
a else f a
forall (f :: * -> *) a. Alternative f => f a
empty
{-# INLINE guardPure #-}

instance Shape V0

instance Shape V1 where
  {-# INLINE shapeToIndex #-}
  {-# INLINE shapeFromIndex #-}
  {-# INLINE shapeIntersect #-}
  {-# INLINE shapeStep #-}
  {-# INLINE shapeInRange #-}
  shapeToIndex :: Layout V1 -> Layout V1 -> Int
shapeToIndex Layout V1
_ (V1 Int
i) = Int
i
  shapeFromIndex :: Layout V1 -> Int -> Layout V1
shapeFromIndex Layout V1
_ Int
i = Int -> Layout V1
forall a. a -> V1 a
V1 Int
i
  shapeIntersect :: Layout V1 -> Layout V1 -> Layout V1
shapeIntersect = Layout V1 -> Layout V1 -> Layout V1
forall a. Ord a => a -> a -> a
min
  shapeStep :: Layout V1 -> Layout V1 -> Maybe (Layout V1)
shapeStep Layout V1
l = (Layout V1 -> Bool) -> Layout V1 -> Maybe (Layout V1)
forall (f :: * -> *) a. Alternative f => (a -> Bool) -> a -> f a
guardPure (Layout V1 -> Layout V1 -> Bool
forall (f :: * -> *). Shape f => Layout f -> Layout f -> Bool
shapeInRange Layout V1
l) (Layout V1 -> Maybe (Layout V1))
-> (Layout V1 -> Layout V1) -> Layout V1 -> Maybe (Layout V1)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Layout V1 -> Layout V1 -> Layout V1
forall a. Num a => a -> a -> a
+Layout V1
1)
  shapeStepBetween :: Layout V1 -> Layout V1 -> Layout V1 -> Maybe (Layout V1)
shapeStepBetween Layout V1
_a Layout V1
b Layout V1
i = (Layout V1 -> Bool) -> Layout V1 -> Maybe (Layout V1)
forall (f :: * -> *) a. Alternative f => (a -> Bool) -> a -> f a
guardPure (Layout V1 -> Layout V1 -> Bool
forall a. Ord a => a -> a -> Bool
> Layout V1
b) Layout V1
i'
    where i' :: Layout V1
i' = Layout V1
i Layout V1 -> Layout V1 -> Layout V1
forall a. Num a => a -> a -> a
+ Layout V1
1
  shapeInRange :: Layout V1 -> Layout V1 -> Bool
shapeInRange Layout V1
m Layout V1
i = Layout V1
i Layout V1 -> Layout V1 -> Bool
forall a. Ord a => a -> a -> Bool
>= Layout V1
0 Bool -> Bool -> Bool
&& Layout V1
i Layout V1 -> Layout V1 -> Bool
forall a. Ord a => a -> a -> Bool
< Layout V1
m

instance Shape V2 where
  shapeToIndex :: Layout V2 -> Layout V2 -> Int
shapeToIndex (V2 Int
x Int
_y) (V2 Int
i Int
j) = Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
xInt -> Int -> Int
forall a. Num a => a -> a -> a
*Int
j
  {-# INLINE shapeToIndex #-}

  shapeFromIndex :: Layout V2 -> Int -> Layout V2
shapeFromIndex (V2 Int
x Int
_y) Int
n = Int -> Int -> Layout V2
forall a. a -> a -> V2 a
V2 Int
i Int
j
    where (Int
j, Int
i) = Int
n Int -> Int -> (Int, Int)
forall a. Integral a => a -> a -> (a, a)
`quotRem` Int
x
  {-# INLINE shapeFromIndex #-}

  shapeStep :: Layout V2 -> Layout V2 -> Maybe (Layout V2)
shapeStep (V2 Int
x Int
y) (V2 Int
i Int
j)
    | Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
x  = Layout V2 -> Maybe (Layout V2)
forall a. a -> Maybe a
Just (Int -> Int -> Layout V2
forall a. a -> a -> V2 a
V2 (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)  Int
j     )
    | Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
y  = Layout V2 -> Maybe (Layout V2)
forall a. a -> Maybe a
Just (Int -> Int -> Layout V2
forall a. a -> a -> V2 a
V2      Int
0  (Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1))
    | Bool
otherwise  = Maybe (Layout V2)
forall a. Maybe a
Nothing
  {-# INLINE shapeStep #-}

  unsafeShapeStep :: Layout V2 -> Layout V2 -> Layout V2
unsafeShapeStep (V2 Int
x Int
_y) (V2 Int
i Int
j)
    | Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
x  = Int -> Int -> Layout V2
forall a. a -> a -> V2 a
V2 (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)  Int
j
    | Bool
otherwise  = Int -> Int -> Layout V2
forall a. a -> a -> V2 a
V2      Int
0  (Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
  {-# INLINE unsafeShapeStep #-}

  shapeStepBetween :: Layout V2 -> Layout V2 -> Layout V2 -> Maybe (Layout V2)
shapeStepBetween (V2 Int
ia Int
_ja) (V2 Int
ib Int
jb) (V2 Int
i Int
j)
    | Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
ib = Layout V2 -> Maybe (Layout V2)
forall a. a -> Maybe a
Just (Int -> Int -> Layout V2
forall a. a -> a -> V2 a
V2 (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)   Int
j     )
    | Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
jb = Layout V2 -> Maybe (Layout V2)
forall a. a -> Maybe a
Just (Int -> Int -> Layout V2
forall a. a -> a -> V2 a
V2      Int
ia  (Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1))
    | Bool
otherwise  = Maybe (Layout V2)
forall a. Maybe a
Nothing
  {-# INLINE shapeStepBetween #-}

instance Shape V3 where
  shapeToIndex :: Layout V3 -> Layout V3 -> Int
shapeToIndex (V3 Int
x Int
y Int
_z) (V3 Int
i Int
j Int
k) = Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
xInt -> Int -> Int
forall a. Num a => a -> a -> a
*(Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
yInt -> Int -> Int
forall a. Num a => a -> a -> a
*Int
k)
  {-# INLINE shapeToIndex #-}

  shapeStep :: Layout V3 -> Layout V3 -> Maybe (Layout V3)
shapeStep (V3 Int
x Int
y Int
z) (V3 Int
i Int
j Int
k)
    | Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
z  = Layout V3 -> Maybe (Layout V3)
forall a. a -> Maybe a
Just (Int -> Int -> Int -> Layout V3
forall a. a -> a -> a -> V3 a
V3 (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)      Int
j       Int
k )
    | Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
y  = Layout V3 -> Maybe (Layout V3)
forall a. a -> Maybe a
Just (Int -> Int -> Int -> Layout V3
forall a. a -> a -> a -> V3 a
V3      Int
0  (Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)      Int
k )
    | Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
x  = Layout V3 -> Maybe (Layout V3)
forall a. a -> Maybe a
Just (Int -> Int -> Int -> Layout V3
forall a. a -> a -> a -> V3 a
V3      Int
0       Int
0  (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1))
    | Bool
otherwise  = Maybe (Layout V3)
forall a. Maybe a
Nothing
  {-# INLINE shapeStep #-}

  shapeStepBetween :: Layout V3 -> Layout V3 -> Layout V3 -> Maybe (Layout V3)
shapeStepBetween (V3 Int
ia Int
ja Int
_ka) (V3 Int
ib Int
jb Int
kb) (V3 Int
i Int
j Int
k)
    | Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
kb  = Layout V3 -> Maybe (Layout V3)
forall a. a -> Maybe a
Just (Int -> Int -> Int -> Layout V3
forall a. a -> a -> a -> V3 a
V3 (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)      Int
j       Int
k )
    | Int
j Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
jb  = Layout V3 -> Maybe (Layout V3)
forall a. a -> Maybe a
Just (Int -> Int -> Int -> Layout V3
forall a. a -> a -> a -> V3 a
V3     Int
ia  (Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)      Int
k )
    | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
ib  = Layout V3 -> Maybe (Layout V3)
forall a. a -> Maybe a
Just (Int -> Int -> Int -> Layout V3
forall a. a -> a -> a -> V3 a
V3     Int
ia      Int
ja  (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1))
    | Bool
otherwise  = Maybe (Layout V3)
forall a. Maybe a
Nothing
  {-# INLINE shapeStepBetween #-}

instance Shape V4 where
  shapeStep :: Layout V4 -> Layout V4 -> Maybe (Layout V4)
shapeStep (V4 Int
x Int
y Int
z Int
w) (V4 Int
i Int
j Int
k Int
l)
    | Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
w  = Layout V4 -> Maybe (Layout V4)
forall a. a -> Maybe a
Just (Int -> Int -> Int -> Int -> Layout V4
forall a. a -> a -> a -> a -> V4 a
V4 (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)      Int
j       Int
k       Int
l )
    | Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
z  = Layout V4 -> Maybe (Layout V4)
forall a. a -> Maybe a
Just (Int -> Int -> Int -> Int -> Layout V4
forall a. a -> a -> a -> a -> V4 a
V4      Int
0  (Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)      Int
k       Int
l )
    | Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
y  = Layout V4 -> Maybe (Layout V4)
forall a. a -> Maybe a
Just (Int -> Int -> Int -> Int -> Layout V4
forall a. a -> a -> a -> a -> V4 a
V4      Int
0       Int
0  (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)      Int
l )
    | Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
x  = Layout V4 -> Maybe (Layout V4)
forall a. a -> Maybe a
Just (Int -> Int -> Int -> Int -> Layout V4
forall a. a -> a -> a -> a -> V4 a
V4      Int
0       Int
0       Int
0  (Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1))
    | Bool
otherwise  = Maybe (Layout V4)
forall a. Maybe a
Nothing
  {-# INLINE shapeStep #-}

  shapeStepBetween :: Layout V4 -> Layout V4 -> Layout V4 -> Maybe (Layout V4)
shapeStepBetween (V4 Int
ia Int
ja Int
ka Int
_la) (V4 Int
ib Int
jb Int
kb Int
lb) (V4 Int
i Int
j Int
k Int
l)
    | Int
l Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
lb  = Layout V4 -> Maybe (Layout V4)
forall a. a -> Maybe a
Just (Int -> Int -> Int -> Int -> Layout V4
forall a. a -> a -> a -> a -> V4 a
V4 (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)      Int
j       Int
k       Int
l )
    | Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
kb  = Layout V4 -> Maybe (Layout V4)
forall a. a -> Maybe a
Just (Int -> Int -> Int -> Int -> Layout V4
forall a. a -> a -> a -> a -> V4 a
V4     Int
ia  (Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)      Int
k       Int
l )
    | Int
j Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
jb  = Layout V4 -> Maybe (Layout V4)
forall a. a -> Maybe a
Just (Int -> Int -> Int -> Int -> Layout V4
forall a. a -> a -> a -> a -> V4 a
V4     Int
ia      Int
ja  (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)      Int
l )
    | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
ib  = Layout V4 -> Maybe (Layout V4)
forall a. a -> Maybe a
Just (Int -> Int -> Int -> Int -> Layout V4
forall a. a -> a -> a -> a -> V4 a
V4     Int
ia      Int
ia      Int
ka  (Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1))
    | Bool
otherwise  = Maybe (Layout V4)
forall a. Maybe a
Nothing
  {-# INLINE shapeStepBetween #-}

-- instance Dim n => Shape (V n)

-- | @'toIndex' l@ and @'fromIndex' l@ form two halfs of an isomorphism.
indexIso :: Shape f => Layout f -> Iso' (f Int) Int
indexIso :: Layout f -> Iso' (Layout f) Int
indexIso Layout f
l = (Layout f -> Int) -> (Int -> Layout f) -> Iso' (Layout f) Int
forall s a b t. (s -> a) -> (b -> t) -> Iso s t a b
iso (Layout f -> Layout f -> Int
forall (f :: * -> *). Shape f => Layout f -> Layout f -> Int
shapeToIndex Layout f
l) (Layout f -> Int -> Layout f
forall (f :: * -> *). Shape f => Layout f -> Int -> Layout f
shapeFromIndex Layout f
l)
{-# INLINE indexIso #-}

------------------------------------------------------------------------
-- HasLayout
------------------------------------------------------------------------

-- | Class of things that have a 'Layout'. This means we can use the
--   same functions for the various different arrays in the library.
class Shape f => HasLayout f a | a -> f where
  -- | Lens onto the  'Layout' of something.
  layout :: Lens' a (Layout f)
  default layout :: (a ~ f Int) => (Layout f -> g (Layout f)) -> a -> g a
  layout = (Layout f -> g (Layout f)) -> a -> g a
forall a. a -> a
id
  {-# INLINE layout #-}

instance i ~ Int => HasLayout V0 (V0 i)
instance i ~ Int => HasLayout V1 (V1 i)
instance i ~ Int => HasLayout V2 (V2 i)
instance i ~ Int => HasLayout V3 (V3 i)
instance i ~ Int => HasLayout V4 (V4 i)

-- | Get the extent of an array.
--
-- @
-- 'extent' :: 'Data.Dense.Base.Array' v f a    -> f 'Int'
-- 'extent' :: 'Data.Dense.Mutable.MArray' v f s a -> f 'Int'
-- 'extent' :: 'Data.Dense.Base.Delayed' f a    -> f 'Int'
-- 'extent' :: 'Data.Dense.Base.Focused' f a    -> f 'Int'
-- @
extent :: HasLayout f a => a -> f Int
extent :: a -> f Int
extent = Getting (f Int) a (f Int) -> a -> f Int
forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view Getting (f Int) a (f Int)
forall (f :: * -> *) a. HasLayout f a => Lens' a (Layout f)
layout
{-# INLINE extent #-}

-- | Get the total number of elements in an array.
--
-- @
-- 'size' :: 'Data.Dense.Base.Array' v f a    -> 'Int'
-- 'size' :: 'Data.Dense.Mutable.MArray' v f s a -> 'Int'
-- 'size' :: 'Data.Dense.Base.Delayed' f a    -> 'Int'
-- 'size' :: 'Data.Dense.Base.Focused' f a    -> 'Int'
-- @
size :: HasLayout f a => a -> Int
size :: a -> Int
size = Layout f -> Int
forall (f :: * -> *). Shape f => Layout f -> Int
shapeSize (Layout f -> Int) -> (a -> Layout f) -> a -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Getting (Layout f) a (Layout f) -> a -> Layout f
forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view Getting (Layout f) a (Layout f)
forall (f :: * -> *) a. HasLayout f a => Lens' a (Layout f)
layout
{-# INLINE size #-}

-- NB: lens already uses indices so we settle for indexes

-- | Indexed fold for all the indexes in the layout.
indexes :: HasLayout f a => IndexedFold Int a (f Int)
indexes :: IndexedFold Int a (f Int)
indexes = (f Int -> f (f Int)) -> a -> f a
forall (f :: * -> *) a. HasLayout f a => Lens' a (Layout f)
layout ((f Int -> f (f Int)) -> a -> f a)
-> (p (f Int) (f (f Int)) -> f Int -> f (f Int))
-> p (f Int) (f (f Int))
-> a
-> f a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. p (f Int) (f (f Int)) -> f Int -> f (f Int)
forall (f :: * -> *).
Shape f =>
IndexedFold Int (Layout f) (Layout f)
shapeIndexes
{-# INLINE indexes #-}

-- | 'indexes' for a 'Shape'.
shapeIndexes :: Shape f => IndexedFold Int (Layout f) (f Int)
shapeIndexes :: IndexedFold Int (Layout f) (Layout f)
shapeIndexes p (Layout f) (f (Layout f))
g Layout f
l = Int -> Maybe (Layout f) -> f (Layout f)
go (Int
0::Int) (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
  go :: Int -> Maybe (Layout f) -> f (Layout f)
go Int
i (Just Layout f
x) = p (Layout f) (f (Layout f)) -> Int -> Layout f -> f (Layout f)
forall i (p :: * -> * -> *) a b.
Indexable i p =>
p a b -> i -> a -> b
indexed p (Layout f) (f (Layout f))
g Int
i Layout f
x f (Layout f) -> f (Layout f) -> f (Layout f)
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> Int -> Maybe (Layout f) -> f (Layout f)
go (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Layout f -> Layout f -> Maybe (Layout f)
forall (f :: * -> *).
Shape f =>
Layout f -> Layout f -> Maybe (Layout f)
shapeStep Layout f
l Layout f
x)
  go Int
_ Maybe (Layout f)
Nothing  = f (Layout f)
forall (f :: * -> *) a. (Contravariant f, Applicative f) => f a
noEffect
{-# INLINE shapeIndexes #-}

-- | Indexed fold starting starting from some point, where the index is
--   the linear index for the original layout.
indexesFrom :: HasLayout f a => f Int -> IndexedFold Int a (f Int)
indexesFrom :: f Int -> IndexedFold Int a (f Int)
indexesFrom f Int
a = (f Int -> f (f Int)) -> a -> f a
forall (f :: * -> *) a. HasLayout f a => Lens' a (Layout f)
layout ((f Int -> f (f Int)) -> a -> f a)
-> (p (f Int) (f (f Int)) -> f Int -> f (f Int))
-> p (f Int) (f (f Int))
-> a
-> f a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. f Int -> IndexedFold Int (f Int) (f Int)
forall (f :: * -> *).
Shape f =>
f Int -> IndexedFold Int (f Int) (f Int)
shapeIndexesFrom f Int
a
{-# INLINE indexesFrom #-}

-- | 'indexesFrom' for a 'Shape'.
shapeIndexesFrom :: Shape f => f Int -> IndexedFold Int (Layout f) (f Int)
shapeIndexesFrom :: f Int -> IndexedFold Int (f Int) (f Int)
shapeIndexesFrom f Int
a p (f Int) (f (f Int))
f f Int
l = f Int -> f Int -> p (f Int) (f (f Int)) -> f Int -> f (f Int)
forall (f :: * -> *).
Shape f =>
f Int -> f Int -> IndexedFold Int (f Int) (f Int)
shapeIndexesBetween f Int
a f Int
l p (f Int) (f (f Int))
f f Int
l
{-# INLINE shapeIndexesFrom #-}

-- | Indexed fold between the two indexes where the index is the linear
--   index for the original layout.
indexesBetween :: HasLayout f a => f Int -> f Int -> IndexedFold Int a (f Int)
indexesBetween :: f Int -> f Int -> IndexedFold Int a (f Int)
indexesBetween f Int
a f Int
b = (f Int -> f (f Int)) -> a -> f a
forall (f :: * -> *) a. HasLayout f a => Lens' a (Layout f)
layout ((f Int -> f (f Int)) -> a -> f a)
-> (p (f Int) (f (f Int)) -> f Int -> f (f Int))
-> p (f Int) (f (f Int))
-> a
-> f a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. f Int -> f Int -> IndexedFold Int (f Int) (f Int)
forall (f :: * -> *).
Shape f =>
f Int -> f Int -> IndexedFold Int (f Int) (f Int)
shapeIndexesBetween f Int
a f Int
b
{-# INLINE indexesBetween #-}

-- | 'indexesBetween' for a 'Shape'.
shapeIndexesBetween :: Shape f => f Int -> f Int -> IndexedFold Int (Layout f) (f Int)
shapeIndexesBetween :: f Int -> f Int -> IndexedFold Int (f Int) (f Int)
shapeIndexesBetween f Int
a f Int
b p (f Int) (f (f Int))
f f Int
l =
  Maybe (f Int) -> f (f Int)
go (if f Int -> f Int -> Bool
forall (f :: * -> *) a. (Eq1 f, Eq a) => f a -> f a -> Bool
eq1 f Int
l f Int
a Bool -> Bool -> Bool
|| Bool -> Bool
not (f Int -> f Int -> Bool
forall (f :: * -> *). Shape f => Layout f -> Layout f -> Bool
shapeInRange f Int
l f Int
b) then Maybe (f Int)
forall a. Maybe a
Nothing else f Int -> Maybe (f Int)
forall a. a -> Maybe a
Just f Int
a) where
    go :: Maybe (f Int) -> f (f Int)
go (Just f Int
x) = p (f Int) (f (f Int)) -> Int -> f Int -> f (f Int)
forall i (p :: * -> * -> *) a b.
Indexable i p =>
p a b -> i -> a -> b
indexed p (f Int) (f (f Int))
f (f Int -> f Int -> Int
forall (f :: * -> *). Shape f => Layout f -> Layout f -> Int
shapeToIndex f Int
l f Int
x) f Int
x f (f Int) -> f (f Int) -> f (f Int)
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> Maybe (f Int) -> f (f Int)
go (f Int -> f Int -> f Int -> Maybe (f Int)
forall (f :: * -> *).
Shape f =>
Layout f -> Layout f -> Layout f -> Maybe (Layout f)
shapeStepBetween f Int
a f Int
b f Int
x)
    go Maybe (f Int)
Nothing  = f (f Int)
forall (f :: * -> *) a. (Contravariant f, Applicative f) => f a
noEffect
{-# INLINE shapeIndexesBetween #-}

------------------------------------------------------------------------
-- Exceptions
------------------------------------------------------------------------

-- Bounds check --------------------------------------------------------

-- | @boundsCheck l i@ performs a bounds check for index @i@ and layout
--   @l@. Throws an 'IndexOutOfBounds' exception when out of range in
--   the form @(i, l)@. This can be caught with the '_IndexOutOfBounds'
--   prism.
--
-- >>> boundsCheck (V2 3 5) (V2 1 4) "in range"
-- "in range"
--
-- >>> boundsCheck (V2 10 20) (V2 10 5) "in bounds"
-- "*** Exception: array index out of range: (V2 10 5, V2 10 20)
--
-- >>> catching _IndexOutOfBounds (boundsCheck (V1 2) (V1 2) (putStrLn "in range")) print
-- "(V1 2, V1 2)"
--
-- The output format is suitable to be read using the '_Show' prism:
--
-- >>> trying (_IndexOutOfBounds . _Show) (boundsCheck (V1 2) (V1 20) (putStrLn "in range")) :: IO (Either (V1 Int, V1 Int) ())
-- Left (V1 20,V1 2)
boundsCheck :: Shape l => Layout l-> l Int -> a -> a
boundsCheck :: Layout l -> Layout l -> a -> a
boundsCheck Layout l
l Layout l
i
  | Layout l -> Layout l -> Bool
forall (f :: * -> *). Shape f => Layout f -> Layout f -> Bool
shapeInRange Layout l
l Layout l
i = a -> a
forall a. a -> a
id
  | Bool
otherwise        = AReview SomeException String -> String -> a -> a
forall b r. AReview SomeException b -> b -> r
throwing AReview SomeException String
forall t. AsArrayException t => Prism' t String
_IndexOutOfBounds (String -> a -> a) -> String -> a -> a
forall a b. (a -> b) -> a -> b
$ String
"(" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Layout l -> String
forall (f :: * -> *). Shape f => f Int -> String
showShape Layout l
i String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
", " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Layout l -> String
forall (f :: * -> *). Shape f => f Int -> String
showShape Layout l
l String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
")"
{-# INLINE boundsCheck #-}

-- Size missmatch ------------------------------------------------------

-- | Thrown when two sizes that should match, don't.
data SizeMissmatch = SizeMissmatch String
  deriving Typeable

instance Exception SizeMissmatch
instance Show SizeMissmatch where
  showsPrec :: Int -> SizeMissmatch -> String -> String
showsPrec Int
_ (SizeMissmatch String
s)
    = String -> String -> String
showString String
"size missmatch"
    (String -> String) -> (String -> String) -> String -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (if Bool -> Bool
not (String -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null String
s) then String -> String -> String
showString String
": " (String -> String) -> (String -> String) -> String -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> String -> String
showString String
s
                       else String -> String
forall a. a -> a
id)

-- | Exception thown from missmatching sizes.
class AsSizeMissmatch t where
  -- | Extract information about an 'SizeMissmatch'.
  --
  -- @
  -- '_SizeMissmatch' :: 'Prism'' 'SizeMissmatch' 'String'
  -- '_SizeMissmatch' :: 'Prism'' 'SomeException' 'String'
  -- @
  _SizeMissmatch :: Prism' t String

instance AsSizeMissmatch SizeMissmatch where
  _SizeMissmatch :: p String (f String) -> p SizeMissmatch (f SizeMissmatch)
_SizeMissmatch = (String -> SizeMissmatch)
-> (SizeMissmatch -> Maybe String) -> Prism' SizeMissmatch String
forall b s a. (b -> s) -> (s -> Maybe a) -> Prism s s a b
prism' String -> SizeMissmatch
SizeMissmatch ((SizeMissmatch -> Maybe String) -> Prism' SizeMissmatch String)
-> (SizeMissmatch -> Maybe String) -> Prism' SizeMissmatch String
forall a b. (a -> b) -> a -> b
$ (\(SizeMissmatch String
s) -> String -> Maybe String
forall a. a -> Maybe a
Just String
s)
  {-# INLINE _SizeMissmatch #-}

instance AsSizeMissmatch SomeException where
  _SizeMissmatch :: p String (f String) -> p SomeException (f SomeException)
_SizeMissmatch = p SizeMissmatch (f SizeMissmatch)
-> p SomeException (f SomeException)
forall a. Exception a => Prism' SomeException a
exception (p SizeMissmatch (f SizeMissmatch)
 -> p SomeException (f SomeException))
-> (p String (f String) -> p SizeMissmatch (f SizeMissmatch))
-> p String (f String)
-> p SomeException (f SomeException)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (forall t. AsSizeMissmatch t => Prism' t String
Prism' SizeMissmatch String
_SizeMissmatch :: Prism' SizeMissmatch String)
  {-# INLINE _SizeMissmatch #-}

-- | Check the sizes are equal. If not, throw 'SizeMissmatch'.
sizeMissmatch :: Int -> Int -> String -> a -> a
sizeMissmatch :: Int -> Int -> String -> a -> a
sizeMissmatch Int
i Int
j String
err
  | Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
j    = a -> a
forall a. a -> a
id
  | Bool
otherwise = AReview SomeException String -> String -> a -> a
forall b r. AReview SomeException b -> b -> r
throwing AReview SomeException String
forall t. AsSizeMissmatch t => Prism' t String
_SizeMissmatch String
err
{-# INLINE sizeMissmatch #-}

-- Utilities -----------------------------------------------------------

-- | Show a shape in the form @VN i1 i2 .. iN@ where @N@ is the 'length'
--   of the shape.
showShape :: Shape f => f Int -> String
showShape :: f Int -> String
showShape f Int
l = String
"V" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show (Getting (Endo (Endo Int)) (f Int) Int -> f Int -> Int
forall s a. Getting (Endo (Endo Int)) s a -> s -> Int
lengthOf Getting (Endo (Endo Int)) (f Int) Int
forall (f :: * -> *) a. Foldable f => IndexedFold Int (f a) a
folded f Int
l) String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
" " String -> String -> String
forall a. [a] -> [a] -> [a]
++ [String] -> String
unwords (Int -> String
forall a. Show a => a -> String
show (Int -> String) -> [Int] -> [String]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> f Int -> [Int]
forall (t :: * -> *) a. Foldable t => t a -> [a]
F.toList f Int
l)