{-# LANGUAGE GADTs               #-}
{-# LANGUAGE LambdaCase          #-}
{-# LANGUAGE RankNTypes          #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeOperators       #-}
-- |
-- Module      : Data.Array.Accelerate.Trafo.Var
-- Copyright   : [2012..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.Trafo.Var
  where

import Data.Array.Accelerate.AST
import Data.Array.Accelerate.AST.Environment
import Data.Array.Accelerate.AST.Idx
import Data.Array.Accelerate.AST.LeftHandSide
import Data.Array.Accelerate.AST.Var
import Data.Array.Accelerate.Representation.Array
import Data.Array.Accelerate.Representation.Type


data DeclareVars s t aenv where
  DeclareVars :: LeftHandSide s t env env'
              -> (env :> env')
              -> (forall env''. env' :> env'' -> Vars s env'' t)
              -> DeclareVars s t env

declareVars :: TupR s t -> DeclareVars s t env
declareVars :: TupR s t -> DeclareVars s t env
declareVars TupR s t
TupRunit
  = LeftHandSide s t env env
-> (env :> env)
-> (forall env''. (env :> env'') -> Vars s env'' t)
-> DeclareVars s t env
forall (s :: * -> *) t env env'.
LeftHandSide s t env env'
-> (env :> env')
-> (forall env''. (env' :> env'') -> Vars s env'' t)
-> DeclareVars s t env
DeclareVars LeftHandSide s t env env
forall env' env v (s :: * -> *).
(env' ~ env, v ~ ()) =>
LeftHandSide s v env env'
LeftHandSideUnit env :> env
forall env. env :> env
weakenId ((forall env''. (env :> env'') -> Vars s env'' t)
 -> DeclareVars s t env)
-> (forall env''. (env :> env'') -> Vars s env'' t)
-> DeclareVars s t env
forall a b. (a -> b) -> a -> b
$ TupR (Var s env'') () -> (env :> env'') -> TupR (Var s env'') ()
forall a b. a -> b -> a
const (TupR (Var s env'') () -> (env :> env'') -> TupR (Var s env'') ())
-> TupR (Var s env'') () -> (env :> env'') -> TupR (Var s env'') ()
forall a b. (a -> b) -> a -> b
$ TupR (Var s env'') ()
forall (s :: * -> *). TupR s ()
TupRunit
declareVars (TupRsingle s t
s)
  = LeftHandSide s t env (env, t)
-> (env :> (env, t))
-> (forall env''. ((env, t) :> env'') -> Vars s env'' t)
-> DeclareVars s t env
forall (s :: * -> *) t env env'.
LeftHandSide s t env env'
-> (env :> env')
-> (forall env''. (env' :> env'') -> Vars s env'' t)
-> DeclareVars s t env
DeclareVars (s t -> LeftHandSide s t env (env, t)
forall (s :: * -> *) v env. s v -> LeftHandSide s v env (env, v)
LeftHandSideSingle s t
s) (((env, t) :> (env, t)) -> env :> (env, t)
forall env t env'. ((env, t) :> env') -> env :> env'
weakenSucc (env, t) :> (env, t)
forall env. env :> env
weakenId) ((forall env''. ((env, t) :> env'') -> Vars s env'' t)
 -> DeclareVars s t env)
-> (forall env''. ((env, t) :> env'') -> Vars s env'' t)
-> DeclareVars s t env
forall a b. (a -> b) -> a -> b
$ \(env, t) :> env''
k -> Var s env'' t -> Vars s env'' t
forall (s :: * -> *) a. s a -> TupR s a
TupRsingle (Var s env'' t -> Vars s env'' t)
-> Var s env'' t -> Vars s env'' t
forall a b. (a -> b) -> a -> b
$ s t -> Idx env'' t -> Var s env'' t
forall (s :: * -> *) env t. s t -> Idx env t -> Var s env t
Var s t
s (Idx env'' t -> Var s env'' t) -> Idx env'' t -> Var s env'' t
forall a b. (a -> b) -> a -> b
$ (env, t) :> env''
k ((env, t) :> env'') -> Idx (env, t) t -> Idx env'' t
forall env env'.
(env :> env') -> forall t'. Idx env t' -> Idx env' t'
>:> Idx (env, t) t
forall env t. Idx (env, t) t
ZeroIdx
declareVars (TupRpair TupR s a
r1 TupR s b
r2)
  | DeclareVars LeftHandSide s a env env'
lhs1 env :> env'
subst1 forall env''. (env' :> env'') -> Vars s env'' a
a1 <- TupR s a -> DeclareVars s a env
forall (s :: * -> *) t env. TupR s t -> DeclareVars s t env
declareVars TupR s a
r1
  , DeclareVars LeftHandSide s b env' env'
lhs2 env' :> env'
subst2 forall env''. (env' :> env'') -> Vars s env'' b
a2 <- TupR s b -> DeclareVars s b env'
forall (s :: * -> *) t env. TupR s t -> DeclareVars s t env
declareVars TupR s b
r2
  = LeftHandSide s (a, b) env env'
-> (env :> env')
-> (forall env''. (env' :> env'') -> Vars s env'' (a, b))
-> DeclareVars s (a, b) env
forall (s :: * -> *) t env env'.
LeftHandSide s t env env'
-> (env :> env')
-> (forall env''. (env' :> env'') -> Vars s env'' t)
-> DeclareVars s t env
DeclareVars (LeftHandSide s a env env'
-> LeftHandSide s b env' env' -> LeftHandSide s (a, b) env env'
forall (s :: * -> *) v1 env env' v2 env''.
LeftHandSide s v1 env env'
-> LeftHandSide s v2 env' env''
-> LeftHandSide s (v1, v2) env env''
LeftHandSidePair LeftHandSide s a env env'
lhs1 LeftHandSide s b env' env'
lhs2) (env' :> env'
subst2 (env' :> env') -> (env :> env') -> env :> env'
forall env2 env3 env1.
(env2 :> env3) -> (env1 :> env2) -> env1 :> env3
.> env :> env'
subst1) ((forall env''. (env' :> env'') -> Vars s env'' (a, b))
 -> DeclareVars s (a, b) env)
-> (forall env''. (env' :> env'') -> Vars s env'' (a, b))
-> DeclareVars s (a, b) env
forall a b. (a -> b) -> a -> b
$ \env' :> env''
k -> (env' :> env'') -> Vars s env'' a
forall env''. (env' :> env'') -> Vars s env'' a
a1 (env' :> env''
k (env' :> env'') -> (env' :> env') -> env' :> env''
forall env2 env3 env1.
(env2 :> env3) -> (env1 :> env2) -> env1 :> env3
.> env' :> env'
subst2) Vars s env'' a -> TupR (Var s env'') b -> Vars s env'' (a, b)
forall (s :: * -> *) a b. TupR s a -> TupR s b -> TupR s (a, b)
`TupRpair` (env' :> env'') -> TupR (Var s env'') b
forall env''. (env' :> env'') -> Vars s env'' b
a2 env' :> env''
k


type InjectAcc  acc = forall env t. PreOpenAcc acc env t -> acc env t
type ExtractAcc acc = forall env t. acc env t -> Maybe (PreOpenAcc acc env t)

avarIn :: InjectAcc acc
       -> ArrayVar aenv a
       -> acc aenv a
avarIn :: InjectAcc acc -> ArrayVar aenv a -> acc aenv a
avarIn InjectAcc acc
inject v :: ArrayVar aenv a
v@(Var ArrayR{} Idx aenv a
_) = PreOpenAcc acc aenv (Array sh e) -> acc aenv (Array sh e)
InjectAcc acc
inject (ArrayVar aenv (Array sh e) -> PreOpenAcc acc aenv (Array sh e)
forall aenv sh e (acc :: * -> * -> *).
ArrayVar aenv (Array sh e) -> PreOpenAcc acc aenv (Array sh e)
Avar ArrayVar aenv a
ArrayVar aenv (Array sh e)
v)

avarsIn :: forall acc aenv arrs.
           InjectAcc acc
        -> ArrayVars aenv arrs
        -> acc aenv arrs
avarsIn :: InjectAcc acc -> ArrayVars aenv arrs -> acc aenv arrs
avarsIn InjectAcc acc
inject = ArrayVars aenv arrs -> acc aenv arrs
forall t. ArrayVars aenv t -> acc aenv t
go
  where
    go :: ArrayVars aenv t -> acc aenv t
    go :: ArrayVars aenv t -> acc aenv t
go ArrayVars aenv t
TupRunit       = PreOpenAcc acc aenv () -> acc aenv ()
InjectAcc acc
inject PreOpenAcc acc aenv ()
forall (acc :: * -> * -> *) aenv. PreOpenAcc acc aenv ()
Anil
    go (TupRsingle Var ArrayR aenv t
v) = InjectAcc acc -> Var ArrayR aenv t -> acc aenv t
forall (acc :: * -> * -> *) aenv a.
InjectAcc acc -> ArrayVar aenv a -> acc aenv a
avarIn InjectAcc acc
inject Var ArrayR aenv t
v
    go (TupRpair TupR (Var ArrayR aenv) a
a TupR (Var ArrayR aenv) b
b) = PreOpenAcc acc aenv (a, b) -> acc aenv (a, b)
InjectAcc acc
inject (TupR (Var ArrayR aenv) a -> acc aenv a
forall t. ArrayVars aenv t -> acc aenv t
go TupR (Var ArrayR aenv) a
a acc aenv a -> acc aenv b -> PreOpenAcc acc aenv (a, b)
forall (acc :: * -> * -> *) aenv as bs.
acc aenv as -> acc aenv bs -> PreOpenAcc acc aenv (as, bs)
`Apair` TupR (Var ArrayR aenv) b -> acc aenv b
forall t. ArrayVars aenv t -> acc aenv t
go TupR (Var ArrayR aenv) b
b)

avarsOut
    :: ExtractAcc acc
    -> PreOpenAcc acc aenv a
    -> Maybe (ArrayVars aenv a)
avarsOut :: ExtractAcc acc -> PreOpenAcc acc aenv a -> Maybe (ArrayVars aenv a)
avarsOut ExtractAcc acc
extract = \case
  PreOpenAcc acc aenv a
Anil   -> TupR (Var ArrayR aenv) () -> Maybe (TupR (Var ArrayR aenv) ())
forall a. a -> Maybe a
Just (TupR (Var ArrayR aenv) () -> Maybe (TupR (Var ArrayR aenv) ()))
-> TupR (Var ArrayR aenv) () -> Maybe (TupR (Var ArrayR aenv) ())
forall a b. (a -> b) -> a -> b
$ TupR (Var ArrayR aenv) ()
forall (s :: * -> *). TupR s ()
TupRunit
  Avar ArrayVar aenv (Array sh e)
v -> TupR (Var ArrayR aenv) (Array sh e)
-> Maybe (TupR (Var ArrayR aenv) (Array sh e))
forall a. a -> Maybe a
Just (TupR (Var ArrayR aenv) (Array sh e)
 -> Maybe (TupR (Var ArrayR aenv) (Array sh e)))
-> TupR (Var ArrayR aenv) (Array sh e)
-> Maybe (TupR (Var ArrayR aenv) (Array sh e))
forall a b. (a -> b) -> a -> b
$ ArrayVar aenv (Array sh e) -> TupR (Var ArrayR aenv) (Array sh e)
forall (s :: * -> *) a. s a -> TupR s a
TupRsingle ArrayVar aenv (Array sh e)
v
  Apair acc aenv as
al acc aenv bs
ar
    | Just PreOpenAcc acc aenv as
pl <- acc aenv as -> Maybe (PreOpenAcc acc aenv as)
ExtractAcc acc
extract acc aenv as
al
    , Just PreOpenAcc acc aenv bs
pr <- acc aenv bs -> Maybe (PreOpenAcc acc aenv bs)
ExtractAcc acc
extract acc aenv bs
ar
    , Just ArrayVars aenv as
as <- ExtractAcc acc
-> PreOpenAcc acc aenv as -> Maybe (ArrayVars aenv as)
forall (acc :: * -> * -> *) aenv a.
ExtractAcc acc -> PreOpenAcc acc aenv a -> Maybe (ArrayVars aenv a)
avarsOut ExtractAcc acc
extract PreOpenAcc acc aenv as
pl
    , Just ArrayVars aenv bs
bs <- ExtractAcc acc
-> PreOpenAcc acc aenv bs -> Maybe (ArrayVars aenv bs)
forall (acc :: * -> * -> *) aenv a.
ExtractAcc acc -> PreOpenAcc acc aenv a -> Maybe (ArrayVars aenv a)
avarsOut ExtractAcc acc
extract PreOpenAcc acc aenv bs
pr
    -> TupR (Var ArrayR aenv) (as, bs)
-> Maybe (TupR (Var ArrayR aenv) (as, bs))
forall a. a -> Maybe a
Just (ArrayVars aenv as
-> ArrayVars aenv bs -> TupR (Var ArrayR aenv) (as, bs)
forall (s :: * -> *) a b. TupR s a -> TupR s b -> TupR s (a, b)
TupRpair ArrayVars aenv as
as ArrayVars aenv bs
bs)
  PreOpenAcc acc aenv a
_ -> Maybe (ArrayVars aenv a)
forall a. Maybe a
Nothing