{-# 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.Simple.Slice as Slice
import qualified Data.Array.Knead.Simple.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 as MultiValue
import qualified LLVM.Extra.Marshal as Marshal

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.MV 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 (Cons slice createSlice deleteSlice) (Array arr createArr deleteArr) =
   Array
      (MultiValue.uncurry $ \paramSlice paramArr ->
         Slice.apply (slice paramSlice) (arr paramArr))
      (Priv.combineCreate createSlice createArr)
      (Priv.combineDelete deleteSlice deleteArr)


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


passAny :: Cubic p rank rank
passAny =
   Cons (const Slice.passAny) (Priv.createPlain $ const ()) Priv.deletePlain

pass ::
   (Unary.Natural rank0, Unary.Natural rank1) =>
   Cubic p rank0 rank1 ->
   Cubic p (Unary.Succ rank0) (Unary.Succ rank1)
pass (Cons slice create delete) = Cons (Slice.pass . slice) create delete

pick ::
   (Unary.Natural rank0, Unary.Natural rank1) =>
   Param.T p Index.Int ->
   Cubic p rank0 rank1 ->
   Cubic p (Unary.Succ rank0) rank1
pick = lift 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 = lift Slice.extrude

lift ::
   (Marshal.MV i) =>
   (Exp i -> Slice.Cubic rank0 rank1 -> Slice.Cubic rank2 rank3) ->
   Param.T p i ->
   Cubic p rank0 rank1 -> Cubic p rank2 rank3
lift f i (Cons slice create delete) =
   Param.withMulti i $ \getI valueI ->
   Cons
      (MultiValue.uncurry $ \slicep ip ->
         f (Expr.lift0 (valueI ip)) (slice slicep))
      (\p -> do
         (ctx, param) <- create p
         return (ctx, (param, getI p)))
      delete

instance Core.Process (T p sh0 sh1) where