{-# 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 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
-}


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
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
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
$c== :: forall sh a. (Eq sh, Eq a) => 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
pforall a. Ord a => a -> a -> Bool
>Int
10) forall a b. (a -> b) -> a -> b
$
         String -> ShowS
showString String
"BoxedArray.fromList " forall b c a. (b -> c) -> (a -> b) -> a -> c
.
         forall a. Show a => Int -> a -> ShowS
showsPrec Int
11 (forall sh a. Array sh a -> sh
shape Array sh a
arr) forall b c a. (b -> c) -> (a -> b) -> a -> c
.
         Char -> ShowS
showChar Char
' ' forall b c a. (b -> c) -> (a -> b) -> a -> c
.
         forall a. Show a => a -> ShowS
shows (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) = forall a. NFData a => a -> ()
rnf (sh
sh, 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 = 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 = forall sh a. C sh => sh -> a -> Array sh a
replicate forall sh. Static sh => sh
Shape.static
   <*> :: forall 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 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 = forall (t :: * -> *) m. (Foldable t, Monoid m) => t m -> m
Fold.fold forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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 = forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
Fold.foldMap a -> m
f forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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 = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
Fold.foldl b -> a -> b
f b
a forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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 = forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
Fold.foldr a -> b -> b
f b
a forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall sh a. Array sh a -> Array a
buffer
   foldl1 :: forall a. (a -> a -> a) -> Array sh a -> a
foldl1 a -> a -> a
f = forall (t :: * -> *) a. Foldable t => (a -> a -> a) -> t a -> a
Fold.foldl1 a -> a -> a
f forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall sh a. Array sh a -> Array a
buffer
   foldr1 :: forall a. (a -> a -> a) -> Array sh a -> a
foldr1 a -> a -> a
f = forall (t :: * -> *) a. Foldable t => (a -> a -> a) -> t a -> a
Fold.foldr1 a -> a -> a
f forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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) = forall sh a. sh -> Array a -> Array sh a
Array sh
sh forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t 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) = forall sh a. sh -> Array a -> Array sh a
Array sh
sh forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (f :: * -> *) a.
(Traversable t, Applicative f) =>
t (f a) -> f (t 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) = forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
liftM (forall sh a. sh -> Array a -> Array sh a
Array sh
sh) forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t 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) = forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
liftM (forall sh a. sh -> Array a -> Array sh a
Array sh
sh) forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t 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) = 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) = 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 = forall a. Array a -> Int -> a
Prim.indexArray Array a
arr forall a b. (a -> b) -> a -> b
$ 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) =
   forall a b. (a -> b) -> [a] -> [b]
List.map (forall a. Array a -> Int -> a
Prim.indexArray Array a
arr) forall a b. (a -> b) -> a -> b
$ forall a. Int -> [a] -> [a]
List.take (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 = forall sh a. sh -> Array a -> Array sh a
Array sh
sh forall a b. (a -> b) -> a -> b
$ forall l. IsList l => Int -> [Item l] -> l
Prim.fromListN (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 = forall l. IsList l => [Item l] -> l
Prim.fromList [a]
xs
   in forall sh a. sh -> Array a -> Array sh a
Array (forall n. n -> ZeroBased n
Shape.ZeroBased forall a b. (a -> b) -> a -> b
$ 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 =
   forall sh a. sh -> Array a -> Array sh a
Array sh
sh forall a b. (a -> b) -> a -> b
$
   forall a. (forall s. ST s a) -> a
ST.runST (forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> m (Array a)
Prim.unsafeFreezeArray  forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (MutableArray (PrimState m) a)
Prim.newArray (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) = forall sh a. sh -> Array a -> Array sh a
Array sh
sh forall a b. (a -> b) -> a -> b
$
   let n :: Int
n = forall sh. C sh => sh -> Int
Shape.size sh
sh
   in forall l. IsList l => Int -> [Item l] -> l
Prim.fromListN Int
n forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
List.map (a -> b
f forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Array a -> Int -> a
Prim.indexArray Array a
arr) forall a b. (a -> b) -> a -> b
$ 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) = forall sh a. sh -> Array a -> Array sh a
Array sh
sha forall a b. (a -> b) -> a -> b
$
   let n :: Int
n = forall sh. C sh => sh -> Int
Shape.size sh
sha
   in forall l. IsList l => Int -> [Item l] -> l
Prim.fromListN Int
n forall a b. (a -> b) -> a -> b
$
      forall a b. (a -> b) -> [a] -> [b]
List.map (\Int
k -> a -> b -> c
f (forall a. Array a -> Int -> a
Prim.indexArray Array a
arra Int
k) (forall a. Array a -> Int -> a
Prim.indexArray Array b
arrb Int
k)) forall a b. (a -> b) -> a -> b
$
      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 = forall sh. C sh => sh -> Int
Shape.size shx
shX in
   let sizeY :: Int
sizeY = forall sh. C sh => sh -> Int
Shape.size shy
shY in
   forall sh a. sh -> Array a -> Array sh a
Array (shx
shXforall sh0 sh1. sh0 -> sh1 -> sh0 ::+ sh1
::+shy
shY) forall a b. (a -> b) -> a -> b
$
   forall a. (forall s. ST s a) -> a
ST.runST (do
      MutableArray s a
arr <-
         forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (MutableArray (PrimState m) a)
Prim.newArray (Int
sizeXforall a. Num a => a -> a -> a
+Int
sizeY)
            (forall a. HasCallStack => String -> a
error String
"Boxed.append: uninitialized element")
      forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a
-> Int -> Array a -> Int -> Int -> m ()
Prim.copyArray MutableArray s a
arr Int
0 Array a
x Int
0 Int
sizeX
      forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a
-> Int -> Array a -> Int -> Int -> m ()
Prim.copyArray MutableArray s a
arr Int
sizeX Array a
y Int
0 Int
sizeY
      forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> m (Array a)
Prim.unsafeFreezeArray MutableArray 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 = forall sh0 sh1 a.
(C sh0, C sh1) =>
Array (sh0 ::+ sh1) a -> Array sh0 a
takeLeft forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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 = forall sh0 sh1 a.
(C sh0, C sh1) =>
Array (sh0 ::+ sh1) a -> Array sh1 a
takeRight forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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 = forall sh0 sh1 a. (sh0 -> sh1) -> Array sh0 a -> Array sh1 a
mapShape (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 =
   forall sh0 sh1 sh2 a.
(C sh0, C sh1, C sh2) =>
Array (sh0 ::+ (sh1 ::+ sh2)) a -> Array sh1 a
takeCenter forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall sh0 sh1 a. (sh0 -> sh1) -> Array sh0 a -> Array sh1 a
mapShape (\(sh0
sh0 ::+ sh1
sh1) -> (Zero
Shape.Zero forall sh0 sh1. sh0 -> sh1 -> sh0 ::+ sh1
::+ sh0
sh0 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 =
   forall sh0 sh1 sh2 a.
(C sh0, C sh1, C sh2) =>
Array (sh0 ::+ (sh1 ::+ sh2)) a -> Array sh1 a
takeCenter forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall sh0 sh1 a. (sh0 -> sh1) -> Array sh0 a -> Array sh1 a
mapShape (\(sh0
sh0 ::+ sh1
sh1) -> (sh0
sh0 forall sh0 sh1. sh0 -> sh1 -> sh0 ::+ sh1
::+ sh1
sh1 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 = (forall sh0 sh1 a.
(C sh0, C sh1) =>
Array (sh0 ::+ sh1) a -> Array sh0 a
takeLeft Array (sh0 ::+ sh1) a
x, 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) =
   forall sh a. sh -> Array a -> Array sh a
Array sh1
sh1 forall a b. (a -> b) -> a -> b
$ forall a. Array a -> Int -> Int -> Array a
Prim.cloneArray Array a
x (forall sh. C sh => sh -> Int
Shape.size sh0
sh0) (forall sh. C sh => sh -> Int
Shape.size sh1
sh1)