{-# LANGUAGE CPP #-}
#include "fusion-phases.h"

-- | Closures.
--   Used when closure converting the source program during vectorisation.
module Data.Array.Parallel.Lifted.Closure 
        ( -- * Closures.
          (:->)(..)
        , ($:)

        -- * Array Closures.
        , PData(..)
        , ($:^), liftedApply

        -- * Closure Construction.
        , closure1,  closure2,  closure3,  closure4,  closure5
        , closure1', closure2', closure3', closure4', closure5')
where
import Data.Array.Parallel.Pretty
import Data.Array.Parallel.PArray.PData.Base
import Data.Array.Parallel.PArray.PData.Unit
import Data.Array.Parallel.PArray.PData.Tuple2
import Data.Array.Parallel.PArray.PData.Tuple3
import Data.Array.Parallel.PArray.PData.Tuple4
import Data.Array.Parallel.PArray.PRepr
import qualified Data.Vector            as V
import GHC.Exts


-- Closures -------------------------------------------------------------------
-- | Define the fixity of the closure type constructor.
infixr 0 :->
infixl 1 $:, $:^

-- | The type of closures.
--   This bundles up:
---   1) the 'vectorised' version of the function that takes an explicit environment
--    2) the 'lifted' version, that works on arrays.
--       The first parameter of the lifted version is the 'lifting context'
--       that gives the length of the arrays being operated on.
--    3) the environment of the closure.
-- 
--   The vectoriser closure-converts the source program so that all functions
--   are expressed in this form.
data (a :-> b)
        = forall env. PA env
        => Clo  (env -> a -> b)
                (Int -> PData env -> PData a -> PData b)
                env

-- | Closure application.
($:) :: (a :-> b) -> a -> b
($:) (Clo fv _fl env) x  = fv env x
{-# INLINE_CLOSURE ($:) #-}


-- Array Closures -------------------------------------------------------------
-- | Arrays of closures (aka array closures)
--   We need to represent arrays of closures when vectorising partial applications.
--
--   For example, consider:
--     @mapP (+) xs   ::  [: Int -> Int :]@
--
--   Representing this an array of thunks doesn't work because we can't evaluate
--   it in a data parallel manner. Instead, we want *one* function applied to many
--   array elements.
-- 
--   Instead, such an array of closures is represented as the vectorised  and
--   lifted versions of (+), along with an environment array xs that contains the
--   partially applied arguments.
--
--     @mapP (+) xs  ==>  AClo plus_v plus_l xs@
--
data instance PData (a :-> b)
        =  forall env. PA env
        => AClo  (env -> a -> b)
                 (Int -> PData env -> PData a -> PData b)
                 (PData env)

data instance PDatas (a :-> b)
        =  forall env. PA env
        => AClos (env -> a -> b)
                 (Int -> PData env -> PData a -> PData b)
                 (PDatas env)


-- | Lifted closure application.
($:^) :: PArray (a :-> b) -> PArray a -> PArray b
PArray n# (AClo _ f es) $:^ PArray _ as 
        = PArray n# (f (I# n#) es as)
{-# INLINE ($:^) #-}


-- | Lifted closure application, taking an explicit lifting context.
liftedApply :: Int -> PData (a :-> b) -> PData a -> PData b
liftedApply n (AClo _ fl envs) as
        = fl n envs as
{-# INLINE_CLOSURE liftedApply #-}


-- Closure Construction -------------------------------------------------------
-- These functions are used for building closure representations of primitive
-- functions. They're used in D.A.P.Lifted.Combinators where we define the 
-- closure converted lifted array combinators that vectorised code uses.

-- | Construct an arity-1 closure,
--   from unlifted and lifted versions of a primitive function.
closure1 
        :: (a -> b)
        -> (Int -> PData a -> PData b)
        -> (a :-> b)

closure1 fv fl  
        = Clo   (\_env -> fv)
                (\n _env -> fl n)
                ()
{-# INLINE_CLOSURE closure1 #-}


-- | Construct an arity-2 closure,
--   from lifted and unlifted versions of a primitive function.
closure2 
        :: forall a b c. PA a
        => (a -> b -> c)
        -> (Int -> PData a -> PData b -> PData c)
        -> (a :-> b :-> c)

closure2 fv fl
 = let  fv_1 _ xa   = Clo fv fl xa
        fl_1 _ _ xs = AClo fv fl xs
        
   in   Clo fv_1 fl_1 ()
{-# INLINE_CLOSURE closure2 #-}


-- | Construct an arity-3 closure
--   from lifted and unlifted versions of a primitive function.
closure3 
        :: forall a b c d. (PA a, PA b)
        => (a -> b -> c -> d)
        -> (Int -> PData a -> PData b -> PData c -> PData d)
        -> (a :-> b :-> c :-> d)
        
closure3 fv fl
 = let  fv_1   _ xa = Clo   fv_2 fl_2 xa
        fl_1 _ _ xs = AClo  fv_2 fl_2 xs

        -----
        fv_2 xa yb   = Clo  fv_3 fl_3 (xa, yb)
        fl_2 _ xs ys = AClo fv_3 fl_3 (PTuple2 xs ys)

        -----
        fv_3 (xa, yb) zc           = fv xa yb zc
        fl_3 n (PTuple2 xs ys) zs  = fl n xs ys zs

   in   Clo fv_1 fl_1 ()
{-# INLINE_CLOSURE closure3 #-}


-- | Construct an arity-4 closure
--   from lifted and unlifted versions of a primitive function.
closure4 
        :: forall a b c d e. (PA a, PA b, PA c)
        => (a -> b -> c -> d -> e)
        -> (Int -> PData a -> PData b -> PData c -> PData d -> PData e)
        -> (a :-> b :-> c :-> d :-> e)
        
closure4 fv fl
 = let  fv_1   _ xa                   = Clo  fv_2 fl_2 xa
        fl_1 _ _ xs                   = AClo fv_2 fl_2 xs

        fv_2   xa yb                  = Clo  fv_3 fl_3 (xa, yb)
        fl_2 _ xs ys                  = AClo fv_3 fl_3 (PTuple2 xs ys)

        fv_3 (xa, yb) zc              = Clo  fv_4 fl_4 (xa, yb, zc)
        fl_3 _ (PTuple2 xs ys) zs     = AClo fv_4 fl_4 (PTuple3 xs ys zs)

        fv_4 (xa, yb, zc) ad          = fv xa yb zc ad
        fl_4 n (PTuple3 xs ys zs) as  = fl n xs ys zs as

   in   Clo fv_1 fl_1 ()
{-# INLINE_CLOSURE closure4 #-}


-- | Construct an arity-5 closure
--   from lifted and unlifted versions of a primitive function.
closure5
        :: forall a b c d e f. (PA a, PA b, PA c, PA d)
        => (a -> b -> c -> d -> e -> f)
        -> (Int -> PData a -> PData b -> PData c -> PData d -> PData e -> PData f)
        -> (a :-> b :-> c :-> d :-> e :-> f)
        
closure5 fv fl
 = let  fv_1   _ xa                     = Clo  fv_2 fl_2 xa
        fl_1 _ _ xs                     = AClo fv_2 fl_2 xs

        fv_2   xa yb                    = Clo  fv_3 fl_3 (xa, yb)
        fl_2 _ xs ys                    = AClo fv_3 fl_3 (PTuple2 xs ys)

        fv_3 (xa, yb) zc                = Clo  fv_4 fl_4 (xa, yb, zc)
        fl_3 _ (PTuple2 xs ys) zs       = AClo fv_4 fl_4 (PTuple3 xs ys zs)

        fv_4 (xa, yb, zc) ad            = Clo  fv_5 fl_5 (xa, yb, zc, ad)
        fl_4 _ (PTuple3 xs ys zs) as    = AClo fv_5 fl_5 (PTuple4 xs ys zs as)

        fv_5 (xa, yb, zc, ad) be        = fv xa yb zc ad be
        fl_5 n (PTuple4 xs ys zs as) bs = fl n xs ys zs as bs

   in   Clo fv_1 fl_1 ()
{-# INLINE_CLOSURE closure5 #-}


-- Closure constructors that take PArrays -------------------------------------
-- These versions are useful when defining prelude functions such as in 
-- D.A.P.Prelude.Int. They let us promote functions that work on PArrays 
-- to closures, while inferring the lifting context from the first argument.

-- | Construct an arity-1 closure.
closure1'
        :: forall a b
        .  (a -> b)
        -> (PArray a -> PArray b)
        -> (a :-> b)

closure1' fv fl 
 = let  {-# INLINE fl' #-}
        fl' (I# n#) pdata
         = case fl (PArray n# pdata) of
                 PArray _ pdata' -> pdata'
   in   closure1 fv fl'
{-# INLINE_CLOSURE closure1' #-}


-- | Construct an arity-2 closure.
closure2'
        :: forall a b c. PA a
        => (a -> b -> c)
        -> (PArray a -> PArray b -> PArray c)
        -> (a :-> b :-> c)

closure2' fv fl 
 = let  {-# INLINE fl' #-}
        fl' (I# n#) !pdata1 !pdata2
         = case fl (PArray n# pdata1) (PArray n# pdata2) of
                 PArray _ pdata' -> pdata'
   in   closure2 fv fl'
{-# INLINE_CLOSURE closure2' #-}


-- | Construct an arity-3 closure.
closure3'
        :: forall a b c d. (PA a, PA b) 
        => (a -> b -> c -> d)
        -> (PArray a -> PArray b -> PArray c -> PArray d)
        -> (a :-> b :-> c :-> d) 

closure3' fv fl 
 = let  {-# INLINE fl' #-}
        fl' (I# n#) !pdata1 !pdata2 !pdata3
         = case fl (PArray n# pdata1) (PArray n# pdata2) (PArray n# pdata3) of
                 PArray _ pdata' -> pdata'
   in   closure3 fv fl'
{-# INLINE_CLOSURE closure3' #-}


-- | Construct an arity-4 closure.
closure4'
        :: forall a b c d e. (PA a, PA b, PA c) 
        => (a -> b -> c -> d -> e)
        -> (PArray a -> PArray b -> PArray c -> PArray d -> PArray e)
        -> (a :-> b :-> c :-> d :-> e) 

closure4' fv fl 
 = let  {-# INLINE fl' #-}
        fl' (I# n#) !pdata1 !pdata2 !pdata3 !pdata4
         = case fl (PArray n# pdata1) (PArray n# pdata2) 
                   (PArray n# pdata3) (PArray n# pdata4) of
                 PArray _ pdata' -> pdata'
   in   closure4 fv fl'
{-# INLINE_CLOSURE closure4' #-}


-- | Construct an arity-5 closure.
closure5'
        :: forall a b c d e f. (PA a, PA b, PA c, PA d) 
        => (a -> b -> c -> d -> e -> f)
        -> (PArray a -> PArray b -> PArray c -> PArray d -> PArray e -> PArray f)
        -> (a :-> b :-> c :-> d :-> e :-> f) 

closure5' fv fl 
 = let  {-# INLINE fl' #-}
        fl' (I# n#) !pdata1 !pdata2 !pdata3 !pdata4 !pdata5
         = case fl (PArray n# pdata1) (PArray n# pdata2) 
                   (PArray n# pdata3) (PArray n# pdata4) 
                   (PArray n# pdata5) of
                 PArray _ pdata' -> pdata'
   in   closure5 fv fl'
{-# INLINE_CLOSURE closure5' #-}


-- PData instance for closures ------------------------------------------------
-- This needs to be here instead of in a module D.A.P.PArray.PData.Closure
-- to break an import loop.
-- We use INLINE_CLOSURE for these bindings instead of INLINE_PDATA because
-- most of the functions return closure constructors, and we want to eliminate
-- these early in the compilation.
--
instance PR (a :-> b) where

  {-# NOINLINE validPR #-}
  validPR (AClo _ _ env)
        = validPA env

  {-# NOINLINE nfPR #-}
  nfPR (AClo fv fl envs)
        = fv `seq` fl `seq` nfPA envs `seq` ()

  -- We can't test functions for equality.
  -- We can't test the environments either, because they're existentially quantified.
  -- Provided the closures have the same type, we just call them similar.
  {-# NOINLINE similarPR #-}
  similarPR _ _
        = True

  {-# NOINLINE coversPR #-}
  coversPR weak (AClo _ _ envs) ix
        = coversPA weak envs ix

  {-# NOINLINE pprpPR #-}
  pprpPR (Clo _ _ env)
        = vcat
        [ text "Clo"
        , pprpPA env ]

  {-# NOINLINE pprpDataPR #-}
  pprpDataPR (AClo _ _ envs)
        = vcat
        [ text "AClo"
        , pprpDataPA envs ]


  -- Constructors -------------------------------
  {-# INLINE_CLOSURE emptyPR #-}
  emptyPR
   = let  die    = error "emptydPR[:->]: no function in empty closure array"
      in  AClo die die (emptyPA :: PData ())

  {-# INLINE_CLOSURE replicatePR #-}
  replicatePR n (Clo fv fl envs)
        = AClo fv fl (replicatePA n envs)

  {-# INLINE_CLOSURE replicatesPR #-}
  replicatesPR lens (AClo fv fl envs)
        = AClo fv fl (replicatesPA lens envs)


  -- Projections --------------------------------
  {-# INLINE_CLOSURE lengthPR #-}
  lengthPR (AClo _ _ envs)
        = lengthPA envs

  {-# INLINE_CLOSURE indexPR #-}
  indexPR (AClo fv fl envs) ix
        = Clo fv fl  $ indexPA envs ix

  {-# INLINE_CLOSURE indexsPR #-}
  indexsPR (AClos fv fl envs) srcixs
        = AClo fv fl $ indexsPA envs srcixs

  {-# INLINE_CLOSURE extractPR #-}
  extractPR (AClo fv fl envs) start len
        = AClo fv fl $ extractPA envs start len

  {-# INLINE_CLOSURE extractssPR #-}
  extractssPR (AClos fv fl envs) ssegd
        = AClo fv fl $ extractssPA envs ssegd

  {-# INLINE_CLOSURE extractvsPR #-}
  extractvsPR (AClos fv fl envs) vsegd
        = AClo fv fl $ extractvsPA envs vsegd


  -- Pack and Combine ---------------------------
  {-# INLINE_CLOSURE packByTagPR #-}
  packByTagPR (AClo fv fl envs) tags tag
        = AClo fv fl $ packByTagPA envs tags tag


  -- Conversions --------------------------------
  {-# NOINLINE toVectorPR #-}
  toVectorPR (AClo fv fl envs)
        = V.map (Clo fv fl) $ toVectorPA envs


  -- PDatas -------------------------------------
  -- When constructing an empty array of closures, we don't know what 
  {-# INLINE_CLOSURE emptydPR #-}
  emptydPR 
   = let die    = error "emptydPR[:->]: no function in empty closure array"
     in  AClos  die die (emptydPA :: PDatas ())

  {-# INLINE_CLOSURE singletondPR #-}
  singletondPR (AClo fv fl env)
        = AClos fv fl $ singletondPA env
        
  {-# INLINE_CLOSURE lengthdPR #-}
  lengthdPR (AClos _ _ env)
        = lengthdPA env
        
  {-# INLINE_CLOSURE indexdPR #-}
  indexdPR (AClos fv fl envs) ix
        = AClo fv fl $ indexdPA envs ix

  {-# NOINLINE toVectordPR #-}
  toVectordPR (AClos fv fl envs)
        = V.map (AClo fv fl) $ toVectordPA envs


  -- Unsupported --------------------------------
  -- To support these operators we'd need to manage closure arrays containing
  -- multiple hetrogenous functions. But this is more work than we care for
  -- right now. Note that the problematic functions are all constructors, and
  -- we can't know that all the parameters contain the same function.
  appendPR      = dieHetroFunctions "appendPR"
  appendsPR     = dieHetroFunctions "appendsPR"
  combine2PR    = dieHetroFunctions "combine2PR"
  fromVectorPR  = dieHetroFunctions "fromVectorPR"
  appenddPR     = dieHetroFunctions "appenddPR"
  fromVectordPR = dieHetroFunctions "fromVectordPR"


dieHetroFunctions :: String -> a
dieHetroFunctions name
 = error $ unlines
   [ "Data.Array.Parallel.Lifted.Closure." ++ name
   , "  Unsupported Array Operation"
   , "  It looks like you're trying to define an array containing multiple"
   , "  hetrogenous functions, or trying to select between multiple arrays"
   , "  of functions in vectorised code. Although we could support this by"
   , "  constructing a new function that selects between them depending on"
   , "  what the array index is, to make that anywhere near efficient is"
   , "  more work than we care to do right now, and we expect this use case"
   , "  to be uncommon. If you want this to work then contact the DPH team"
   , "  and ask what you can do to help." ]


-- PRepr Instance -------------------------------------------------------------
-- This needs to be here instead of in D.A.P.PRepr.Instances 
-- to break an import loop.
--
type instance PRepr (a :-> b) 
        = a :-> b

instance (PA a, PA b) => PA (a :-> b) where
  toPRepr       = id
  fromPRepr     = id
  toArrPRepr    = id
  fromArrPRepr  = id
  toArrPReprs   = id
  fromArrPReprs = id