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

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

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

   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 System.IO.Unsafe (unsafePerformIO)
import Foreign.Storable (Storable)
import Foreign.ForeignPtr (withForeignPtr)

import Control.Monad.ST (runST)

import qualified Data.Map as Map
import qualified Data.Set as Set
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.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
>>>
>>> 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


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

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)


infixl 9 !

(!) :: (Shape.Indexed sh, Storable a) => Array sh a -> Shape.Index sh -> a
(!) Array sh a
arr Index sh
ix = (forall s. ST s a) -> 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
   Array (ST s) sh a -> Index sh -> ST s a
forall (m :: * -> *) sh a.
(PrimMonad m, Indexed sh, Storable a) =>
Array m sh a -> Index sh -> m a
MutArray.read 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 -> Array sh c
forall a. HasCallStack => String -> a
error String
"zipWith: 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)



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