{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
module Data.Array.Comfort.Boxed.Unchecked (
   Array(..),
   reshape,
   mapShape,
   (!),
   toList,
   fromList,
   vectorFromList,
   replicate,
   map,
   zipWith,

   append,
   take, drop,
   takeLeft, takeRight, split,
   takeCenter,
   ) where

import qualified Data.Array.Comfort.Shape as Shape
import qualified Data.Primitive.Array as Prim
import Data.Array.Comfort.Shape ((::+)((::+)))

-- FixMe: In GHC-7.4.2 there is no instance PrimMonad (Lazy.ST s)
-- import qualified Control.Monad.ST.Lazy as ST
import qualified Control.Monad.ST.Strict as ST
import Control.Monad (liftM)
import Control.Applicative (Applicative, pure, (<*>), (<$>))
import Control.DeepSeq (NFData, rnf)

import qualified Data.Traversable as Trav
import qualified Data.Foldable as Fold
import qualified Data.List as List
import Prelude hiding (map, zipWith, replicate, take, drop)


{- $setup
>>> import qualified Data.Array.Comfort.Boxed as Array
>>> import qualified Data.Array.Comfort.Shape as Shape
>>> import Data.Array.Comfort.Boxed (Array, (!))
>>> import Data.Tuple.HT (swap)
>>> import Control.Applicative ((<$>))
>>>
>>> import qualified Test.QuickCheck as QC
>>>
>>> type ShapeInt = Shape.ZeroBased Int
>>>
>>> genArray :: QC.Gen (Array ShapeInt Char)
>>> genArray = Array.vectorFromList <$> QC.arbitrary
>>>
>>> newtype ArrayChar = ArrayChar (Array ShapeInt Char)
>>>    deriving (Show)
>>>
>>> instance QC.Arbitrary ArrayChar where
>>>    arbitrary = fmap ArrayChar genArray
>>>
>>>
>>> transpose ::
>>>    (Shape.Indexed sh0, Shape.Indexed sh1) =>
>>>    Array (sh0,sh1) a -> Array (sh1,sh0) a
>>> transpose a =
>>>    fmap (\(i,j) -> a!(j,i)) $ Array.indices $ swap $ Array.shape a
-}


data Array sh a =
   Array {
      forall sh a. Array sh a -> sh
shape :: sh,
      forall sh a. Array sh a -> Array a
buffer :: Prim.Array a
   } deriving (Array sh a -> Array sh a -> Bool
(Array sh a -> Array sh a -> Bool)
-> (Array sh a -> Array sh a -> Bool) -> Eq (Array sh a)
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
forall sh a. (Eq sh, Eq a) => Array sh a -> Array sh a -> Bool
$c== :: forall sh a. (Eq sh, Eq a) => Array sh a -> Array sh a -> Bool
== :: Array sh a -> Array sh a -> Bool
$c/= :: forall sh a. (Eq sh, Eq a) => Array sh a -> Array sh a -> Bool
/= :: Array sh a -> Array sh a -> Bool
Eq)

instance (Shape.C sh, Show sh, Show a) => Show (Array sh a) where
   showsPrec :: Int -> Array sh a -> ShowS
showsPrec Int
p Array sh a
arr =
      Bool -> ShowS -> ShowS
showParen (Int
pInt -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>Int
10) (ShowS -> ShowS) -> ShowS -> ShowS
forall a b. (a -> b) -> a -> b
$
         String -> ShowS
showString String
"BoxedArray.fromList " ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
.
         Int -> sh -> ShowS
forall a. Show a => Int -> a -> ShowS
showsPrec Int
11 (Array sh a -> sh
forall sh a. Array sh a -> sh
shape Array sh a
arr) ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
.
         Char -> ShowS
showChar Char
' ' ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
.
         [a] -> ShowS
forall a. Show a => a -> ShowS
shows (Array sh a -> [a]
forall sh a. C sh => Array sh a -> [a]
toList Array sh a
arr)


instance (Shape.C sh, NFData sh, NFData a) => NFData (Array sh a) where
   rnf :: Array sh a -> ()
rnf a :: Array sh a
a@(Array sh
sh Array a
_arr) = (sh, [a]) -> ()
forall a. NFData a => a -> ()
rnf (sh
sh, Array sh a -> [a]
forall sh a. C sh => Array sh a -> [a]
toList Array sh a
a)

instance (Shape.C sh) => Functor (Array sh) where
   fmap :: forall a b. (a -> b) -> Array sh a -> Array sh b
fmap = (a -> b) -> Array sh a -> Array sh b
forall sh a b. C sh => (a -> b) -> Array sh a -> Array sh b
map

{- |
We must restrict 'Applicative' to 'Shape.Static' because of 'pure'.
Because the shape is static, we do not need a size check in '(<*>)'.
-}
instance (Shape.Static sh) => Applicative (Array sh) where
   pure :: forall a. a -> Array sh a
pure = sh -> a -> Array sh a
forall sh a. C sh => sh -> a -> Array sh a
replicate sh
forall sh. Static sh => sh
Shape.static
   <*> :: forall a b. Array sh (a -> b) -> Array sh a -> Array sh b
(<*>) = ((a -> b) -> a -> b)
-> Array sh (a -> b) -> Array sh a -> Array sh b
forall sh a b c.
C sh =>
(a -> b -> c) -> Array sh a -> Array sh b -> Array sh c
zipWith (a -> b) -> a -> b
forall a b. (a -> b) -> a -> b
($)

instance (Shape.C sh) => Fold.Foldable (Array sh) where
   fold :: forall m. Monoid m => Array sh m -> m
fold = Array m -> m
forall m. Monoid m => Array m -> m
forall (t :: * -> *) m. (Foldable t, Monoid m) => t m -> m
Fold.fold (Array m -> m) -> (Array sh m -> Array m) -> Array sh m -> m
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array sh m -> Array m
forall sh a. Array sh a -> Array a
buffer
   foldMap :: forall m a. Monoid m => (a -> m) -> Array sh a -> m
foldMap a -> m
f = (a -> m) -> Array a -> m
forall m a. Monoid m => (a -> m) -> Array a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
Fold.foldMap a -> m
f (Array a -> m) -> (Array sh a -> Array a) -> Array sh a -> m
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array sh a -> Array a
forall sh a. Array sh a -> Array a
buffer
   foldl :: forall b a. (b -> a -> b) -> b -> Array sh a -> b
foldl b -> a -> b
f b
a = (b -> a -> b) -> b -> Array a -> b
forall b a. (b -> a -> b) -> b -> Array a -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
Fold.foldl b -> a -> b
f b
a (Array a -> b) -> (Array sh a -> Array a) -> Array sh a -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array sh a -> Array a
forall sh a. Array sh a -> Array a
buffer
   foldr :: forall a b. (a -> b -> b) -> b -> Array sh a -> b
foldr a -> b -> b
f b
a = (a -> b -> b) -> b -> Array a -> b
forall a b. (a -> b -> b) -> b -> Array a -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
Fold.foldr a -> b -> b
f b
a (Array a -> b) -> (Array sh a -> Array a) -> Array sh a -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array sh a -> Array a
forall sh a. Array sh a -> Array a
buffer
   foldl1 :: forall a. (a -> a -> a) -> Array sh a -> a
foldl1 a -> a -> a
f = (a -> a -> a) -> Array a -> a
forall a. (a -> a -> a) -> Array a -> a
forall (t :: * -> *) a. Foldable t => (a -> a -> a) -> t a -> a
Fold.foldl1 a -> a -> a
f (Array a -> a) -> (Array sh a -> Array a) -> Array sh a -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array sh a -> Array a
forall sh a. Array sh a -> Array a
buffer
   foldr1 :: forall a. (a -> a -> a) -> Array sh a -> a
foldr1 a -> a -> a
f = (a -> a -> a) -> Array a -> a
forall a. (a -> a -> a) -> Array a -> a
forall (t :: * -> *) a. Foldable t => (a -> a -> a) -> t a -> a
Fold.foldr1 a -> a -> a
f (Array a -> a) -> (Array sh a -> Array a) -> Array sh a -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array sh a -> Array a
forall sh a. Array sh a -> Array a
buffer

instance (Shape.C sh) => Trav.Traversable (Array sh) where
   traverse :: forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Array sh a -> f (Array sh b)
traverse a -> f b
f (Array sh
sh Array a
arr) = sh -> Array b -> Array sh b
forall sh a. sh -> Array a -> Array sh a
Array sh
sh (Array b -> Array sh b) -> f (Array b) -> f (Array sh b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (a -> f b) -> Array a -> f (Array b)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Array a -> f (Array b)
Trav.traverse a -> f b
f Array a
arr
   sequenceA :: forall (f :: * -> *) a.
Applicative f =>
Array sh (f a) -> f (Array sh a)
sequenceA (Array sh
sh Array (f a)
arr) = sh -> Array a -> Array sh a
forall sh a. sh -> Array a -> Array sh a
Array sh
sh (Array a -> Array sh a) -> f (Array a) -> f (Array sh a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Array (f a) -> f (Array a)
forall (t :: * -> *) (f :: * -> *) a.
(Traversable t, Applicative f) =>
t (f a) -> f (t a)
forall (f :: * -> *) a. Applicative f => Array (f a) -> f (Array a)
Trav.sequenceA Array (f a)
arr
   mapM :: forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Array sh a -> m (Array sh b)
mapM a -> m b
f (Array sh
sh Array a
arr) = (Array b -> Array sh b) -> m (Array b) -> m (Array sh b)
forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
liftM (sh -> Array b -> Array sh b
forall sh a. sh -> Array a -> Array sh a
Array sh
sh) (m (Array b) -> m (Array sh b)) -> m (Array b) -> m (Array sh b)
forall a b. (a -> b) -> a -> b
$ (a -> m b) -> Array a -> m (Array b)
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Array a -> m (Array b)
Trav.mapM a -> m b
f Array a
arr
   sequence :: forall (m :: * -> *) a. Monad m => Array sh (m a) -> m (Array sh a)
sequence (Array sh
sh Array (m a)
arr) = (Array a -> Array sh a) -> m (Array a) -> m (Array sh a)
forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
liftM (sh -> Array a -> Array sh a
forall sh a. sh -> Array a -> Array sh a
Array sh
sh) (m (Array a) -> m (Array sh a)) -> m (Array a) -> m (Array sh a)
forall a b. (a -> b) -> a -> b
$ Array (m a) -> m (Array a)
forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
forall (m :: * -> *) a. Monad m => Array (m a) -> m (Array a)
Trav.sequence Array (m a)
arr


reshape :: sh1 -> Array sh0 a -> Array sh1 a
reshape :: forall sh1 sh0 a. sh1 -> Array sh0 a -> Array sh1 a
reshape sh1
sh (Array sh0
_ Array a
arr) = sh1 -> Array a -> Array sh1 a
forall sh a. sh -> Array a -> Array sh a
Array sh1
sh Array a
arr

mapShape :: (sh0 -> sh1) -> Array sh0 a -> Array sh1 a
mapShape :: forall sh0 sh1 a. (sh0 -> sh1) -> Array sh0 a -> Array sh1 a
mapShape sh0 -> sh1
f (Array sh0
sh Array a
arr) = sh1 -> Array a -> Array sh1 a
forall sh a. sh -> Array a -> Array sh a
Array (sh0 -> sh1
f sh0
sh) Array a
arr


infixl 9 !

(!) :: (Shape.Indexed sh) => Array sh a -> Shape.Index sh -> a
! :: forall sh a. Indexed sh => Array sh a -> Index sh -> a
(!) (Array sh
sh Array a
arr) Index sh
ix = Array a -> Int -> a
forall a. Array a -> Int -> a
Prim.indexArray Array a
arr (Int -> a) -> Int -> a
forall a b. (a -> b) -> a -> b
$ sh -> Index sh -> Int
forall sh. Indexed sh => sh -> Index sh -> Int
Shape.uncheckedOffset sh
sh Index sh
ix

toList :: (Shape.C sh) => Array sh a -> [a]
toList :: forall sh a. C sh => Array sh a -> [a]
toList (Array sh
sh Array a
arr) =
   (Int -> a) -> [Int] -> [a]
forall a b. (a -> b) -> [a] -> [b]
List.map (Array a -> Int -> a
forall a. Array a -> Int -> a
Prim.indexArray Array a
arr) ([Int] -> [a]) -> [Int] -> [a]
forall a b. (a -> b) -> a -> b
$ Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
List.take (sh -> Int
forall sh. C sh => sh -> Int
Shape.size sh
sh) [Int
0..]

fromList :: (Shape.C sh) => sh -> [a] -> Array sh a
fromList :: forall sh a. C sh => sh -> [a] -> Array sh a
fromList sh
sh [a]
xs = sh -> Array a -> Array sh a
forall sh a. sh -> Array a -> Array sh a
Array sh
sh (Array a -> Array sh a) -> Array a -> Array sh a
forall a b. (a -> b) -> a -> b
$ Int -> [a] -> Array a
forall a. Int -> [a] -> Array a
Prim.arrayFromListN (sh -> Int
forall sh. C sh => sh -> Int
Shape.size sh
sh) [a]
xs

vectorFromList :: [a] -> Array (Shape.ZeroBased Int) a
vectorFromList :: forall a. [a] -> Array (ZeroBased Int) a
vectorFromList [a]
xs =
   let arr :: Array a
arr = [a] -> Array a
forall a. [a] -> Array a
Prim.arrayFromList [a]
xs
   in ZeroBased Int -> Array a -> Array (ZeroBased Int) a
forall sh a. sh -> Array a -> Array sh a
Array (Int -> ZeroBased Int
forall n. n -> ZeroBased n
Shape.ZeroBased (Int -> ZeroBased Int) -> Int -> ZeroBased Int
forall a b. (a -> b) -> a -> b
$ Array a -> Int
forall a. Array a -> Int
Prim.sizeofArray Array a
arr) Array a
arr

replicate :: (Shape.C sh) => sh -> a -> Array sh a
replicate :: forall sh a. C sh => sh -> a -> Array sh a
replicate sh
sh a
a =
   sh -> Array a -> Array sh a
forall sh a. sh -> Array a -> Array sh a
Array sh
sh (Array a -> Array sh a) -> Array a -> Array sh a
forall a b. (a -> b) -> a -> b
$
   (forall s. ST s (Array a)) -> Array a
forall a. (forall s. ST s a) -> a
ST.runST (MutableArray s a -> ST s (Array a)
MutableArray (PrimState (ST s)) a -> ST s (Array a)
forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> m (Array a)
Prim.unsafeFreezeArray  (MutableArray s a -> ST s (Array a))
-> ST s (MutableArray s a) -> ST s (Array a)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Int -> a -> ST s (MutableArray (PrimState (ST s)) a)
forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (MutableArray (PrimState m) a)
Prim.newArray (sh -> Int
forall sh. C sh => sh -> Int
Shape.size sh
sh) a
a)

map :: (Shape.C sh) => (a -> b) -> Array sh a -> Array sh b
map :: forall sh a b. C sh => (a -> b) -> Array sh a -> Array sh b
map a -> b
f (Array sh
sh Array a
arr) = sh -> Array b -> Array sh b
forall sh a. sh -> Array a -> Array sh a
Array sh
sh (Array b -> Array sh b) -> Array b -> Array sh b
forall a b. (a -> b) -> a -> b
$
   let n :: Int
n = sh -> Int
forall sh. C sh => sh -> Int
Shape.size sh
sh
   in Int -> [b] -> Array b
forall a. Int -> [a] -> Array a
Prim.arrayFromListN Int
n ([b] -> Array b) -> [b] -> Array b
forall a b. (a -> b) -> a -> b
$ (Int -> b) -> [Int] -> [b]
forall a b. (a -> b) -> [a] -> [b]
List.map (a -> b
f (a -> b) -> (Int -> a) -> Int -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array a -> Int -> a
forall a. Array a -> Int -> a
Prim.indexArray Array a
arr) ([Int] -> [b]) -> [Int] -> [b]
forall a b. (a -> b) -> a -> b
$ Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
List.take Int
n [Int
0..]

zipWith ::
   (Shape.C sh) => (a -> b -> c) -> Array sh a -> Array sh b -> Array sh c
zipWith :: forall sh a b c.
C sh =>
(a -> b -> c) -> Array sh a -> Array sh b -> Array sh c
zipWith a -> b -> c
f (Array sh
sha Array a
arra) (Array sh
_shb Array b
arrb) = sh -> Array c -> Array sh c
forall sh a. sh -> Array a -> Array sh a
Array sh
sha (Array c -> Array sh c) -> Array c -> Array sh c
forall a b. (a -> b) -> a -> b
$
   let n :: Int
n = sh -> Int
forall sh. C sh => sh -> Int
Shape.size sh
sha
   in Int -> [c] -> Array c
forall a. Int -> [a] -> Array a
Prim.arrayFromListN Int
n ([c] -> Array c) -> [c] -> Array c
forall a b. (a -> b) -> a -> b
$
      (Int -> c) -> [Int] -> [c]
forall a b. (a -> b) -> [a] -> [b]
List.map (\Int
k -> a -> b -> c
f (Array a -> Int -> a
forall a. Array a -> Int -> a
Prim.indexArray Array a
arra Int
k) (Array b -> Int -> b
forall a. Array a -> Int -> a
Prim.indexArray Array b
arrb Int
k)) ([Int] -> [c]) -> [Int] -> [c]
forall a b. (a -> b) -> a -> b
$
      Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
List.take Int
n [Int
0..]



infixr 5 `append`

append ::
   (Shape.C shx, Shape.C shy) =>
   Array shx a -> Array shy a -> Array (shx::+shy) a
append :: forall shx shy a.
(C shx, C shy) =>
Array shx a -> Array shy a -> Array (shx ::+ shy) a
append (Array shx
shX Array a
x) (Array shy
shY Array a
y) =
   let sizeX :: Int
sizeX = shx -> Int
forall sh. C sh => sh -> Int
Shape.size shx
shX in
   let sizeY :: Int
sizeY = shy -> Int
forall sh. C sh => sh -> Int
Shape.size shy
shY in
   (shx ::+ shy) -> Array a -> Array (shx ::+ shy) a
forall sh a. sh -> Array a -> Array sh a
Array (shx
shXshx -> shy -> shx ::+ shy
forall sh0 sh1. sh0 -> sh1 -> sh0 ::+ sh1
::+shy
shY) (Array a -> Array (shx ::+ shy) a)
-> Array a -> Array (shx ::+ shy) a
forall a b. (a -> b) -> a -> b
$
   (forall s. ST s (Array a)) -> Array a
forall a. (forall s. ST s a) -> a
ST.runST (do
      MutableArray s a
arr <-
         Int -> a -> ST s (MutableArray (PrimState (ST s)) a)
forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (MutableArray (PrimState m) a)
Prim.newArray (Int
sizeXInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
sizeY)
            (String -> a
forall a. HasCallStack => String -> a
error String
"Boxed.append: uninitialized element")
      MutableArray (PrimState (ST s)) a
-> Int -> Array a -> Int -> Int -> ST s ()
forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a
-> Int -> Array a -> Int -> Int -> m ()
Prim.copyArray MutableArray s a
MutableArray (PrimState (ST s)) a
arr Int
0 Array a
x Int
0 Int
sizeX
      MutableArray (PrimState (ST s)) a
-> Int -> Array a -> Int -> Int -> ST s ()
forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a
-> Int -> Array a -> Int -> Int -> m ()
Prim.copyArray MutableArray s a
MutableArray (PrimState (ST s)) a
arr Int
sizeX Array a
y Int
0 Int
sizeY
      MutableArray (PrimState (ST s)) a -> ST s (Array a)
forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> m (Array a)
Prim.unsafeFreezeArray MutableArray s a
MutableArray (PrimState (ST s)) a
arr)

{- |
prop> \(QC.NonNegative n) (ArrayChar x)  ->  x == Array.mapShape (Shape.ZeroBased . Shape.size) (Array.append (Array.take n x) (Array.drop n x))
-}
take, drop ::
   (Integral n) =>
   n -> Array (Shape.ZeroBased n) a -> Array (Shape.ZeroBased n) a
take :: forall n a.
Integral n =>
n -> Array (ZeroBased n) a -> Array (ZeroBased n) a
take n
n = Array (ZeroBased n ::+ ZeroBased n) a -> Array (ZeroBased n) a
forall sh0 sh1 a.
(C sh0, C sh1) =>
Array (sh0 ::+ sh1) a -> Array sh0 a
takeLeft (Array (ZeroBased n ::+ ZeroBased n) a -> Array (ZeroBased n) a)
-> (Array (ZeroBased n) a -> Array (ZeroBased n ::+ ZeroBased n) a)
-> Array (ZeroBased n) a
-> Array (ZeroBased n) a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. n -> Array (ZeroBased n) a -> Array (ZeroBased n ::+ ZeroBased n) a
forall n a.
Integral n =>
n -> Array (ZeroBased n) a -> Array (ZeroBased n ::+ ZeroBased n) a
splitN n
n
drop :: forall n a.
Integral n =>
n -> Array (ZeroBased n) a -> Array (ZeroBased n) a
drop n
n = Array (ZeroBased n ::+ ZeroBased n) a -> Array (ZeroBased n) a
forall sh0 sh1 a.
(C sh0, C sh1) =>
Array (sh0 ::+ sh1) a -> Array sh1 a
takeRight (Array (ZeroBased n ::+ ZeroBased n) a -> Array (ZeroBased n) a)
-> (Array (ZeroBased n) a -> Array (ZeroBased n ::+ ZeroBased n) a)
-> Array (ZeroBased n) a
-> Array (ZeroBased n) a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. n -> Array (ZeroBased n) a -> Array (ZeroBased n ::+ ZeroBased n) a
forall n a.
Integral n =>
n -> Array (ZeroBased n) a -> Array (ZeroBased n ::+ ZeroBased n) a
splitN n
n

splitN ::
   (Integral n) =>
   n -> Array (Shape.ZeroBased n) a ->
   Array (Shape.ZeroBased n ::+ Shape.ZeroBased n) a
splitN :: forall n a.
Integral n =>
n -> Array (ZeroBased n) a -> Array (ZeroBased n ::+ ZeroBased n) a
splitN n
n = (ZeroBased n -> ZeroBased n ::+ ZeroBased n)
-> Array (ZeroBased n) a -> Array (ZeroBased n ::+ ZeroBased n) a
forall sh0 sh1 a. (sh0 -> sh1) -> Array sh0 a -> Array sh1 a
mapShape (n -> ZeroBased n -> ZeroBased n ::+ ZeroBased n
forall n. Real n => n -> ZeroBased n -> ZeroBased n ::+ ZeroBased n
Shape.zeroBasedSplit n
n)

{- |
prop> \(ArrayChar x) (ArrayChar y) -> let xy = Array.append x y in x == Array.takeLeft xy  &&  y == Array.takeRight xy
-}
takeLeft ::
   (Shape.C sh0, Shape.C sh1) =>
   Array (sh0::+sh1) a -> Array sh0 a
takeLeft :: forall sh0 sh1 a.
(C sh0, C sh1) =>
Array (sh0 ::+ sh1) a -> Array sh0 a
takeLeft =
   Array (Zero ::+ (sh0 ::+ sh1)) a -> Array sh0 a
forall sh0 sh1 sh2 a.
(C sh0, C sh1, C sh2) =>
Array (sh0 ::+ (sh1 ::+ sh2)) a -> Array sh1 a
takeCenter (Array (Zero ::+ (sh0 ::+ sh1)) a -> Array sh0 a)
-> (Array (sh0 ::+ sh1) a -> Array (Zero ::+ (sh0 ::+ sh1)) a)
-> Array (sh0 ::+ sh1) a
-> Array sh0 a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((sh0 ::+ sh1) -> Zero ::+ (sh0 ::+ sh1))
-> Array (sh0 ::+ sh1) a -> Array (Zero ::+ (sh0 ::+ sh1)) a
forall sh0 sh1 a. (sh0 -> sh1) -> Array sh0 a -> Array sh1 a
mapShape (\(sh0
sh0 ::+ sh1
sh1) -> (Zero
Shape.Zero Zero -> (sh0 ::+ sh1) -> Zero ::+ (sh0 ::+ sh1)
forall sh0 sh1. sh0 -> sh1 -> sh0 ::+ sh1
::+ sh0
sh0 sh0 -> sh1 -> sh0 ::+ sh1
forall sh0 sh1. sh0 -> sh1 -> sh0 ::+ sh1
::+ sh1
sh1))

takeRight ::
   (Shape.C sh0, Shape.C sh1) =>
   Array (sh0::+sh1) a -> Array sh1 a
takeRight :: forall sh0 sh1 a.
(C sh0, C sh1) =>
Array (sh0 ::+ sh1) a -> Array sh1 a
takeRight =
   Array (sh0 ::+ (sh1 ::+ Zero)) a -> Array sh1 a
forall sh0 sh1 sh2 a.
(C sh0, C sh1, C sh2) =>
Array (sh0 ::+ (sh1 ::+ sh2)) a -> Array sh1 a
takeCenter (Array (sh0 ::+ (sh1 ::+ Zero)) a -> Array sh1 a)
-> (Array (sh0 ::+ sh1) a -> Array (sh0 ::+ (sh1 ::+ Zero)) a)
-> Array (sh0 ::+ sh1) a
-> Array sh1 a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((sh0 ::+ sh1) -> sh0 ::+ (sh1 ::+ Zero))
-> Array (sh0 ::+ sh1) a -> Array (sh0 ::+ (sh1 ::+ Zero)) a
forall sh0 sh1 a. (sh0 -> sh1) -> Array sh0 a -> Array sh1 a
mapShape (\(sh0
sh0 ::+ sh1
sh1) -> (sh0
sh0 sh0 -> (sh1 ::+ Zero) -> sh0 ::+ (sh1 ::+ Zero)
forall sh0 sh1. sh0 -> sh1 -> sh0 ::+ sh1
::+ sh1
sh1 sh1 -> Zero -> sh1 ::+ Zero
forall sh0 sh1. sh0 -> sh1 -> sh0 ::+ sh1
::+ Zero
Shape.Zero))

split ::
   (Shape.C sh0, Shape.C sh1) =>
   Array (sh0::+sh1) a -> (Array sh0 a, Array sh1 a)
split :: forall sh0 sh1 a.
(C sh0, C sh1) =>
Array (sh0 ::+ sh1) a -> (Array sh0 a, Array sh1 a)
split Array (sh0 ::+ sh1) a
x = (Array (sh0 ::+ sh1) a -> Array sh0 a
forall sh0 sh1 a.
(C sh0, C sh1) =>
Array (sh0 ::+ sh1) a -> Array sh0 a
takeLeft Array (sh0 ::+ sh1) a
x, Array (sh0 ::+ sh1) a -> Array sh1 a
forall sh0 sh1 a.
(C sh0, C sh1) =>
Array (sh0 ::+ sh1) a -> Array sh1 a
takeRight Array (sh0 ::+ sh1) a
x)

{- |
prop> \(ArrayChar x) (ArrayChar y) (ArrayChar z) -> let xyz = Array.append x $ Array.append y z in y == Array.takeCenter xyz
-}
takeCenter ::
   (Shape.C sh0, Shape.C sh1, Shape.C sh2) =>
   Array (sh0::+sh1::+sh2) a -> Array sh1 a
takeCenter :: forall sh0 sh1 sh2 a.
(C sh0, C sh1, C sh2) =>
Array (sh0 ::+ (sh1 ::+ sh2)) a -> Array sh1 a
takeCenter (Array (sh0
sh0::+sh1
sh1::+sh2
_sh2) Array a
x) =
   sh1 -> Array a -> Array sh1 a
forall sh a. sh -> Array a -> Array sh a
Array sh1
sh1 (Array a -> Array sh1 a) -> Array a -> Array sh1 a
forall a b. (a -> b) -> a -> b
$ Array a -> Int -> Int -> Array a
forall a. Array a -> Int -> Int -> Array a
Prim.cloneArray Array a
x (sh0 -> Int
forall sh. C sh => sh -> Int
Shape.size sh0
sh0) (sh1 -> Int
forall sh. C sh => sh -> Int
Shape.size sh1
sh1)