{-# LANGUAGE ConstraintKinds  #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GADTs            #-}
{-# LANGUAGE TypeOperators    #-}
{-# OPTIONS_HADDOCK hide #-}
-- |
-- Module      : Data.Array.Accelerate.LLVM.AST
-- Copyright   : [2017..2020] The Accelerate Team
-- License     : BSD3
--
-- Maintainer  : Trevor L. McDonell <trevor.mcdonell@gmail.com>
-- Stability   : experimental
-- Portability : non-portable (GHC extensions)
--

module Data.Array.Accelerate.LLVM.AST (

  DelayedOpenAcc(..),
  PreOpenAccCommand(..),
  PreOpenAccSkeleton(..),
  UnzipIdx(..),
  HasInitialValue,

) where

import Data.Array.Accelerate.LLVM.Execute.Async

import Data.Array.Accelerate.AST                                    ( PreOpenAfun(..), HasArraysR(..), ArrayVar, ALeftHandSide, Exp, Direction, PrimBool, arrayR )
import Data.Array.Accelerate.AST.Idx
import Data.Array.Accelerate.AST.Var
import Data.Array.Accelerate.Representation.Array
import Data.Array.Accelerate.Representation.Shape
import Data.Array.Accelerate.Representation.Type
import Data.Array.Accelerate.Type


-- | Non-computational array program operations, parameterised over array
-- variables represented as de Bruijn indices.
--
data PreOpenAccCommand acc arch aenv a where

  Avar        :: ArrayVar                   aenv arrs
              -> PreOpenAccCommand acc arch aenv arrs

  Alet        :: ALeftHandSide bnd aenv aenv'
              -> acc                   arch aenv  bnd
              -> acc                   arch aenv' body
              -> PreOpenAccCommand acc arch aenv  body

  Alloc       :: ArrayR (Array sh e)
              -> Exp                        aenv sh
              -> PreOpenAccCommand acc arch aenv (Array sh e)

  Use         :: ArrayR (Array sh e)
              -> Array sh e
              -> PreOpenAccCommand acc arch aenv (Array sh e)

  Unit        :: TypeR e
              -> Exp                        aenv e
              -> PreOpenAccCommand acc arch aenv (Scalar e)

  Apair       :: acc                   arch aenv arrs1
              -> acc                   arch aenv arrs2
              -> PreOpenAccCommand acc arch aenv (arrs1, arrs2)

  Anil        :: PreOpenAccCommand acc arch aenv ()

  Apply       :: ArraysR bs
              -> PreOpenAfun      (acc arch) aenv (as -> bs)
              -> acc                   arch  aenv as
              -> PreOpenAccCommand acc arch  aenv bs

  Aforeign    :: ArraysR bs
              -> String
              -> (as -> Par arch (FutureR arch bs))
              -> acc                   arch aenv as
              -> PreOpenAccCommand acc arch aenv bs

  Acond       :: Exp                         aenv PrimBool
              -> acc                   arch  aenv arrs
              -> acc                   arch  aenv arrs
              -> PreOpenAccCommand acc arch  aenv arrs

  Awhile      :: PreOpenAfun      (acc arch) aenv (arrs -> Scalar PrimBool)
              -> PreOpenAfun      (acc arch) aenv (arrs -> arrs)
              -> acc                   arch  aenv arrs
              -> PreOpenAccCommand acc arch  aenv arrs

  Reshape     :: ShapeR  sh
              -> Exp                         aenv sh
              -> ArrayVar                    aenv (Array sh' e)
              -> PreOpenAccCommand acc arch  aenv (Array sh  e)

  Unzip       :: UnzipIdx tup e
              -> ArrayVar                   aenv (Array sh tup)
              -> PreOpenAccCommand acc arch aenv (Array sh e)


-- | Collective array computations parameterised over array variables
-- represented as de Bruijn indices.
--
data PreOpenAccSkeleton acc arch aenv a where

  -- Producers. The only way these terms can appear in the AST is if they
  -- are applied to a manifest array.
  --
  Map         :: TypeR b
              -> acc                    arch aenv (Array sh a)
              -> PreOpenAccSkeleton acc arch aenv (Array sh b)

  Generate    :: ArrayR (Array sh e)
              -> Exp                         aenv sh
              -> PreOpenAccSkeleton acc arch aenv (Array sh e)

  Transform   :: ArrayR (Array sh' b)
              -> Exp                         aenv sh'
              -> acc                    arch aenv (Array sh  a)
              -> PreOpenAccSkeleton acc arch aenv (Array sh' b)

  Backpermute :: ShapeR sh'
              -> Exp                         aenv sh'
              -> acc                    arch aenv (Array sh  e)
              -> PreOpenAccSkeleton acc arch aenv (Array sh' e)

  -- Consumers. These may have been applied to either manifest or delayed
  -- array data.
  --
  Fold        :: HasInitialValue
              -> DelayedOpenAcc     acc arch aenv (Array (sh, Int) e)
              -> PreOpenAccSkeleton acc arch aenv (Array sh e)

  FoldSeg     :: IntegralType i
              -> HasInitialValue
              -> DelayedOpenAcc     acc arch aenv (Array (sh, Int) e)
              -> DelayedOpenAcc     acc arch aenv (Segments i)
              -> PreOpenAccSkeleton acc arch aenv (Array (sh, Int) e)

  Scan        :: Direction
              -> HasInitialValue
              -> DelayedOpenAcc     acc arch aenv (Array (sh, Int) e)
              -> PreOpenAccSkeleton acc arch aenv (Array (sh, Int) e)

  Scan'       :: Direction
              -> DelayedOpenAcc     acc arch aenv (Array (sh, Int) e)
              -> PreOpenAccSkeleton acc arch aenv (Array (sh, Int) e, Array sh e)

  Permute     :: acc                    arch aenv (Array sh' e)     -- target array (default values)
              -> DelayedOpenAcc     acc arch aenv (Array sh  e)     -- source values
              -> PreOpenAccSkeleton acc arch aenv (Array sh' e)

  Stencil1    :: TypeR b
              -> sh                                                 -- stencil offset/halo size
              -> DelayedOpenAcc     acc arch aenv (Array sh a)
              -> PreOpenAccSkeleton acc arch aenv (Array sh b)

  Stencil2    :: TypeR c
              -> sh                                                 -- stencil offset/halo size
              -> DelayedOpenAcc     acc arch aenv (Array sh a)
              -> DelayedOpenAcc     acc arch aenv (Array sh b)
              -> PreOpenAccSkeleton acc arch aenv (Array sh c)

data UnzipIdx a b where
  UnzipId   ::                                   UnzipIdx a a
  UnzipPrj  :: PairIdx a b   -> UnzipIdx b c ->  UnzipIdx a c
  UnzipUnit ::                                   UnzipIdx a ()
  UnzipPair :: UnzipIdx a b1 -> UnzipIdx a b2 -> UnzipIdx a (b1, b2)

-- Denotes whether the fold or scan has an initial value.
-- When False, this is a fold1 or scan1.
--
type HasInitialValue = Bool

-- | Representation for array arguments.
--
-- If the argument is a delayed array (that is, it was fused into its
-- consumer) we only need to keep track of the extent of the argument. If
-- the argument is a manifest array, we recurse into the subterm.
--
data DelayedOpenAcc acc arch aenv a where
  Delayed     :: ArrayR (Array sh e)
              -> Exp aenv sh
              -> DelayedOpenAcc acc arch aenv (Array sh e)

  Manifest    :: ArraysR (Array sh e)
              -> acc arch aenv (Array sh e)
              -> DelayedOpenAcc acc arch aenv (Array sh e)

instance HasArraysR (acc arch) => HasArraysR (PreOpenAccCommand acc arch) where
  {-# INLINEABLE arraysR #-}
  arraysR (Avar (Var repr _))                   = TupRsingle repr
  arraysR (Alet _ _ a)                          = arraysR a
  arraysR (Alloc repr _)                        = TupRsingle repr
  arraysR (Use repr _)                          = TupRsingle repr
  arraysR (Unit tp _)                           = TupRsingle $ ArrayR ShapeRz tp
  arraysR (Apair a1 a2)                         = arraysR a1 `TupRpair` arraysR a2
  arraysR Anil                                  = TupRunit
  arraysR (Apply repr _ _)                      = repr
  arraysR (Aforeign repr _ _ _)                 = repr
  arraysR (Acond _ a1 _)                        = arraysR a1
  arraysR (Awhile _ _ a)                        = arraysR a
  arraysR (Reshape shr _ (Var (ArrayR _ tp) _)) = TupRsingle $ ArrayR shr tp
  arraysR (Unzip idx (Var (ArrayR shr tp) _))   = TupRsingle $ ArrayR shr $ go idx tp
    where
      go :: UnzipIdx a b -> TypeR a -> TypeR b
      go UnzipId                    t              = t
      go (UnzipPrj PairIdxLeft ix)  (TupRpair t _) = go ix t
      go (UnzipPrj PairIdxRight ix) (TupRpair _ t) = go ix t
      go UnzipUnit                  _              = TupRunit
      go (UnzipPair ix1 ix2)        t              = go ix1 t `TupRpair` go ix2 t
      go _                          _              = error "Time enough for life to unfold all the precious things life has in store."

instance HasArraysR (acc arch) => HasArraysR (PreOpenAccSkeleton acc arch) where
  {-# INLINEABLE arraysR #-}
  arraysR (Map tp a)               = let ArrayR shr _ = arrayR a
                                     in  TupRsingle $ ArrayR shr tp
  arraysR (Generate repr _)        = TupRsingle repr
  arraysR (Transform repr _ _)     = TupRsingle repr
  arraysR (Backpermute shr _ a)    = TupRsingle $ ArrayR shr $ arrayRtype $ arrayR a
  arraysR (Fold _ a)               = let ArrayR (ShapeRsnoc shr) tp = arrayR a
                                     in  TupRsingle $ ArrayR shr tp
  arraysR (FoldSeg _ _ a _)        = arraysR a
  arraysR (Scan _ _ a)             = arraysR a
  arraysR (Scan' _ a)              = let ArrayR (ShapeRsnoc shr) tp = arrayR a
                                     in  TupRsingle (ArrayR (ShapeRsnoc shr) tp) `TupRpair` TupRsingle (ArrayR shr tp)
  arraysR (Permute a _)            = arraysR a
  arraysR (Stencil1 tp _ a)        = let ArrayR shr _ = arrayR a
                                     in  TupRsingle $ ArrayR shr tp
  arraysR (Stencil2 tp _ a _)      = let ArrayR shr _ = arrayR a
                                     in  TupRsingle $ ArrayR shr tp

instance HasArraysR (acc arch) => HasArraysR (DelayedOpenAcc acc arch) where
  {-# INLINEABLE arraysR #-}
  arraysR (Delayed  repr _) = TupRsingle repr
  arraysR (Manifest repr _) = repr