{-# LANGUAGE QuasiQuotes #-}

-- | C code generation for functions.
module Futhark.CodeGen.Backends.GenericC.Fun
  ( compileFun,
    compileVoidFun,
    module Futhark.CodeGen.Backends.GenericC.Monad,
    module Futhark.CodeGen.Backends.GenericC.Code,
  )
where

import Control.Monad
import Futhark.CodeGen.Backends.GenericC.Code
import Futhark.CodeGen.Backends.GenericC.Monad
import Futhark.CodeGen.ImpCode
import Futhark.MonadFreshNames
import Language.C.Quote.OpenCL qualified as C
import Language.C.Syntax qualified as C

compileFunBody :: [C.Exp] -> [Param] -> Code op -> CompilerM op s ()
compileFunBody :: forall op s. [Exp] -> [Param] -> Code op -> CompilerM op s ()
compileFunBody [Exp]
output_ptrs [Param]
outputs Code op
code = do
  (Param -> CompilerM op s ()) -> [Param] -> CompilerM op s ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ Param -> CompilerM op s ()
forall {op} {s}. Param -> CompilerM op s ()
declareOutput [Param]
outputs
  Code op -> CompilerM op s ()
forall op s. Code op -> CompilerM op s ()
compileCode Code op
code
  (Exp -> Param -> CompilerM op s ())
-> [Exp] -> [Param] -> CompilerM op s ()
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m ()
zipWithM_ Exp -> Param -> CompilerM op s ()
forall {a} {op} {s}. ToExp a => a -> Param -> CompilerM op s ()
setRetVal' [Exp]
output_ptrs [Param]
outputs
  where
    declareOutput :: Param -> CompilerM op s ()
declareOutput (MemParam VName
name Space
space) =
      VName -> Space -> CompilerM op s ()
forall op s. VName -> Space -> CompilerM op s ()
declMem VName
name Space
space
    declareOutput (ScalarParam VName
name PrimType
pt) = do
      let ctp :: Type
ctp = PrimType -> Type
primTypeToCType PrimType
pt
      InitGroup -> CompilerM op s ()
forall op s. InitGroup -> CompilerM op s ()
decl [C.cdecl|$ty:ctp $id:name;|]

    setRetVal' :: a -> Param -> CompilerM op s ()
setRetVal' a
p (MemParam VName
name Space
space) =
      -- It is required that the memory block is already initialised
      -- (although it may be NULL).
      Exp -> VName -> Space -> CompilerM op s ()
forall a b op s.
(ToExp a, ToExp b) =>
a -> b -> Space -> CompilerM op s ()
setMem [C.cexp|*$exp:p|] VName
name Space
space
    setRetVal' a
p (ScalarParam VName
name PrimType
_) =
      Stm -> CompilerM op s ()
forall op s. Stm -> CompilerM op s ()
stm [C.cstm|*$exp:p = $id:name;|]

compileInput :: Param -> CompilerM op s C.Param
compileInput :: forall op s. Param -> CompilerM op s Param
compileInput (ScalarParam VName
name PrimType
bt) = do
  let ctp :: Type
ctp = PrimType -> Type
primTypeToCType PrimType
bt
  Param -> CompilerM op s Param
forall a. a -> CompilerM op s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [C.cparam|$ty:ctp $id:name|]
compileInput (MemParam VName
name Space
space) = do
  Type
ty <- VName -> Space -> CompilerM op s Type
forall op s. VName -> Space -> CompilerM op s Type
memToCType VName
name Space
space
  Param -> CompilerM op s Param
forall a. a -> CompilerM op s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [C.cparam|$ty:ty $id:name|]

compileOutput :: Param -> CompilerM op s (C.Param, C.Exp)
compileOutput :: forall op s. Param -> CompilerM op s (Param, Exp)
compileOutput (ScalarParam VName
name PrimType
bt) = do
  let ctp :: Type
ctp = PrimType -> Type
primTypeToCType PrimType
bt
  VName
p_name <- [Char] -> CompilerM op s VName
forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
newVName ([Char] -> CompilerM op s VName) -> [Char] -> CompilerM op s VName
forall a b. (a -> b) -> a -> b
$ [Char]
"out_" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ VName -> [Char]
baseString VName
name
  (Param, Exp) -> CompilerM op s (Param, Exp)
forall a. a -> CompilerM op s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([C.cparam|$ty:ctp *$id:p_name|], [C.cexp|$id:p_name|])
compileOutput (MemParam VName
name Space
space) = do
  Type
ty <- VName -> Space -> CompilerM op s Type
forall op s. VName -> Space -> CompilerM op s Type
memToCType VName
name Space
space
  VName
p_name <- [Char] -> CompilerM op s VName
forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
newVName ([Char] -> CompilerM op s VName) -> [Char] -> CompilerM op s VName
forall a b. (a -> b) -> a -> b
$ VName -> [Char]
baseString VName
name [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"_p"
  (Param, Exp) -> CompilerM op s (Param, Exp)
forall a. a -> CompilerM op s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([C.cparam|$ty:ty *$id:p_name|], [C.cexp|$id:p_name|])

compileFun :: [C.BlockItem] -> [C.Param] -> (Name, Function op) -> CompilerM op s (C.Definition, C.Func)
compileFun :: forall op s.
[BlockItem]
-> [Param]
-> (Name, Function op)
-> CompilerM op s (Definition, Func)
compileFun [BlockItem]
get_constants [Param]
extra (Name
fname, func :: Function op
func@(Function Maybe EntryPoint
_ [Param]
outputs [Param]
inputs Code op
body)) = CompilerM op s (Definition, Func)
-> CompilerM op s (Definition, Func)
forall op s a. CompilerM op s a -> CompilerM op s a
inNewFunction (CompilerM op s (Definition, Func)
 -> CompilerM op s (Definition, Func))
-> CompilerM op s (Definition, Func)
-> CompilerM op s (Definition, Func)
forall a b. (a -> b) -> a -> b
$ do
  ([Param]
outparams, [Exp]
out_ptrs) <- (Param -> CompilerM op s (Param, Exp))
-> [Param] -> CompilerM op s ([Param], [Exp])
forall (m :: * -> *) a b c.
Applicative m =>
(a -> m (b, c)) -> [a] -> m ([b], [c])
mapAndUnzipM Param -> CompilerM op s (Param, Exp)
forall op s. Param -> CompilerM op s (Param, Exp)
compileOutput [Param]
outputs
  [Param]
inparams <- (Param -> CompilerM op s Param)
-> [Param] -> CompilerM op s [Param]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM Param -> CompilerM op s Param
forall op s. Param -> CompilerM op s Param
compileInput [Param]
inputs

  Map VName Space
-> ([BlockItem] -> [Stm] -> CompilerM op s (Definition, Func))
-> CompilerM op s (Definition, Func)
forall op s a.
Map VName Space
-> ([BlockItem] -> [Stm] -> CompilerM op s a) -> CompilerM op s a
cachingMemory (Function op -> Map VName Space
forall a. Function a -> Map VName Space
lexicalMemoryUsage Function op
func) (([BlockItem] -> [Stm] -> CompilerM op s (Definition, Func))
 -> CompilerM op s (Definition, Func))
-> ([BlockItem] -> [Stm] -> CompilerM op s (Definition, Func))
-> CompilerM op s (Definition, Func)
forall a b. (a -> b) -> a -> b
$ \[BlockItem]
decl_cached [Stm]
free_cached -> do
    [BlockItem]
body' <- CompilerM op s () -> CompilerM op s [BlockItem]
forall op s. CompilerM op s () -> CompilerM op s [BlockItem]
collect (CompilerM op s () -> CompilerM op s [BlockItem])
-> CompilerM op s () -> CompilerM op s [BlockItem]
forall a b. (a -> b) -> a -> b
$ [Exp] -> [Param] -> Code op -> CompilerM op s ()
forall op s. [Exp] -> [Param] -> Code op -> CompilerM op s ()
compileFunBody [Exp]
out_ptrs [Param]
outputs Code op
body
    [BlockItem]
decl_mem <- CompilerM op s [BlockItem]
forall op s. CompilerM op s [BlockItem]
declAllocatedMem
    [BlockItem]
free_mem <- CompilerM op s [BlockItem]
forall op s. CompilerM op s [BlockItem]
freeAllocatedMem
    let futhark_function :: DeclSpec
futhark_function =
          [Storage] -> [TypeQual] -> TypeSpec -> SrcLoc -> DeclSpec
C.DeclSpec [] [[Char] -> SrcLoc -> TypeQual
C.EscTypeQual [Char]
"FUTHARK_FUN_ATTR" SrcLoc
forall a. Monoid a => a
mempty] (Maybe Sign -> SrcLoc -> TypeSpec
C.Tint Maybe Sign
forall a. Maybe a
Nothing SrcLoc
forall a. Monoid a => a
mempty) SrcLoc
forall a. Monoid a => a
mempty

    (Definition, Func) -> CompilerM op s (Definition, Func)
forall a. a -> CompilerM op s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
      ( [C.cedecl|$spec:futhark_function $id:(funName fname)($params:extra, $params:outparams, $params:inparams);|],
        [C.cfun|$spec:futhark_function $id:(funName fname)($params:extra, $params:outparams, $params:inparams) {
               $stms:ignores
               int err = 0;
               $items:decl_cached
               $items:decl_mem
               $items:get_constants
               $items:body'
              cleanup:
               {
               $stms:free_cached
               $items:free_mem
               }
               return err;
  }|]
      )
  where
    -- Ignore all the boilerplate parameters, just in case we don't
    -- actually need to use them.
    ignores :: [Stm]
ignores = [[C.cstm|(void)$id:p;|] | C.Param (Just Id
p) DeclSpec
_ Decl
_ SrcLoc
_ <- [Param]
extra]

-- | Generate code for a function that returns void (meaning it cannot
-- fail) and has no extra parameters (meaning it cannot allocate
-- memory non-lexxical or do anything fancy).
compileVoidFun :: [C.BlockItem] -> (Name, Function op) -> CompilerM op s (C.Definition, C.Func)
compileVoidFun :: forall op s.
[BlockItem]
-> (Name, Function op) -> CompilerM op s (Definition, Func)
compileVoidFun [BlockItem]
get_constants (Name
fname, func :: Function op
func@(Function Maybe EntryPoint
_ [Param]
outputs [Param]
inputs Code op
body)) = CompilerM op s (Definition, Func)
-> CompilerM op s (Definition, Func)
forall op s a. CompilerM op s a -> CompilerM op s a
inNewFunction (CompilerM op s (Definition, Func)
 -> CompilerM op s (Definition, Func))
-> CompilerM op s (Definition, Func)
-> CompilerM op s (Definition, Func)
forall a b. (a -> b) -> a -> b
$ do
  ([Param]
outparams, [Exp]
out_ptrs) <- (Param -> CompilerM op s (Param, Exp))
-> [Param] -> CompilerM op s ([Param], [Exp])
forall (m :: * -> *) a b c.
Applicative m =>
(a -> m (b, c)) -> [a] -> m ([b], [c])
mapAndUnzipM Param -> CompilerM op s (Param, Exp)
forall op s. Param -> CompilerM op s (Param, Exp)
compileOutput [Param]
outputs
  [Param]
inparams <- (Param -> CompilerM op s Param)
-> [Param] -> CompilerM op s [Param]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM Param -> CompilerM op s Param
forall op s. Param -> CompilerM op s Param
compileInput [Param]
inputs

  Map VName Space
-> ([BlockItem] -> [Stm] -> CompilerM op s (Definition, Func))
-> CompilerM op s (Definition, Func)
forall op s a.
Map VName Space
-> ([BlockItem] -> [Stm] -> CompilerM op s a) -> CompilerM op s a
cachingMemory (Function op -> Map VName Space
forall a. Function a -> Map VName Space
lexicalMemoryUsage Function op
func) (([BlockItem] -> [Stm] -> CompilerM op s (Definition, Func))
 -> CompilerM op s (Definition, Func))
-> ([BlockItem] -> [Stm] -> CompilerM op s (Definition, Func))
-> CompilerM op s (Definition, Func)
forall a b. (a -> b) -> a -> b
$ \[BlockItem]
decl_cached [Stm]
free_cached -> do
    [BlockItem]
body' <- CompilerM op s () -> CompilerM op s [BlockItem]
forall op s. CompilerM op s () -> CompilerM op s [BlockItem]
collect (CompilerM op s () -> CompilerM op s [BlockItem])
-> CompilerM op s () -> CompilerM op s [BlockItem]
forall a b. (a -> b) -> a -> b
$ [Exp] -> [Param] -> Code op -> CompilerM op s ()
forall op s. [Exp] -> [Param] -> Code op -> CompilerM op s ()
compileFunBody [Exp]
out_ptrs [Param]
outputs Code op
body
    let futhark_function :: DeclSpec
futhark_function =
          [Storage] -> [TypeQual] -> TypeSpec -> SrcLoc -> DeclSpec
C.DeclSpec [] [[Char] -> SrcLoc -> TypeQual
C.EscTypeQual [Char]
"FUTHARK_FUN_ATTR" SrcLoc
forall a. Monoid a => a
mempty] (SrcLoc -> TypeSpec
C.Tvoid SrcLoc
forall a. Monoid a => a
mempty) SrcLoc
forall a. Monoid a => a
mempty

    (Definition, Func) -> CompilerM op s (Definition, Func)
forall a. a -> CompilerM op s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
      ( [C.cedecl|$spec:futhark_function $id:(funName fname)($params:outparams, $params:inparams);|],
        [C.cfun|$spec:futhark_function $id:(funName fname)($params:outparams, $params:inparams) {
               $items:decl_cached
               $items:get_constants
               $items:body'
               $stms:free_cached
               }|]
      )