module Data.Array.Comfort.Boxed (
   Array,
   shape,
   reshape,
   mapShape,
   accessMaybe, (!),
   Array.toList,
   Array.fromList,
   Array.vectorFromList,
   toAssociations,
   fromMap,
   toMap,
   fromTuple,
   toTuple,
   fromRecord,
   toRecord,
   fromContainer,
   toContainer,
   indices,
   Array.replicate,
   cartesian,

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

   pick,
   Array.append,
   Array.take, Array.drop,
   Array.takeLeft, Array.takeRight, Array.split,
   Array.takeCenter,
   ) where

import qualified Data.Array.Comfort.Boxed.Unchecked as Array
import qualified Data.Array.Comfort.Container as Container
import qualified Data.Array.Comfort.Check as Check
import qualified Data.Array.Comfort.Shape.Tuple as TupleShape
import qualified Data.Array.Comfort.Shape as Shape
import Data.Array.Comfort.Boxed.Unchecked (Array(Array))

import qualified Data.Primitive.Array as Prim

import qualified Control.Monad.Primitive as PrimM
import qualified Control.Monad.Trans.State as MS
import Control.Monad.ST (runST)
import Control.Applicative (liftA2, (<$>))

import qualified Data.Foldable as Fold
import qualified Data.Map as Map
import qualified Data.Set as Set
import Data.Map (Map)
import Data.Set (Set)
import Data.Traversable (Traversable, traverse)
import Data.Foldable (forM_)
import Data.Either.HT (maybeRight)

import Prelude hiding (zipWith, replicate)


{- $setup
>>> import qualified Data.Array.Comfort.Boxed as Array
>>> import qualified Data.Array.Comfort.Shape as Shape
>>> import Data.Array.Comfort.Boxed (Array, (!))
>>>
>>> import qualified Test.QuickCheck as QC
>>>
>>> type ShapeInt = Shape.ZeroBased Int
>>>
>>> genArray2 :: QC.Gen (Array (ShapeInt,ShapeInt) Char)
>>> genArray2 = do
>>>    xs <- QC.arbitrary
>>>    let n = length xs
>>>    (k,m) <-
>>>       if n == 0
>>>          then QC.elements [(,) 0, flip (,) 0] <*> QC.choose (1,20)
>>>          else fmap (\m -> (div n m, m)) $ QC.choose (1,n)
>>>    return $
>>>       Array.fromList (Shape.ZeroBased k, Shape.ZeroBased m) $ take (k*m) xs
-}


shape :: Array.Array sh a -> sh
shape :: forall sh a. 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 :: forall sh0 sh1 a.
(C sh0, C sh1) =>
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
"Boxed" 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 :: forall sh0 sh1 a.
(C sh0, C sh1) =>
(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


indices :: (Shape.Indexed sh) => sh -> Array.Array sh (Shape.Index sh)
indices :: forall sh. Indexed sh => sh -> Array sh (Index sh)
indices sh
sh = sh -> [Index sh] -> Array sh (Index sh)
forall sh a. C sh => sh -> [a] -> Array sh a
Array.fromList sh
sh ([Index sh] -> Array sh (Index sh))
-> [Index sh] -> Array sh (Index sh)
forall a b. (a -> b) -> a -> b
$ sh -> [Index sh]
forall sh. Indexed sh => sh -> [Index sh]
Shape.indices sh
sh

fromMap :: (Ord k) => Map k a -> Array (Set k) a
fromMap :: forall k a. Ord k => Map k a -> Array (Set k) a
fromMap Map k a
m = Set k -> [a] -> Array (Set k) a
forall sh a. C sh => sh -> [a] -> Array sh a
Array.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) => Array (Set k) a -> Map k a
toMap :: forall k a. Ord k => 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 => Array sh a -> [a]
Array.toList Array (Set k) a
arr)

fromTuple ::
   (TupleShape.NestedTuple tuple) =>
   Shape.DataTuple tuple a -> Array (Shape.NestedTuple ixtype tuple) a
fromTuple :: forall tuple a ixtype.
NestedTuple tuple =>
DataTuple tuple a -> Array (NestedTuple ixtype tuple) a
fromTuple DataTuple tuple a
tuple =
   case State Element (tuple, [a]) -> Element -> (tuple, [a])
forall s a. State s a -> s -> a
MS.evalState (DataTuple tuple a -> State Element (tuple, [a])
forall a. DataTuple tuple a -> State Element (tuple, [a])
forall shape a.
NestedTuple shape =>
DataTuple shape a -> State Element (shape, [a])
TupleShape.decons DataTuple tuple a
tuple) (Int -> Element
Shape.Element Int
0) of
      (tuple
sh, [a]
xs) -> NestedTuple ixtype tuple
-> [a] -> Array (NestedTuple ixtype tuple) a
forall sh a. C sh => sh -> [a] -> Array sh a
Array.fromList (tuple -> NestedTuple ixtype tuple
forall ixtype tuple. tuple -> NestedTuple ixtype tuple
Shape.NestedTuple tuple
sh) [a]
xs

toTuple ::
   (TupleShape.NestedTuple tuple) =>
   Array (Shape.NestedTuple ixtype tuple) a -> Shape.DataTuple tuple a
toTuple :: forall tuple ixtype a.
NestedTuple tuple =>
Array (NestedTuple ixtype tuple) a -> DataTuple tuple a
toTuple Array (NestedTuple ixtype tuple) a
arr =
   State [a] (DataTuple tuple a) -> [a] -> DataTuple tuple a
forall s a. State s a -> s -> a
MS.evalState
      (tuple -> State [a] (DataTuple tuple a)
forall shape a.
ElementTuple shape =>
shape -> State [a] (DataTuple shape a)
TupleShape.cons (tuple -> State [a] (DataTuple tuple a))
-> tuple -> State [a] (DataTuple tuple a)
forall a b. (a -> b) -> a -> b
$ NestedTuple ixtype tuple -> tuple
forall ixtype tuple. NestedTuple ixtype tuple -> tuple
Shape.getNestedTuple (NestedTuple ixtype tuple -> tuple)
-> NestedTuple ixtype tuple -> tuple
forall a b. (a -> b) -> a -> b
$ Array (NestedTuple ixtype tuple) a -> NestedTuple ixtype tuple
forall sh a. Array sh a -> sh
shape Array (NestedTuple ixtype tuple) a
arr)
      (Array (NestedTuple ixtype tuple) a -> [a]
forall sh a. C sh => Array sh a -> [a]
Array.toList Array (NestedTuple ixtype tuple) a
arr)

fromRecord ::
   (Traversable f) =>
   f a -> Array (Shape.Record f) a
fromRecord :: forall (f :: * -> *) a. Traversable f => f a -> Array (Record f) a
fromRecord f a
xs =
   Record f -> [a] -> Array (Record f) a
forall sh a. C sh => sh -> [a] -> Array sh a
Array.fromList
      (f Element -> Record f
forall (f :: * -> *). f Element -> Record f
Shape.Record (f Element -> Record f) -> f Element -> Record f
forall a b. (a -> b) -> a -> b
$ (State Element (f Element) -> Element -> f Element)
-> Element -> State Element (f Element) -> f Element
forall a b c. (a -> b -> c) -> b -> a -> c
flip State Element (f Element) -> Element -> f Element
forall s a. State s a -> s -> a
MS.evalState (Int -> Element
Shape.Element Int
0) (State Element (f Element) -> f Element)
-> State Element (f Element) -> f Element
forall a b. (a -> b) -> a -> b
$
       (a -> StateT Element Identity Element)
-> f a -> State Element (f Element)
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) -> f a -> f (f b)
traverse (StateT Element Identity Element
-> a -> StateT Element Identity Element
forall a b. a -> b -> a
const StateT Element Identity Element
TupleShape.next) f a
xs)
      (f a -> [a]
forall a. f a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
Fold.toList f a
xs)

toRecord ::
   (Traversable f) =>
   Array (Shape.Record f) a -> f a
toRecord :: forall (f :: * -> *) a. Traversable f => Array (Record f) a -> f a
toRecord Array (Record f) a
arr =
   State [a] (f a) -> [a] -> f a
forall s a. State s a -> s -> a
MS.evalState
      ((Element -> StateT [a] Identity a) -> f Element -> State [a] (f a)
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) -> f a -> f (f b)
traverse (StateT [a] Identity a -> Element -> StateT [a] Identity a
forall a b. a -> b -> a
const StateT [a] Identity a
forall a. State [a] a
TupleShape.get) (f Element -> State [a] (f a)) -> f Element -> State [a] (f a)
forall a b. (a -> b) -> a -> b
$
       (\(Shape.Record f Element
record) -> f Element
record) (Record f -> f Element) -> Record f -> f Element
forall a b. (a -> b) -> a -> b
$ Array (Record f) a -> Record f
forall sh a. Array sh a -> sh
shape Array (Record f) a
arr)
      (Array (Record f) a -> [a]
forall sh a. C sh => Array sh a -> [a]
Array.toList Array (Record f) a
arr)

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

toContainer :: (Container.C f) => Array (Container.Shape f) a -> f a
toContainer :: forall (f :: * -> *) a. C f => Array (Shape f) a -> f a
toContainer Array (Shape f) a
arr = Shape f -> [a] -> f a
forall a. 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 => Array sh a -> [a]
Array.toList Array (Shape f) 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 a
arr =
   (String -> a) -> (a -> a) -> Either String a -> a
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (String -> a
forall a. HasCallStack => String -> a
error (String -> a) -> (String -> String) -> String -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (String
"Array.Comfort.Boxed.!: " String -> String -> String
forall a. [a] -> [a] -> [a]
++)) a -> a
forall a. a -> a
id (Either String a -> a)
-> (Index sh -> Either String a) -> Index sh -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array sh a -> Index sh -> Either String a
forall sh a.
Indexed sh =>
Array sh a -> Index sh -> Either String a
accessEither Array sh a
arr

accessMaybe :: (Shape.Indexed sh) => Array sh a -> Shape.Index sh -> Maybe a
accessMaybe :: forall sh a. Indexed sh => Array sh a -> Index sh -> Maybe a
accessMaybe Array sh a
arr = Either String a -> Maybe a
forall a b. Either a b -> Maybe b
maybeRight (Either String a -> Maybe a)
-> (Index sh -> Either String a) -> Index sh -> Maybe a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array sh a -> Index sh -> Either String a
forall sh a.
Indexed sh =>
Array sh a -> Index sh -> Either String a
accessEither Array sh a
arr

accessEither ::
   (Shape.Indexed sh) => Array sh a -> Shape.Index sh -> Either String a
accessEither :: forall sh a.
Indexed sh =>
Array sh a -> Index sh -> Either String a
accessEither (Array sh
sh Array a
arr) Index sh
ix =
   (Int -> a) -> Either String Int -> Either String a
forall a b. (a -> b) -> Either String a -> Either String b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Array a -> Int -> a
forall a. Array a -> Int -> a
Prim.indexArray Array a
arr) (Either String Int -> Either String a)
-> Either String Int -> Either String a
forall a b. (a -> b) -> a -> b
$ Result Checked Int -> Either String Int
forall a. Result Checked a -> Either String a
Shape.getChecked (Result Checked Int -> Either String Int)
-> Result Checked Int -> Either String Int
forall a b. (a -> b) -> a -> b
$ sh -> Index sh -> Result Checked Int
forall sh check.
(Indexed sh, Checking check) =>
sh -> Index sh -> Result check Int
forall check. Checking check => sh -> Index sh -> Result check Int
Shape.unifiedOffset sh
sh Index sh
ix


zipWith ::
   (Shape.C sh, Eq sh) =>
   (a -> b -> c) -> Array sh a -> Array sh b -> Array sh c
zipWith :: forall sh a b c.
(C sh, Eq sh) =>
(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 =>
(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) => Array sh a -> [(Shape.Index sh, a)] -> Array sh a
// :: forall sh a.
Indexed sh =>
Array sh a -> [(Index sh, a)] -> Array sh a
(//) (Array sh
sh Array 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
   MutableArray s a
marr <- Array a -> Int -> Int -> ST s (MutableArray (PrimState (ST s)) a)
forall (m :: * -> *) a.
PrimMonad m =>
Array a -> Int -> Int -> m (MutableArray (PrimState m) a)
Prim.thawArray Array a
arr Int
0 (sh -> Int
forall sh. C sh => sh -> Int
Shape.size sh
sh)
   [(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
ix,a
a) -> MutableArray (PrimState (ST s)) a -> Int -> a -> ST s ()
forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> Int -> a -> m ()
Prim.writeArray MutableArray s a
MutableArray (PrimState (ST s)) a
marr (sh -> Index sh -> Int
forall sh. Indexed sh => sh -> Index sh -> Int
Shape.offset sh
sh Index sh
ix) 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) -> ST s (Array a) -> ST s (Array sh a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> 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
marr)

accumulate ::
   (Shape.Indexed sh) =>
   (a -> b -> a) -> Array sh a -> [(Shape.Index sh, b)] -> Array sh a
accumulate :: forall sh a b.
Indexed sh =>
(a -> b -> a) -> Array sh a -> [(Index sh, b)] -> Array sh a
accumulate a -> b -> a
f (Array sh
sh Array 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
   MutableArray s a
marr <- Array a -> Int -> Int -> ST s (MutableArray (PrimState (ST s)) a)
forall (m :: * -> *) a.
PrimMonad m =>
Array a -> Int -> Int -> m (MutableArray (PrimState m) a)
Prim.thawArray Array a
arr Int
0 (sh -> Int
forall sh. C sh => sh -> Int
Shape.size sh
sh)
   [(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) -> MutableArray (PrimState (ST s)) a -> Int -> (a -> a) -> ST s ()
forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> Int -> (a -> a) -> m ()
updateArray MutableArray s a
MutableArray (PrimState (ST s)) a
marr (sh -> Index sh -> Int
forall sh. Indexed sh => sh -> Index sh -> Int
Shape.offset sh
sh 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
   sh -> Array a -> Array sh a
forall sh a. sh -> Array a -> Array sh a
Array sh
sh (Array a -> Array sh a) -> ST s (Array a) -> ST s (Array sh a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> 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
marr)

updateArray ::
   PrimM.PrimMonad m =>
   Prim.MutableArray (PrimM.PrimState m) a -> Int -> (a -> a) -> m ()
updateArray :: forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> Int -> (a -> a) -> m ()
updateArray MutableArray (PrimState m) a
marr Int
k a -> a
f = MutableArray (PrimState m) a -> Int -> a -> m ()
forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> Int -> a -> m ()
Prim.writeArray MutableArray (PrimState m) a
marr Int
k (a -> m ()) -> (a -> a) -> a -> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> a
f (a -> m ()) -> m a -> m ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< MutableArray (PrimState m) a -> Int -> m a
forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> Int -> m a
Prim.readArray MutableArray (PrimState m) a
marr Int
k

toAssociations :: (Shape.Indexed sh) => Array sh a -> [(Shape.Index sh, a)]
toAssociations :: forall sh a. Indexed sh => 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 => Array sh a -> [a]
Array.toList Array sh a
arr)

fromAssociations ::
   (Shape.Indexed sh) => a -> sh -> [(Shape.Index sh, a)] -> Array sh a
fromAssociations :: forall sh a. Indexed sh => 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
   MutableArray s a
marr <- 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
   [(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
ix,a
x) -> MutableArray (PrimState (ST s)) a -> Int -> a -> ST s ()
forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> Int -> a -> m ()
Prim.writeArray MutableArray s a
MutableArray (PrimState (ST s)) a
marr (sh -> Index sh -> Int
forall sh. Indexed sh => sh -> Index sh -> Int
Shape.offset sh
sh Index sh
ix) a
x
   sh -> Array a -> Array sh a
forall sh a. sh -> Array a -> Array sh a
Array sh
sh (Array a -> Array sh a) -> ST s (Array a) -> ST s (Array sh a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> 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
marr)



{- |
prop> :{
   QC.forAll genArray2 $ \xs ->
   let shape = Array.shape xs in
   Shape.size shape > 0   QC.==>
   QC.forAll (QC.elements $ Shape.indices shape) $ \(ix0,ix1) ->
      Array.pick xs ix0 ! ix1 == xs!(ix0,ix1)
:}
-}
pick ::
   (Shape.Indexed sh0, Shape.C sh1) =>
   Array (sh0,sh1) a -> Shape.Index sh0 -> Array sh1 a
pick :: forall sh0 sh1 a.
(Indexed sh0, C sh1) =>
Array (sh0, sh1) a -> Index sh0 -> Array sh1 a
pick (Array (sh0
sh0,sh1
sh1) Array a
x) Index sh0
ix0 =
   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
$
   let k :: Int
k = sh1 -> Int
forall sh. C sh => sh -> Int
Shape.size sh1
sh1
   in Array a -> Int -> Int -> Array a
forall a. Array a -> Int -> Int -> Array a
Prim.cloneArray Array a
x (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
* Int
k) Int
k


cartesian ::
   (Shape.C sh0, Shape.C sh1) =>
   Array sh0 a -> Array sh1 b -> Array (sh0,sh1) (a,b)
cartesian :: forall sh0 sh1 a b.
(C sh0, C sh1) =>
Array sh0 a -> Array sh1 b -> Array (sh0, sh1) (a, b)
cartesian Array sh0 a
a Array sh1 b
b =
   (sh0, sh1) -> [(a, b)] -> Array (sh0, sh1) (a, b)
forall sh a. C sh => sh -> [a] -> Array sh a
Array.fromList (Array sh0 a -> sh0
forall sh a. Array sh a -> sh
shape Array sh0 a
a, Array sh1 b -> sh1
forall sh a. Array sh a -> sh
shape Array sh1 b
b) ([(a, b)] -> Array (sh0, sh1) (a, b))
-> [(a, b)] -> Array (sh0, sh1) (a, b)
forall a b. (a -> b) -> a -> b
$
      (a -> b -> (a, b)) -> [a] -> [b] -> [(a, b)]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 (,) (Array sh0 a -> [a]
forall sh a. C sh => Array sh a -> [a]
Array.toList Array sh0 a
a) (Array sh1 b -> [b]
forall sh a. C sh => Array sh a -> [a]
Array.toList Array sh1 b
b)