{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE TypeOperators #-}
module Data.Array.Knead.Parameterized.Slice (
   T,
   apply,
   Cubic,
   passAny,
   pass,
   pick,
   extrude,
   (Core.$:.),
   ) where

import qualified Data.Array.Knead.Parameterized.Private as Priv
import Data.Array.Knead.Parameterized.Private (Array(Array), )

import qualified Data.Array.Knead.Symbolic.Slice as Slice
import qualified Data.Array.Knead.Symbolic.Private as Core

import qualified Data.Array.Knead.Shape.Cubic.Int as Index
import qualified Data.Array.Knead.Shape.Cubic as Cubic
import qualified Data.Array.Knead.Shape as Shape
import qualified Data.Array.Knead.Expression as Expr
import Data.Array.Knead.Expression (Exp, )

import qualified LLVM.DSL.Parameter as Param

import qualified LLVM.Extra.Multi.Value.Marshal as Marshal
import qualified LLVM.Extra.Multi.Value as MultiValue

import qualified Type.Data.Num.Unary as Unary


{-
This wrapper data type is pretty much the same as Parameterized.Array
but there seems to be no benefit from using the same data structure for it.
-}
data T p sh0 sh1 =
   forall parameter context.
   (Marshal.C parameter) =>
   Cons {
      ()
_core :: MultiValue.T parameter -> Slice.T sh0 sh1,
      ()
_createContext :: p -> IO (context, parameter),
      ()
_deleteContext :: context -> IO ()
   }

apply ::
   (Shape.C sh0, Shape.C sh1, MultiValue.C a) =>
   T p sh0 sh1 ->
   Array p sh0 a ->
   Array p sh1 a
apply :: forall sh0 sh1 a p.
(C sh0, C sh1, C a) =>
T p sh0 sh1 -> Array p sh0 a -> Array p sh1 a
apply (Cons T parameter -> T sh0 sh1
slice p -> IO (context, parameter)
createSlice context -> IO ()
deleteSlice) (Array T parameter -> Array sh0 a
arr p -> IO (context, parameter)
createArr context -> IO ()
deleteArr) =
   (T (parameter, parameter) -> Array sh1 a)
-> (p -> IO ((context, context), (parameter, parameter)))
-> ((context, context) -> IO ())
-> Array p sh1 a
forall p sh a parameter context.
C parameter =>
(T parameter -> Array sh a)
-> (p -> IO (context, parameter))
-> (context -> IO ())
-> Array p sh a
Array
      ((T parameter -> T parameter -> Array sh1 a)
-> T (parameter, parameter) -> Array sh1 a
forall a b c. (T a -> T b -> c) -> T (a, b) -> c
MultiValue.uncurry ((T parameter -> T parameter -> Array sh1 a)
 -> T (parameter, parameter) -> Array sh1 a)
-> (T parameter -> T parameter -> Array sh1 a)
-> T (parameter, parameter)
-> Array sh1 a
forall a b. (a -> b) -> a -> b
$ \T parameter
paramSlice T parameter
paramArr ->
         T sh0 sh1 -> Array sh0 a -> Array sh1 a
forall (array :: * -> * -> *) sh0 sh1 a.
(C array, C sh0, C sh1, C a) =>
T sh0 sh1 -> array sh0 a -> array sh1 a
Slice.apply (T parameter -> T sh0 sh1
slice T parameter
paramSlice) (T parameter -> Array sh0 a
arr T parameter
paramArr))
      ((p -> IO (context, parameter))
-> (p -> IO (context, parameter))
-> p
-> IO ((context, context), (parameter, parameter))
forall (m :: * -> *) p ctxA paramA ctxB paramB.
Monad m =>
(p -> m (ctxA, paramA))
-> (p -> m (ctxB, paramB))
-> p
-> m ((ctxA, ctxB), (paramA, paramB))
Priv.combineCreate p -> IO (context, parameter)
createSlice p -> IO (context, parameter)
createArr)
      ((context -> IO ())
-> (context -> IO ()) -> (context, context) -> IO ()
forall (m :: * -> *) ctxA ctxB.
Monad m =>
(ctxA -> m ()) -> (ctxB -> m ()) -> (ctxA, ctxB) -> m ()
Priv.combineDelete context -> IO ()
deleteSlice context -> IO ()
deleteArr)


type Cubic p rank0 rank1 = T p (Cubic.Shape rank0) (Cubic.Shape rank1)


passAny :: Cubic p rank rank
passAny :: forall p rank. Cubic p rank rank
passAny =
   (T () -> T (Shape rank) (Shape rank))
-> (p -> IO ((), ()))
-> (() -> IO ())
-> T p (Shape rank) (Shape rank)
forall p sh0 sh1 parameter context.
C parameter =>
(T parameter -> T sh0 sh1)
-> (p -> IO (context, parameter))
-> (context -> IO ())
-> T p sh0 sh1
Cons (T (Shape rank) (Shape rank) -> T () -> T (Shape rank) (Shape rank)
forall a b. a -> b -> a
const T (Shape rank) (Shape rank)
forall rank. Cubic rank rank
Slice.passAny) ((p -> ()) -> p -> IO ((), ())
forall (m :: * -> *) p pl. Monad m => (p -> pl) -> p -> m ((), pl)
Priv.createPlain ((p -> ()) -> p -> IO ((), ())) -> (p -> ()) -> p -> IO ((), ())
forall a b. (a -> b) -> a -> b
$ () -> p -> ()
forall a b. a -> b -> a
const ()) () -> IO ()
forall (m :: * -> *). Monad m => () -> m ()
Priv.deletePlain

pass ::
   (Unary.Natural rank0, Unary.Natural rank1) =>
   Cubic p rank0 rank1 ->
   Cubic p (Unary.Succ rank0) (Unary.Succ rank1)
pass :: forall rank0 rank1 p.
(Natural rank0, Natural rank1) =>
Cubic p rank0 rank1 -> Cubic p (Succ rank0) (Succ rank1)
pass (Cons T parameter -> T (Shape rank0) (Shape rank1)
slice p -> IO (context, parameter)
create context -> IO ()
delete) = (T parameter -> T (Shape (Succ rank0)) (Shape (Succ rank1)))
-> (p -> IO (context, parameter))
-> (context -> IO ())
-> T p (Shape (Succ rank0)) (Shape (Succ rank1))
forall p sh0 sh1 parameter context.
C parameter =>
(T parameter -> T sh0 sh1)
-> (p -> IO (context, parameter))
-> (context -> IO ())
-> T p sh0 sh1
Cons (T (Shape rank0) (Shape rank1)
-> T (Shape (Succ rank0)) (Shape (Succ rank1))
forall rank0 rank1.
(Natural rank0, Natural rank1) =>
Cubic rank0 rank1 -> Cubic (Succ rank0) (Succ rank1)
Slice.pass (T (Shape rank0) (Shape rank1)
 -> T (Shape (Succ rank0)) (Shape (Succ rank1)))
-> (T parameter -> T (Shape rank0) (Shape rank1))
-> T parameter
-> T (Shape (Succ rank0)) (Shape (Succ rank1))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. T parameter -> T (Shape rank0) (Shape rank1)
slice) p -> IO (context, parameter)
create context -> IO ()
delete

pick ::
   (Unary.Natural rank0, Unary.Natural rank1) =>
   Param.T p Index.Int ->
   Cubic p rank0 rank1 ->
   Cubic p (Unary.Succ rank0) rank1
pick :: forall rank0 rank1 p.
(Natural rank0, Natural rank1) =>
T p Int -> Cubic p rank0 rank1 -> Cubic p (Succ rank0) rank1
pick = (Exp Int -> Cubic rank0 rank1 -> Cubic (Succ rank0) rank1)
-> T p Int -> Cubic p rank0 rank1 -> Cubic p (Succ rank0) rank1
forall i rank0 rank1 rank2 rank3 p.
C i =>
(Exp i -> Cubic rank0 rank1 -> Cubic rank2 rank3)
-> T p i -> Cubic p rank0 rank1 -> Cubic p rank2 rank3
lift Exp Int -> Cubic rank0 rank1 -> Cubic (Succ rank0) rank1
forall rank0 rank1.
(Natural rank0, Natural rank1) =>
Exp Int -> Cubic rank0 rank1 -> Cubic (Succ rank0) rank1
Slice.pick

extrude ::
   (Unary.Natural rank0, Unary.Natural rank1) =>
   Param.T p Index.Int ->
   Cubic p rank0 rank1 ->
   Cubic p rank0 (Unary.Succ rank1)
extrude :: forall rank0 rank1 p.
(Natural rank0, Natural rank1) =>
T p Int -> Cubic p rank0 rank1 -> Cubic p rank0 (Succ rank1)
extrude = (Exp Int -> Cubic rank0 rank1 -> Cubic rank0 (Succ rank1))
-> T p Int -> Cubic p rank0 rank1 -> Cubic p rank0 (Succ rank1)
forall i rank0 rank1 rank2 rank3 p.
C i =>
(Exp i -> Cubic rank0 rank1 -> Cubic rank2 rank3)
-> T p i -> Cubic p rank0 rank1 -> Cubic p rank2 rank3
lift Exp Int -> Cubic rank0 rank1 -> Cubic rank0 (Succ rank1)
forall rank0 rank1.
(Natural rank0, Natural rank1) =>
Exp Int -> Cubic rank0 rank1 -> Cubic rank0 (Succ rank1)
Slice.extrude

lift ::
   (Marshal.C i) =>
   (Exp i -> Slice.Cubic rank0 rank1 -> Slice.Cubic rank2 rank3) ->
   Param.T p i ->
   Cubic p rank0 rank1 -> Cubic p rank2 rank3
lift :: forall i rank0 rank1 rank2 rank3 p.
C i =>
(Exp i -> Cubic rank0 rank1 -> Cubic rank2 rank3)
-> T p i -> Cubic p rank0 rank1 -> Cubic p rank2 rank3
lift Exp i -> Cubic rank0 rank1 -> Cubic rank2 rank3
f T p i
i (Cons T parameter -> Cubic rank0 rank1
slice p -> IO (context, parameter)
create context -> IO ()
delete) =
   T p i
-> (forall parameters.
    C parameters =>
    (p -> parameters)
    -> (T parameters -> T i)
    -> T p (T ShapeTag rank2) (T ShapeTag rank3))
-> T p (T ShapeTag rank2) (T ShapeTag rank3)
forall b p a.
C b =>
T p b
-> (forall parameters.
    C parameters =>
    (p -> parameters) -> (T parameters -> T b) -> a)
-> a
Param.withMulti T p i
i ((forall parameters.
  C parameters =>
  (p -> parameters)
  -> (T parameters -> T i)
  -> T p (T ShapeTag rank2) (T ShapeTag rank3))
 -> T p (T ShapeTag rank2) (T ShapeTag rank3))
-> (forall parameters.
    C parameters =>
    (p -> parameters)
    -> (T parameters -> T i)
    -> T p (T ShapeTag rank2) (T ShapeTag rank3))
-> T p (T ShapeTag rank2) (T ShapeTag rank3)
forall a b. (a -> b) -> a -> b
$ \p -> parameters
getI T parameters -> T i
valueI ->
   (T (parameter, parameters) -> Cubic rank2 rank3)
-> (p -> IO (context, (parameter, parameters)))
-> (context -> IO ())
-> T p (T ShapeTag rank2) (T ShapeTag rank3)
forall p sh0 sh1 parameter context.
C parameter =>
(T parameter -> T sh0 sh1)
-> (p -> IO (context, parameter))
-> (context -> IO ())
-> T p sh0 sh1
Cons
      ((T parameter -> T parameters -> Cubic rank2 rank3)
-> T (parameter, parameters) -> Cubic rank2 rank3
forall a b c. (T a -> T b -> c) -> T (a, b) -> c
MultiValue.uncurry ((T parameter -> T parameters -> Cubic rank2 rank3)
 -> T (parameter, parameters) -> Cubic rank2 rank3)
-> (T parameter -> T parameters -> Cubic rank2 rank3)
-> T (parameter, parameters)
-> Cubic rank2 rank3
forall a b. (a -> b) -> a -> b
$ \T parameter
slicep T parameters
ip ->
         Exp i -> Cubic rank0 rank1 -> Cubic rank2 rank3
f (T i -> Exp i
forall a. T a -> Exp a
forall (val :: * -> *) a. Value val => T a -> val a
Expr.lift0 (T parameters -> T i
valueI T parameters
ip)) (T parameter -> Cubic rank0 rank1
slice T parameter
slicep))
      (\p
p -> do
         (context
ctx, parameter
param) <- p -> IO (context, parameter)
create p
p
         (context, (parameter, parameters))
-> IO (context, (parameter, parameters))
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (context
ctx, (parameter
param, p -> parameters
getI p
p)))
      context -> IO ()
delete

instance Core.Process (T p sh0 sh1) where