{- |
Generate and apply index maps.
This unifies the @replicate@ and @slice@ functions of the @accelerate@ package.
However the structure of slicing and replicating cannot depend on parameters.
If you need that, you must use 'ShapeDep.backpermute' and friends.
-}
{-
Some notes on the design choice:

Instead of the shallow embedding implemented by the 'T' type,
we could maintain a symbolic representation of the Slice and Replicate pattern,
like the accelerate package does.
We actually used that representation in former versions.
It has however some drawbacks:

* We need additional type functions that map from the pattern
  to the source and the target shape and we need a proof,
  that the images of these type functions are actually shapes.
  This worked already, but was rather cumbersome.

* We need a way to store and pass this pattern through the Parameter handler.
  This yields new problems:
  We need a wrapper type for wrapping Index, Shape, Slice, Replicate, Fold patterns.
  Then the question is whether we use one Wrap type with a phantom parameter
  or whether we define a Wrap type for every pattern type.
  That is, the options are to write either

  > Wrap Shape (Z:.Int:.Int)

  or

  > Shape (Z:.Int:.Int)

  The first one seems to save us many duplicate instances of
  Storable, MultiValue etc.
  and it allows us easily to reuse the (:.) for all kinds of patterns.
  However, we need a way to restrict the element type of the (:.)-list elements.
  We can define that using variable ConstraintKinds,
  but e.g. we are not able to add a Storable superclass constraint
  to the instance Storable (Wrap constr).
  That is, we are left with the second option
  and had to define a lot of similar Storable, MultiValue instances.
-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
module Data.Array.Knead.Simple.Slice (
   T,
   Linear,
   apply,
   passAny,
   pass,
   pick,
   extrude,
   (Core.$:.),
   ) where

import qualified Data.Array.Knead.Simple.ShapeDependent as ShapeDep
import qualified Data.Array.Knead.Simple.Private as Core

import qualified Data.Array.Knead.Index.Linear as Linear
import qualified Data.Array.Knead.Index.Nested.Shape as Shape
import qualified Data.Array.Knead.Expression as Expr
import Data.Array.Knead.Index.Linear ((#:.), (:.)((:.)), )
import Data.Array.Knead.Expression (Exp, )

import qualified LLVM.Extra.Multi.Value as MultiValue
import LLVM.Extra.Multi.Value (atom, )

import Prelude hiding (zipWith, zipWith3, zip, zip3, replicate, )



{-
This data type is almost identical to Core.Array.
The only difference is,
that the shape @sh1@ in T can depend on another shape @sh0@.
-}
data T sh0 sh1 =
   forall ix0 ix1.
   (Shape.Index sh0 ~ ix0, Shape.Index sh1 ~ ix1) =>
   Cons
      (Exp sh0 -> Exp sh1)
      (Exp ix1 -> Exp ix0)

{- |
This is essentially a 'ShapeDep.backpermute'.
-}
apply ::
   (Core.C array, Shape.C sh0, Shape.C sh1, MultiValue.C a) =>
   T sh0 sh1 ->
   array sh0 a ->
   array sh1 a
apply (Cons fsh fix) =
   ShapeDep.backpermute fsh fix


type Linear sh0 sh1 = T (Linear.Shape sh0) (Linear.Shape sh1)

{- |
Like @Any@ in @accelerate@.
-}
passAny :: Linear sh sh
passAny = Cons id id

{- |
Like @All@ in @accelerate@.
-}
pass ::
   Linear sh0 sh1 ->
   Linear (sh0:.i) (sh1:.i)
pass (Cons fsh fix) =
   Cons
      (Expr.modify (Linear.shape (atom:.atom)) $ \(sh:.s) -> fsh sh :. s)
      (Expr.modify (Linear.index (atom:.atom)) $ \(ix:.i) -> fix ix :. i)

{- |
Like @Int@ in @accelerate/slice@.
-}
pick ::
   Exp i ->
   Linear sh0 sh1 ->
   Linear (sh0:.i) sh1
pick i (Cons fsh fix) =
   Cons
      (fsh . Linear.tail)
      (\ix -> fix ix #:. i)

{- |
Like @Int@ in @accelerate/replicate@.
-}
extrude ::
   Exp i ->
   Linear sh0 sh1 ->
   Linear sh0 (sh1:.i)
extrude n (Cons fsh fix) =
   Cons
      (\sh -> fsh sh #:. n)
      (fix . Linear.tail)


instance Core.Process (T sh0 sh1) where