{-# 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 ()
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)
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
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
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