{-# LANGUAGE TypeFamilies #-}
module Data.Array.Comfort.Boxed.Strict.Unchecked where

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.ST.Strict as STStrict
import qualified Control.Monad.Trans.Class as MT
import qualified Control.Monad.Trans.State as MS

import Prelude hiding (map, zipWith)


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. (forall s. ST s a) -> a
STStrict.runST (forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall (m :: * -> *) a. Monad m => Array a -> Int -> m a
Prim.indexArrayM Array a
arr) forall a b. (a -> b) -> a -> b
$ forall a. Int -> [a] -> [a]
take (forall sh. C sh => sh -> Int
Shape.size sh
sh) [Int
0..])

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
$ forall a b. (a -> b) -> Array a -> Array b
Prim.mapArray' a -> b
f Array a
arr

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
$
   forall a. (forall s. ST s a) -> a
STStrict.runST
      (forall a b c. (a -> b -> c) -> b -> a -> c
flip forall (m :: * -> *) s a. Monad m => StateT s m a -> s -> m a
MS.evalStateT Int
0 forall a b. (a -> b) -> a -> b
$
       forall (m :: * -> *) a b.
PrimMonad m =>
(a -> m b) -> Array a -> m (Array b)
Prim.traverseArrayP
         (\a
a -> do
            Int
k <- forall (m :: * -> *) s. Monad m => StateT s m s
MS.get
            b
b <- forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
MT.lift forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. Monad m => Array a -> Int -> m a
Prim.indexArrayM Array b
arrb Int
k
            forall (m :: * -> *) s. Monad m => s -> StateT s m ()
MS.put (Int
kforall a. Num a => a -> a -> a
+Int
1)
            forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ a -> b -> c
f a
a b
b)
         Array a
arra)