{-# LANGUAGE TypeFamilies #-}
module Data.Array.Comfort.Storable (
   Array,
   shape,
   reshape,
   mapShape,

   accessMaybe, (!),
   Array.toList,
   Array.vectorFromList,
   toAssociations,
   fromList,
   fromMap, toMap,
   fromContainer,
   toContainer,
   sample,
   fromBoxed,
   toBoxed,
   fromStorableVector,
   toStorableVector,

   Array.map,
   Array.mapWithIndex,
   zipWith,
   (//),
   accumulate,
   fromAssociations,

   pick,
   toRowArray,
   fromRowArray,
   Array.singleton,
   Array.append,
   Array.take, Array.drop,
   Array.takeLeft, Array.takeRight, Array.split,
   Array.takeCenter,

   Array.sum, Array.product,
   minimum, argMinimum,
   maximum, argMaximum,
   limits,
   Array.foldl,
   foldl1,
   foldMap,
   ) where

import qualified Data.Array.Comfort.Storable.Mutable.Unchecked as MutArrayNC
import qualified Data.Array.Comfort.Storable.Mutable as MutArray
import qualified Data.Array.Comfort.Storable.Unchecked as Array
import qualified Data.Array.Comfort.Storable.Memory as Memory
import qualified Data.Array.Comfort.Container as Container
import qualified Data.Array.Comfort.Boxed as BoxedArray
import qualified Data.Array.Comfort.Check as Check
import qualified Data.Array.Comfort.Shape as Shape
import Data.Array.Comfort.Storable.Unchecked (Array(Array))
import Data.Array.Comfort.Shape ((::+)((::+)))

import System.IO.Unsafe (unsafePerformIO)
import Foreign.Marshal.Array (copyArray, advancePtr)
import Foreign.Storable (Storable)
import Foreign.ForeignPtr (withForeignPtr)

import Control.Monad.ST (runST)

import qualified Data.StorableVector.Base as SVB
import qualified Data.Map as Map
import qualified Data.Set as Set
import qualified Data.Traversable as Trav
import qualified Data.Foldable as Fold
import qualified Data.List as List
import qualified Data.Tuple.Strict as StrictTuple
import Data.Map (Map)
import Data.Set (Set)
import Data.Foldable (forM_)
import Data.Maybe (fromMaybe)
import Data.Semigroup
         (Semigroup, (<>), Min(Min,getMin), Max(Max,getMax), Arg(Arg))

import Prelude2010 hiding (map, zipWith, foldl1, minimum, maximum)
import Prelude ()


{- $setup
>>> import qualified Data.Array.Comfort.Storable as Array
>>> import qualified Data.Array.Comfort.Shape as Shape
>>> import Data.Array.Comfort.Storable (Array, (!))
>>>
>>> import qualified Test.QuickCheck as QC
>>> import Test.ChasingBottoms.IsBottom (isBottom)
>>>
>>> import Control.Applicative ((<$>), (<*>))
>>>
>>> import Data.Word (Word16)
>>>
>>> type ShapeInt = Shape.ZeroBased Int
>>>
>>> genArray :: QC.Gen (Array ShapeInt Word16)
>>> genArray = Array.vectorFromList <$> QC.arbitrary
>>>
>>> genArray2 :: QC.Gen (Array (ShapeInt,ShapeInt) Word16)
>>> genArray2 = do
>>>    xs <- QC.arbitrary
>>>    let n = length xs
>>>    (k,m) <-
>>>       if n == 0
>>>          then QC.elements [(,) 0, flip (,) 0] <*> QC.choose (1,n)
>>>          else fmap (\m -> (div n m, m)) $ QC.choose (1,n)
>>>    return $ Array.fromList (Shape.ZeroBased k, Shape.ZeroBased m) xs
>>>
>>> genNonEmptyArray2 :: QC.Gen (Array (ShapeInt,ShapeInt) Word16)
>>> genNonEmptyArray2 = do
>>>    xs <- QC.getNonEmpty <$> QC.arbitrary
>>>    let n = length xs
>>>    m <- QC.choose (1,n)
>>>    return $ Array.fromList (Shape.ZeroBased (div n m), Shape.ZeroBased m) xs
>>>
>>> infix 4 ==?
>>> (==?) :: a -> a -> (a,a)
>>> (==?) = (,)
>>>
>>> forAllNonEmpty :: (Eq b) => (Array ShapeInt Word16 -> (b,b)) -> QC.Property
>>> forAllNonEmpty f =
>>>    QC.forAll genArray $ \xs ->
>>>    case f xs of
>>>       (resultArray,resultList) ->
>>>          if Array.shape xs == Shape.ZeroBased 0
>>>             then isBottom resultArray
>>>             else resultArray == resultList
-}


shape :: Array sh a -> sh
shape :: Array sh a -> sh
shape = Array sh a -> sh
forall sh a. Array sh a -> sh
Array.shape

reshape :: (Shape.C sh0, Shape.C sh1) => sh1 -> Array sh0 a -> Array sh1 a
reshape :: sh1 -> Array sh0 a -> Array sh1 a
reshape = String
-> (Array sh0 a -> sh0)
-> (sh1 -> Array sh0 a -> Array sh1 a)
-> sh1
-> Array sh0 a
-> Array sh1 a
forall sh0 sh1 array0 array1.
(C sh0, C sh1) =>
String
-> (array0 -> sh0)
-> (sh1 -> array0 -> array1)
-> sh1
-> array0
-> array1
Check.reshape String
"Storable" Array sh0 a -> sh0
forall sh a. Array sh a -> sh
shape sh1 -> Array sh0 a -> Array sh1 a
forall sh1 sh0 a. sh1 -> Array sh0 a -> Array sh1 a
Array.reshape

mapShape ::
   (Shape.C sh0, Shape.C sh1) => (sh0 -> sh1) -> Array sh0 a -> Array sh1 a
mapShape :: (sh0 -> sh1) -> Array sh0 a -> Array sh1 a
mapShape sh0 -> sh1
f Array sh0 a
arr = sh1 -> Array sh0 a -> Array sh1 a
forall sh0 sh1 a.
(C sh0, C sh1) =>
sh1 -> Array sh0 a -> Array sh1 a
reshape (sh0 -> sh1
f (sh0 -> sh1) -> sh0 -> sh1
forall a b. (a -> b) -> a -> b
$ Array sh0 a -> sh0
forall sh a. Array sh a -> sh
shape Array sh0 a
arr) Array sh0 a
arr


{- |
>>> Array.fromList (Shape.ZeroBased (5::Int)) ['a'..]
StorableArray.fromList (ZeroBased {zeroBasedSize = 5}) "abcde"
-}
fromList :: (Shape.C sh, Storable a) => sh -> [a] -> Array sh a
fromList :: sh -> [a] -> Array sh a
fromList sh
sh [a]
arr = (forall s. ST s (Array sh a)) -> Array sh a
forall a. (forall s. ST s a) -> a
runST (Array (ST s) sh a -> ST s (Array sh a)
forall (m :: * -> *) sh a.
(PrimMonad m, C sh, Storable a) =>
Array m sh a -> m (Array sh a)
MutArrayNC.unsafeFreeze (Array (ST s) sh a -> ST s (Array sh a))
-> ST s (Array (ST s) sh a) -> ST s (Array sh a)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< sh -> [a] -> ST s (Array (ST s) sh a)
forall (m :: * -> *) sh a.
(PrimMonad m, C sh, Storable a) =>
sh -> [a] -> m (Array m sh a)
MutArray.fromList sh
sh [a]
arr)

fromMap :: (Ord k, Storable a) => Map k a -> Array (Set k) a
fromMap :: Map k a -> Array (Set k) a
fromMap Map k a
m = Set k -> [a] -> Array (Set k) a
forall sh a. (C sh, Storable a) => sh -> [a] -> Array sh a
fromList (Map k a -> Set k
forall k a. Map k a -> Set k
Map.keysSet Map k a
m) (Map k a -> [a]
forall k a. Map k a -> [a]
Map.elems Map k a
m)

toMap :: (Ord k, Storable a) => Array (Set k) a -> Map k a
toMap :: Array (Set k) a -> Map k a
toMap Array (Set k) a
arr = [(k, a)] -> Map k a
forall k a. Eq k => [(k, a)] -> Map k a
Map.fromAscList ([(k, a)] -> Map k a) -> [(k, a)] -> Map k a
forall a b. (a -> b) -> a -> b
$ [k] -> [a] -> [(k, a)]
forall a b. [a] -> [b] -> [(a, b)]
zip (Set k -> [k]
forall a. Set a -> [a]
Set.toAscList (Set k -> [k]) -> Set k -> [k]
forall a b. (a -> b) -> a -> b
$ Array (Set k) a -> Set k
forall sh a. Array sh a -> sh
shape Array (Set k) a
arr) (Array (Set k) a -> [a]
forall sh a. (C sh, Storable a) => Array sh a -> [a]
Array.toList Array (Set k) a
arr)

fromContainer ::
   (Container.C f, Storable a) => f a -> Array (Container.Shape f) a
fromContainer :: f a -> Array (Shape f) a
fromContainer f a
xs = Shape f -> [a] -> Array (Shape f) a
forall sh a. (C sh, Storable a) => sh -> [a] -> Array sh a
fromList (f a -> Shape f
forall (f :: * -> *) a. C f => f a -> Shape f
Container.toShape f a
xs) (f a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
Fold.toList f a
xs)

toContainer ::
   (Container.C f, Storable a) => Array (Container.Shape f) a -> f a
toContainer :: Array (Shape f) a -> f a
toContainer Array (Shape f) a
arr = Shape f -> [a] -> f a
forall (f :: * -> *) a. C f => Shape f -> [a] -> f a
Container.fromList (Array (Shape f) a -> Shape f
forall sh a. Array sh a -> sh
Array.shape Array (Shape f) a
arr) (Array (Shape f) a -> [a]
forall sh a. (C sh, Storable a) => Array sh a -> [a]
Array.toList Array (Shape f) a
arr)

sample ::
   (Shape.Indexed sh, Storable a) => sh -> (Shape.Index sh -> a) -> Array sh a
sample :: sh -> (Index sh -> a) -> Array sh a
sample sh
sh Index sh -> a
f = sh -> [a] -> Array sh a
forall sh a. (C sh, Storable a) => sh -> [a] -> Array sh a
Array.fromList sh
sh ([a] -> Array sh a) -> [a] -> Array sh a
forall a b. (a -> b) -> a -> b
$ (Index sh -> a) -> [Index sh] -> [a]
forall a b. (a -> b) -> [a] -> [b]
List.map Index sh -> a
f ([Index sh] -> [a]) -> [Index sh] -> [a]
forall a b. (a -> b) -> a -> b
$ sh -> [Index sh]
forall sh. Indexed sh => sh -> [Index sh]
Shape.indices sh
sh


fromBoxed :: (Shape.C sh, Storable a) => BoxedArray.Array sh a -> Array sh a
fromBoxed :: Array sh a -> Array sh a
fromBoxed Array sh a
arr = sh -> [a] -> Array sh a
forall sh a. (C sh, Storable a) => sh -> [a] -> Array sh a
Array.fromList (Array sh a -> sh
forall sh a. Array sh a -> sh
BoxedArray.shape Array sh a
arr) ([a] -> Array sh a) -> [a] -> Array sh a
forall a b. (a -> b) -> a -> b
$ Array sh a -> [a]
forall sh a. C sh => Array sh a -> [a]
BoxedArray.toList Array sh a
arr

toBoxed :: (Shape.C sh, Storable a) => Array sh a -> BoxedArray.Array sh a
toBoxed :: Array sh a -> Array sh a
toBoxed Array sh a
arr = sh -> [a] -> Array sh a
forall sh a. C sh => sh -> [a] -> Array sh a
BoxedArray.fromList (Array sh a -> sh
forall sh a. Array sh a -> sh
Array.shape Array sh a
arr) ([a] -> Array sh a) -> [a] -> Array sh a
forall a b. (a -> b) -> a -> b
$ Array sh a -> [a]
forall sh a. (C sh, Storable a) => Array sh a -> [a]
Array.toList Array sh a
arr


fromStorableVector ::
   (Storable a) => SVB.Vector a -> Array (Shape.ZeroBased Int) a
fromStorableVector :: Vector a -> Array (ZeroBased Int) a
fromStorableVector Vector a
xs =
   case Vector a -> (ForeignPtr a, Int, Int)
forall a. Vector a -> (ForeignPtr a, Int, Int)
SVB.toForeignPtr Vector a
xs of
      (ForeignPtr a
fptr,Int
0,Int
n) -> ZeroBased Int -> ForeignPtr a -> Array (ZeroBased Int) a
forall sh a. sh -> ForeignPtr a -> Array sh a
Array (Int -> ZeroBased Int
forall n. n -> ZeroBased n
Shape.ZeroBased Int
n) ForeignPtr a
fptr
      (ForeignPtr a
fptr,Int
s,Int
n) ->
         Array (ZeroBased Int ::+ ZeroBased Int) a
-> Array (ZeroBased Int) a
forall sh0 sh1 a.
(C sh0, C sh1, Storable a) =>
Array (sh0 ::+ sh1) a -> Array sh1 a
Array.takeRight (Array (ZeroBased Int ::+ ZeroBased Int) a
 -> Array (ZeroBased Int) a)
-> Array (ZeroBased Int ::+ ZeroBased Int) a
-> Array (ZeroBased Int) a
forall a b. (a -> b) -> a -> b
$
         (ZeroBased Int ::+ ZeroBased Int)
-> ForeignPtr a -> Array (ZeroBased Int ::+ ZeroBased Int) a
forall sh a. sh -> ForeignPtr a -> Array sh a
Array (Int -> ZeroBased Int
forall n. n -> ZeroBased n
Shape.ZeroBased Int
s ZeroBased Int -> ZeroBased Int -> ZeroBased Int ::+ ZeroBased Int
forall sh0 sh1. sh0 -> sh1 -> sh0 ::+ sh1
::+ Int -> ZeroBased Int
forall n. n -> ZeroBased n
Shape.ZeroBased Int
n) ForeignPtr a
fptr

toStorableVector :: (Shape.C sh, Storable a) => Array sh a -> SVB.Vector a
toStorableVector :: Array sh a -> Vector a
toStorableVector (Array sh
sh ForeignPtr a
fptr) =
   ForeignPtr a -> Int -> Vector a
forall a. ForeignPtr a -> Int -> Vector a
SVB.fromForeignPtr ForeignPtr a
fptr (Int -> Vector a) -> Int -> Vector a
forall a b. (a -> b) -> a -> b
$ sh -> Int
forall sh. C sh => sh -> Int
Shape.size sh
sh


toAssociations ::
   (Shape.Indexed sh, Storable a) => Array sh a -> [(Shape.Index sh, a)]
toAssociations :: Array sh a -> [(Index sh, a)]
toAssociations Array sh a
arr = [Index sh] -> [a] -> [(Index sh, a)]
forall a b. [a] -> [b] -> [(a, b)]
zip (sh -> [Index sh]
forall sh. Indexed sh => sh -> [Index sh]
Shape.indices (sh -> [Index sh]) -> sh -> [Index sh]
forall a b. (a -> b) -> a -> b
$ Array sh a -> sh
forall sh a. Array sh a -> sh
shape Array sh a
arr) (Array sh a -> [a]
forall sh a. (C sh, Storable a) => Array sh a -> [a]
Array.toList Array sh a
arr)


errorArray :: String -> String -> a
errorArray :: String -> String -> a
errorArray String
name String
msg =
   String -> a
forall a. HasCallStack => String -> a
error (String
"Array.Comfort.Storable." String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
name String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
": " String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
msg)

infixl 9 !

(!) :: (Shape.Indexed sh, Storable a) => Array sh a -> Shape.Index sh -> a
(!) Array sh a
arr Index sh
ix =
   a -> Maybe a -> a
forall a. a -> Maybe a -> a
fromMaybe (String -> String -> a
forall a. String -> String -> a
errorArray String
"!" String
"index out of bounds") (Maybe a -> a) -> Maybe a -> a
forall a b. (a -> b) -> a -> b
$
   Array sh a -> Index sh -> Maybe a
forall sh a.
(Indexed sh, Storable a) =>
Array sh a -> Index sh -> Maybe a
accessMaybe Array sh a
arr Index sh
ix

accessMaybe ::
   (Shape.Indexed sh, Storable a) => Array sh a -> Shape.Index sh -> Maybe a
accessMaybe :: Array sh a -> Index sh -> Maybe a
accessMaybe Array sh a
arr Index sh
ix = (forall s. ST s (Maybe a)) -> Maybe a
forall a. (forall s. ST s a) -> a
runST (do
   Array (ST s) sh a
marr <- Array sh a -> ST s (Array (ST s) sh a)
forall (m :: * -> *) sh a.
(PrimMonad m, C sh, Storable a) =>
Array sh a -> m (Array m sh a)
MutArrayNC.unsafeThaw Array sh a
arr
   Maybe (ST s a) -> ST s (Maybe a)
forall (t :: * -> *) (f :: * -> *) a.
(Traversable t, Applicative f) =>
t (f a) -> f (t a)
Trav.sequenceA (Maybe (ST s a) -> ST s (Maybe a))
-> Maybe (ST s a) -> ST s (Maybe a)
forall a b. (a -> b) -> a -> b
$ Array (ST s) sh a -> Index sh -> Maybe (ST s a)
forall (m :: * -> *) sh a.
(PrimMonad m, Indexed sh, Storable a) =>
Array m sh a -> Index sh -> Maybe (m a)
MutArray.readMaybe Array (ST s) sh a
marr Index sh
ix)


zipWith ::
   (Shape.C sh, Eq sh, Storable a, Storable b, Storable c) =>
   (a -> b -> c) -> Array sh a -> Array sh b -> Array sh c
zipWith :: (a -> b -> c) -> Array sh a -> Array sh b -> Array sh c
zipWith a -> b -> c
f Array sh a
a Array sh b
b =
   if Array sh a -> sh
forall sh a. Array sh a -> sh
shape Array sh a
a sh -> sh -> Bool
forall a. Eq a => a -> a -> Bool
== Array sh b -> sh
forall sh a. Array sh a -> sh
shape Array sh b
b
      then (a -> b -> c) -> Array sh a -> Array sh b -> Array sh c
forall sh a b c.
(C sh, Storable a, Storable b, Storable c) =>
(a -> b -> c) -> Array sh a -> Array sh b -> Array sh c
Array.zipWith a -> b -> c
f Array sh a
a Array sh b
b
      else String -> String -> Array sh c
forall a. String -> String -> a
errorArray String
"zipWith" String
"shapes mismatch"

(//) ::
   (Shape.Indexed sh, Storable a) =>
   Array sh a -> [(Shape.Index sh, a)] -> Array sh a
// :: Array sh a -> [(Index sh, a)] -> Array sh a
(//) Array sh a
arr [(Index sh, a)]
xs = (forall s. ST s (Array sh a)) -> Array sh a
forall a. (forall s. ST s a) -> a
runST (do
   Array (ST s) sh a
marr <- Array sh a -> ST s (Array (ST s) sh a)
forall (m :: * -> *) sh a.
(PrimMonad m, C sh, Storable a) =>
Array sh a -> m (Array m sh a)
MutArray.thaw Array sh a
arr
   [(Index sh, a)] -> ((Index sh, a) -> ST s ()) -> ST s ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [(Index sh, a)]
xs (((Index sh, a) -> ST s ()) -> ST s ())
-> ((Index sh, a) -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ (Index sh -> a -> ST s ()) -> (Index sh, a) -> ST s ()
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry ((Index sh -> a -> ST s ()) -> (Index sh, a) -> ST s ())
-> (Index sh -> a -> ST s ()) -> (Index sh, a) -> ST s ()
forall a b. (a -> b) -> a -> b
$ Array (ST s) sh a -> Index sh -> a -> ST s ()
forall (m :: * -> *) sh a.
(PrimMonad m, Indexed sh, Storable a) =>
Array m sh a -> Index sh -> a -> m ()
MutArray.write Array (ST s) sh a
marr
   Array (ST s) sh a -> ST s (Array sh a)
forall (m :: * -> *) sh a.
(PrimMonad m, C sh, Storable a) =>
Array m sh a -> m (Array sh a)
MutArrayNC.unsafeFreeze Array (ST s) sh a
marr)

accumulate ::
   (Shape.Indexed sh, Storable a) =>
   (a -> b -> a) -> Array sh a -> [(Shape.Index sh, b)] -> Array sh a
accumulate :: (a -> b -> a) -> Array sh a -> [(Index sh, b)] -> Array sh a
accumulate a -> b -> a
f Array sh a
arr [(Index sh, b)]
xs = (forall s. ST s (Array sh a)) -> Array sh a
forall a. (forall s. ST s a) -> a
runST (do
   Array (ST s) sh a
marr <- Array sh a -> ST s (Array (ST s) sh a)
forall (m :: * -> *) sh a.
(PrimMonad m, C sh, Storable a) =>
Array sh a -> m (Array m sh a)
MutArray.thaw Array sh a
arr
   [(Index sh, b)] -> ((Index sh, b) -> ST s ()) -> ST s ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [(Index sh, b)]
xs (((Index sh, b) -> ST s ()) -> ST s ())
-> ((Index sh, b) -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \(Index sh
ix,b
b) -> Array (ST s) sh a -> Index sh -> (a -> a) -> ST s ()
forall (m :: * -> *) sh a.
(PrimMonad m, Indexed sh, Storable a) =>
Array m sh a -> Index sh -> (a -> a) -> m ()
MutArray.update Array (ST s) sh a
marr Index sh
ix ((a -> a) -> ST s ()) -> (a -> a) -> ST s ()
forall a b. (a -> b) -> a -> b
$ (a -> b -> a) -> b -> a -> a
forall a b c. (a -> b -> c) -> b -> a -> c
flip a -> b -> a
f b
b
   Array (ST s) sh a -> ST s (Array sh a)
forall (m :: * -> *) sh a.
(PrimMonad m, C sh, Storable a) =>
Array m sh a -> m (Array sh a)
MutArrayNC.unsafeFreeze Array (ST s) sh a
marr)

fromAssociations ::
   (Shape.Indexed sh, Storable a) =>
   a -> sh -> [(Shape.Index sh, a)] -> Array sh a
fromAssociations :: a -> sh -> [(Index sh, a)] -> Array sh a
fromAssociations a
a sh
sh [(Index sh, a)]
xs = (forall s. ST s (Array sh a)) -> Array sh a
forall a. (forall s. ST s a) -> a
runST (do
   Array (ST s) sh a
marr <- sh -> a -> ST s (Array (ST s) sh a)
forall (m :: * -> *) sh a.
(PrimMonad m, C sh, Storable a) =>
sh -> a -> m (Array m sh a)
MutArray.new sh
sh a
a
   [(Index sh, a)] -> ((Index sh, a) -> ST s ()) -> ST s ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [(Index sh, a)]
xs (((Index sh, a) -> ST s ()) -> ST s ())
-> ((Index sh, a) -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ (Index sh -> a -> ST s ()) -> (Index sh, a) -> ST s ()
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry ((Index sh -> a -> ST s ()) -> (Index sh, a) -> ST s ())
-> (Index sh -> a -> ST s ()) -> (Index sh, a) -> ST s ()
forall a b. (a -> b) -> a -> b
$ Array (ST s) sh a -> Index sh -> a -> ST s ()
forall (m :: * -> *) sh a.
(PrimMonad m, Indexed sh, Storable a) =>
Array m sh a -> Index sh -> a -> m ()
MutArray.write Array (ST s) sh a
marr
   Array (ST s) sh a -> ST s (Array sh a)
forall (m :: * -> *) sh a.
(PrimMonad m, C sh, Storable a) =>
Array m sh a -> m (Array sh a)
MutArrayNC.unsafeFreeze Array (ST s) sh a
marr)


{- |
prop> QC.forAll genNonEmptyArray2 $ \xs -> QC.forAll (QC.elements $ Shape.indices $ Array.shape xs) $ \(ix0,ix1) -> Array.pick xs ix0 ! ix1 == xs!(ix0,ix1)
-}
pick ::
   (Shape.Indexed sh0, Shape.C sh1, Storable a) =>
   Array (sh0,sh1) a -> Shape.Index sh0 -> Array sh1 a
pick :: Array (sh0, sh1) a -> Index sh0 -> Array sh1 a
pick (Array (sh0
sh0,sh1
sh1) ForeignPtr a
x) Index sh0
ix0 =
   sh1 -> (Int -> Ptr a -> IO ()) -> Array sh1 a
forall sh a.
(C sh, Storable a) =>
sh -> (Int -> Ptr a -> IO ()) -> Array sh a
Array.unsafeCreateWithSize sh1
sh1 ((Int -> Ptr a -> IO ()) -> Array sh1 a)
-> (Int -> Ptr a -> IO ()) -> Array sh1 a
forall a b. (a -> b) -> a -> b
$ \Int
k Ptr a
yPtr ->
   ForeignPtr a -> (Ptr a -> IO ()) -> IO ()
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr a
x ((Ptr a -> IO ()) -> IO ()) -> (Ptr a -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Ptr a
xPtr ->
      Ptr a -> Ptr a -> Int -> IO ()
forall a. Storable a => Ptr a -> Ptr a -> Int -> IO ()
copyArray Ptr a
yPtr (Ptr a -> Int -> Ptr a
forall a. Storable a => Ptr a -> Int -> Ptr a
advancePtr Ptr a
xPtr (sh0 -> Index sh0 -> Int
forall sh. Indexed sh => sh -> Index sh -> Int
Shape.offset sh0
sh0 Index sh0
ix0 Int -> Int -> Int
forall a. Num a => a -> a -> a
* sh1 -> Int
forall sh. C sh => sh -> Int
Shape.size sh1
sh1)) Int
k

toRowArray ::
   (Shape.Indexed sh0, Shape.C sh1, Storable a) =>
   Array (sh0,sh1) a -> BoxedArray.Array sh0 (Array sh1 a)
toRowArray :: Array (sh0, sh1) a -> Array sh0 (Array sh1 a)
toRowArray Array (sh0, sh1) a
x = (Index sh0 -> Array sh1 a)
-> Array sh0 (Index sh0) -> Array sh0 (Array sh1 a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Array (sh0, sh1) a -> Index sh0 -> Array sh1 a
forall sh0 sh1 a.
(Indexed sh0, C sh1, Storable a) =>
Array (sh0, sh1) a -> Index sh0 -> Array sh1 a
pick Array (sh0, sh1) a
x) (Array sh0 (Index sh0) -> Array sh0 (Array sh1 a))
-> Array sh0 (Index sh0) -> Array sh0 (Array sh1 a)
forall a b. (a -> b) -> a -> b
$ sh0 -> Array sh0 (Index sh0)
forall sh. Indexed sh => sh -> Array sh (Index sh)
BoxedArray.indices (sh0 -> Array sh0 (Index sh0)) -> sh0 -> Array sh0 (Index sh0)
forall a b. (a -> b) -> a -> b
$ (sh0, sh1) -> sh0
forall a b. (a, b) -> a
fst ((sh0, sh1) -> sh0) -> (sh0, sh1) -> sh0
forall a b. (a -> b) -> a -> b
$ Array (sh0, sh1) a -> (sh0, sh1)
forall sh a. Array sh a -> sh
Array.shape Array (sh0, sh1) a
x

{- |
It is a checked error if a row width differs from the result array width.

prop> QC.forAll genArray2 $ \xs -> xs == Array.fromRowArray (snd $ Array.shape xs) (Array.toRowArray xs)
-}
fromRowArray ::
   (Shape.C sh0, Shape.C sh1, Eq sh1, Storable a) =>
   sh1 -> BoxedArray.Array sh0 (Array sh1 a) -> Array (sh0,sh1) a
fromRowArray :: sh1 -> Array sh0 (Array sh1 a) -> Array (sh0, sh1) a
fromRowArray sh1
sh1 Array sh0 (Array sh1 a)
x =
   (sh0, sh1) -> (Ptr a -> IO ()) -> Array (sh0, sh1) a
forall sh a.
(C sh, Storable a) =>
sh -> (Ptr a -> IO ()) -> Array sh a
Array.unsafeCreate (Array sh0 (Array sh1 a) -> sh0
forall sh a. Array sh a -> sh
BoxedArray.shape Array sh0 (Array sh1 a)
x, sh1
sh1) ((Ptr a -> IO ()) -> Array (sh0, sh1) a)
-> (Ptr a -> IO ()) -> Array (sh0, sh1) a
forall a b. (a -> b) -> a -> b
$ \Ptr a
yPtr ->
   let k :: Int
k = sh1 -> Int
forall sh. C sh => sh -> Int
Shape.size sh1
sh1 in
   [(Int, Array sh1 a)] -> ((Int, Array sh1 a) -> IO ()) -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([Int] -> [Array sh1 a] -> [(Int, Array sh1 a)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0,Int
k..] (Array sh0 (Array sh1 a) -> [Array sh1 a]
forall sh a. C sh => Array sh a -> [a]
BoxedArray.toList Array sh0 (Array sh1 a)
x)) (((Int, Array sh1 a) -> IO ()) -> IO ())
-> ((Int, Array sh1 a) -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \(Int
j, Array sh1
sh1i ForeignPtr a
row) ->
   if sh1
sh1 sh1 -> sh1 -> Bool
forall a. Eq a => a -> a -> Bool
== sh1
sh1i
      then ForeignPtr a -> (Ptr a -> IO ()) -> IO ()
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr a
row ((Ptr a -> IO ()) -> IO ()) -> (Ptr a -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Ptr a
xPtr -> Ptr a -> Ptr a -> Int -> IO ()
forall a. Storable a => Ptr a -> Ptr a -> Int -> IO ()
copyArray (Ptr a -> Int -> Ptr a
forall a. Storable a => Ptr a -> Int -> Ptr a
advancePtr Ptr a
yPtr Int
j) Ptr a
xPtr Int
k
      else String -> String -> IO ()
forall a. String -> String -> a
errorArray String
"fromRowArray" String
"mismatching row width"


{- |
It is a checked error if the vector is empty.

prop> forAllNonEmpty $ \xs -> Array.minimum xs ==? minimum (Array.toList xs)
-}
minimum :: (Shape.C sh, Storable a, Ord a) => Array sh a -> a
minimum :: Array sh a -> a
minimum = (a -> a -> a) -> Array sh a -> a
forall sh a. (C sh, Storable a) => (a -> a -> a) -> Array sh a -> a
foldl1 a -> a -> a
forall a. Ord a => a -> a -> a
min

{- |
It is a checked error if the vector is empty.

prop> forAllNonEmpty $ \xs -> Array.maximum xs ==? maximum (Array.toList xs)
-}
maximum :: (Shape.C sh, Storable a, Ord a) => Array sh a -> a
maximum :: Array sh a -> a
maximum = (a -> a -> a) -> Array sh a -> a
forall sh a. (C sh, Storable a) => (a -> a -> a) -> Array sh a -> a
foldl1 a -> a -> a
forall a. Ord a => a -> a -> a
max

{-# INLINE foldl1 #-}
foldl1 :: (Shape.C sh, Storable a) => (a -> a -> a) -> Array sh a -> a
foldl1 :: (a -> a -> a) -> Array sh a -> a
foldl1 a -> a -> a
op (Array sh
sh ForeignPtr a
x) = IO a -> a
forall a. IO a -> a
unsafePerformIO (IO a -> a) -> IO a -> a
forall a b. (a -> b) -> a -> b
$
   ForeignPtr a -> (Ptr a -> IO a) -> IO a
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr a
x ((Ptr a -> IO a) -> IO a) -> (Ptr a -> IO a) -> IO a
forall a b. (a -> b) -> a -> b
$ \Ptr a
xPtr ->
      (Int -> a -> a) -> (a -> a -> a) -> Int -> Ptr a -> Int -> IO a
forall a b.
Storable a =>
(Int -> a -> b) -> (b -> b -> b) -> Int -> Ptr a -> Int -> IO b
Memory.foldl1 ((a -> a) -> Int -> a -> a
forall a b. a -> b -> a
const a -> a
forall a. a -> a
id) a -> a -> a
op (sh -> Int
forall sh. C sh => sh -> Int
Shape.size sh
sh) Ptr a
xPtr Int
1

{- |
prop> forAllNonEmpty $ \xs -> Array.limits xs ==? (Array.minimum xs, Array.maximum xs)
-}
limits :: (Shape.C sh, Storable a, Ord a) => Array sh a -> (a,a)
limits :: Array sh a -> (a, a)
limits = (Min a -> a, Max a -> a) -> (Min a, Max a) -> (a, a)
forall a c b d. (a -> c, b -> d) -> (a, b) -> (c, d)
StrictTuple.mapPair (Min a -> a
forall a. Min a -> a
getMin, Max a -> a
forall a. Max a -> a
getMax) ((Min a, Max a) -> (a, a))
-> (Array sh a -> (Min a, Max a)) -> Array sh a -> (a, a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> (Min a, Max a)) -> Array sh a -> (Min a, Max a)
forall sh a m.
(C sh, Storable a, Ord a, Semigroup m) =>
(a -> m) -> Array sh a -> m
foldMap (\a
x -> (a -> Min a
forall a. a -> Min a
Min a
x, a -> Max a
forall a. a -> Max a
Max a
x))

{-# INLINE foldMap #-}
foldMap ::
   (Shape.C sh, Storable a, Ord a, Semigroup m) => (a -> m) -> Array sh a -> m
foldMap :: (a -> m) -> Array sh a -> m
foldMap a -> m
f (Array sh
sh ForeignPtr a
x) = IO m -> m
forall a. IO a -> a
unsafePerformIO (IO m -> m) -> IO m -> m
forall a b. (a -> b) -> a -> b
$
   ForeignPtr a -> (Ptr a -> IO m) -> IO m
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr a
x ((Ptr a -> IO m) -> IO m) -> (Ptr a -> IO m) -> IO m
forall a b. (a -> b) -> a -> b
$ \Ptr a
xPtr ->
      (Int -> a -> m) -> (m -> m -> m) -> Int -> Ptr a -> Int -> IO m
forall a b.
Storable a =>
(Int -> a -> b) -> (b -> b -> b) -> Int -> Ptr a -> Int -> IO b
Memory.foldl1 ((a -> m) -> Int -> a -> m
forall a b. a -> b -> a
const a -> m
f) m -> m -> m
forall a. Semigroup a => a -> a -> a
(<>) (sh -> Int
forall sh. C sh => sh -> Int
Shape.size sh
sh) Ptr a
xPtr Int
1


argMinimum, argMaximum ::
   (Shape.InvIndexed sh, Shape.Index sh ~ ix, Storable a, Ord a) =>
   Array sh a -> (ix,a)
argMinimum :: Array sh a -> (ix, a)
argMinimum Array sh a
xs = Array sh a -> Arg a Int -> (Index sh, a)
forall sh a.
InvIndexed sh =>
Array sh a -> Arg a Int -> (Index sh, a)
unArg Array sh a
xs (Arg a Int -> (Index sh, a)) -> Arg a Int -> (Index sh, a)
forall a b. (a -> b) -> a -> b
$ Min (Arg a Int) -> Arg a Int
forall a. Min a -> a
getMin (Min (Arg a Int) -> Arg a Int) -> Min (Arg a Int) -> Arg a Int
forall a b. (a -> b) -> a -> b
$ (Int -> a -> Min (Arg a Int)) -> Array sh a -> Min (Arg a Int)
forall sh a m.
(C sh, Storable a, Semigroup m) =>
(Int -> a -> m) -> Array sh a -> m
foldMapWithIndex (\Int
k a
x -> Arg a Int -> Min (Arg a Int)
forall a. a -> Min a
Min (a -> Int -> Arg a Int
forall a b. a -> b -> Arg a b
Arg a
x Int
k)) Array sh a
xs
argMaximum :: Array sh a -> (ix, a)
argMaximum Array sh a
xs = Array sh a -> Arg a Int -> (Index sh, a)
forall sh a.
InvIndexed sh =>
Array sh a -> Arg a Int -> (Index sh, a)
unArg Array sh a
xs (Arg a Int -> (Index sh, a)) -> Arg a Int -> (Index sh, a)
forall a b. (a -> b) -> a -> b
$ Max (Arg a Int) -> Arg a Int
forall a. Max a -> a
getMax (Max (Arg a Int) -> Arg a Int) -> Max (Arg a Int) -> Arg a Int
forall a b. (a -> b) -> a -> b
$ (Int -> a -> Max (Arg a Int)) -> Array sh a -> Max (Arg a Int)
forall sh a m.
(C sh, Storable a, Semigroup m) =>
(Int -> a -> m) -> Array sh a -> m
foldMapWithIndex (\Int
k a
x -> Arg a Int -> Max (Arg a Int)
forall a. a -> Max a
Max (a -> Int -> Arg a Int
forall a b. a -> b -> Arg a b
Arg a
x Int
k)) Array sh a
xs

unArg ::
   (Shape.InvIndexed sh) => Array sh a -> Arg a Int -> (Shape.Index sh, a)
unArg :: Array sh a -> Arg a Int -> (Index sh, a)
unArg Array sh a
xs (Arg a
x Int
k) = (sh -> Int -> Index sh
forall sh. InvIndexed sh => sh -> Int -> Index sh
Shape.indexFromOffset (Array sh a -> sh
forall sh a. Array sh a -> sh
Array.shape Array sh a
xs) Int
k, a
x)

{-# INLINE foldMapWithIndex #-}
foldMapWithIndex ::
   (Shape.C sh, Storable a, Semigroup m) => (Int -> a -> m) -> Array sh a -> m
foldMapWithIndex :: (Int -> a -> m) -> Array sh a -> m
foldMapWithIndex Int -> a -> m
f (Array sh
sh ForeignPtr a
x) = IO m -> m
forall a. IO a -> a
unsafePerformIO (IO m -> m) -> IO m -> m
forall a b. (a -> b) -> a -> b
$
   ForeignPtr a -> (Ptr a -> IO m) -> IO m
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr a
x ((Ptr a -> IO m) -> IO m) -> (Ptr a -> IO m) -> IO m
forall a b. (a -> b) -> a -> b
$ \Ptr a
xPtr -> (Int -> a -> m) -> (m -> m -> m) -> Int -> Ptr a -> Int -> IO m
forall a b.
Storable a =>
(Int -> a -> b) -> (b -> b -> b) -> Int -> Ptr a -> Int -> IO b
Memory.foldl1 Int -> a -> m
f m -> m -> m
forall a. Semigroup a => a -> a -> a
(<>) (sh -> Int
forall sh. C sh => sh -> Int
Shape.size sh
sh) Ptr a
xPtr Int
1