{-# LANGUAGE Rank2Types #-}
module Data.Array.Comfort.Storable.Unchecked.Creator where

import qualified Data.Array.Comfort.Shape.SubSize as SubSize
import qualified Data.Array.Comfort.Shape as Shape
import Data.Array.Comfort.Storable.Private (Array(Array))

import Foreign.Storable (Storable, )
import Foreign.Ptr (Ptr, )

import qualified Foreign.Marshal.Array.Guarded as Alloc
import qualified Control.Monad.Trans.Cont as MC
import Control.Monad.Primitive (PrimMonad, unsafeIOToPrim)

import Data.Biapplicative (Biapplicative(bipure, (<<*>>)))
import Data.Bifunctor (Bifunctor(bimap))

import Data.Tuple.HT (mapFst)


newtype Creator arr ptr = Creator (forall a. (ptr -> IO a) -> IO (arr, a))

liftIO :: IO ptr -> Creator () ptr
liftIO :: forall ptr. IO ptr -> Creator () ptr
liftIO IO ptr
act = (forall a. (ptr -> IO a) -> IO ((), a)) -> Creator () ptr
forall arr ptr.
(forall a. (ptr -> IO a) -> IO (arr, a)) -> Creator arr ptr
Creator ((forall a. (ptr -> IO a) -> IO ((), a)) -> Creator () ptr)
-> (forall a. (ptr -> IO a) -> IO ((), a)) -> Creator () ptr
forall a b. (a -> b) -> a -> b
$ \ptr -> IO a
f -> (a -> ((), a)) -> IO a -> IO ((), a)
forall a b. (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((,) ()) (IO a -> IO ((), a)) -> IO a -> IO ((), a)
forall a b. (a -> b) -> a -> b
$ ptr -> IO a
f (ptr -> IO a) -> IO ptr -> IO a
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< IO ptr
act

liftContT :: (forall a. MC.ContT a IO ptr) -> Creator () ptr
liftContT :: forall ptr. (forall a. ContT a IO ptr) -> Creator () ptr
liftContT forall a. ContT a IO ptr
act = (forall a. (ptr -> IO a) -> IO ((), a)) -> Creator () ptr
forall arr ptr.
(forall a. (ptr -> IO a) -> IO (arr, a)) -> Creator arr ptr
Creator ((forall a. (ptr -> IO a) -> IO ((), a)) -> Creator () ptr)
-> (forall a. (ptr -> IO a) -> IO ((), a)) -> Creator () ptr
forall a b. (a -> b) -> a -> b
$ \ptr -> IO a
f -> (a -> ((), a)) -> IO a -> IO ((), a)
forall a b. (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((,) ()) (IO a -> IO ((), a)) -> IO a -> IO ((), a)
forall a b. (a -> b) -> a -> b
$ ContT a IO ptr -> (ptr -> IO a) -> IO a
forall {k} (r :: k) (m :: k -> *) a.
ContT r m a -> (a -> m r) -> m r
MC.runContT ContT a IO ptr
forall a. ContT a IO ptr
act ptr -> IO a
f

instance Functor (Creator arr) where
   fmap :: forall a b. (a -> b) -> Creator arr a -> Creator arr b
fmap a -> b
g (Creator forall a. (a -> IO a) -> IO (arr, a)
act) = (forall a. (b -> IO a) -> IO (arr, a)) -> Creator arr b
forall arr ptr.
(forall a. (ptr -> IO a) -> IO (arr, a)) -> Creator arr ptr
Creator ((forall a. (b -> IO a) -> IO (arr, a)) -> Creator arr b)
-> (forall a. (b -> IO a) -> IO (arr, a)) -> Creator arr b
forall a b. (a -> b) -> a -> b
$ \b -> IO a
f -> (a -> IO a) -> IO (arr, a)
forall a. (a -> IO a) -> IO (arr, a)
act (b -> IO a
f (b -> IO a) -> (a -> b) -> a -> IO a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> b
g)


pair ::
   Creator arr0 ptr0 -> Creator arr1 ptr1 ->
   Creator (arr0,arr1) (ptr0,ptr1)
pair :: forall arr0 ptr0 arr1 ptr1.
Creator arr0 ptr0
-> Creator arr1 ptr1 -> Creator (arr0, arr1) (ptr0, ptr1)
pair (Creator forall a. (ptr0 -> IO a) -> IO (arr0, a)
act0) (Creator forall a. (ptr1 -> IO a) -> IO (arr1, a)
act1) =
   (forall a. ((ptr0, ptr1) -> IO a) -> IO ((arr0, arr1), a))
-> Creator (arr0, arr1) (ptr0, ptr1)
forall arr ptr.
(forall a. (ptr -> IO a) -> IO (arr, a)) -> Creator arr ptr
Creator ((forall a. ((ptr0, ptr1) -> IO a) -> IO ((arr0, arr1), a))
 -> Creator (arr0, arr1) (ptr0, ptr1))
-> (forall a. ((ptr0, ptr1) -> IO a) -> IO ((arr0, arr1), a))
-> Creator (arr0, arr1) (ptr0, ptr1)
forall a b. (a -> b) -> a -> b
$ \(ptr0, ptr1) -> IO a
f ->
      ((arr0, (arr1, a)) -> ((arr0, arr1), a))
-> IO (arr0, (arr1, a)) -> IO ((arr0, arr1), a)
forall a b. (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\(arr0
arr0,(arr1
arr1,a
a)) -> ((arr0
arr0,arr1
arr1),a
a)) (IO (arr0, (arr1, a)) -> IO ((arr0, arr1), a))
-> IO (arr0, (arr1, a)) -> IO ((arr0, arr1), a)
forall a b. (a -> b) -> a -> b
$
      (ptr0 -> IO (arr1, a)) -> IO (arr0, (arr1, a))
forall a. (ptr0 -> IO a) -> IO (arr0, a)
act0 ((ptr0 -> IO (arr1, a)) -> IO (arr0, (arr1, a)))
-> (ptr0 -> IO (arr1, a)) -> IO (arr0, (arr1, a))
forall a b. (a -> b) -> a -> b
$ \ptr0
ptr0 ->
      (ptr1 -> IO a) -> IO (arr1, a)
forall a. (ptr1 -> IO a) -> IO (arr1, a)
act1 ((ptr1 -> IO a) -> IO (arr1, a)) -> (ptr1 -> IO a) -> IO (arr1, a)
forall a b. (a -> b) -> a -> b
$ \ptr1
ptr1 ->
         (ptr0, ptr1) -> IO a
f (ptr0
ptr0, ptr1
ptr1)

instance Bifunctor Creator where
   bimap :: forall a b c d. (a -> b) -> (c -> d) -> Creator a c -> Creator b d
bimap a -> b
g c -> d
h (Creator forall a. (c -> IO a) -> IO (a, a)
act) = (forall a. (d -> IO a) -> IO (b, a)) -> Creator b d
forall arr ptr.
(forall a. (ptr -> IO a) -> IO (arr, a)) -> Creator arr ptr
Creator ((forall a. (d -> IO a) -> IO (b, a)) -> Creator b d)
-> (forall a. (d -> IO a) -> IO (b, a)) -> Creator b d
forall a b. (a -> b) -> a -> b
$ \d -> IO a
f -> ((a, a) -> (b, a)) -> IO (a, a) -> IO (b, a)
forall a b. (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((a -> b) -> (a, a) -> (b, a)
forall a c b. (a -> c) -> (a, b) -> (c, b)
mapFst a -> b
g) (IO (a, a) -> IO (b, a)) -> IO (a, a) -> IO (b, a)
forall a b. (a -> b) -> a -> b
$ (c -> IO a) -> IO (a, a)
forall a. (c -> IO a) -> IO (a, a)
act (d -> IO a
f (d -> IO a) -> (c -> d) -> c -> IO a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. c -> d
h)

instance Biapplicative Creator where
   bipure :: forall a b. a -> b -> Creator a b
bipure a
a b
b = (forall a. (b -> IO a) -> IO (a, a)) -> Creator a b
forall arr ptr.
(forall a. (ptr -> IO a) -> IO (arr, a)) -> Creator arr ptr
Creator ((forall a. (b -> IO a) -> IO (a, a)) -> Creator a b)
-> (forall a. (b -> IO a) -> IO (a, a)) -> Creator a b
forall a b. (a -> b) -> a -> b
$ \b -> IO a
f -> (a -> (a, a)) -> IO a -> IO (a, a)
forall a b. (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((,) a
a) (IO a -> IO (a, a)) -> IO a -> IO (a, a)
forall a b. (a -> b) -> a -> b
$ b -> IO a
f b
b
   Creator (a -> b) (c -> d)
creator0 <<*>> :: forall a b c d.
Creator (a -> b) (c -> d) -> Creator a c -> Creator b d
<<*>> Creator a c
creator1 =
      ((a -> b, a) -> b)
-> ((c -> d, c) -> d)
-> Creator (a -> b, a) (c -> d, c)
-> Creator b d
forall a b c d. (a -> b) -> (c -> d) -> Creator a c -> Creator b d
forall (p :: * -> * -> *) a b c d.
Bifunctor p =>
(a -> b) -> (c -> d) -> p a c -> p b d
bimap (((a -> b) -> a -> b) -> (a -> b, a) -> b
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry (a -> b) -> a -> b
forall a. a -> a
id) (((c -> d) -> c -> d) -> (c -> d, c) -> d
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry (c -> d) -> c -> d
forall a. a -> a
id) (Creator (a -> b, a) (c -> d, c) -> Creator b d)
-> Creator (a -> b, a) (c -> d, c) -> Creator b d
forall a b. (a -> b) -> a -> b
$ Creator (a -> b) (c -> d)
-> Creator a c -> Creator (a -> b, a) (c -> d, c)
forall arr0 ptr0 arr1 ptr1.
Creator arr0 ptr0
-> Creator arr1 ptr1 -> Creator (arr0, arr1) (ptr0, ptr1)
pair Creator (a -> b) (c -> d)
creator0 Creator a c
creator1

unsafeRun :: (PrimMonad m) => Creator arr ptr -> (ptr -> IO ()) -> m arr
unsafeRun :: forall (m :: * -> *) arr ptr.
PrimMonad m =>
Creator arr ptr -> (ptr -> IO ()) -> m arr
unsafeRun (Creator forall a. (ptr -> IO a) -> IO (arr, a)
act) ptr -> IO ()
f = IO arr -> m arr
forall (m :: * -> *) a. PrimMonad m => IO a -> m a
unsafeIOToPrim (IO arr -> m arr) -> IO arr -> m arr
forall a b. (a -> b) -> a -> b
$ ((arr, ()) -> arr) -> IO (arr, ()) -> IO arr
forall a b. (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\(arr
arr,()) -> arr
arr) (IO (arr, ()) -> IO arr) -> IO (arr, ()) -> IO arr
forall a b. (a -> b) -> a -> b
$ (ptr -> IO ()) -> IO (arr, ())
forall a. (ptr -> IO a) -> IO (arr, a)
act ptr -> IO ()
f

unsafeRunWithResult ::
   (PrimMonad m) => Creator arr ptr -> (ptr -> IO b) -> m (arr, b)
unsafeRunWithResult :: forall (m :: * -> *) arr ptr b.
PrimMonad m =>
Creator arr ptr -> (ptr -> IO b) -> m (arr, b)
unsafeRunWithResult (Creator forall a. (ptr -> IO a) -> IO (arr, a)
act) ptr -> IO b
f = IO (arr, b) -> m (arr, b)
forall (m :: * -> *) a. PrimMonad m => IO a -> m a
unsafeIOToPrim (IO (arr, b) -> m (arr, b)) -> IO (arr, b) -> m (arr, b)
forall a b. (a -> b) -> a -> b
$ (ptr -> IO b) -> IO (arr, b)
forall a. (ptr -> IO a) -> IO (arr, a)
act ptr -> IO b
f

{-# INLINE create #-}
create ::
   (Shape.C sh, Storable a) =>
   sh -> Creator (Array sh a) (Ptr a)
create :: forall sh a.
(C sh, Storable a) =>
sh -> Creator (Array sh a) (Ptr a)
create sh
sh = ((Int, Ptr a) -> Ptr a)
-> Creator (Array sh a) (Int, Ptr a)
-> Creator (Array sh a) (Ptr a)
forall a b.
(a -> b) -> Creator (Array sh a) a -> Creator (Array sh a) b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Int, Ptr a) -> Ptr a
forall a b. (a, b) -> b
snd (Creator (Array sh a) (Int, Ptr a) -> Creator (Array sh a) (Ptr a))
-> Creator (Array sh a) (Int, Ptr a)
-> Creator (Array sh a) (Ptr a)
forall a b. (a -> b) -> a -> b
$ sh -> Creator (Array sh a) (Int, Ptr a)
forall sh a.
(C sh, Storable a) =>
sh -> Creator (Array sh a) (Int, Ptr a)
createWithSize sh
sh

{-# INLINE createWithSize #-}
createWithSize ::
   (Shape.C sh, Storable a) =>
   sh -> Creator (Array sh a) (Int, Ptr a)
createWithSize :: forall sh a.
(C sh, Storable a) =>
sh -> Creator (Array sh a) (Int, Ptr a)
createWithSize sh
sh =
   let size :: Int
size = sh -> Int
forall sh. C sh => sh -> Int
Shape.size sh
sh
   in (forall a. ((Int, Ptr a) -> IO a) -> IO (Array sh a, a))
-> Creator (Array sh a) (Int, Ptr a)
forall arr ptr.
(forall a. (ptr -> IO a) -> IO (arr, a)) -> Creator arr ptr
Creator ((forall a. ((Int, Ptr a) -> IO a) -> IO (Array sh a, a))
 -> Creator (Array sh a) (Int, Ptr a))
-> (forall a. ((Int, Ptr a) -> IO a) -> IO (Array sh a, a))
-> Creator (Array sh a) (Int, Ptr a)
forall a b. (a -> b) -> a -> b
$ \(Int, Ptr a) -> IO a
f ->
         ((ForeignPtr a, a) -> (Array sh a, a))
-> IO (ForeignPtr a, a) -> IO (Array sh a, a)
forall a b. (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((ForeignPtr a -> Array sh a)
-> (ForeignPtr a, a) -> (Array sh a, a)
forall a c b. (a -> c) -> (a, b) -> (c, b)
mapFst (sh -> ForeignPtr a -> Array sh a
forall sh a. sh -> ForeignPtr a -> Array sh a
Array sh
sh)) (IO (ForeignPtr a, a) -> IO (Array sh a, a))
-> IO (ForeignPtr a, a) -> IO (Array sh a, a)
forall a b. (a -> b) -> a -> b
$ Int -> (Ptr a -> IO a) -> IO (ForeignPtr a, a)
forall a b.
Storable a =>
Int -> (Ptr a -> IO b) -> IO (ForeignPtr a, b)
Alloc.create Int
size ((Ptr a -> IO a) -> IO (ForeignPtr a, a))
-> (Ptr a -> IO a) -> IO (ForeignPtr a, a)
forall a b. (a -> b) -> a -> b
$ ((Int, Ptr a) -> IO a) -> Int -> Ptr a -> IO a
forall a b c. ((a, b) -> c) -> a -> b -> c
curry (Int, Ptr a) -> IO a
f Int
size

{-# INLINE createWithSizes #-}
createWithSizes ::
   (Shape.C sh, Storable a) =>
   SubSize.T sh nsize -> sh -> Creator (Array sh a) (nsize, Ptr a)
createWithSizes :: forall sh a nsize.
(C sh, Storable a) =>
T sh nsize -> sh -> Creator (Array sh a) (nsize, Ptr a)
createWithSizes (SubSize.Cons sh -> (Int, nsize)
subSize) sh
sh =
   let (Int
size, nsize
subSizes) = sh -> (Int, nsize)
subSize sh
sh
   in (forall a. ((nsize, Ptr a) -> IO a) -> IO (Array sh a, a))
-> Creator (Array sh a) (nsize, Ptr a)
forall arr ptr.
(forall a. (ptr -> IO a) -> IO (arr, a)) -> Creator arr ptr
Creator ((forall a. ((nsize, Ptr a) -> IO a) -> IO (Array sh a, a))
 -> Creator (Array sh a) (nsize, Ptr a))
-> (forall a. ((nsize, Ptr a) -> IO a) -> IO (Array sh a, a))
-> Creator (Array sh a) (nsize, Ptr a)
forall a b. (a -> b) -> a -> b
$ \(nsize, Ptr a) -> IO a
f ->
         ((ForeignPtr a, a) -> (Array sh a, a))
-> IO (ForeignPtr a, a) -> IO (Array sh a, a)
forall a b. (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((ForeignPtr a -> Array sh a)
-> (ForeignPtr a, a) -> (Array sh a, a)
forall a c b. (a -> c) -> (a, b) -> (c, b)
mapFst (sh -> ForeignPtr a -> Array sh a
forall sh a. sh -> ForeignPtr a -> Array sh a
Array sh
sh)) (IO (ForeignPtr a, a) -> IO (Array sh a, a))
-> IO (ForeignPtr a, a) -> IO (Array sh a, a)
forall a b. (a -> b) -> a -> b
$ Int -> (Ptr a -> IO a) -> IO (ForeignPtr a, a)
forall a b.
Storable a =>
Int -> (Ptr a -> IO b) -> IO (ForeignPtr a, b)
Alloc.create Int
size ((Ptr a -> IO a) -> IO (ForeignPtr a, a))
-> (Ptr a -> IO a) -> IO (ForeignPtr a, a)
forall a b. (a -> b) -> a -> b
$ ((nsize, Ptr a) -> IO a) -> nsize -> Ptr a -> IO a
forall a b c. ((a, b) -> c) -> a -> b -> c
curry (nsize, Ptr a) -> IO a
f nsize
subSizes