{-# LANGUAGE DataKinds,
FlexibleContexts,
GADTs,
RankNTypes,
KindSignatures #-}
module Language.Hakaru.CodeGen.Types
( buildDeclaration
, buildDeclaration'
, buildPtrDeclaration
, typeDeclaration
, typePtrDeclaration
, typeName
, arrayDeclaration
, arrayStruct
, arraySize
, arrayData
, arrayPtrSize
, arrayPtrData
, mdataDeclaration
, mdataPtrDeclaration
, mdataStruct
, mdataStruct'
, mdataWeight
, mdataSample
, mdataPtrWeight
, mdataPtrSample
, datumDeclaration
, datumStruct
, datumSum
, datumProd
, datumFst
, datumSnd
, datumIndex
, functionDef
, closureStructure
, buildType
, castTo
, castToPtrOf
, callStruct
, buildStruct
, buildUnion
, binaryOp
) where
import Control.Monad.State
import Language.Hakaru.Syntax.ABT
import Language.Hakaru.Syntax.AST
import Language.Hakaru.Syntax.IClasses
import Language.Hakaru.Types.DataKind
import Language.Hakaru.Types.HClasses
import Language.Hakaru.Types.Sing
import Language.Hakaru.CodeGen.AST
import Language.Hakaru.CodeGen.Libs
buildDeclaration :: CTypeSpec -> Ident -> CDecl
buildDeclaration ctyp ident =
CDecl [ CTypeSpec ctyp ]
[( CDeclr Nothing (CDDeclrIdent ident)
, Nothing)]
buildDeclaration' :: [CTypeSpec] -> Ident -> CDecl
buildDeclaration' specs ident =
CDecl (fmap CTypeSpec specs)
[( CDeclr Nothing (CDDeclrIdent ident)
, Nothing)]
buildPtrDeclaration :: CTypeSpec -> Ident -> CDecl
buildPtrDeclaration ctyp ident =
CDecl [ CTypeSpec ctyp ]
[( CDeclr (Just $ CPtrDeclr []) (CDDeclrIdent ident)
, Nothing)]
typeDeclaration :: Sing (a :: Hakaru) -> Ident -> CDecl
typeDeclaration typ ident =
CDecl (fmap CTypeSpec $ buildType typ)
[( CDeclr Nothing (CDDeclrIdent ident)
, Nothing)]
typePtrDeclaration :: Sing (a :: Hakaru) -> Ident -> CDecl
typePtrDeclaration typ ident =
CDecl (fmap CTypeSpec $ buildType typ)
[( CDeclr (Just $ CPtrDeclr [])
(CDDeclrIdent ident)
, Nothing)]
typeName :: Sing (a :: Hakaru) -> String
typeName SInt = "int"
typeName SNat = "nat"
typeName SReal = "real"
typeName SProb = "prob"
typeName (SArray t) = "array_" ++ typeName t
typeName (SMeasure t) = "mdata_" ++ typeName t
typeName f@(SFun _ _) = error $ "typeName{SFun} doen't make sense: unknown context for {" ++ show f ++ "}"
typeName (SData _ t) = "dat_" ++ datumSumName t
where datumSumName :: Sing (a :: [[HakaruFun]]) -> String
datumSumName SVoid = "V"
datumSumName (SPlus p s) = datumProdName p ++ datumSumName s
datumProdName :: Sing (a :: [HakaruFun]) -> String
datumProdName SDone = "D"
datumProdName (SEt x p) = datumPrimName x ++ datumProdName p
datumPrimName :: Sing (a :: HakaruFun) -> String
datumPrimName SIdent = "I"
datumPrimName (SKonst s) = "K" ++ typeName s
arrayStruct :: Sing (a :: Hakaru) -> CExtDecl
arrayStruct t = CDeclExt (CDecl [CTypeSpec $ arrayStruct' t] [])
arrayStruct' :: Sing (a :: Hakaru) -> CTypeSpec
arrayStruct' t = aStruct
where aSize = buildDeclaration' [CUnsigned,CInt] (Ident "size")
aData = typePtrDeclaration t (Ident "data")
aStruct = buildStruct (Just . Ident . typeName . SArray $ t) [aSize,aData]
arrayDeclaration
:: Sing (a :: Hakaru)
-> Ident
-> CDecl
arrayDeclaration = buildDeclaration . callStruct . typeName . SArray
arraySize :: CExpr -> CExpr
arraySize e = CMember e (Ident "size") True
arrayData :: CExpr -> CExpr
arrayData e = CMember e (Ident "data") True
arrayPtrSize :: CExpr -> CExpr
arrayPtrSize e = CMember e (Ident "size") False
arrayPtrData :: CExpr -> CExpr
arrayPtrData e = CMember e (Ident "data") False
mdataStruct :: Sing (a :: Hakaru) -> CExtDecl
mdataStruct t = CDeclExt (CDecl [CTypeSpec $ mdataStruct' t] [])
mdataStruct' :: Sing (a :: Hakaru) -> CTypeSpec
mdataStruct' t = mdStruct
where weight = buildDeclaration CDouble (Ident "weight")
sample = typeDeclaration t (Ident "sample")
mdStruct = buildStruct (Just . Ident . typeName . SMeasure $ t) [weight,sample]
mdataDeclaration
:: Sing (a :: Hakaru)
-> Ident
-> CDecl
mdataDeclaration = buildDeclaration . callStruct . typeName . SMeasure
mdataPtrDeclaration
:: Sing (a :: Hakaru)
-> Ident
-> CDecl
mdataPtrDeclaration = buildPtrDeclaration . callStruct . typeName . SMeasure
mdataWeight :: CExpr -> CExpr
mdataWeight d = CMember d (Ident "weight") True
mdataSample :: CExpr -> CExpr
mdataSample d = CMember d (Ident "sample") True
mdataPtrWeight :: CExpr -> CExpr
mdataPtrWeight d = CMember d (Ident "weight") False
mdataPtrSample :: CExpr -> CExpr
mdataPtrSample d = CMember d (Ident "sample") False
datumStruct :: (Sing (HData' t)) -> CExtDecl
datumStruct dat@(SData _ typ)
= CDeclExt $ datumSum dat typ (Ident (typeName dat))
datumDeclaration
:: (Sing (HData' t))
-> Ident
-> CDecl
datumDeclaration = buildDeclaration . callStruct . typeName
datumSum
:: Sing (HData' t)
-> Sing (a :: [[HakaruFun]])
-> Ident
-> CDecl
datumSum dat funs ident =
let declrs = fst $ runState (datumSum' dat funs) cNameStream
union = buildDeclaration (buildUnion declrs) (Ident "sum")
ind = buildDeclaration CInt (Ident "index")
struct = buildStruct (Just ident) $ case declrs of
[] -> [ind]
_ -> [ind,union]
in CDecl [ CTypeSpec struct ] []
datumSum'
:: Sing (HData' t)
-> Sing (a :: [[HakaruFun]])
-> State [String] [CDecl]
datumSum' _ SVoid = return []
datumSum' dat (SPlus prod rest) =
do (name:names) <- get
put names
let ident = Ident name
mdecl = datumProd dat prod ident
rest' <- datumSum' dat rest
case mdecl of
Nothing -> return rest'
Just d -> return $ [d] ++ rest'
datumProd
:: Sing (HData' t)
-> Sing (a :: [HakaruFun])
-> Ident
-> Maybe CDecl
datumProd _ SDone _ = Nothing
datumProd dat funs ident =
let declrs = fst $ runState (datumProd' dat funs) cNameStream
in Just $ buildDeclaration (buildStruct Nothing $ declrs) ident
datumProd'
:: Sing (HData' t)
-> Sing (a :: [HakaruFun])
-> State [String] [CDecl]
datumProd' _ SDone = return []
datumProd' dat (SEt x ps) =
do x' <- datumPrim dat x
ps' <- datumProd' dat ps
return $ x' ++ ps'
datumPrim
:: Sing (HData' t)
-> Sing (a :: HakaruFun)
-> State [String] [CDecl]
datumPrim dat prim =
do (name:names) <- get
put names
let ident = Ident name
decl = case prim of
SIdent -> datumDeclaration dat ident
(SKonst k) -> typeDeclaration k ident
return [decl]
datumFst :: CExpr -> CExpr
datumFst x = x ... "sum" ... "a" ... "a"
datumSnd :: CExpr -> CExpr
datumSnd x = x ... "sum" ... "a" ... "b"
datumIndex :: CExpr -> CExpr
datumIndex x = x ... "index"
functionDef
:: Sing (a :: Hakaru)
-> Ident
-> [CDecl]
-> [CDecl]
-> [CStat]
-> CFunDef
functionDef typ ident argDecls internalDecls stmts =
CFunDef (fmap CTypeSpec $ buildType typ)
(CDeclr Nothing (CDDeclrIdent ident))
argDecls
(CCompound ((fmap CBlockDecl internalDecls)
++ (fmap CBlockStat stmts)))
closureStructure
:: forall (a :: Hakaru) xs
. [SomeVariable (KindOf a)]
-> List1 Variable (xs :: [Hakaru])
-> Ident
-> Sing a
-> CExtDecl
closureStructure fvs as i@(Ident name) typ = CDeclExt $
(CDecl [CTypeSpec $ (buildStruct (Just i) (codePtr:(declFvs cNameStream fvs)))]
[])
where declFvs _ [] = []
declFvs (n:ns) ((SomeVariable (Variable _ _ typ')):as') =
typeDeclaration typ' (Ident n) : declFvs ns as'
declFvc [] (_:_) = error "Ran out of identifiers but still had some types to assign"
codePtr = CDecl (fmap CTypeSpec . buildType $ typ)
[(CDeclr Nothing
(CDDeclrFun
(CDDeclrRec
(CDeclr (Just . CPtrDeclr $ [])
(CDDeclrIdent . Ident $ "_code_ptr")))
([callStruct name]:(varTypes as)))
,Nothing)]
varTypes :: List1 Variable (xs :: [Hakaru]) -> [[CTypeSpec]]
varTypes = foldMap11 (\(Variable _ _ typ') -> [buildType typ'])
buildType :: Sing (a :: Hakaru) -> [CTypeSpec]
buildType SInt = [CInt]
buildType SNat = [CUnsigned, CInt]
buildType SProb = [CDouble]
buildType SReal = [CDouble]
buildType (SMeasure x) = [callStruct . typeName . SMeasure $ x]
buildType (SArray t) = [callStruct . typeName . SArray $ t]
buildType (SFun _ x) = buildType $ x
buildType d@(SData _ _) = [callStruct . typeName $ d]
castTo :: [CTypeSpec] -> CExpr -> CExpr
castTo t = CCast (CTypeName t False)
castToPtrOf :: [CTypeSpec] -> CExpr -> CExpr
castToPtrOf t = CCast (CTypeName t True)
buildStruct :: Maybe Ident -> [CDecl] -> CTypeSpec
buildStruct mi decls =
CSUType (CSUSpec CStructTag mi decls)
callStruct :: String -> CTypeSpec
callStruct name =
CSUType (CSUSpec CStructTag (Just (Ident name)) [])
buildUnion :: [CDecl] -> CTypeSpec
buildUnion decls =
CSUType (CSUSpec CUnionTag Nothing decls)
binaryOp :: NaryOp a -> CExpr -> CExpr -> CExpr
binaryOp (Sum HSemiring_Prob) a b = CBinary CAddOp (expE a) (expE b)
binaryOp (Prod HSemiring_Prob) a b = CBinary CAddOp a b
binaryOp (Sum _) a b = CBinary CAddOp a b
binaryOp (Prod _) a b = CBinary CMulOp a b
binaryOp And a b = CUnary CNegOp (CBinary CEqOp a b)
binaryOp Or a b = CBinary CAndOp a b
binaryOp Xor a b = CBinary CLorOp a b
binaryOp x _ _ = error $ "TODO: binaryOp " ++ show x