{-# LANGUAGE CPP,
             BangPatterns,
             DataKinds,
             FlexibleContexts,
             FlexibleInstances,
             GADTs,
             KindSignatures,
             PolyKinds,
             StandaloneDeriving,
             TypeOperators,
             RankNTypes        #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}

----------------------------------------------------------------
--                                                    2016.07.01
-- |
-- Module      :  Language.Hakaru.CodeGen.CodeGenMonad
-- Copyright   :  Copyright (c) 2016 the Hakaru team
-- License     :  BSD3
-- Maintainer  :  zsulliva@indiana.edu
-- Stability   :  experimental
-- Portability :  GHC-only
--
--   This module provides a monad for C code generation as well
-- as some useful helper functions for manipulating it
----------------------------------------------------------------


module Language.Hakaru.CodeGen.CodeGenMonad
  ( CodeGen
  , CG(..)
  , runCodeGen
  , runCodeGenBlock
  , runCodeGenWith
  , emptyCG

  -- codegen effects
  , declare
  , declare'
  , assign
  , putStat
  , putExprStat
  , extDeclare
  , extDeclareTypes

  , funCG
  , whenPar
  , parDo
  , seqDo

  , reserveIdent
  , genIdent
  , genIdent'

  -- Hakaru specific
  , createIdent
  , createIdent'
  , lookupIdent

  -- control mechanisms
  , ifCG
  , whileCG
  , doWhileCG
  , forCG
  , reductionCG
  , codeBlockCG

  -- memory
  , putMallocStat
  ) where

import Control.Monad.State.Strict

#if __GLASGOW_HASKELL__ < 710
import Data.Monoid (Monoid(..))
import Control.Applicative ((<$>))
#endif

import Language.Hakaru.Syntax.ABT hiding (var)
import Language.Hakaru.Types.DataKind
import Language.Hakaru.Types.Sing
import Language.Hakaru.CodeGen.Types
import Language.Hakaru.CodeGen.AST
import Language.Hakaru.CodeGen.Libs

import Data.Number.Nat (fromNat)
import qualified Data.IntMap.Strict as IM
import qualified Data.Text          as T
import qualified Data.Set           as S

-- CG after "codegen", holds the state of a codegen computation
data CG = CG
  { CG -> [String]
freshNames    :: [String]     -- ^ fresh names for code generations
  , CG -> Set String
reservedNames :: S.Set String -- ^ reserve names during code generations
  , CG -> [CExtDecl]
extDecls      :: [CExtDecl]   -- ^ total external declarations
  , CG -> [CDecl]
declarations  :: [CDecl]      -- ^ declarations in local block
  , CG -> [CStat]
statements    :: [CStat]      -- ^ statements can include assignments as well as other side-effects
  , CG -> Env
varEnv        :: Env          -- ^ mapping between Hakaru vars and codegeneration vars
  , CG -> Bool
managedMem    :: Bool         -- ^ garbage collected block
  , CG -> Bool
sharedMem     :: Bool         -- ^ shared memory supported block (OpenMP)
  , CG -> Bool
simd          :: Bool         -- ^ support single instruction multiple data instructions  (OpenMP)
  , CG -> Bool
distributed   :: Bool         -- ^ distributed supported block
  , CG -> Bool
logProbs      :: Bool         -- ^ true by default, but might not matter in some situations
  }

emptyCG :: CG
emptyCG :: CG
emptyCG = [String]
-> Set String
-> [CExtDecl]
-> [CDecl]
-> [CStat]
-> Env
-> Bool
-> Bool
-> Bool
-> Bool
-> Bool
-> CG
CG [String]
cNameStream Set String
forall a. Monoid a => a
mempty [CExtDecl]
forall a. Monoid a => a
mempty [] [] Env
emptyEnv Bool
False Bool
False Bool
False Bool
False Bool
True

type CodeGen = State CG

runCodeGen :: CodeGen a -> ([CExtDecl],[CDecl], [CStat])
runCodeGen :: CodeGen a -> ([CExtDecl], [CDecl], [CStat])
runCodeGen CodeGen a
m =
  let (a
_, CG
cg) = CodeGen a -> CG -> (a, CG)
forall s a. State s a -> s -> (a, s)
runState CodeGen a
m CG
emptyCG
  in  ( [CExtDecl] -> [CExtDecl]
forall a. [a] -> [a]
reverse ([CExtDecl] -> [CExtDecl]) -> [CExtDecl] -> [CExtDecl]
forall a b. (a -> b) -> a -> b
$ CG -> [CExtDecl]
extDecls     CG
cg
      , [CDecl] -> [CDecl]
forall a. [a] -> [a]
reverse ([CDecl] -> [CDecl]) -> [CDecl] -> [CDecl]
forall a b. (a -> b) -> a -> b
$ CG -> [CDecl]
declarations CG
cg
      , [CStat] -> [CStat]
forall a. [a] -> [a]
reverse ([CStat] -> [CStat]) -> [CStat] -> [CStat]
forall a b. (a -> b) -> a -> b
$ CG -> [CStat]
statements   CG
cg )


runCodeGenBlock :: CodeGen a -> CodeGen CStat
runCodeGenBlock :: CodeGen a -> CodeGen CStat
runCodeGenBlock CodeGen a
m =
  do CG
cg <- StateT CG Identity CG
forall s (m :: * -> *). MonadState s m => m s
get
     let (a
_,CG
cg') = CodeGen a -> CG -> (a, CG)
forall s a. State s a -> s -> (a, s)
runState CodeGen a
m (CG -> (a, CG)) -> CG -> (a, CG)
forall a b. (a -> b) -> a -> b
$ CG
cg { statements :: [CStat]
statements = []
                                   , declarations :: [CDecl]
declarations = [] }
     CG -> StateT CG Identity ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put (CG -> StateT CG Identity ()) -> CG -> StateT CG Identity ()
forall a b. (a -> b) -> a -> b
$ CG
cg' { statements :: [CStat]
statements   = CG -> [CStat]
statements CG
cg
               , declarations :: [CDecl]
declarations = CG -> [CDecl]
declarations CG
cg' [CDecl] -> [CDecl] -> [CDecl]
forall a. [a] -> [a] -> [a]
++ CG -> [CDecl]
declarations CG
cg
               }
     CStat -> CodeGen CStat
forall (m :: * -> *) a. Monad m => a -> m a
return (CStat -> CodeGen CStat) -> (CG -> CStat) -> CG -> CodeGen CStat
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [CCompoundBlockItem] -> CStat
CCompound ([CCompoundBlockItem] -> CStat)
-> (CG -> [CCompoundBlockItem]) -> CG -> CStat
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (CStat -> CCompoundBlockItem) -> [CStat] -> [CCompoundBlockItem]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap CStat -> CCompoundBlockItem
CBlockStat ([CStat] -> [CCompoundBlockItem])
-> (CG -> [CStat]) -> CG -> [CCompoundBlockItem]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [CStat] -> [CStat]
forall a. [a] -> [a]
reverse ([CStat] -> [CStat]) -> (CG -> [CStat]) -> CG -> [CStat]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CG -> [CStat]
statements (CG -> CodeGen CStat) -> CG -> CodeGen CStat
forall a b. (a -> b) -> a -> b
$ CG
cg'

runCodeGenWith :: CodeGen a -> CG -> [CExtDecl]
runCodeGenWith :: CodeGen a -> CG -> [CExtDecl]
runCodeGenWith CodeGen a
cg CG
start = let (a
_,CG
cg') = CodeGen a -> CG -> (a, CG)
forall s a. State s a -> s -> (a, s)
runState CodeGen a
cg CG
start in [CExtDecl] -> [CExtDecl]
forall a. [a] -> [a]
reverse ([CExtDecl] -> [CExtDecl]) -> [CExtDecl] -> [CExtDecl]
forall a b. (a -> b) -> a -> b
$ CG -> [CExtDecl]
extDecls CG
cg'

--------------------------------------------------------------------------------

whenPar :: CodeGen () -> CodeGen ()
whenPar :: StateT CG Identity () -> StateT CG Identity ()
whenPar StateT CG Identity ()
m = (CG -> Bool
sharedMem (CG -> Bool) -> StateT CG Identity CG -> StateT CG Identity Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> StateT CG Identity CG
forall s (m :: * -> *). MonadState s m => m s
get) StateT CG Identity Bool
-> (Bool -> StateT CG Identity ()) -> StateT CG Identity ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (\Bool
b -> Bool -> StateT CG Identity () -> StateT CG Identity ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
b StateT CG Identity ()
m)

parDo :: CodeGen a -> CodeGen a
parDo :: CodeGen a -> CodeGen a
parDo CodeGen a
m = do
  CG
cg <- StateT CG Identity CG
forall s (m :: * -> *). MonadState s m => m s
get
  CG -> StateT CG Identity ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put (CG
cg { sharedMem :: Bool
sharedMem = Bool
True } )
  a
a <- CodeGen a
m
  CG
cg' <- StateT CG Identity CG
forall s (m :: * -> *). MonadState s m => m s
get
  CG -> StateT CG Identity ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put (CG
cg' { sharedMem :: Bool
sharedMem = CG -> Bool
sharedMem CG
cg } )
  a -> CodeGen a
forall (m :: * -> *) a. Monad m => a -> m a
return a
a

seqDo :: CodeGen a -> CodeGen a
seqDo :: CodeGen a -> CodeGen a
seqDo CodeGen a
m = do
  CG
cg <- StateT CG Identity CG
forall s (m :: * -> *). MonadState s m => m s
get
  CG -> StateT CG Identity ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put (CG
cg { sharedMem :: Bool
sharedMem = Bool
False } )
  a
a <- CodeGen a
m
  CG
cg' <- StateT CG Identity CG
forall s (m :: * -> *). MonadState s m => m s
get
  CG -> StateT CG Identity ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put (CG
cg' { sharedMem :: Bool
sharedMem = CG -> Bool
sharedMem CG
cg } )
  a -> CodeGen a
forall (m :: * -> *) a. Monad m => a -> m a
return a
a

--------------------------------------------------------------------------------

reserveIdent :: String -> CodeGen Ident
reserveIdent :: String -> CodeGen Ident
reserveIdent String
s = do
  StateT CG Identity CG
forall s (m :: * -> *). MonadState s m => m s
get StateT CG Identity CG
-> (CG -> StateT CG Identity ()) -> StateT CG Identity ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \CG
cg -> CG -> StateT CG Identity ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put (CG -> StateT CG Identity ()) -> CG -> StateT CG Identity ()
forall a b. (a -> b) -> a -> b
$ CG
cg { reservedNames :: Set String
reservedNames = String
s String -> Set String -> Set String
forall a. Ord a => a -> Set a -> Set a
`S.insert` CG -> Set String
reservedNames CG
cg }
  Ident -> CodeGen Ident
forall (m :: * -> *) a. Monad m => a -> m a
return (String -> Ident
Ident String
s)


genIdent :: CodeGen Ident
genIdent :: CodeGen Ident
genIdent = String -> CodeGen Ident
genIdent' String
""

genIdent' :: String -> CodeGen Ident
genIdent' :: String -> CodeGen Ident
genIdent' String
s =
  do CG
cg <- StateT CG Identity CG
forall s (m :: * -> *). MonadState s m => m s
get
     let ([String]
freshNs,String
name) = [String] -> Set String -> ([String], String)
pullName (CG -> [String]
freshNames CG
cg) (CG -> Set String
reservedNames CG
cg)
     CG -> StateT CG Identity ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put (CG -> StateT CG Identity ()) -> CG -> StateT CG Identity ()
forall a b. (a -> b) -> a -> b
$ CG
cg { freshNames :: [String]
freshNames = [String]
freshNs }
     Ident -> CodeGen Ident
forall (m :: * -> *) a. Monad m => a -> m a
return (Ident -> CodeGen Ident) -> Ident -> CodeGen Ident
forall a b. (a -> b) -> a -> b
$ String -> Ident
Ident String
name
  where pullName :: [String] -> S.Set String -> ([String],String)
        pullName :: [String] -> Set String -> ([String], String)
pullName (String
n:[String]
names) Set String
reserved =
          let name :: String
name = String
s String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"_" String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
n in
          if String -> Set String -> Bool
forall a. Ord a => a -> Set a -> Bool
S.member String
name Set String
reserved
          then let ([String]
names',String
out) = [String] -> Set String -> ([String], String)
pullName [String]
names Set String
reserved
               in  (String
nString -> [String] -> [String]
forall a. a -> [a] -> [a]
:[String]
names',String
out)
          else ([String]
names,String
name)
        pullName [String]
_ Set String
_ = String -> ([String], String)
forall a. HasCallStack => String -> a
error String
"should not happen, names is infinite"



createIdent :: Variable (a :: Hakaru) -> CodeGen Ident
createIdent :: Variable a -> CodeGen Ident
createIdent = String -> Variable a -> CodeGen Ident
forall (a :: Hakaru). String -> Variable a -> CodeGen Ident
createIdent' String
""

createIdent' :: String -> Variable (a :: Hakaru) -> CodeGen Ident
createIdent' :: String -> Variable a -> CodeGen Ident
createIdent' String
s var :: Variable a
var@(Variable Text
name Nat
_ Sing a
_) =
  do !CG
cg <- StateT CG Identity CG
forall s (m :: * -> *). MonadState s m => m s
get
     let ident :: Ident
ident = String -> Ident
Ident (String -> Ident) -> String -> Ident
forall a b. (a -> b) -> a -> b
$ [String] -> String
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [(Char -> String) -> String -> String
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Char -> String
toAscii (String -> String) -> (Text -> String) -> Text -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> String
T.unpack (Text -> String) -> Text -> String
forall a b. (a -> b) -> a -> b
$ Text
name
                                ,String
"_",String
s,String
"_",[String] -> String
forall a. [a] -> a
head ([String] -> String) -> [String] -> String
forall a b. (a -> b) -> a -> b
$ CG -> [String]
freshNames CG
cg ]
         env' :: Env
env'  = Variable a -> Ident -> Env -> Env
forall (a :: Hakaru). Variable a -> Ident -> Env -> Env
updateEnv Variable a
var Ident
ident (CG -> Env
varEnv CG
cg)
     CG -> StateT CG Identity ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put (CG -> StateT CG Identity ()) -> CG -> StateT CG Identity ()
forall a b. (a -> b) -> a -> b
$! CG
cg { freshNames :: [String]
freshNames = [String] -> [String]
forall a. [a] -> [a]
tail ([String] -> [String]) -> [String] -> [String]
forall a b. (a -> b) -> a -> b
$ CG -> [String]
freshNames CG
cg
               , varEnv :: Env
varEnv     = Env
env' }
     Ident -> CodeGen Ident
forall (m :: * -> *) a. Monad m => a -> m a
return Ident
ident
  where toAscii :: Char -> String
toAscii Char
c = let num :: Int
num = Char -> Int
forall a. Enum a => a -> Int
fromEnum Char
c in
                    if Int
num Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
48 Bool -> Bool -> Bool
|| Int
num Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
122
                    then String
"u" String -> String -> String
forall a. [a] -> [a] -> [a]
++ (Int -> String
forall a. Show a => a -> String
show Int
num)
                    else [Char
c]

lookupIdent :: Variable (a :: Hakaru) -> CodeGen Ident
lookupIdent :: Variable a -> CodeGen Ident
lookupIdent Variable a
var =
  do !CG
cg <- StateT CG Identity CG
forall s (m :: * -> *). MonadState s m => m s
get
     let !env :: Env
env = CG -> Env
varEnv CG
cg
     case Variable a -> Env -> Maybe Ident
forall (a :: Hakaru). Variable a -> Env -> Maybe Ident
lookupVar Variable a
var Env
env of
       Maybe Ident
Nothing -> String -> CodeGen Ident
forall a. HasCallStack => String -> a
error (String -> CodeGen Ident) -> String -> CodeGen Ident
forall a b. (a -> b) -> a -> b
$ String
"lookupIdent: var not found --" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Variable a -> String
forall a. Show a => a -> String
show Variable a
var
       Just Ident
i  -> Ident -> CodeGen Ident
forall (m :: * -> *) a. Monad m => a -> m a
return Ident
i

-- | types like SData and SMeasure are impure in that they will produce extra
--   code in the CodeGenMonad while literal types SReal, SInt, SNat, and SProb
--   do not
declare :: Sing (a :: Hakaru) -> Ident -> CodeGen ()
declare :: Sing a -> Ident -> StateT CG Identity ()
declare Sing a
SInt  = CDecl -> StateT CG Identity ()
declare' (CDecl -> StateT CG Identity ())
-> (Ident -> CDecl) -> Ident -> StateT CG Identity ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Sing 'HInt -> Ident -> CDecl
forall (a :: Hakaru). Sing a -> Ident -> CDecl
typeDeclaration Sing 'HInt
SInt
declare Sing a
SNat  = CDecl -> StateT CG Identity ()
declare' (CDecl -> StateT CG Identity ())
-> (Ident -> CDecl) -> Ident -> StateT CG Identity ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Sing 'HNat -> Ident -> CDecl
forall (a :: Hakaru). Sing a -> Ident -> CDecl
typeDeclaration Sing 'HNat
SNat
declare Sing a
SProb = CDecl -> StateT CG Identity ()
declare' (CDecl -> StateT CG Identity ())
-> (Ident -> CDecl) -> Ident -> StateT CG Identity ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Sing 'HProb -> Ident -> CDecl
forall (a :: Hakaru). Sing a -> Ident -> CDecl
typeDeclaration Sing 'HProb
SProb
declare Sing a
SReal = CDecl -> StateT CG Identity ()
declare' (CDecl -> StateT CG Identity ())
-> (Ident -> CDecl) -> Ident -> StateT CG Identity ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Sing 'HReal -> Ident -> CDecl
forall (a :: Hakaru). Sing a -> Ident -> CDecl
typeDeclaration Sing 'HReal
SReal
declare m :: Sing a
m@(SMeasure t) = \Ident
i ->
  Sing a -> StateT CG Identity ()
forall (a :: Hakaru). Sing a -> StateT CG Identity ()
extDeclareTypes Sing a
m StateT CG Identity ()
-> StateT CG Identity () -> StateT CG Identity ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> (CDecl -> StateT CG Identity ()
declare' (CDecl -> StateT CG Identity ()) -> CDecl -> StateT CG Identity ()
forall a b. (a -> b) -> a -> b
$ Sing a -> Ident -> CDecl
forall (a :: Hakaru). Sing a -> Ident -> CDecl
mdataDeclaration Sing a
t Ident
i)

declare a :: Sing a
a@(SArray t) = \Ident
i ->
  Sing a -> StateT CG Identity ()
forall (a :: Hakaru). Sing a -> StateT CG Identity ()
extDeclareTypes Sing a
a StateT CG Identity ()
-> StateT CG Identity () -> StateT CG Identity ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> (CDecl -> StateT CG Identity ()
declare' (CDecl -> StateT CG Identity ()) -> CDecl -> StateT CG Identity ()
forall a b. (a -> b) -> a -> b
$ Sing a -> Ident -> CDecl
forall (a :: Hakaru). Sing a -> Ident -> CDecl
arrayDeclaration Sing a
t Ident
i)

declare d :: Sing a
d@(SData _ _)  = \Ident
i ->
  Sing a -> StateT CG Identity ()
forall (a :: Hakaru). Sing a -> StateT CG Identity ()
extDeclareTypes Sing a
d StateT CG Identity ()
-> StateT CG Identity () -> StateT CG Identity ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> (CDecl -> StateT CG Identity ()
declare' (CDecl -> StateT CG Identity ()) -> CDecl -> StateT CG Identity ()
forall a b. (a -> b) -> a -> b
$ Sing (HData' t) -> Ident -> CDecl
forall (t :: HakaruCon). Sing (HData' t) -> Ident -> CDecl
datumDeclaration Sing a
Sing (HData' t)
d Ident
i)

declare f :: Sing a
f@(SFun _ _) = \Ident
_ ->
  Sing a -> StateT CG Identity ()
forall (a :: Hakaru). Sing a -> StateT CG Identity ()
extDeclareTypes Sing a
f StateT CG Identity ()
-> StateT CG Identity () -> StateT CG Identity ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> () -> StateT CG Identity ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
  -- this currently avoids declaration if the type is a lambda, this is hacky

-- | for types that contain subtypes we need to recursively traverse them and
--   build up a list of external type declarations.
--   For example: Measure (Array Nat) will need to have structures for arrays
--   declared before the top level type
extDeclareTypes :: Sing (a :: Hakaru) -> CodeGen ()
extDeclareTypes :: Sing a -> StateT CG Identity ()
extDeclareTypes Sing a
SInt          = () -> StateT CG Identity ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
extDeclareTypes Sing a
SNat          = () -> StateT CG Identity ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
extDeclareTypes Sing a
SReal         = () -> StateT CG Identity ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
extDeclareTypes Sing a
SProb         = () -> StateT CG Identity ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
extDeclareTypes (SMeasure i)  = Sing a -> StateT CG Identity ()
forall (a :: Hakaru). Sing a -> StateT CG Identity ()
extDeclareTypes Sing a
i StateT CG Identity ()
-> StateT CG Identity () -> StateT CG Identity ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> CExtDecl -> StateT CG Identity ()
extDeclare (Sing a -> CExtDecl
forall (a :: Hakaru). Sing a -> CExtDecl
mdataStruct Sing a
i)
extDeclareTypes (SArray i)    = Sing a -> StateT CG Identity ()
forall (a :: Hakaru). Sing a -> StateT CG Identity ()
extDeclareTypes Sing a
i StateT CG Identity ()
-> StateT CG Identity () -> StateT CG Identity ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> CExtDecl -> StateT CG Identity ()
extDeclare (Sing a -> CExtDecl
forall (a :: Hakaru). Sing a -> CExtDecl
arrayStruct Sing a
i)
extDeclareTypes (SFun x y)    = Sing a -> StateT CG Identity ()
forall (a :: Hakaru). Sing a -> StateT CG Identity ()
extDeclareTypes Sing a
x StateT CG Identity ()
-> StateT CG Identity () -> StateT CG Identity ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Sing b -> StateT CG Identity ()
forall (a :: Hakaru). Sing a -> StateT CG Identity ()
extDeclareTypes Sing b
y
extDeclareTypes d :: Sing a
d@(SData _ i) = Sing (Code t) -> StateT CG Identity ()
forall (a :: [[HakaruFun]]). Sing a -> StateT CG Identity ()
extDeclDatum Sing (Code t)
i    StateT CG Identity ()
-> StateT CG Identity () -> StateT CG Identity ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> CExtDecl -> StateT CG Identity ()
extDeclare (Sing (HData' t) -> CExtDecl
forall (t :: HakaruCon). Sing (HData' t) -> CExtDecl
datumStruct Sing a
Sing (HData' t)
d)
  where extDeclDatum :: Sing (a :: [[HakaruFun]]) -> CodeGen ()
        extDeclDatum :: Sing a -> StateT CG Identity ()
extDeclDatum Sing a
SVoid       = () -> StateT CG Identity ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
        extDeclDatum (SPlus p s) = Sing xss -> StateT CG Identity ()
forall (a :: [[HakaruFun]]). Sing a -> StateT CG Identity ()
extDeclDatum Sing xss
s StateT CG Identity ()
-> StateT CG Identity () -> StateT CG Identity ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Sing xs -> StateT CG Identity ()
forall (a :: [HakaruFun]). Sing a -> StateT CG Identity ()
datumProdTypes Sing xs
p

        datumProdTypes :: Sing (a :: [HakaruFun]) -> CodeGen ()
        datumProdTypes :: Sing a -> StateT CG Identity ()
datumProdTypes Sing a
SDone     = () -> StateT CG Identity ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
        datumProdTypes (SEt x p) = Sing xs -> StateT CG Identity ()
forall (a :: [HakaruFun]). Sing a -> StateT CG Identity ()
datumProdTypes Sing xs
p StateT CG Identity ()
-> StateT CG Identity () -> StateT CG Identity ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Sing x -> StateT CG Identity ()
forall (a :: HakaruFun). Sing a -> StateT CG Identity ()
datumPrimTypes Sing x
x

        datumPrimTypes :: Sing (a :: HakaruFun) -> CodeGen ()
        datumPrimTypes :: Sing a -> StateT CG Identity ()
datumPrimTypes Sing a
SIdent     = () -> StateT CG Identity ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
        datumPrimTypes (SKonst s) = Sing a -> StateT CG Identity ()
forall (a :: Hakaru). Sing a -> StateT CG Identity ()
extDeclareTypes Sing a
s

declare' :: CDecl -> CodeGen ()
declare' :: CDecl -> StateT CG Identity ()
declare' CDecl
d = do CG
cg <- StateT CG Identity CG
forall s (m :: * -> *). MonadState s m => m s
get
                CG -> StateT CG Identity ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put (CG -> StateT CG Identity ()) -> CG -> StateT CG Identity ()
forall a b. (a -> b) -> a -> b
$ CG
cg { declarations :: [CDecl]
declarations = CDecl
dCDecl -> [CDecl] -> [CDecl]
forall a. a -> [a] -> [a]
:(CG -> [CDecl]
declarations CG
cg) }

putStat :: CStat -> CodeGen ()
putStat :: CStat -> StateT CG Identity ()
putStat CStat
s = do CG
cg <- StateT CG Identity CG
forall s (m :: * -> *). MonadState s m => m s
get
               CG -> StateT CG Identity ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put (CG -> StateT CG Identity ()) -> CG -> StateT CG Identity ()
forall a b. (a -> b) -> a -> b
$ CG
cg { statements :: [CStat]
statements = CStat
sCStat -> [CStat] -> [CStat]
forall a. a -> [a] -> [a]
:(CG -> [CStat]
statements CG
cg) }

putExprStat :: CExpr -> CodeGen ()
putExprStat :: CExpr -> StateT CG Identity ()
putExprStat = CStat -> StateT CG Identity ()
putStat (CStat -> StateT CG Identity ())
-> (CExpr -> CStat) -> CExpr -> StateT CG Identity ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Maybe CExpr -> CStat
CExpr (Maybe CExpr -> CStat) -> (CExpr -> Maybe CExpr) -> CExpr -> CStat
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CExpr -> Maybe CExpr
forall a. a -> Maybe a
Just

assign :: Ident -> CExpr -> CodeGen ()
assign :: Ident -> CExpr -> StateT CG Identity ()
assign Ident
ident CExpr
e = CStat -> StateT CG Identity ()
putStat (CStat -> StateT CG Identity ())
-> (CExpr -> CStat) -> CExpr -> StateT CG Identity ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Maybe CExpr -> CStat
CExpr (Maybe CExpr -> CStat) -> (CExpr -> Maybe CExpr) -> CExpr -> CStat
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CExpr -> Maybe CExpr
forall a. a -> Maybe a
Just (CExpr -> StateT CG Identity ()) -> CExpr -> StateT CG Identity ()
forall a b. (a -> b) -> a -> b
$ (Ident -> CExpr
CVar Ident
ident CExpr -> CExpr -> CExpr
.=. CExpr
e)


extDeclare :: CExtDecl -> CodeGen ()
extDeclare :: CExtDecl -> StateT CG Identity ()
extDeclare CExtDecl
d = do CG
cg <- StateT CG Identity CG
forall s (m :: * -> *). MonadState s m => m s
get
                  let extds :: [CExtDecl]
extds = CG -> [CExtDecl]
extDecls CG
cg
                      extds' :: [CExtDecl]
extds' = if CExtDecl -> [CExtDecl] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
elem CExtDecl
d [CExtDecl]
extds
                               then [CExtDecl]
extds
                               else CExtDecl
dCExtDecl -> [CExtDecl] -> [CExtDecl]
forall a. a -> [a] -> [a]
:[CExtDecl]
extds
                  CG -> StateT CG Identity ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put (CG -> StateT CG Identity ()) -> CG -> StateT CG Identity ()
forall a b. (a -> b) -> a -> b
$ CG
cg { extDecls :: [CExtDecl]
extDecls = [CExtDecl]
extds' }

---------
-- ENV --
---------

newtype Env = Env (IM.IntMap Ident)
  deriving Int -> Env -> String -> String
[Env] -> String -> String
Env -> String
(Int -> Env -> String -> String)
-> (Env -> String) -> ([Env] -> String -> String) -> Show Env
forall a.
(Int -> a -> String -> String)
-> (a -> String) -> ([a] -> String -> String) -> Show a
showList :: [Env] -> String -> String
$cshowList :: [Env] -> String -> String
show :: Env -> String
$cshow :: Env -> String
showsPrec :: Int -> Env -> String -> String
$cshowsPrec :: Int -> Env -> String -> String
Show

emptyEnv :: Env
emptyEnv :: Env
emptyEnv = IntMap Ident -> Env
Env IntMap Ident
forall a. IntMap a
IM.empty

updateEnv :: Variable (a :: Hakaru) -> Ident -> Env -> Env
updateEnv :: Variable a -> Ident -> Env -> Env
updateEnv (Variable Text
_ Nat
nat Sing a
_) Ident
ident (Env IntMap Ident
env) =
  IntMap Ident -> Env
Env (IntMap Ident -> Env) -> IntMap Ident -> Env
forall a b. (a -> b) -> a -> b
$! Int -> Ident -> IntMap Ident -> IntMap Ident
forall a. Int -> a -> IntMap a -> IntMap a
IM.insert (Nat -> Int
fromNat Nat
nat) Ident
ident IntMap Ident
env

lookupVar :: Variable (a :: Hakaru) -> Env -> Maybe Ident
lookupVar :: Variable a -> Env -> Maybe Ident
lookupVar (Variable Text
_ Nat
nat Sing a
_) (Env IntMap Ident
env) =
  Int -> IntMap Ident -> Maybe Ident
forall a. Int -> IntMap a -> Maybe a
IM.lookup (Nat -> Int
fromNat Nat
nat) IntMap Ident
env

--------------------------------------------------------------------------------
--                      Control Flow and Code Blocks                          --
--------------------------------------------------------------------------------
{-
Monadic operations funCG, ifCG, whileCG, forCG, reductionCG, and codeBlockCG all
generate compound C statements (several declarations and statements surrounded
by '{' '}'). It is important that these code blocks float external functions and
imports to the top of the generated C file AND keep a set of the variable
declarations local to the block of code.
-}

funCG :: [CTypeSpec] -> Ident -> [CDecl] -> CodeGen () -> CodeGen ()
funCG :: [CTypeSpec]
-> Ident
-> [CDecl]
-> StateT CG Identity ()
-> StateT CG Identity ()
funCG [CTypeSpec]
ts Ident
ident [CDecl]
args StateT CG Identity ()
m =
  do CG
cg <- StateT CG Identity CG
forall s (m :: * -> *). MonadState s m => m s
get
     let (()
_,CG
cg') = StateT CG Identity () -> CG -> ((), CG)
forall s a. State s a -> s -> (a, s)
runState StateT CG Identity ()
m (CG -> ((), CG)) -> CG -> ((), CG)
forall a b. (a -> b) -> a -> b
$ CG
cg { statements :: [CStat]
statements   = []
                                   , declarations :: [CDecl]
declarations = []
                                   , freshNames :: [String]
freshNames   = [String]
cNameStream }
     let decls :: [CDecl]
decls = [CDecl] -> [CDecl]
forall a. [a] -> [a]
reverse ([CDecl] -> [CDecl]) -> (CG -> [CDecl]) -> CG -> [CDecl]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CG -> [CDecl]
declarations (CG -> [CDecl]) -> CG -> [CDecl]
forall a b. (a -> b) -> a -> b
$ CG
cg'
         stmts :: [CStat]
stmts = [CStat] -> [CStat]
forall a. [a] -> [a]
reverse ([CStat] -> [CStat]) -> (CG -> [CStat]) -> CG -> [CStat]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CG -> [CStat]
statements   (CG -> [CStat]) -> CG -> [CStat]
forall a b. (a -> b) -> a -> b
$ CG
cg'
     -- reset local statements and declarations
     CG -> StateT CG Identity ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put (CG -> StateT CG Identity ()) -> CG -> StateT CG Identity ()
forall a b. (a -> b) -> a -> b
$ CG
cg' { statements :: [CStat]
statements   = CG -> [CStat]
statements CG
cg
               , declarations :: [CDecl]
declarations = CG -> [CDecl]
declarations CG
cg
               , freshNames :: [String]
freshNames   = CG -> [String]
freshNames CG
cg }
     CExtDecl -> StateT CG Identity ()
extDeclare (CExtDecl -> StateT CG Identity ())
-> (CFunDef -> CExtDecl) -> CFunDef -> StateT CG Identity ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CFunDef -> CExtDecl
CFunDefExt (CFunDef -> StateT CG Identity ())
-> CFunDef -> StateT CG Identity ()
forall a b. (a -> b) -> a -> b
$
       [CDeclSpec] -> CDeclr -> [CDecl] -> CStat -> CFunDef
CFunDef ((CTypeSpec -> CDeclSpec) -> [CTypeSpec] -> [CDeclSpec]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap CTypeSpec -> CDeclSpec
CTypeSpec [CTypeSpec]
ts)
               (Maybe CPtrDeclr -> CDirectDeclr -> CDeclr
CDeclr Maybe CPtrDeclr
forall a. Maybe a
Nothing (Ident -> CDirectDeclr
CDDeclrIdent Ident
ident))
               [CDecl]
args
               ([CCompoundBlockItem] -> CStat
CCompound (((CDecl -> CCompoundBlockItem) -> [CDecl] -> [CCompoundBlockItem]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap CDecl -> CCompoundBlockItem
CBlockDecl [CDecl]
decls) [CCompoundBlockItem]
-> [CCompoundBlockItem] -> [CCompoundBlockItem]
forall a. [a] -> [a] -> [a]
++ ((CStat -> CCompoundBlockItem) -> [CStat] -> [CCompoundBlockItem]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap CStat -> CCompoundBlockItem
CBlockStat [CStat]
stmts)))

ifCG :: CExpr -> CodeGen () -> CodeGen () -> CodeGen ()
ifCG :: CExpr
-> StateT CG Identity ()
-> StateT CG Identity ()
-> StateT CG Identity ()
ifCG CExpr
bE StateT CG Identity ()
m1 StateT CG Identity ()
m2 =
  do CG
cg <- StateT CG Identity CG
forall s (m :: * -> *). MonadState s m => m s
get
     let (()
_,CG
cg') = StateT CG Identity () -> CG -> ((), CG)
forall s a. State s a -> s -> (a, s)
runState StateT CG Identity ()
m1 (CG -> ((), CG)) -> CG -> ((), CG)
forall a b. (a -> b) -> a -> b
$ CG
cg { statements :: [CStat]
statements   = []
                                    , declarations :: [CDecl]
declarations = [] }
         (()
_,CG
cg'') = StateT CG Identity () -> CG -> ((), CG)
forall s a. State s a -> s -> (a, s)
runState StateT CG Identity ()
m2 (CG -> ((), CG)) -> CG -> ((), CG)
forall a b. (a -> b) -> a -> b
$ CG
cg' { statements :: [CStat]
statements   = []
                                      , declarations :: [CDecl]
declarations = [] }
         thnBlock :: [CCompoundBlockItem]
thnBlock =  ((CDecl -> CCompoundBlockItem) -> [CDecl] -> [CCompoundBlockItem]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap CDecl -> CCompoundBlockItem
CBlockDecl ([CDecl] -> [CDecl]
forall a. [a] -> [a]
reverse ([CDecl] -> [CDecl]) -> [CDecl] -> [CDecl]
forall a b. (a -> b) -> a -> b
$ CG -> [CDecl]
declarations CG
cg'))
                  [CCompoundBlockItem]
-> [CCompoundBlockItem] -> [CCompoundBlockItem]
forall a. [a] -> [a] -> [a]
++ ((CStat -> CCompoundBlockItem) -> [CStat] -> [CCompoundBlockItem]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap CStat -> CCompoundBlockItem
CBlockStat ([CStat] -> [CStat]
forall a. [a] -> [a]
reverse ([CStat] -> [CStat]) -> [CStat] -> [CStat]
forall a b. (a -> b) -> a -> b
$ CG -> [CStat]
statements CG
cg'))
         elsBlock :: [CCompoundBlockItem]
elsBlock =  ((CDecl -> CCompoundBlockItem) -> [CDecl] -> [CCompoundBlockItem]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap CDecl -> CCompoundBlockItem
CBlockDecl ([CDecl] -> [CDecl]
forall a. [a] -> [a]
reverse ([CDecl] -> [CDecl]) -> [CDecl] -> [CDecl]
forall a b. (a -> b) -> a -> b
$ CG -> [CDecl]
declarations CG
cg'')
                  [CCompoundBlockItem]
-> [CCompoundBlockItem] -> [CCompoundBlockItem]
forall a. [a] -> [a] -> [a]
++ ((CStat -> CCompoundBlockItem) -> [CStat] -> [CCompoundBlockItem]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap CStat -> CCompoundBlockItem
CBlockStat ([CStat] -> [CStat]
forall a. [a] -> [a]
reverse ([CStat] -> [CStat]) -> [CStat] -> [CStat]
forall a b. (a -> b) -> a -> b
$ CG -> [CStat]
statements CG
cg'')))
     CG -> StateT CG Identity ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put (CG -> StateT CG Identity ()) -> CG -> StateT CG Identity ()
forall a b. (a -> b) -> a -> b
$ CG
cg'' { statements :: [CStat]
statements = CG -> [CStat]
statements CG
cg
                , declarations :: [CDecl]
declarations = CG -> [CDecl]
declarations CG
cg }
     CStat -> StateT CG Identity ()
putStat (CStat -> StateT CG Identity ()) -> CStat -> StateT CG Identity ()
forall a b. (a -> b) -> a -> b
$ CExpr -> CStat -> Maybe CStat -> CStat
CIf CExpr
bE
                   ([CCompoundBlockItem] -> CStat
CCompound [CCompoundBlockItem]
thnBlock)
                   (case [CCompoundBlockItem]
elsBlock of
                      [] -> Maybe CStat
forall a. Maybe a
Nothing
                      [CCompoundBlockItem]
_  -> CStat -> Maybe CStat
forall a. a -> Maybe a
Just (CStat -> Maybe CStat)
-> ([CCompoundBlockItem] -> CStat)
-> [CCompoundBlockItem]
-> Maybe CStat
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [CCompoundBlockItem] -> CStat
CCompound ([CCompoundBlockItem] -> Maybe CStat)
-> [CCompoundBlockItem] -> Maybe CStat
forall a b. (a -> b) -> a -> b
$ [CCompoundBlockItem]
elsBlock)

whileCG' :: Bool -> CExpr -> CodeGen () -> CodeGen ()
whileCG' :: Bool -> CExpr -> StateT CG Identity () -> StateT CG Identity ()
whileCG' Bool
isDoWhile CExpr
bE StateT CG Identity ()
m =
  do CG
cg <- StateT CG Identity CG
forall s (m :: * -> *). MonadState s m => m s
get
     let (()
_,CG
cg') = StateT CG Identity () -> CG -> ((), CG)
forall s a. State s a -> s -> (a, s)
runState StateT CG Identity ()
m (CG -> ((), CG)) -> CG -> ((), CG)
forall a b. (a -> b) -> a -> b
$ CG
cg { statements :: [CStat]
statements = []
                                   , declarations :: [CDecl]
declarations = [] }
     CG -> StateT CG Identity ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put (CG -> StateT CG Identity ()) -> CG -> StateT CG Identity ()
forall a b. (a -> b) -> a -> b
$ CG
cg' { statements :: [CStat]
statements = CG -> [CStat]
statements CG
cg
               , declarations :: [CDecl]
declarations = CG -> [CDecl]
declarations CG
cg }
     CStat -> StateT CG Identity ()
putStat (CStat -> StateT CG Identity ()) -> CStat -> StateT CG Identity ()
forall a b. (a -> b) -> a -> b
$ CExpr -> CStat -> Bool -> CStat
CWhile CExpr
bE
                      ([CCompoundBlockItem] -> CStat
CCompound ([CCompoundBlockItem] -> CStat) -> [CCompoundBlockItem] -> CStat
forall a b. (a -> b) -> a -> b
$ ((CDecl -> CCompoundBlockItem) -> [CDecl] -> [CCompoundBlockItem]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap CDecl -> CCompoundBlockItem
CBlockDecl ([CDecl] -> [CDecl]
forall a. [a] -> [a]
reverse ([CDecl] -> [CDecl]) -> [CDecl] -> [CDecl]
forall a b. (a -> b) -> a -> b
$ CG -> [CDecl]
declarations CG
cg')
                                [CCompoundBlockItem]
-> [CCompoundBlockItem] -> [CCompoundBlockItem]
forall a. [a] -> [a] -> [a]
++ ((CStat -> CCompoundBlockItem) -> [CStat] -> [CCompoundBlockItem]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap CStat -> CCompoundBlockItem
CBlockStat ([CStat] -> [CStat]
forall a. [a] -> [a]
reverse ([CStat] -> [CStat]) -> [CStat] -> [CStat]
forall a b. (a -> b) -> a -> b
$ CG -> [CStat]
statements CG
cg'))))
                      Bool
isDoWhile
whileCG :: CExpr -> CodeGen () -> CodeGen ()
whileCG :: CExpr -> StateT CG Identity () -> StateT CG Identity ()
whileCG = Bool -> CExpr -> StateT CG Identity () -> StateT CG Identity ()
whileCG' Bool
False

doWhileCG :: CExpr -> CodeGen () -> CodeGen ()
doWhileCG :: CExpr -> StateT CG Identity () -> StateT CG Identity ()
doWhileCG = Bool -> CExpr -> StateT CG Identity () -> StateT CG Identity ()
whileCG' Bool
True

-- forCG and reductionCG both create C for loops, their difference lies in the
-- parallel code they generate
forCG
  :: CExpr
  -> CExpr
  -> CExpr
  -> CodeGen ()
  -> CodeGen ()
forCG :: CExpr
-> CExpr -> CExpr -> StateT CG Identity () -> StateT CG Identity ()
forCG CExpr
iter CExpr
cond CExpr
inc StateT CG Identity ()
body =
  do CG
cg <- StateT CG Identity CG
forall s (m :: * -> *). MonadState s m => m s
get
     let (()
_,CG
cg') = StateT CG Identity () -> CG -> ((), CG)
forall s a. State s a -> s -> (a, s)
runState StateT CG Identity ()
body (CG -> ((), CG)) -> CG -> ((), CG)
forall a b. (a -> b) -> a -> b
$ CG
cg { statements :: [CStat]
statements = []
                                      , declarations :: [CDecl]
declarations = []
                                      , sharedMem :: Bool
sharedMem = Bool
False }
     CG -> StateT CG Identity ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put (CG -> StateT CG Identity ()) -> CG -> StateT CG Identity ()
forall a b. (a -> b) -> a -> b
$ CG
cg' { statements :: [CStat]
statements   = CG -> [CStat]
statements CG
cg
               , declarations :: [CDecl]
declarations = CG -> [CDecl]
declarations CG
cg
               , sharedMem :: Bool
sharedMem    = CG -> Bool
sharedMem CG
cg } -- only use pragmas at the top level
     StateT CG Identity () -> StateT CG Identity ()
whenPar (StateT CG Identity () -> StateT CG Identity ())
-> (OMP -> StateT CG Identity ()) -> OMP -> StateT CG Identity ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CStat -> StateT CG Identity ()
putStat (CStat -> StateT CG Identity ())
-> (OMP -> CStat) -> OMP -> StateT CG Identity ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Preprocessor -> CStat
CPPStat (Preprocessor -> CStat) -> (OMP -> Preprocessor) -> OMP -> CStat
forall b c a. (b -> c) -> (a -> b) -> a -> c
. OMP -> Preprocessor
ompToPP (OMP -> StateT CG Identity ()) -> OMP -> StateT CG Identity ()
forall a b. (a -> b) -> a -> b
$ Directive -> OMP
OMP ([Directive] -> Directive
Parallel [Directive
For])
     CStat -> StateT CG Identity ()
putStat (CStat -> StateT CG Identity ()) -> CStat -> StateT CG Identity ()
forall a b. (a -> b) -> a -> b
$ Maybe CExpr -> Maybe CExpr -> Maybe CExpr -> CStat -> CStat
CFor (CExpr -> Maybe CExpr
forall a. a -> Maybe a
Just CExpr
iter)
                    (CExpr -> Maybe CExpr
forall a. a -> Maybe a
Just CExpr
cond)
                    (CExpr -> Maybe CExpr
forall a. a -> Maybe a
Just CExpr
inc)
                    ([CCompoundBlockItem] -> CStat
CCompound ([CCompoundBlockItem] -> CStat) -> [CCompoundBlockItem] -> CStat
forall a b. (a -> b) -> a -> b
$  ((CDecl -> CCompoundBlockItem) -> [CDecl] -> [CCompoundBlockItem]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap CDecl -> CCompoundBlockItem
CBlockDecl ([CDecl] -> [CDecl]
forall a. [a] -> [a]
reverse ([CDecl] -> [CDecl]) -> [CDecl] -> [CDecl]
forall a b. (a -> b) -> a -> b
$ CG -> [CDecl]
declarations CG
cg')
                               [CCompoundBlockItem]
-> [CCompoundBlockItem] -> [CCompoundBlockItem]
forall a. [a] -> [a] -> [a]
++ ((CStat -> CCompoundBlockItem) -> [CStat] -> [CCompoundBlockItem]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap CStat -> CCompoundBlockItem
CBlockStat ([CStat] -> [CStat]
forall a. [a] -> [a]
reverse ([CStat] -> [CStat]) -> [CStat] -> [CStat]
forall a b. (a -> b) -> a -> b
$ CG -> [CStat]
statements CG
cg'))))

{-
The operation for a reduction is either a builtin binary op (which is a built
in OpenMP reducer),

OR, it must be specified for a given Hakaru type. This will generate fuctions
for the monoidal operations mempty and mappend, use these to generate OpenMP
reduction declarations, and then outfit a for loop with the pragma calling the
reduction.
-}
reductionCG
  :: Either CBinaryOp
            ( Sing (a :: Hakaru)             -- type of reduction sections
            , CExpr -> CodeGen ()            -- monoidal unit
            , CExpr -> CExpr -> CodeGen () ) -- monoidal multiplication
  -> CExpr       -- accumulator var
  -> CExpr       -- iteration var
  -> CExpr       -- iteration condition
  -> CExpr       -- iteration increment
  -> CodeGen ()  -- body of the loop
  -> CodeGen ()
reductionCG :: Either
  CBinaryOp
  (Sing a, CExpr -> StateT CG Identity (),
   CExpr -> CExpr -> StateT CG Identity ())
-> CExpr
-> CExpr
-> CExpr
-> CExpr
-> StateT CG Identity ()
-> StateT CG Identity ()
reductionCG Either
  CBinaryOp
  (Sing a, CExpr -> StateT CG Identity (),
   CExpr -> CExpr -> StateT CG Identity ())
op CExpr
acc CExpr
iter CExpr
cond CExpr
inc StateT CG Identity ()
body =
  do CG
cg <- StateT CG Identity CG
forall s (m :: * -> *). MonadState s m => m s
get
     let (()
_,CG
cg') = StateT CG Identity () -> CG -> ((), CG)
forall s a. State s a -> s -> (a, s)
runState StateT CG Identity ()
body (CG -> ((), CG)) -> CG -> ((), CG)
forall a b. (a -> b) -> a -> b
$ CG
cg { statements :: [CStat]
statements   = []
                                      , declarations :: [CDecl]
declarations = []
                                      , sharedMem :: Bool
sharedMem    = Bool
False } -- only use pragmas at the top level
     CG -> StateT CG Identity ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put (CG -> StateT CG Identity ()) -> CG -> StateT CG Identity ()
forall a b. (a -> b) -> a -> b
$ CG
cg' { statements :: [CStat]
statements   = CG -> [CStat]
statements CG
cg
               , declarations :: [CDecl]
declarations = CG -> [CDecl]
declarations CG
cg
               , sharedMem :: Bool
sharedMem    = CG -> Bool
sharedMem CG
cg }
     StateT CG Identity () -> StateT CG Identity ()
whenPar (StateT CG Identity () -> StateT CG Identity ())
-> StateT CG Identity () -> StateT CG Identity ()
forall a b. (a -> b) -> a -> b
$
       case Either
  CBinaryOp
  (Sing a, CExpr -> StateT CG Identity (),
   CExpr -> CExpr -> StateT CG Identity ())
op of
         Left CBinaryOp
binop ->
           CStat -> StateT CG Identity ()
putStat (CStat -> StateT CG Identity ())
-> (OMP -> CStat) -> OMP -> StateT CG Identity ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Preprocessor -> CStat
CPPStat (Preprocessor -> CStat) -> (OMP -> Preprocessor) -> OMP -> CStat
forall b c a. (b -> c) -> (a -> b) -> a -> c
. OMP -> Preprocessor
ompToPP (OMP -> StateT CG Identity ()) -> OMP -> StateT CG Identity ()
forall a b. (a -> b) -> a -> b
$
             Directive -> OMP
OMP ([Directive] -> Directive
Parallel [Directive
For,Either CBinaryOp Ident -> [CExpr] -> Directive
Reduction (CBinaryOp -> Either CBinaryOp Ident
forall a b. a -> Either a b
Left CBinaryOp
binop) [CExpr
acc]])
         Right (Sing a
typ,CExpr -> StateT CG Identity ()
unit,CExpr -> CExpr -> StateT CG Identity ()
mul) ->
           do { Ident
redId <- Sing a
-> (CExpr -> StateT CG Identity ())
-> (CExpr -> CExpr -> StateT CG Identity ())
-> CodeGen Ident
forall (a :: Hakaru).
Sing a
-> (CExpr -> StateT CG Identity ())
-> (CExpr -> CExpr -> StateT CG Identity ())
-> CodeGen Ident
declareReductionCG Sing a
typ CExpr -> StateT CG Identity ()
unit CExpr -> CExpr -> StateT CG Identity ()
mul
              ; CStat -> StateT CG Identity ()
putStat (CStat -> StateT CG Identity ())
-> (OMP -> CStat) -> OMP -> StateT CG Identity ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Preprocessor -> CStat
CPPStat (Preprocessor -> CStat) -> (OMP -> Preprocessor) -> OMP -> CStat
forall b c a. (b -> c) -> (a -> b) -> a -> c
. OMP -> Preprocessor
ompToPP (OMP -> StateT CG Identity ()) -> OMP -> StateT CG Identity ()
forall a b. (a -> b) -> a -> b
$
                  Directive -> OMP
OMP ([Directive] -> Directive
Parallel [Directive
For,Either CBinaryOp Ident -> [CExpr] -> Directive
Reduction (Ident -> Either CBinaryOp Ident
forall a b. b -> Either a b
Right Ident
redId) [CExpr
acc]]) }
     CStat -> StateT CG Identity ()
putStat (CStat -> StateT CG Identity ()) -> CStat -> StateT CG Identity ()
forall a b. (a -> b) -> a -> b
$ Maybe CExpr -> Maybe CExpr -> Maybe CExpr -> CStat -> CStat
CFor (CExpr -> Maybe CExpr
forall a. a -> Maybe a
Just CExpr
iter)
                    (CExpr -> Maybe CExpr
forall a. a -> Maybe a
Just CExpr
cond)
                    (CExpr -> Maybe CExpr
forall a. a -> Maybe a
Just CExpr
inc)
                    ([CCompoundBlockItem] -> CStat
CCompound ([CCompoundBlockItem] -> CStat) -> [CCompoundBlockItem] -> CStat
forall a b. (a -> b) -> a -> b
$  ((CDecl -> CCompoundBlockItem) -> [CDecl] -> [CCompoundBlockItem]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap CDecl -> CCompoundBlockItem
CBlockDecl ([CDecl] -> [CDecl]
forall a. [a] -> [a]
reverse ([CDecl] -> [CDecl]) -> [CDecl] -> [CDecl]
forall a b. (a -> b) -> a -> b
$ CG -> [CDecl]
declarations CG
cg')
                               [CCompoundBlockItem]
-> [CCompoundBlockItem] -> [CCompoundBlockItem]
forall a. [a] -> [a] -> [a]
++ ((CStat -> CCompoundBlockItem) -> [CStat] -> [CCompoundBlockItem]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap CStat -> CCompoundBlockItem
CBlockStat ([CStat] -> [CStat]
forall a. [a] -> [a]
reverse ([CStat] -> [CStat]) -> [CStat] -> [CStat]
forall a b. (a -> b) -> a -> b
$ CG -> [CStat]
statements CG
cg'))))

-- given a monoid for a Hakaru type, generate the appropriate openMP reduction
-- declaration and return its unique identifier
declareReductionCG
  :: Sing (a :: Hakaru)
  -> (CExpr -> CodeGen ())
  -> (CExpr -> CExpr -> CodeGen ())
  -> CodeGen Ident
declareReductionCG :: Sing a
-> (CExpr -> StateT CG Identity ())
-> (CExpr -> CExpr -> StateT CG Identity ())
-> CodeGen Ident
declareReductionCG Sing a
typ CExpr -> StateT CG Identity ()
unit CExpr -> CExpr -> StateT CG Identity ()
mul =
  do Ident
redId <- String -> CodeGen Ident
genIdent' String
"red"
     Ident
unitId <- String -> CodeGen Ident
genIdent' String
"unit"
     Ident
mulId <- String -> CodeGen Ident
genIdent' String
"mul"
     let declType :: Ident -> CDecl
declType = Sing a -> Ident -> CDecl
forall (a :: Hakaru). Sing a -> Ident -> CDecl
typePtrDeclaration Sing a
typ

     Ident
inId <- String -> CodeGen Ident
genIdent' String
"in"
     [CTypeSpec]
-> Ident
-> [CDecl]
-> StateT CG Identity ()
-> StateT CG Identity ()
funCG [CTypeSpec
CVoid] Ident
unitId [Ident -> CDecl
declType Ident
inId] (StateT CG Identity () -> StateT CG Identity ())
-> StateT CG Identity () -> StateT CG Identity ()
forall a b. (a -> b) -> a -> b
$
       CExpr -> StateT CG Identity ()
unit (CExpr -> StateT CG Identity ())
-> (Ident -> CExpr) -> Ident -> StateT CG Identity ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ident -> CExpr
CVar (Ident -> StateT CG Identity ()) -> Ident -> StateT CG Identity ()
forall a b. (a -> b) -> a -> b
$ Ident
inId

     Ident
outId <- String -> CodeGen Ident
genIdent' String
"out"
     Ident
in2Id <- String -> CodeGen Ident
genIdent' String
"in"
     [CTypeSpec]
-> Ident
-> [CDecl]
-> StateT CG Identity ()
-> StateT CG Identity ()
funCG [CTypeSpec
CVoid] Ident
mulId [Ident -> CDecl
declType Ident
outId,Ident -> CDecl
declType Ident
in2Id] (StateT CG Identity () -> StateT CG Identity ())
-> StateT CG Identity () -> StateT CG Identity ()
forall a b. (a -> b) -> a -> b
$
         CExpr -> CExpr -> StateT CG Identity ()
mul (Ident -> CExpr
CVar Ident
outId) (Ident -> CExpr
CVar Ident
in2Id)

     let typ' :: CTypeSpec
typ' = case Sing a -> [CTypeSpec]
forall (a :: Hakaru). Sing a -> [CTypeSpec]
buildType Sing a
typ of
                  (CTypeSpec
x:[CTypeSpec]
_) -> CTypeSpec
x
                  [CTypeSpec]
_ -> String -> CTypeSpec
forall a. HasCallStack => String -> a
error (String -> CTypeSpec) -> String -> CTypeSpec
forall a b. (a -> b) -> a -> b
$ String
"buildType{" String -> String -> String
forall a. [a] -> [a] -> [a]
++ (Sing a -> String
forall a. Show a => a -> String
show Sing a
typ) String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"}"
     CStat -> StateT CG Identity ()
putStat (CStat -> StateT CG Identity ())
-> (OMP -> CStat) -> OMP -> StateT CG Identity ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Preprocessor -> CStat
CPPStat (Preprocessor -> CStat) -> (OMP -> Preprocessor) -> OMP -> CStat
forall b c a. (b -> c) -> (a -> b) -> a -> c
. OMP -> Preprocessor
ompToPP (OMP -> StateT CG Identity ()) -> OMP -> StateT CG Identity ()
forall a b. (a -> b) -> a -> b
$
       Directive -> OMP
OMP (Ident -> CTypeSpec -> CExpr -> CExpr -> Directive
DeclareRed Ident
redId
                       CTypeSpec
typ'
                       (CExpr -> [CExpr] -> CExpr
CCall (Ident -> CExpr
CVar Ident
mulId)
                              ((String -> CExpr) -> [String] -> [CExpr]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (CExpr -> CExpr
address (CExpr -> CExpr) -> (String -> CExpr) -> String -> CExpr
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ident -> CExpr
CVar (Ident -> CExpr) -> (String -> Ident) -> String -> CExpr
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> Ident
Ident)
                                    [String
"omp_in",String
"omp_out"]))
                       (CExpr -> [CExpr] -> CExpr
CCall (Ident -> CExpr
CVar Ident
unitId)
                              [CExpr -> CExpr
address (CExpr -> CExpr) -> (String -> CExpr) -> String -> CExpr
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ident -> CExpr
CVar (Ident -> CExpr) -> (String -> Ident) -> String -> CExpr
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> Ident
Ident (String -> CExpr) -> String -> CExpr
forall a b. (a -> b) -> a -> b
$ String
"omp_priv"]))
     Ident -> CodeGen Ident
forall (m :: * -> *) a. Monad m => a -> m a
return Ident
redId


-- not control flow, but like these it creates a block with local variables
codeBlockCG :: CodeGen () -> CodeGen ()
codeBlockCG :: StateT CG Identity () -> StateT CG Identity ()
codeBlockCG StateT CG Identity ()
body =
  do CG
cg <- StateT CG Identity CG
forall s (m :: * -> *). MonadState s m => m s
get
     let (()
_,CG
cg') = StateT CG Identity () -> CG -> ((), CG)
forall s a. State s a -> s -> (a, s)
runState StateT CG Identity ()
body (CG -> ((), CG)) -> CG -> ((), CG)
forall a b. (a -> b) -> a -> b
$ CG
cg { statements :: [CStat]
statements = []
                                      , declarations :: [CDecl]
declarations = [] }
     CG -> StateT CG Identity ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put (CG -> StateT CG Identity ()) -> CG -> StateT CG Identity ()
forall a b. (a -> b) -> a -> b
$ CG
cg' { statements :: [CStat]
statements = CG -> [CStat]
statements CG
cg
               , declarations :: [CDecl]
declarations = CG -> [CDecl]
declarations CG
cg }
     CStat -> StateT CG Identity ()
putStat (CStat -> StateT CG Identity ()) -> CStat -> StateT CG Identity ()
forall a b. (a -> b) -> a -> b
$ ([CCompoundBlockItem] -> CStat
CCompound ([CCompoundBlockItem] -> CStat) -> [CCompoundBlockItem] -> CStat
forall a b. (a -> b) -> a -> b
$  ((CDecl -> CCompoundBlockItem) -> [CDecl] -> [CCompoundBlockItem]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap CDecl -> CCompoundBlockItem
CBlockDecl ([CDecl] -> [CDecl]
forall a. [a] -> [a]
reverse ([CDecl] -> [CDecl]) -> [CDecl] -> [CDecl]
forall a b. (a -> b) -> a -> b
$ CG -> [CDecl]
declarations CG
cg')
                          [CCompoundBlockItem]
-> [CCompoundBlockItem] -> [CCompoundBlockItem]
forall a. [a] -> [a] -> [a]
++ ((CStat -> CCompoundBlockItem) -> [CStat] -> [CCompoundBlockItem]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap CStat -> CCompoundBlockItem
CBlockStat ([CStat] -> [CStat]
forall a. [a] -> [a]
reverse ([CStat] -> [CStat]) -> [CStat] -> [CStat]
forall a b. (a -> b) -> a -> b
$ CG -> [CStat]
statements CG
cg'))))



--------------------------------------------------------------------------------
-- ^ Takes a cexpression for the location and size and a hakaru type, and
--   generates a statement for allocating the memory
putMallocStat :: CExpr -> CExpr -> Sing (a :: Hakaru) -> CodeGen ()
putMallocStat :: CExpr -> CExpr -> Sing a -> StateT CG Identity ()
putMallocStat CExpr
loc CExpr
size Sing a
typ = do
  Bool
isManagedMem <- CG -> Bool
managedMem (CG -> Bool) -> StateT CG Identity CG -> StateT CG Identity Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> StateT CG Identity CG
forall s (m :: * -> *). MonadState s m => m s
get
  let malloc' :: CExpr -> CExpr
malloc' = if Bool
isManagedMem then CExpr -> CExpr
gcMalloc else CExpr -> CExpr
mallocE
      typ' :: [CTypeSpec]
typ' = Sing a -> [CTypeSpec]
forall (a :: Hakaru). Sing a -> [CTypeSpec]
buildType Sing a
typ
  CExpr -> StateT CG Identity ()
putExprStat (CExpr -> StateT CG Identity ()) -> CExpr -> StateT CG Identity ()
forall a b. (a -> b) -> a -> b
$   CExpr
loc
              CExpr -> CExpr -> CExpr
.=. ( CTypeName -> CExpr -> CExpr
CCast ([CTypeSpec] -> Bool -> CTypeName
CTypeName [CTypeSpec]
typ' Bool
True)
                  (CExpr -> CExpr) -> CExpr -> CExpr
forall a b. (a -> b) -> a -> b
$ CExpr -> CExpr
malloc' (CExpr
size CExpr -> CExpr -> CExpr
.*. (CTypeName -> CExpr
CSizeOfType ([CTypeSpec] -> Bool -> CTypeName
CTypeName [CTypeSpec]
typ' Bool
False))))