content: {-# LANGUAGE QuasiQuotes #-}
  {-# OPTIONS_GHC -fno-warn-orphans #-}
  -- |
  -- Module      : Data.Array.Accelerate.CUDA.CodeGen.Monad
  -- Copyright   : [2008..2014] Manuel M T Chakravarty, Gabriele Keller
  --               [2009..2014] Trevor L. McDonell
  -- License     : BSD3
  --
  -- Maintainer  : Trevor L. McDonell <tmcdonell@cse.unsw.edu.au>
  -- Stability   : experimental
  -- Portability : non-portable (GHC extensions)
  --
  
  module Data.Array.Accelerate.CUDA.CodeGen.Monad (
  
    CUDA, Gen, AccST(..), ExpST(..),
    runCUDA, runCGM, evalCGM, execCGM, pushEnv, getEnv, fresh, bind, use,
  
  ) where
  
  import Prelude                                          hiding ( exp )
  import Data.HashSet                                     ( HashSet )
  import Data.Hashable
  import Control.Monad
  import Control.Monad.State.Strict
  import Language.C.Quote.CUDA
  import qualified Language.C                             as C
  import qualified Data.HashSet                           as Set
  
  import Data.Array.Accelerate.Trafo
  import Data.Array.Accelerate.CUDA.CodeGen.Type
  
  instance Hashable C.Exp where
    hashWithSalt salt = hashWithSalt salt . show
  
  
  -- The state of the code generator monad. The outer monad is used to generate
  -- fresh variable names and collect any headers required for foreign functions.
  -- The inner is used to collect local bindings (for let-bound expressions) and
  -- to record when previously bound variables are later used for dead-code
  -- analysis.
  --
  -- This separation is required so that names are unique across all generated
  -- code fragments of a skeleton.
  --
  type CUDA       = State  AccST
  type Gen        = StateT ExpST CUDA
  
  data AccST = AccST
    { counter     :: {-# UNPACK #-} !Int
    , headers     :: !(HashSet String)
    }
  
  data ExpST = ExpST
    -- A stack of (typically) declarations that must be evaluated before
    -- computation of the main expression. These are typically introduced by
    -- let-bindings. The list is kept in reverse order, with the last to be
    -- evaluated (newest addition) at the front of the list.
    --
    { localBindings       :: [C.BlockItem]
  
    -- A set of the Var's that we know have been used in the expression. With this
    -- we can do def-use analysis for simple dead-code elimination.
    --
    , usedTerms           :: !(HashSet C.Exp)
    }
  
  
  -- Run the code generator with a fresh environment, returning the result and
  -- final state.
  --
  runCUDA :: CUDA a -> (a, AccST)
  runCUDA a = runState a (AccST 0 Set.empty)
  
  runCGM :: Gen a -> CUDA (a, ExpST)
  runCGM a = runStateT a (ExpST [] Set.empty)
  
  evalCGM :: Gen a -> CUDA a
  evalCGM = fmap fst . runCGM
  
  execCGM :: Gen a -> CUDA ExpST
  execCGM = fmap snd . runCGM
  
  
  -- Create new binding points for the C expressions associated with the given AST
  -- term, unless the term is itself a variable.
  --
  pushEnv :: DelayedOpenExp env aenv t -> [C.Exp] -> Gen [C.Exp]
  pushEnv exp cs =
    let tys               = expType exp
        zipWithM' xs ys f = zipWithM f xs ys
    in
    zipWithM' tys cs $ \ty c ->
      case c of
        C.Var{}   -> return c
        C.Const{} -> return c
        _         -> bind ty c
  
  
  -- Return the local environment code, consisting of a list of initialisation
  -- declarations and statements. During construction, these are introduced to the
  -- front of the list, so reverse to get in execution order.
  --
  getEnv :: Gen [C.BlockItem]
  getEnv = reverse `fmap` gets localBindings
  
  -- Generate a fresh variable name
  --
  fresh :: CUDA String
  fresh = state $ \s -> let n = counter s in ('v':show n, s { counter = n+1 })
  
  -- Add an expression of given type to the environment and return the (new,
  -- unique) binding name that can be used in place of the thing just bound.
  --
  bind :: C.Type -> C.Exp -> Gen C.Exp
  bind t e =
    case e of
      C.Var{}     -> return e
      _           -> do
        name <- lift fresh
        modify (\st -> st { localBindings = [citem| const $ty:t $id:name = $exp:e;|] : localBindings st })
        return [cexp| $id:name |]
  
  
  -- Add an expression to the set marking that it will be used to generate the
  -- output value(s). If the term exists in the reverse let-map, add that binding
  -- instead.
  --
  use :: C.Exp -> Gen C.Exp
  use e@(C.Var{}) = do
    modify (\st -> st { usedTerms = Set.insert e (usedTerms st) })
    return e
  --
  use e           = return e
  
  