{-# LANGUAGE GADTs #-}
{-# OPTIONS_HADDOCK hide #-}
-- |
-- Module      : Data.Array.Accelerate.LLVM.CodeGen.Environment
-- Copyright   : [2015..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.CodeGen.Environment
  where

import Data.IntMap                                              ( IntMap )
import Data.String
import Text.Printf
import qualified Data.IntMap                                    as IM

import Data.Array.Accelerate.AST                                ( ArrayVar )
import Data.Array.Accelerate.AST.Idx                            ( Idx(..), idxToInt )
import Data.Array.Accelerate.AST.Var                            ( Var(..) )
import Data.Array.Accelerate.Error                              ( internalError )
import Data.Array.Accelerate.Representation.Array               ( Array, ArrayR(..) )

import Data.Array.Accelerate.LLVM.CodeGen.IR

import LLVM.AST.Type.Name

import GHC.Stack


-- Scalar environment
-- ==================

-- | An environment for local scalar expression bindings, encoded at the value
-- level as a heterogenous snoc list, and on the type level as nested tuples.
--
data Val env where
  Empty ::                          Val ()
  Push  :: Val env -> Operands t -> Val (env, t)

-- | Projection of a value from the valuation environment using a de Bruijn
-- index.
--
prj :: Idx env t -> Val env -> Operands t
prj :: Idx env t -> Val env -> Operands t
prj Idx env t
ZeroIdx      (Push Val env
_   Operands t
v) = Operands t
Operands t
v
prj (SuccIdx Idx env1 t
ix) (Push Val env
val Operands t
_) = Idx env1 t -> Val env1 -> Operands t
forall env t. Idx env t -> Val env -> Operands t
prj Idx env1 t
ix Val env1
Val env
val


-- Array environment
-- =================

-- | A mapping between the environment index of a free array variable and the
-- Name of that array to be used in the generated code.
--
-- This simply compresses the array indices into a continuous range, rather than
-- directly using the integer equivalent of the de Bruijn index. Thus, the
-- result is still sensitive to the order of let bindings, but not of any
-- intermediate (unused) free array variables.
--
type Gamma aenv = IntMap (Label, Idx' aenv)

data Idx' aenv where
  Idx' :: ArrayR (Array sh e) -> Idx aenv (Array sh e) -> Idx' aenv

-- Projection of a value from the array environment using a de Bruijn index.
-- This returns a pair of operands to access the shape and array data
-- respectively.
--
aprj :: HasCallStack => Idx aenv t -> Gamma aenv -> Name t
aprj :: Idx aenv t -> Gamma aenv -> Name t
aprj Idx aenv t
ix Gamma aenv
aenv =
  case Key -> Gamma aenv -> Maybe (Label, Idx' aenv)
forall a. Key -> IntMap a -> Maybe a
IM.lookup (Idx aenv t -> Key
forall env t. Idx env t -> Key
idxToInt Idx aenv t
ix) Gamma aenv
aenv of
    Maybe (Label, Idx' aenv)
Nothing             -> String -> Name t
forall a. HasCallStack => String -> a
internalError String
"free variable not registered"
    Just (Label ShortByteString
n,Idx' aenv
_)    -> ShortByteString -> Name t
forall a. ShortByteString -> Name a
Name ShortByteString
n


-- | Construct the array environment index, will be used by code generation to
-- map free array variable indices to names in the generated code.
--
makeGamma :: IntMap (Idx' aenv) -> Gamma aenv
makeGamma :: IntMap (Idx' aenv) -> Gamma aenv
makeGamma = (Key, Gamma aenv) -> Gamma aenv
forall a b. (a, b) -> b
snd ((Key, Gamma aenv) -> Gamma aenv)
-> (IntMap (Idx' aenv) -> (Key, Gamma aenv))
-> IntMap (Idx' aenv)
-> Gamma aenv
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Key -> Idx' aenv -> (Key, (Label, Idx' aenv)))
-> Key -> IntMap (Idx' aenv) -> (Key, Gamma aenv)
forall a b c. (a -> b -> (a, c)) -> a -> IntMap b -> (a, IntMap c)
IM.mapAccum (\Key
n Idx' aenv
ix -> (Key
nKey -> Key -> Key
forall a. Num a => a -> a -> a
+Key
1, Key -> Idx' aenv -> (Label, Idx' aenv)
forall aenv. Key -> Idx' aenv -> (Label, Idx' aenv)
toAval Key
n Idx' aenv
ix)) Key
0
  where
    toAval :: Int -> Idx' aenv -> (Label, Idx' aenv)
    toAval :: Key -> Idx' aenv -> (Label, Idx' aenv)
toAval Key
n Idx' aenv
ix = (String -> Label
forall a. IsString a => String -> a
fromString (String -> Key -> String
forall r. PrintfType r => String -> r
printf String
"fv%d" Key
n), Idx' aenv
ix)

-- | A free variable
--
freevar :: ArrayVar aenv a -> IntMap (Idx' aenv)
freevar :: ArrayVar aenv a -> IntMap (Idx' aenv)
freevar (Var repr :: ArrayR a
repr@ArrayR{} Idx aenv a
ix) = Key -> Idx' aenv -> IntMap (Idx' aenv)
forall a. Key -> a -> IntMap a
IM.singleton (Idx aenv a -> Key
forall env t. Idx env t -> Key
idxToInt Idx aenv a
ix) (ArrayR (Array sh e) -> Idx aenv (Array sh e) -> Idx' aenv
forall sh e aenv.
ArrayR (Array sh e) -> Idx aenv (Array sh e) -> Idx' aenv
Idx' ArrayR a
ArrayR (Array sh e)
repr Idx aenv a
Idx aenv (Array sh e)
ix)