{-# LANGUAGE CPP,
             BangPatterns,
             DataKinds,
             FlexibleContexts,
             GADTs,
             KindSignatures,
             ScopedTypeVariables,
             RankNTypes,
             TypeOperators #-}

----------------------------------------------------------------
--                                                    2016.06.23
-- |
-- Module      :  Language.Hakaru.CodeGen.Flatten
-- Copyright   :  Copyright (c) 2016 the Hakaru team
-- License     :  BSD3
-- Maintainer  :  zsulliva@indiana.edu
-- Stability   :  experimental
-- Portability :  GHC-only
--
--   Flatten takes Hakaru ABTs and C vars and returns a CStatement
-- assigning the var to the flattened ABT.
--
----------------------------------------------------------------


module Language.Hakaru.CodeGen.Flatten
  ( flattenABT
  , flattenVar
  , flattenTerm
  , flattenWithName
  , flattenWithName'
  , localVar
  , localVar'
  , opComment
  ) where

import Language.Hakaru.CodeGen.CodeGenMonad
import Language.Hakaru.CodeGen.AST
import Language.Hakaru.CodeGen.Libs
import Language.Hakaru.CodeGen.Types

import Language.Hakaru.Syntax.AST
import Language.Hakaru.Syntax.ABT
import Language.Hakaru.Syntax.TypeOf
import Language.Hakaru.Syntax.Datum hiding (Ident)
import Language.Hakaru.Syntax.Reducer
import Language.Hakaru.Syntax.IClasses
import qualified Language.Hakaru.Syntax.Prelude as HKP
import Language.Hakaru.Types.DataKind
import Language.Hakaru.Types.HClasses
import Language.Hakaru.Types.Coercion
import Language.Hakaru.Types.Sing

import           Control.Monad.State.Strict
import           Data.Number.Natural
import           Data.Ratio
import qualified Data.List.NonEmpty as NE
import qualified Data.Sequence      as S
import qualified Data.Foldable      as F
import qualified Data.Traversable   as T


#if __GLASGOW_HASKELL__ < 710
import           Control.Applicative (pure)
import           Control.Monad (replicateM)
import           Data.Functor
import           Data.Monoid        hiding (Product,Sum)
#endif


opComment :: String -> CStat
opComment :: String -> CStat
opComment String
opStr = String -> CStat
CComment (String -> CStat) -> String -> CStat
forall a b. (a -> b) -> a -> b
$ [String] -> String
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [String
space,String
" ",String
opStr,String
" ",String
space]
  where size :: Int
size  = (Int
50 Int -> Int -> Int
forall a. Num a => a -> a -> a
- (String -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length String
opStr)) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
8
        space :: String
space = Int -> Char -> String
forall a. Int -> a -> [a]
replicate Int
size Char
'-'

--------------------------------------------------------------------------------
--                                 Top Level                                  --
--------------------------------------------------------------------------------
{-

flattening an ABT will produce a continuation that takes a CExpr representing
a location where the value of the ABT should be stored. Return type of the
the continuation is CodeGen Bool, where the computed bool is whether or not
there is a Reject inside the ABT. Therefore it is only needed when computing
mochastic values

-}

localVar :: Sing (a :: Hakaru) -> CodeGen CExpr
localVar :: Sing a -> CodeGen CExpr
localVar Sing a
typ = Sing a -> String -> CodeGen CExpr
forall (a :: Hakaru). Sing a -> String -> CodeGen CExpr
localVar' Sing a
typ String
""

localVar' :: Sing (a :: Hakaru) -> String -> CodeGen CExpr
localVar' :: Sing a -> String -> CodeGen CExpr
localVar' Sing a
typ String
s =
  do Ident
eId <- String -> CodeGen Ident
genIdent' String
s
     Sing a -> Ident -> CodeGen ()
forall (a :: Hakaru). Sing a -> Ident -> CodeGen ()
declare Sing a
typ Ident
eId
     CExpr -> CodeGen CExpr
forall (m :: * -> *) a. Monad m => a -> m a
return (Ident -> CExpr
CVar Ident
eId)


flattenWithName'
  :: ABT Term abt
  => abt '[] a
  -> String
  -> CodeGen CExpr
flattenWithName' :: abt '[] a -> String -> CodeGen CExpr
flattenWithName' abt '[] a
abt String
hint = do
  Ident
ident <- String -> CodeGen Ident
genIdent' String
hint
  Sing a -> Ident -> CodeGen ()
forall (a :: Hakaru). Sing a -> Ident -> CodeGen ()
declare (abt '[] a -> Sing a
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> Sing a
typeOf abt '[] a
abt) Ident
ident
  let cvar :: CExpr
cvar = Ident -> CExpr
CVar Ident
ident
  abt '[] a -> CExpr -> CodeGen ()
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> CExpr -> CodeGen ()
flattenABT abt '[] a
abt CExpr
cvar
  CExpr -> CodeGen CExpr
forall (m :: * -> *) a. Monad m => a -> m a
return CExpr
cvar

flattenWithName
  :: ABT Term abt
  => abt '[] a
  -> CodeGen CExpr
flattenWithName :: abt '[] a -> CodeGen CExpr
flattenWithName abt '[] a
abt = abt '[] a -> String -> CodeGen CExpr
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> String -> CodeGen CExpr
flattenWithName' abt '[] a
abt String
""

flattenABT
  :: ABT Term abt
  => abt '[] a
  -> (CExpr -> CodeGen ())
flattenABT :: abt '[] a -> CExpr -> CodeGen ()
flattenABT abt '[] a
abt = abt '[] a
-> (Variable a -> CExpr -> CodeGen ())
-> (Term abt a -> CExpr -> CodeGen ())
-> CExpr
-> CodeGen ()
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k) r.
ABT syn abt =>
abt '[] a -> (Variable a -> r) -> (syn abt a -> r) -> r
caseVarSyn abt '[] a
abt Variable a -> CExpr -> CodeGen ()
forall (a :: Hakaru). Variable a -> CExpr -> CodeGen ()
flattenVar Term abt a -> CExpr -> CodeGen ()
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
Term abt a -> CExpr -> CodeGen ()
flattenTerm

-- note that variables will find their values in the state of the CodeGen monad
flattenVar
  :: Variable (a :: Hakaru)
  -> (CExpr -> CodeGen ())
flattenVar :: Variable a -> CExpr -> CodeGen ()
flattenVar Variable a
v = \CExpr
loc ->
  do CExpr
v' <- Ident -> CExpr
CVar (Ident -> CExpr) -> CodeGen Ident -> CodeGen CExpr
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Variable a -> CodeGen Ident
forall (a :: Hakaru). Variable a -> CodeGen Ident
lookupIdent Variable a
v
     CExpr -> CodeGen ()
putExprStat (CExpr -> CodeGen ()) -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ CExpr
loc CExpr -> CExpr -> CExpr
.=. CExpr
v'

flattenTerm
  :: ABT Term abt
  => Term abt a
  -> (CExpr -> CodeGen ())
-- SCon can contain mochastic terms
flattenTerm :: Term abt a -> CExpr -> CodeGen ()
flattenTerm (SCon args a
x :$ SArgs abt args
ys)         = SCon args a -> SArgs abt args -> CExpr -> CodeGen ()
forall (abt :: [Hakaru] -> Hakaru -> *)
       (args :: [([Hakaru], Hakaru)]) (a :: Hakaru).
ABT Term abt =>
SCon args a -> SArgs abt args -> CExpr -> CodeGen ()
flattenSCon SCon args a
x SArgs abt args
ys

flattenTerm (NaryOp_ NaryOp a
t Seq (abt '[] a)
s)     = NaryOp a -> Seq (abt '[] a) -> CExpr -> CodeGen ()
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
NaryOp a -> Seq (abt '[] a) -> CExpr -> CodeGen ()
flattenNAryOp NaryOp a
t Seq (abt '[] a)
s
flattenTerm (Literal_ Literal a
x)      = Literal a -> CExpr -> CodeGen ()
forall (a :: Hakaru). Literal a -> CExpr -> CodeGen ()
flattenLit Literal a
x
flattenTerm (Empty_ Sing ('HArray a)
_)        = String -> CExpr -> CodeGen ()
forall a. HasCallStack => String -> a
error String
"TODO: flattenTerm{Empty}"

flattenTerm (Datum_ Datum (abt '[]) (HData' t)
d)        = Datum (abt '[]) (HData' t) -> CExpr -> CodeGen ()
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: HakaruCon).
ABT Term abt =>
Datum (abt '[]) (HData' a) -> CExpr -> CodeGen ()
flattenDatum Datum (abt '[]) (HData' t)
d
flattenTerm (Case_ abt '[] a
c [Branch a abt a]
bs)      = abt '[] a -> [Branch a abt a] -> CExpr -> CodeGen ()
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru)
       (b :: Hakaru).
ABT Term abt =>
abt '[] a -> [Branch a abt b] -> CExpr -> CodeGen ()
flattenCase abt '[] a
c [Branch a abt a]
bs

flattenTerm (Bucket abt '[] 'HNat
b abt '[] 'HNat
e Reducer abt '[] a
rs)   = abt '[] 'HNat
-> abt '[] 'HNat -> Reducer abt '[] a -> CExpr -> CodeGen ()
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] 'HNat
-> abt '[] 'HNat -> Reducer abt '[] a -> CExpr -> CodeGen ()
flattenBucket abt '[] 'HNat
b abt '[] 'HNat
e Reducer abt '[] a
rs

flattenTerm (Array_ abt '[] 'HNat
s abt '[ 'HNat] a
e)      = abt '[] 'HNat -> abt '[ 'HNat] a -> CExpr -> CodeGen ()
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] 'HNat -> abt '[ 'HNat] a -> CExpr -> CodeGen ()
flattenArray abt '[] 'HNat
s abt '[ 'HNat] a
e
flattenTerm (ArrayLiteral_ [abt '[] a]
s) = [abt '[] a] -> CExpr -> CodeGen ()
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
[abt '[] a] -> CExpr -> CodeGen ()
flattenArrayLiteral [abt '[] a]
s


---------------------
-- Mochastic Terms --
---------------------
flattenTerm (Superpose_ NonEmpty (abt '[] 'HProb, abt '[] ('HMeasure a))
wes)  = NonEmpty (abt '[] 'HProb, abt '[] ('HMeasure a))
-> CExpr -> CodeGen ()
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
NonEmpty (abt '[] 'HProb, abt '[] ('HMeasure a))
-> CExpr -> CodeGen ()
flattenSuperpose NonEmpty (abt '[] 'HProb, abt '[] ('HMeasure a))
wes
flattenTerm (Reject_ Sing ('HMeasure a)
_)       = \CExpr
loc -> CExpr -> CodeGen ()
putExprStat (CExpr -> CExpr
mdataWeight CExpr
loc CExpr -> CExpr -> CExpr
.=. (Integer -> CExpr
intE Integer
0))


--------------------------------------------------------------------------------
--                                  SCon                                      --
--------------------------------------------------------------------------------

flattenSCon
  :: ( ABT Term abt )
  => SCon args a
  -> SArgs abt args
  -> (CExpr -> CodeGen ())

flattenSCon :: SCon args a -> SArgs abt args -> CExpr -> CodeGen ()
flattenSCon SCon args a
Let_ =
  \(abt vars a
expr :* abt vars a
body :* SArgs abt args
End) ->
    \CExpr
loc -> do
      abt '[a] a -> (Variable a -> abt '[] a -> CodeGen ()) -> CodeGen ()
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (x :: k) (xs :: [k]) (a :: k) r.
ABT syn abt =>
abt (x : xs) a -> (Variable x -> abt xs a -> r) -> r
caseBind abt vars a
abt '[a] a
body ((Variable a -> abt '[] a -> CodeGen ()) -> CodeGen ())
-> (Variable a -> abt '[] a -> CodeGen ()) -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ \v :: Variable a
v@(Variable Text
_ Nat
_ Sing a
typ) abt '[] a
body'->
        do Ident
ident <- Variable a -> CodeGen Ident
forall (a :: Hakaru). Variable a -> CodeGen Ident
createIdent Variable a
v
           Sing a -> Ident -> CodeGen ()
forall (a :: Hakaru). Sing a -> Ident -> CodeGen ()
declare Sing a
typ Ident
ident
           abt '[] a -> CExpr -> CodeGen ()
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> CExpr -> CodeGen ()
flattenABT abt vars a
abt '[] a
expr (Ident -> CExpr
CVar Ident
ident)
           abt '[] a -> CExpr -> CodeGen ()
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> CExpr -> CodeGen ()
flattenABT abt '[] a
body' CExpr
loc

-- Lambdas produce functions and then return a function label exprssion
flattenSCon SCon args a
Lam_ =
  \(abt vars a
body :* SArgs abt args
End) ->
    \CExpr
loc ->
      abt '[a] a
-> (forall (ys :: [Hakaru]) (b :: Hakaru).
    List1 Variable ys -> abt '[] b -> CodeGen ())
-> CodeGen ()
forall (abt :: [Hakaru] -> Hakaru -> *) (x :: Hakaru) (a :: Hakaru)
       r.
ABT Term abt =>
abt '[x] a
-> (forall (ys :: [Hakaru]) (b :: Hakaru).
    List1 Variable ys -> abt '[] b -> r)
-> r
coalesceLambda abt vars a
abt '[a] a
body ((forall (ys :: [Hakaru]) (b :: Hakaru).
  List1 Variable ys -> abt '[] b -> CodeGen ())
 -> CodeGen ())
-> (forall (ys :: [Hakaru]) (b :: Hakaru).
    List1 Variable ys -> abt '[] b -> CodeGen ())
-> CodeGen ()
forall a b. (a -> b) -> a -> b
$ \List1 Variable ys
args abt '[] b
body' ->
        let freevars :: [SomeVariable (KindOf b)]
freevars = VarSet (KindOf b) -> [SomeVariable (KindOf b)]
forall k (kproxy :: KProxy k).
VarSet kproxy -> [SomeVariable kproxy]
fromVarSet (VarSet (KindOf b) -> [SomeVariable (KindOf b)])
-> (abt '[] b -> VarSet (KindOf b))
-> abt '[] b
-> [SomeVariable (KindOf b)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. abt '[] b -> VarSet (KindOf b)
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (xs :: [k]) (a :: k).
ABT syn abt =>
abt xs a -> VarSet (KindOf a)
freeVars (abt '[] b -> [SomeVariable (KindOf b)])
-> abt '[] b -> [SomeVariable (KindOf b)]
forall a b. (a -> b) -> a -> b
$ abt '[] b
body'
            retTyp :: Sing b
retTyp   = abt '[] b -> Sing b
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> Sing a
typeOf abt '[] b
body'
        in do { -- create code block and closure structure
                [CDecl]
args' <- [StateT CG Identity CDecl] -> StateT CG Identity [CDecl]
forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
sequence ([StateT CG Identity CDecl] -> StateT CG Identity [CDecl])
-> (List1 Variable ys -> [StateT CG Identity CDecl])
-> List1 Variable ys
-> StateT CG Identity [CDecl]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (forall (i :: Hakaru). Variable i -> [StateT CG Identity CDecl])
-> List1 Variable ys -> [StateT CG Identity CDecl]
forall k1 k2 (f :: (k1 -> *) -> k2 -> *) m (a :: k1 -> *)
       (j :: k2).
(Foldable11 f, Monoid m) =>
(forall (i :: k1). a i -> m) -> f a j -> m
foldMap11 forall (i :: Hakaru). Variable i -> [StateT CG Identity CDecl]
argDecl (List1 Variable ys -> StateT CG Identity [CDecl])
-> List1 Variable ys -> StateT CG Identity [CDecl]
forall a b. (a -> b) -> a -> b
$ List1 Variable ys
args
              ; Ident
envId <- String -> CodeGen Ident
genIdent' String
"env"
              ; Ident
fnId  <- String -> CodeGen Ident
genIdent' String
"fn"
              ; closDataId :: Ident
closDataId@(Ident String
clos_n) <- String -> CodeGen Ident
genIdent' String
"clos_data"
              ; CExtDecl -> CodeGen ()
extDeclare ([SomeVariable (KindOf b)]
-> List1 Variable ys -> Ident -> Sing b -> CExtDecl
forall (a :: Hakaru) (xs :: [Hakaru]).
[SomeVariable (KindOf b)]
-> List1 Variable xs -> Ident -> Sing a -> CExtDecl
closureStructure [SomeVariable (KindOf b)]
freevars List1 Variable ys
args Ident
closDataId Sing b
retTyp)
              ; [CTypeSpec] -> Ident -> [CDecl] -> CodeGen () -> CodeGen ()
funCG (Sing b -> [CTypeSpec]
forall (a :: Hakaru). Sing a -> [CTypeSpec]
buildType Sing b
retTyp)
                      Ident
fnId
                      ((CTypeSpec -> Ident -> CDecl
buildDeclaration (String -> CTypeSpec
callStruct String
clos_n) Ident
envId)CDecl -> [CDecl] -> [CDecl]
forall a. a -> [a] -> [a]
:[CDecl]
args') (CodeGen () -> CodeGen ()) -> CodeGen () -> CodeGen ()
forall a b. (a -> b) -> a -> b
$
                  do { CStat -> CodeGen ()
putStat (String -> CStat
opComment String
"Begin Unpack Closure")
                       -- need to re-declare variables in functions scope
                     ; (SomeVariable (KindOf b) -> CodeGen ())
-> [SomeVariable (KindOf b)] -> CodeGen ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (\(SomeVariable v :: Variable a
v@(Variable Text
_ Nat
_ Sing a
typ)) ->
                               Variable a -> CodeGen Ident
forall (a :: Hakaru). Variable a -> CodeGen Ident
lookupIdent Variable a
v CodeGen Ident -> (Ident -> CodeGen ()) -> CodeGen ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Sing a -> Ident -> CodeGen ()
forall (a :: Hakaru). Sing a -> Ident -> CodeGen ()
declare Sing a
typ)
                             [SomeVariable (KindOf b)]
freevars
                     ; CExpr -> [String] -> [SomeVariable (KindOf b)] -> CodeGen ()
forall (a :: Hakaru).
CExpr -> [String] -> [SomeVariable (KindOf b)] -> CodeGen ()
unpackClosure (Ident -> CExpr
CVar Ident
envId) [String]
cNameStream [SomeVariable (KindOf b)]
freevars
                     ; CStat -> CodeGen ()
putStat (String -> CStat
opComment String
"End Unpack Closure")
                     ; CExpr
x <- abt '[] b -> CodeGen CExpr
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> CodeGen CExpr
flattenWithName abt '[] b
body'
                     ; CStat -> CodeGen ()
putStat (CStat -> CodeGen ()) -> (CExpr -> CStat) -> CExpr -> CodeGen ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Maybe CExpr -> CStat
CReturn (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 -> CodeGen ()) -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ CExpr
x }

                -- create the closure object
              ; Ident
closureId <- String -> CodeGen Ident
genIdent' String
"closure"
              ; CDecl -> CodeGen ()
declare' (CDecl -> CodeGen ()) -> (Ident -> CDecl) -> Ident -> CodeGen ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CTypeSpec -> Ident -> CDecl
buildDeclaration (String -> CTypeSpec
callStruct String
clos_n) (Ident -> CodeGen ()) -> Ident -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ Ident
closureId
              ; CStat -> CodeGen ()
putStat (String -> CStat
opComment String
"Begin Pack Closure")
              ; CExpr -> CodeGen ()
putExprStat (CExpr -> CodeGen ()) -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ ((Ident -> CExpr
CVar Ident
closureId) CExpr -> String -> CExpr
... String
"_code_ptr") CExpr -> CExpr -> CExpr
.=. (CExpr -> CExpr
address (Ident -> CExpr
CVar Ident
fnId))
              ; CExpr -> [String] -> [SomeVariable (KindOf b)] -> CodeGen ()
forall (a :: Hakaru).
CExpr -> [String] -> [SomeVariable (KindOf b)] -> CodeGen ()
packClosure (Ident -> CExpr
CVar Ident
closureId) [String]
cNameStream [SomeVariable (KindOf b)]
freevars
              ; CStat -> CodeGen ()
putStat (String -> CStat
opComment String
"End Pack Closure")
              ; CExpr -> CodeGen ()
putExprStat (CExpr -> CodeGen ()) -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ CExpr
loc CExpr -> CExpr -> CExpr
.=. (Ident -> CExpr
CVar Ident
closureId) }

  where -- collapses nested lambdas of one argument to lambdas that take a list
        -- arguments
        coalesceLambda
          :: ( ABT Term abt )
          => abt '[x] a
          -> (forall (ys :: [Hakaru]) b. List1 Variable ys -> abt '[] b -> r)
          -> r
        coalesceLambda :: abt '[x] a
-> (forall (ys :: [Hakaru]) (b :: Hakaru).
    List1 Variable ys -> abt '[] b -> r)
-> r
coalesceLambda abt '[x] a
abt forall (ys :: [Hakaru]) (b :: Hakaru).
List1 Variable ys -> abt '[] b -> r
k =
          abt '[x] a -> (Variable x -> abt '[] a -> r) -> r
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (x :: k) (xs :: [k]) (a :: k) r.
ABT syn abt =>
abt (x : xs) a -> (Variable x -> abt xs a -> r) -> r
caseBind abt '[x] a
abt ((Variable x -> abt '[] a -> r) -> r)
-> (Variable x -> abt '[] a -> r) -> r
forall a b. (a -> b) -> a -> b
$ \Variable x
v abt '[] a
abt' ->
            abt '[] a -> (Variable a -> r) -> (Term abt a -> r) -> r
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k) r.
ABT syn abt =>
abt '[] a -> (Variable a -> r) -> (syn abt a -> r) -> r
caseVarSyn abt '[] a
abt' (r -> Variable a -> r
forall a b. a -> b -> a
const (List1 Variable '[x] -> abt '[] a -> r
forall (ys :: [Hakaru]) (b :: Hakaru).
List1 Variable ys -> abt '[] b -> r
k (Variable x -> List1 Variable '[] -> List1 Variable '[x]
forall a (a :: a -> *) (x :: a) (xs :: [a]).
a x -> List1 a xs -> List1 a (x : xs)
Cons1 Variable x
v List1 Variable '[]
forall k (a :: k -> *). List1 a '[]
Nil1) abt '[] a
abt')) ((Term abt a -> r) -> r) -> (Term abt a -> r) -> r
forall a b. (a -> b) -> a -> b
$ \Term abt a
term ->
              case Term abt a
term of
                (SCon args a
Lam_ :$ abt vars a
body :* SArgs abt args
End) ->
                  abt '[a] a
-> (forall (ys :: [Hakaru]) (b :: Hakaru).
    List1 Variable ys -> abt '[] b -> r)
-> r
forall (abt :: [Hakaru] -> Hakaru -> *) (x :: Hakaru) (a :: Hakaru)
       r.
ABT Term abt =>
abt '[x] a
-> (forall (ys :: [Hakaru]) (b :: Hakaru).
    List1 Variable ys -> abt '[] b -> r)
-> r
coalesceLambda abt vars a
abt '[a] a
body ((forall (ys :: [Hakaru]) (b :: Hakaru).
  List1 Variable ys -> abt '[] b -> r)
 -> r)
-> (forall (ys :: [Hakaru]) (b :: Hakaru).
    List1 Variable ys -> abt '[] b -> r)
-> r
forall a b. (a -> b) -> a -> b
$ \List1 Variable ys
vars abt '[] b
abt'' -> List1 Variable (x : ys) -> abt '[] b -> r
forall (ys :: [Hakaru]) (b :: Hakaru).
List1 Variable ys -> abt '[] b -> r
k (Variable x -> List1 Variable ys -> List1 Variable (x : ys)
forall a (a :: a -> *) (x :: a) (xs :: [a]).
a x -> List1 a xs -> List1 a (x : xs)
Cons1 Variable x
v List1 Variable ys
vars) abt '[] b
abt''
                Term abt a
_ -> List1 Variable '[x] -> abt '[] a -> r
forall (ys :: [Hakaru]) (b :: Hakaru).
List1 Variable ys -> abt '[] b -> r
k (Variable x -> List1 Variable '[] -> List1 Variable '[x]
forall a (a :: a -> *) (x :: a) (xs :: [a]).
a x -> List1 a xs -> List1 a (x : xs)
Cons1 Variable x
v List1 Variable '[]
forall k (a :: k -> *). List1 a '[]
Nil1) abt '[] a
abt'

        argDecl :: Variable (a :: Hakaru) -> [CodeGen CDecl]
        argDecl :: Variable a -> [StateT CG Identity CDecl]
argDecl v :: Variable a
v@(Variable Text
_ Nat
_ Sing a
typ) =
          [do { Ident
ident <- Variable a -> CodeGen Ident
forall (a :: Hakaru). Variable a -> CodeGen Ident
createIdent Variable a
v ; CDecl -> StateT CG Identity CDecl
forall (m :: * -> *) a. Monad m => a -> m a
return (Sing a -> Ident -> CDecl
forall (a :: Hakaru). Sing a -> Ident -> CDecl
typeDeclaration Sing a
typ Ident
ident) }]

        -- captures the environment variables in closure object
        packClosure, unpackClosure
          :: CExpr
          -> [String]
          -> [SomeVariable (KindOf (a :: Hakaru))]
          -> CodeGen ()
        packClosure :: CExpr -> [String] -> [SomeVariable (KindOf b)] -> CodeGen ()
packClosure CExpr
_ [String]
_      [] = () -> CodeGen ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
        packClosure CExpr
c (String
n:[String]
ns) ((SomeVariable Variable a
a):[SomeVariable (KindOf b)]
as) =
          do { CExpr
a' <- Ident -> CExpr
CVar (Ident -> CExpr) -> CodeGen Ident -> CodeGen CExpr
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Variable a -> CodeGen Ident
forall (a :: Hakaru). Variable a -> CodeGen Ident
lookupIdent Variable a
a
             ; CExpr -> CodeGen ()
putExprStat (CExpr -> CodeGen ()) -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ CExpr
c CExpr -> String -> CExpr
... String
n CExpr -> CExpr -> CExpr
.=. CExpr
a'
             ; CExpr -> [String] -> [SomeVariable (KindOf b)] -> CodeGen ()
forall (a :: Hakaru).
CExpr -> [String] -> [SomeVariable (KindOf b)] -> CodeGen ()
packClosure CExpr
c [String]
ns [SomeVariable (KindOf b)]
as }
        packClosure CExpr
_ [String]
_ [SomeVariable (KindOf b)]
_ = String -> CodeGen ()
forall a. HasCallStack => String -> a
error String
"this isn't possible"

        unpackClosure :: CExpr -> [String] -> [SomeVariable (KindOf b)] -> CodeGen ()
unpackClosure CExpr
_ [String]
_      [] = () -> CodeGen ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
        unpackClosure CExpr
c (String
n:[String]
ns) ((SomeVariable Variable a
a):[SomeVariable (KindOf b)]
as) =
          do { CExpr
a' <- Ident -> CExpr
CVar (Ident -> CExpr) -> CodeGen Ident -> CodeGen CExpr
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Variable a -> CodeGen Ident
forall (a :: Hakaru). Variable a -> CodeGen Ident
lookupIdent Variable a
a
             ; CExpr -> CodeGen ()
putExprStat (CExpr -> CodeGen ()) -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ CExpr
a' CExpr -> CExpr -> CExpr
.=. CExpr
c CExpr -> String -> CExpr
... String
n
             ; CExpr -> [String] -> [SomeVariable (KindOf b)] -> CodeGen ()
forall (a :: Hakaru).
CExpr -> [String] -> [SomeVariable (KindOf b)] -> CodeGen ()
unpackClosure CExpr
c [String]
ns [SomeVariable (KindOf b)]
as }
        unpackClosure CExpr
_ [String]
_ [SomeVariable (KindOf b)]
_ = String -> CodeGen ()
forall a. HasCallStack => String -> a
error String
"this isn't possible"

flattenSCon SCon args a
App_  =
 \(abt vars a
fun :* abt vars a
arg :* SArgs abt args
End) ->
   \CExpr
loc ->
     do { CExpr
closE <- abt '[] a -> String -> CodeGen CExpr
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> String -> CodeGen CExpr
flattenWithName' abt vars a
abt '[] a
fun String
"closure"
        ; CExpr
paramE <- abt '[] a -> String -> CodeGen CExpr
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> String -> CodeGen CExpr
flattenWithName' abt vars a
abt '[] a
arg String
"param"
        ; CExpr -> CodeGen ()
putExprStat (CExpr -> CodeGen ()) -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ CExpr
loc CExpr -> CExpr -> CExpr
.=. CExpr -> [CExpr] -> CExpr
CCall (CExpr -> CExpr
indirect (CExpr
closE CExpr -> String -> CExpr
... String
"_code_ptr"))
                                      [CExpr
closE,CExpr
paramE] }

flattenSCon (PrimOp_ PrimOp typs a
op) = PrimOp typs a -> SArgs abt args -> CExpr -> CodeGen ()
forall (abt :: [Hakaru] -> Hakaru -> *) (typs :: [Hakaru])
       (args :: [([Hakaru], Hakaru)]) (a :: Hakaru).
(ABT Term abt, typs ~ UnLCs args, args ~ LCs typs) =>
PrimOp typs a -> SArgs abt args -> CExpr -> CodeGen ()
flattenPrimOp PrimOp typs a
op

flattenSCon (ArrayOp_ ArrayOp typs a
op) = ArrayOp typs a -> SArgs abt args -> CExpr -> CodeGen ()
forall (abt :: [Hakaru] -> Hakaru -> *) (typs :: [Hakaru])
       (args :: [([Hakaru], Hakaru)]) (a :: Hakaru).
(ABT Term abt, typs ~ UnLCs args, args ~ LCs typs) =>
ArrayOp typs a -> SArgs abt args -> CExpr -> CodeGen ()
flattenArrayOp ArrayOp typs a
op

flattenSCon (Summate HDiscrete a
_ HSemiring a
sr) =
  \(abt vars a
lo :* abt vars a
hi :* abt vars a
body :* SArgs abt args
End) ->
    \CExpr
loc ->
      let  semiTyp :: Sing a
semiTyp = HSemiring a -> Sing a
forall (a :: Hakaru). HSemiring a -> Sing a
sing_HSemiring HSemiring a
sr in
        do CExpr
loE <- abt '[] a -> String -> CodeGen CExpr
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> String -> CodeGen CExpr
flattenWithName' abt vars a
abt '[] a
lo String
"lo"
           CExpr
hiE <- abt '[] a -> String -> CodeGen CExpr
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> String -> CodeGen CExpr
flattenWithName' abt vars a
abt '[] a
hi String
"hi"

           CStat -> CodeGen ()
putStat (CStat -> CodeGen ()) -> CStat -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ String -> CStat
opComment String
"Begin Summate"

           case Sing a
semiTyp of
             -- special prob branch
             Sing a
SProb -> do
               Ident
summateArrId <- String -> CodeGen Ident
genIdent' String
"summate_arr"
               Sing ('HArray 'HProb) -> Ident -> CodeGen ()
forall (a :: Hakaru). Sing a -> Ident -> CodeGen ()
declare (Sing 'HProb -> Sing ('HArray 'HProb)
forall (a :: Hakaru). Sing a -> Sing ('HArray a)
SArray Sing 'HProb
SProb) Ident
summateArrId
               let summateArrE :: CExpr
summateArrE = Ident -> CExpr
CVar Ident
summateArrId
               CExpr -> CodeGen ()
putExprStat (CExpr -> CodeGen ()) -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ CExpr -> CExpr
arraySize CExpr
summateArrE CExpr -> CExpr -> CExpr
.=. (CExpr
hiE CExpr -> CExpr -> CExpr
.-. CExpr
loE)
               CExpr -> CodeGen ()
putExprStat (CExpr -> CodeGen ()) -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ CExpr -> CExpr
arrayData CExpr
summateArrE
                     CExpr -> CExpr -> CExpr
.=. ([CTypeSpec] -> CExpr -> CExpr
castToPtrOf [CTypeSpec
CDouble]
                           (CExpr -> CExpr
mallocE ((CExpr -> CExpr
arraySize CExpr
summateArrE) CExpr -> CExpr -> CExpr
.*.
                                    (CTypeName -> CExpr
CSizeOfType ([CTypeSpec] -> Bool -> CTypeName
CTypeName [CTypeSpec
CDouble] Bool
False)))))
               abt '[a] a -> CExpr -> CExpr -> CodeGen ()
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru)
       (b :: Hakaru).
ABT Term abt =>
abt '[a] b -> CExpr -> CExpr -> CodeGen ()
lseSummateArrayCG abt vars a
abt '[a] a
body CExpr
summateArrE CExpr
loc
               CExpr -> CodeGen ()
putExprStat (CExpr -> CodeGen ()) -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ CExpr -> CExpr
freeE (CExpr -> CExpr
arrayData CExpr
summateArrE)

             Sing a
_ ->
               abt '[a] a -> (Variable a -> abt '[] a -> CodeGen ()) -> CodeGen ()
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (x :: k) (xs :: [k]) (a :: k) r.
ABT syn abt =>
abt (x : xs) a -> (Variable x -> abt xs a -> r) -> r
caseBind abt vars a
abt '[a] a
body ((Variable a -> abt '[] a -> CodeGen ()) -> CodeGen ())
-> (Variable a -> abt '[] a -> CodeGen ()) -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ \Variable a
v abt '[] a
body' -> do
                 Ident
iterI <- Variable a -> CodeGen Ident
forall (a :: Hakaru). Variable a -> CodeGen Ident
createIdent Variable a
v
                 Sing 'HNat -> Ident -> CodeGen ()
forall (a :: Hakaru). Sing a -> Ident -> CodeGen ()
declare Sing 'HNat
SNat Ident
iterI

                 Ident
accI <- String -> CodeGen Ident
genIdent' String
"acc"
                 Sing a -> Ident -> CodeGen ()
forall (a :: Hakaru). Sing a -> Ident -> CodeGen ()
declare Sing a
semiTyp Ident
accI
                 Ident -> CExpr -> CodeGen ()
assign Ident
accI (case Sing a
semiTyp of
                                Sing a
SReal -> Float -> CExpr
floatE Float
0
                                Sing a
_     -> Integer -> CExpr
intE Integer
0)

                 let accVar :: CExpr
accVar  = Ident -> CExpr
CVar Ident
accI
                     iterVar :: CExpr
iterVar = Ident -> CExpr
CVar Ident
iterI

                 Either
  CBinaryOp
  (Sing Any, CExpr -> CodeGen (), CExpr -> CExpr -> CodeGen ())
-> CExpr -> CExpr -> CExpr -> CExpr -> CodeGen () -> CodeGen ()
forall (a :: Hakaru).
Either
  CBinaryOp
  (Sing a, CExpr -> CodeGen (), CExpr -> CExpr -> CodeGen ())
-> CExpr -> CExpr -> CExpr -> CExpr -> CodeGen () -> CodeGen ()
reductionCG (CBinaryOp
-> Either
     CBinaryOp
     (Sing Any, CExpr -> CodeGen (), CExpr -> CExpr -> CodeGen ())
forall a b. a -> Either a b
Left CBinaryOp
CAddOp)
                             CExpr
accVar
                             (CExpr
iterVar CExpr -> CExpr -> CExpr
.=. CExpr
loE)
                             (CExpr
iterVar CExpr -> CExpr -> CExpr
.<. CExpr
hiE)
                             (CUnaryOp -> CExpr -> CExpr
CUnary CUnaryOp
CPostIncOp CExpr
iterVar) (CodeGen () -> CodeGen ()) -> CodeGen () -> CodeGen ()
forall a b. (a -> b) -> a -> b
$
                   (CExpr -> CodeGen ()
putExprStat (CExpr -> CodeGen ()) -> (CExpr -> CExpr) -> CExpr -> CodeGen ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (CExpr
accVar CExpr -> CExpr -> CExpr
.+=.) (CExpr -> CodeGen ()) -> CodeGen CExpr -> CodeGen ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< abt '[] a -> CodeGen CExpr
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> CodeGen CExpr
flattenWithName abt '[] a
body')
                 CExpr -> CodeGen ()
putExprStat (CExpr -> CodeGen ()) -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ CExpr
loc CExpr -> CExpr -> CExpr
.=. CExpr
accVar

           CStat -> CodeGen ()
putStat (CStat -> CodeGen ()) -> CStat -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ String -> CStat
opComment String
"End Summate"



flattenSCon (Product HDiscrete a
_ HSemiring a
sr) =
  \(abt vars a
lo :* abt vars a
hi :* abt vars a
body :* SArgs abt args
End) ->
    \CExpr
loc ->
      let  semiTyp :: Sing a
semiTyp = HSemiring a -> Sing a
forall (a :: Hakaru). HSemiring a -> Sing a
sing_HSemiring HSemiring a
sr in
        do CExpr
loE <- abt '[] a -> String -> CodeGen CExpr
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> String -> CodeGen CExpr
flattenWithName' abt vars a
abt '[] a
lo String
"lo"
           CExpr
hiE <- abt '[] a -> String -> CodeGen CExpr
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> String -> CodeGen CExpr
flattenWithName' abt vars a
abt '[] a
hi String
"hi"

           CStat -> CodeGen ()
putStat (CStat -> CodeGen ()) -> CStat -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ String -> CStat
opComment String
"Begin Product"

           case Sing a
semiTyp of
           -- special prob branch
             Sing a
SProb -> abt '[a] a -> CExpr -> CExpr -> CExpr -> CodeGen ()
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru)
       (b :: Hakaru).
ABT Term abt =>
abt '[a] b -> CExpr -> CExpr -> CExpr -> CodeGen ()
kahanSummationCG abt vars a
abt '[a] a
body CExpr
loE CExpr
hiE CExpr
loc

             Sing a
_ ->
               abt '[a] a -> (Variable a -> abt '[] a -> CodeGen ()) -> CodeGen ()
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (x :: k) (xs :: [k]) (a :: k) r.
ABT syn abt =>
abt (x : xs) a -> (Variable x -> abt xs a -> r) -> r
caseBind abt vars a
abt '[a] a
body ((Variable a -> abt '[] a -> CodeGen ()) -> CodeGen ())
-> (Variable a -> abt '[] a -> CodeGen ()) -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ \Variable a
v abt '[] a
body' -> do
                 Ident
iterI <- Variable a -> CodeGen Ident
forall (a :: Hakaru). Variable a -> CodeGen Ident
createIdent Variable a
v
                 Sing 'HNat -> Ident -> CodeGen ()
forall (a :: Hakaru). Sing a -> Ident -> CodeGen ()
declare Sing 'HNat
SNat Ident
iterI

                 Ident
accI <- String -> CodeGen Ident
genIdent' String
"acc"
                 Sing a -> Ident -> CodeGen ()
forall (a :: Hakaru). Sing a -> Ident -> CodeGen ()
declare Sing a
semiTyp Ident
accI
                 Ident -> CExpr -> CodeGen ()
assign Ident
accI (case Sing a
semiTyp of
                                Sing a
SReal -> Float -> CExpr
floatE Float
1
                                Sing a
_     -> Integer -> CExpr
intE Integer
1)

                 let accVar :: CExpr
accVar  = Ident -> CExpr
CVar Ident
accI
                     iterVar :: CExpr
iterVar = Ident -> CExpr
CVar Ident
iterI

                 Either
  CBinaryOp
  (Sing Any, CExpr -> CodeGen (), CExpr -> CExpr -> CodeGen ())
-> CExpr -> CExpr -> CExpr -> CExpr -> CodeGen () -> CodeGen ()
forall (a :: Hakaru).
Either
  CBinaryOp
  (Sing a, CExpr -> CodeGen (), CExpr -> CExpr -> CodeGen ())
-> CExpr -> CExpr -> CExpr -> CExpr -> CodeGen () -> CodeGen ()
reductionCG (CBinaryOp
-> Either
     CBinaryOp
     (Sing Any, CExpr -> CodeGen (), CExpr -> CExpr -> CodeGen ())
forall a b. a -> Either a b
Left CBinaryOp
CMulOp)
                             CExpr
accVar
                             (CExpr
iterVar CExpr -> CExpr -> CExpr
.=. CExpr
loE)
                             (CExpr
iterVar CExpr -> CExpr -> CExpr
.<. CExpr
hiE)
                             (CUnaryOp -> CExpr -> CExpr
CUnary CUnaryOp
CPostIncOp CExpr
iterVar) (CodeGen () -> CodeGen ()) -> CodeGen () -> CodeGen ()
forall a b. (a -> b) -> a -> b
$
                    (CExpr -> CodeGen ()
putExprStat (CExpr -> CodeGen ()) -> (CExpr -> CExpr) -> CExpr -> CodeGen ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (CExpr
accVar CExpr -> CExpr -> CExpr
.*=.) (CExpr -> CodeGen ()) -> CodeGen CExpr -> CodeGen ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< abt '[] a -> CodeGen CExpr
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> CodeGen CExpr
flattenWithName abt '[] a
body')
                 CExpr -> CodeGen ()
putExprStat (CExpr
loc CExpr -> CExpr -> CExpr
.=. CExpr
accVar)
           CStat -> CodeGen ()
putStat (CStat -> CodeGen ()) -> CStat -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ String -> CStat
opComment String
"End Product"


--------------------
-- SCon Coercions --
--------------------
{- Helpers found by searching "Coercion Helpers" -}

flattenSCon (CoerceTo_ Coercion a a
ctyp) =
  \(abt vars a
e :* SArgs abt args
End) ->
    \CExpr
loc ->
      do CExpr
eE <- abt '[] a -> CodeGen CExpr
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> CodeGen CExpr
flattenWithName abt vars a
abt '[] a
e
         CExpr
cE <- Coercion a a -> CExpr -> CodeGen CExpr
forall (a :: Hakaru) (b :: Hakaru).
Coercion a b -> CExpr -> CodeGen CExpr
coerceToCG Coercion a a
ctyp CExpr
eE
         CExpr -> CodeGen ()
putExprStat (CExpr -> CodeGen ()) -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ CExpr
loc CExpr -> CExpr -> CExpr
.=. CExpr
cE

flattenSCon (UnsafeFrom_ Coercion a b
ctyp) =
  \(abt vars a
e :* SArgs abt args
End) ->
    \CExpr
loc ->
      do CExpr
eE <- abt '[] a -> CodeGen CExpr
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> CodeGen CExpr
flattenWithName abt vars a
abt '[] a
e
         CExpr
cE <- Coercion a b -> CExpr -> CodeGen CExpr
forall (a :: Hakaru) (b :: Hakaru).
Coercion a b -> CExpr -> CodeGen CExpr
coerceFromCG Coercion a b
ctyp CExpr
eE
         CExpr -> CodeGen ()
putExprStat (CExpr -> CodeGen ()) -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ CExpr
loc CExpr -> CExpr -> CExpr
.=. CExpr
cE

-----------------------------------
-- SCons in the Stochastic Monad --
-----------------------------------

flattenSCon (MeasureOp_ MeasureOp typs a
op) = MeasureOp typs a -> SArgs abt args -> CExpr -> CodeGen ()
forall (abt :: [Hakaru] -> Hakaru -> *) (typs :: [Hakaru])
       (args :: [([Hakaru], Hakaru)]) (a :: Hakaru).
(ABT Term abt, typs ~ UnLCs args, args ~ LCs typs) =>
MeasureOp typs a -> SArgs abt args -> CExpr -> CodeGen ()
flattenMeasureOp MeasureOp typs a
op

flattenSCon SCon args a
Dirac           =
  \(abt vars a
e :* SArgs abt args
End) ->
    \CExpr
loc ->
      do CExpr
sE <- abt '[] a -> String -> CodeGen CExpr
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> String -> CodeGen CExpr
flattenWithName' abt vars a
abt '[] a
e String
"samp"
         CExpr -> CodeGen ()
putExprStat (CExpr -> CodeGen ()) -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ CExpr -> CExpr
mdataWeight CExpr
loc CExpr -> CExpr -> CExpr
.=. (Float -> CExpr
floatE Float
0)
         CExpr -> CodeGen ()
putExprStat (CExpr -> CodeGen ()) -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ CExpr -> CExpr
mdataSample CExpr
loc CExpr -> CExpr -> CExpr
.=. CExpr
sE

flattenSCon SCon args a
MBind           =
  \(abt vars a
ma :* abt vars a
b :* SArgs abt args
End) ->
    \CExpr
loc ->
      abt '[a] a -> (Variable a -> abt '[] a -> CodeGen ()) -> CodeGen ()
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (x :: k) (xs :: [k]) (a :: k) r.
ABT syn abt =>
abt (x : xs) a -> (Variable x -> abt xs a -> r) -> r
caseBind abt vars a
abt '[a] a
b ((Variable a -> abt '[] a -> CodeGen ()) -> CodeGen ())
-> (Variable a -> abt '[] a -> CodeGen ()) -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ \v :: Variable a
v@(Variable Text
_ Nat
_ Sing a
typ) abt '[] a
mb ->
        do -- first
           CExpr
mE <- abt '[] a -> String -> CodeGen CExpr
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> String -> CodeGen CExpr
flattenWithName' abt vars a
abt '[] a
ma String
"m"

           -- assign that sample to var
           Ident
vId <- Variable a -> CodeGen Ident
forall (a :: Hakaru). Variable a -> CodeGen Ident
createIdent Variable a
v
           Sing a -> Ident -> CodeGen ()
forall (a :: Hakaru). Sing a -> Ident -> CodeGen ()
declare Sing a
typ Ident
vId
           Ident -> CExpr -> CodeGen ()
assign Ident
vId (CExpr -> CExpr
mdataSample CExpr
mE)
           abt '[] a -> CExpr -> CodeGen ()
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> CExpr -> CodeGen ()
flattenABT abt '[] a
mb CExpr
loc
           CExpr -> CodeGen ()
putExprStat (CExpr -> CodeGen ()) -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ CExpr -> CExpr
mdataWeight CExpr
loc CExpr -> CExpr -> CExpr
.+=. (CExpr -> CExpr
mdataWeight CExpr
mE)

-- for now plats make use of a global sample
flattenSCon SCon args a
Plate           =
  \(abt vars a
size :* abt vars a
b :* SArgs abt args
End) ->
    \CExpr
loc ->
      abt '[ 'HNat] a
-> (Variable 'HNat -> abt '[] a -> CodeGen ()) -> CodeGen ()
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (x :: k) (xs :: [k]) (a :: k) r.
ABT syn abt =>
abt (x : xs) a -> (Variable x -> abt xs a -> r) -> r
caseBind abt vars a
abt '[ 'HNat] a
b ((Variable 'HNat -> abt '[] a -> CodeGen ()) -> CodeGen ())
-> (Variable 'HNat -> abt '[] a -> CodeGen ()) -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ \Variable 'HNat
v abt '[] a
body ->
        do CExpr
sizeE <- abt '[] a -> String -> CodeGen CExpr
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> String -> CodeGen CExpr
flattenWithName' abt vars a
abt '[] a
size String
"s"
           Bool
isMM <- 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
           Bool -> CodeGen () -> CodeGen ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Bool -> Bool
not Bool
isMM) (String -> CodeGen ()
forall a. HasCallStack => String -> a
error String
"plate will leak memory without the '-g' flag and boehm-gc")
           CExpr -> CodeGen ()
putExprStat (CExpr -> CodeGen ()) -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ (CExpr -> CExpr
arraySize (CExpr -> CExpr) -> (CExpr -> CExpr) -> CExpr -> CExpr
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CExpr -> CExpr
mdataSample (CExpr -> CExpr) -> CExpr -> CExpr
forall a b. (a -> b) -> a -> b
$ CExpr
loc) CExpr -> CExpr -> CExpr
.=. CExpr
sizeE
           CExpr -> CExpr -> Sing a -> CodeGen ()
forall (a :: Hakaru). CExpr -> CExpr -> Sing a -> CodeGen ()
putMallocStat (CExpr -> CExpr
arrayData (CExpr -> CExpr) -> (CExpr -> CExpr) -> CExpr -> CExpr
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CExpr -> CExpr
mdataSample (CExpr -> CExpr) -> CExpr -> CExpr
forall a b. (a -> b) -> a -> b
$ CExpr
loc) CExpr
sizeE (abt '[] a -> Sing a
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> Sing a
typeOf abt '[] a
body)

           Ident
weightId <- String -> CodeGen Ident
genIdent' String
"w"
           Sing 'HProb -> Ident -> CodeGen ()
forall (a :: Hakaru). Sing a -> Ident -> CodeGen ()
declare Sing 'HProb
SProb Ident
weightId
           let weightE :: CExpr
weightE = Ident -> CExpr
CVar Ident
weightId
           Ident -> CExpr -> CodeGen ()
assign Ident
weightId (Float -> CExpr
floatE Float
0)

           Ident
itId <- Variable 'HNat -> CodeGen Ident
forall (a :: Hakaru). Variable a -> CodeGen Ident
createIdent Variable 'HNat
v
           Sing 'HNat -> Ident -> CodeGen ()
forall (a :: Hakaru). Sing a -> Ident -> CodeGen ()
declare Sing 'HNat
SNat Ident
itId
           let itE :: CExpr
itE = Ident -> CExpr
CVar Ident
itId
               currInd :: CExpr
currInd  = CExpr -> CExpr -> CExpr
index (CExpr -> CExpr
arrayData (CExpr -> CExpr) -> (CExpr -> CExpr) -> CExpr -> CExpr
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CExpr -> CExpr
mdataSample (CExpr -> CExpr) -> CExpr -> CExpr
forall a b. (a -> b) -> a -> b
$ CExpr
loc) CExpr
itE

           Ident
sampId <- String -> CodeGen Ident
genIdent' String
"samp"
           Sing a -> Ident -> CodeGen ()
forall (a :: Hakaru). Sing a -> Ident -> CodeGen ()
declare (abt '[] a -> Sing a
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> Sing a
typeOf (abt '[] a -> Sing a) -> abt '[] a -> Sing a
forall a b. (a -> b) -> a -> b
$ abt '[] a
body) Ident
sampId
           let sampE :: CExpr
sampE = Ident -> CExpr
CVar Ident
sampId

           Either
  CBinaryOp
  (Sing Any, CExpr -> CodeGen (), CExpr -> CExpr -> CodeGen ())
-> CExpr -> CExpr -> CExpr -> CExpr -> CodeGen () -> CodeGen ()
forall (a :: Hakaru).
Either
  CBinaryOp
  (Sing a, CExpr -> CodeGen (), CExpr -> CExpr -> CodeGen ())
-> CExpr -> CExpr -> CExpr -> CExpr -> CodeGen () -> CodeGen ()
reductionCG (CBinaryOp
-> Either
     CBinaryOp
     (Sing Any, CExpr -> CodeGen (), CExpr -> CExpr -> CodeGen ())
forall a b. a -> Either a b
Left CBinaryOp
CAddOp)
                       CExpr
weightE
                       (CExpr
itE CExpr -> CExpr -> CExpr
.=. (Integer -> CExpr
intE Integer
0))
                       (CExpr
itE CExpr -> CExpr -> CExpr
.<. CExpr
sizeE)
                       (CUnaryOp -> CExpr -> CExpr
CUnary CUnaryOp
CPostIncOp CExpr
itE)
                       (do abt '[] a -> CExpr -> CodeGen ()
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> CExpr -> CodeGen ()
flattenABT abt '[] a
body CExpr
sampE
                           CExpr -> CodeGen ()
putExprStat (CExpr
currInd CExpr -> CExpr -> CExpr
.=. (CExpr -> CExpr
mdataSample CExpr
sampE))
                           CExpr -> CodeGen ()
putExprStat (CExpr
weightE CExpr -> CExpr -> CExpr
.+=. (CExpr -> CExpr
mdataWeight CExpr
sampE)))

           CExpr -> CodeGen ()
putExprStat (CExpr -> CodeGen ()) -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ CExpr -> CExpr
mdataWeight CExpr
loc CExpr -> CExpr -> CExpr
.=. CExpr
weightE

-----------------------------------
-- SCon's that arent implemented --
-----------------------------------

flattenSCon SCon args a
x               = \SArgs abt args
_ -> \CExpr
_ -> String -> CodeGen ()
forall a. HasCallStack => String -> a
error (String -> CodeGen ()) -> String -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ String
"TODO: flattenSCon: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ SCon args a -> String
forall a. Show a => a -> String
show SCon args a
x



--------------------------------------------------------------------------------
--                                 NaryOps                                    --
--------------------------------------------------------------------------------

flattenNAryOp :: ABT Term abt
              => NaryOp a
              -> S.Seq (abt '[] a)
              -> (CExpr -> CodeGen ())
flattenNAryOp :: NaryOp a -> Seq (abt '[] a) -> CExpr -> CodeGen ()
flattenNAryOp NaryOp a
op Seq (abt '[] a)
args =
  \CExpr
loc ->
    do Seq CExpr
es <- (abt '[] a -> CodeGen CExpr)
-> Seq (abt '[] a) -> StateT CG Identity (Seq CExpr)
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
T.mapM abt '[] a -> CodeGen CExpr
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> CodeGen CExpr
flattenWithName Seq (abt '[] a)
args
       case NaryOp a
op of
         NaryOp a
And -> NaryOp a -> Seq CExpr -> CExpr -> CodeGen ()
forall (a :: Hakaru). NaryOp a -> Seq CExpr -> CExpr -> CodeGen ()
boolNaryOp NaryOp a
op Seq CExpr
es CExpr
loc
         NaryOp a
Or  -> NaryOp a -> Seq CExpr -> CExpr -> CodeGen ()
forall (a :: Hakaru). NaryOp a -> Seq CExpr -> CExpr -> CodeGen ()
boolNaryOp NaryOp a
op Seq CExpr
es CExpr
loc
         NaryOp a
Xor -> NaryOp a -> Seq CExpr -> CExpr -> CodeGen ()
forall (a :: Hakaru). NaryOp a -> Seq CExpr -> CExpr -> CodeGen ()
boolNaryOp NaryOp a
op Seq CExpr
es CExpr
loc
         NaryOp a
Iff -> NaryOp a -> Seq CExpr -> CExpr -> CodeGen ()
forall (a :: Hakaru). NaryOp a -> Seq CExpr -> CExpr -> CodeGen ()
boolNaryOp NaryOp a
op Seq CExpr
es CExpr
loc
         (Sum HSemiring a
HSemiring_Prob) -> Seq CExpr -> CExpr -> CodeGen ()
logSumExpCG Seq CExpr
es CExpr
loc
         NaryOp a
_ -> let opE :: CExpr
opE = (CExpr -> CExpr -> CExpr) -> CExpr -> Seq CExpr -> CExpr
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
F.foldr (NaryOp a -> CExpr -> CExpr -> CExpr
forall (a :: Hakaru). NaryOp a -> CExpr -> CExpr -> CExpr
binaryOp NaryOp a
op) (Seq CExpr -> Int -> CExpr
forall a. Seq a -> Int -> a
S.index Seq CExpr
es Int
0) (Int -> Seq CExpr -> Seq CExpr
forall a. Int -> Seq a -> Seq a
S.drop Int
1 Seq CExpr
es)
              in  CExpr -> CodeGen ()
putExprStat (CExpr
loc CExpr -> CExpr -> CExpr
.=. CExpr
opE)

  where boolNaryOp :: NaryOp a -> Seq CExpr -> CExpr -> CodeGen ()
boolNaryOp NaryOp a
op' Seq CExpr
es' CExpr
loc' =
          let indexOf :: CExpr -> CExpr
indexOf CExpr
x = CExpr -> Ident -> Bool -> CExpr
CMember CExpr
x (String -> Ident
Ident String
"index") Bool
True
              es'' :: Seq CExpr
es''      = (CExpr -> CExpr) -> Seq CExpr -> Seq CExpr
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap CExpr -> CExpr
indexOf Seq CExpr
es'
              expr :: CExpr
expr      = (CExpr -> CExpr -> CExpr) -> CExpr -> Seq CExpr -> CExpr
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
F.foldr (NaryOp a -> CExpr -> CExpr -> CExpr
forall (a :: Hakaru). NaryOp a -> CExpr -> CExpr -> CExpr
binaryOp NaryOp a
op')
                                  (Seq CExpr -> Int -> CExpr
forall a. Seq a -> Int -> a
S.index Seq CExpr
es'' Int
0)
                                  (Int -> Seq CExpr -> Seq CExpr
forall a. Int -> Seq a -> Seq a
S.drop Int
1 Seq CExpr
es'')
          in  CExpr -> CodeGen ()
putExprStat ((CExpr -> CExpr
indexOf CExpr
loc') CExpr -> CExpr -> CExpr
.=. CExpr
expr)


--------------------------------------------------------------------------------
--                                  Literals                                  --
--------------------------------------------------------------------------------

flattenLit
  :: Literal a
  -> (CExpr -> CodeGen ())
flattenLit :: Literal a -> CExpr -> CodeGen ()
flattenLit Literal a
lit =
  \CExpr
loc ->
    case Literal a
lit of
      (LNat Natural
x)  -> CExpr -> CodeGen ()
putExprStat (CExpr -> CodeGen ()) -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ CExpr
loc CExpr -> CExpr -> CExpr
.=. (Integer -> CExpr
intE (Integer -> CExpr) -> Integer -> CExpr
forall a b. (a -> b) -> a -> b
$ Natural -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral Natural
x)
      (LInt Integer
x)  -> CExpr -> CodeGen ()
putExprStat (CExpr -> CodeGen ()) -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ CExpr
loc CExpr -> CExpr -> CExpr
.=. (Integer -> CExpr
intE Integer
x)
      (LReal Rational
x) -> CExpr -> CodeGen ()
putExprStat (CExpr -> CodeGen ()) -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ CExpr
loc CExpr -> CExpr -> CExpr
.=. (Float -> CExpr
floatE (Float -> CExpr) -> Float -> CExpr
forall a b. (a -> b) -> a -> b
$ Rational -> Float
forall a. Fractional a => Rational -> a
fromRational Rational
x)
      (LProb NonNegativeRational
x) -> let rat :: Rational
rat = NonNegativeRational -> Rational
fromNonNegativeRational NonNegativeRational
x
                       x' :: Float
x'  = (Integer -> Float
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Integer -> Float) -> Integer -> Float
forall a b. (a -> b) -> a -> b
$ Rational -> Integer
forall a. Ratio a -> a
numerator Rational
rat)
                           Float -> Float -> Float
forall a. Fractional a => a -> a -> a
/ (Integer -> Float
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Integer -> Float) -> Integer -> Float
forall a b. (a -> b) -> a -> b
$ Rational -> Integer
forall a. Ratio a -> a
denominator Rational
rat)
                       xE :: CExpr
xE  = CExpr -> CExpr
log1pE (Float -> CExpr
floatE Float
x' CExpr -> CExpr -> CExpr
.-. Integer -> CExpr
intE Integer
1)
                   in CExpr -> CodeGen ()
putExprStat (CExpr
loc CExpr -> CExpr -> CExpr
.=. CExpr
xE)

--------------------------------------------------------------------------------
--                                Array and ArrayOps                          --
--------------------------------------------------------------------------------


flattenArray
  :: (ABT Term abt)
  => (abt '[] 'HNat)
  -> (abt '[ 'HNat ] a)
  -> (CExpr -> CodeGen ())
flattenArray :: abt '[] 'HNat -> abt '[ 'HNat] a -> CExpr -> CodeGen ()
flattenArray abt '[] 'HNat
arity abt '[ 'HNat] a
body =
  \CExpr
loc ->
    abt '[ 'HNat] a
-> (Variable 'HNat -> abt '[] a -> CodeGen ()) -> CodeGen ()
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (x :: k) (xs :: [k]) (a :: k) r.
ABT syn abt =>
abt (x : xs) a -> (Variable x -> abt xs a -> r) -> r
caseBind abt '[ 'HNat] a
body ((Variable 'HNat -> abt '[] a -> CodeGen ()) -> CodeGen ())
-> (Variable 'HNat -> abt '[] a -> CodeGen ()) -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ \Variable 'HNat
v abt '[] a
body' -> do
      let arityE :: CExpr
arityE = CExpr -> CExpr
arraySize CExpr
loc
          dataE :: CExpr
dataE  = CExpr -> CExpr
arrayData CExpr
loc
          typ :: Sing a
typ    = abt '[] a -> Sing a
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> Sing a
typeOf abt '[] a
body'

      abt '[] 'HNat -> CExpr -> CodeGen ()
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> CExpr -> CodeGen ()
flattenABT abt '[] 'HNat
arity CExpr
arityE

      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
      CExpr -> CodeGen ()
putExprStat (CExpr -> CodeGen ()) -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$   CExpr
dataE
                  CExpr -> CExpr -> CExpr
.=. (CTypeName -> CExpr -> CExpr
CCast ([CTypeSpec] -> Bool -> CTypeName
CTypeName (Sing a -> [CTypeSpec]
forall (a :: Hakaru). Sing a -> [CTypeSpec]
buildType Sing a
typ) Bool
True)
                             (CExpr -> CExpr
malloc' (CExpr
arityE CExpr -> CExpr -> CExpr
.*. (CTypeName -> CExpr
CSizeOfType ([CTypeSpec] -> Bool -> CTypeName
CTypeName (Sing a -> [CTypeSpec]
forall (a :: Hakaru). Sing a -> [CTypeSpec]
buildType Sing a
typ) Bool
False)))))

      Ident
itId  <- Variable 'HNat -> CodeGen Ident
forall (a :: Hakaru). Variable a -> CodeGen Ident
createIdent Variable 'HNat
v
      Sing 'HNat -> Ident -> CodeGen ()
forall (a :: Hakaru). Sing a -> Ident -> CodeGen ()
declare Sing 'HNat
SNat Ident
itId
      let itE :: CExpr
itE     = Ident -> CExpr
CVar Ident
itId
          currInd :: CExpr
currInd = CExpr -> CExpr -> CExpr
index CExpr
dataE CExpr
itE

      CStat -> CodeGen ()
putStat (CStat -> CodeGen ()) -> CStat -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ String -> CStat
opComment String
"Begin Array"
      CExpr -> CExpr -> CExpr -> CodeGen () -> CodeGen ()
forCG (CExpr
itE CExpr -> CExpr -> CExpr
.=. (Integer -> CExpr
intE Integer
0))
            (CExpr
itE CExpr -> CExpr -> CExpr
.<. CExpr
arityE)
            (CUnaryOp -> CExpr -> CExpr
CUnary CUnaryOp
CPostIncOp CExpr
itE)
            (abt '[] a -> CExpr -> CodeGen ()
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> CExpr -> CodeGen ()
flattenABT abt '[] a
body' CExpr
currInd)
      CStat -> CodeGen ()
putStat (CStat -> CodeGen ()) -> CStat -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ String -> CStat
opComment String
"End Array"

flattenArrayLiteral
  :: ( ABT Term abt )
  => [abt '[] a]
  -> (CExpr -> CodeGen ())
flattenArrayLiteral :: [abt '[] a] -> CExpr -> CodeGen ()
flattenArrayLiteral [abt '[] a]
es =
  \CExpr
loc -> do
    Ident
arrId <- CodeGen Ident
genIdent
    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 arity :: Integer
arity = Int -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Integer) -> ([abt '[] a] -> Int) -> [abt '[] a] -> Integer
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [abt '[] a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([abt '[] a] -> Integer) -> [abt '[] a] -> Integer
forall a b. (a -> b) -> a -> b
$ [abt '[] a]
es
        typ :: Sing a
typ   = abt '[] a -> Sing a
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> Sing a
typeOf (abt '[] a -> Sing a)
-> ([abt '[] a] -> abt '[] a) -> [abt '[] a] -> Sing a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [abt '[] a] -> abt '[] a
forall a. [a] -> a
head ([abt '[] a] -> Sing a) -> [abt '[] a] -> Sing a
forall a b. (a -> b) -> a -> b
$ [abt '[] a]
es
        arrE :: CExpr
arrE = Ident -> CExpr
CVar Ident
arrId
        malloc' :: CExpr -> CExpr
malloc' = if Bool
isManagedMem then CExpr -> CExpr
gcMalloc else CExpr -> CExpr
mallocE

    Sing ('HArray a) -> Ident -> CodeGen ()
forall (a :: Hakaru). Sing a -> Ident -> CodeGen ()
declare (Sing a -> Sing ('HArray a)
forall (a :: Hakaru). Sing a -> Sing ('HArray a)
SArray Sing a
typ) Ident
arrId
    CExpr -> CodeGen ()
putExprStat (CExpr -> CodeGen ()) -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$   (CExpr -> CExpr
arrayData CExpr
arrE)
                CExpr -> CExpr -> CExpr
.=. (CTypeName -> CExpr -> CExpr
CCast ([CTypeSpec] -> Bool -> CTypeName
CTypeName (Sing a -> [CTypeSpec]
forall (a :: Hakaru). Sing a -> [CTypeSpec]
buildType Sing a
typ) Bool
True)
                           (CExpr -> CExpr
malloc' ((Integer -> CExpr
intE Integer
arity) CExpr -> CExpr -> CExpr
.*. (CTypeName -> CExpr
CSizeOfType ([CTypeSpec] -> Bool -> CTypeName
CTypeName (Sing a -> [CTypeSpec]
forall (a :: Hakaru). Sing a -> [CTypeSpec]
buildType Sing a
typ) Bool
False)))))

    CExpr -> CodeGen ()
putExprStat (CExpr -> CodeGen ()) -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ CExpr -> CExpr
arraySize CExpr
arrE CExpr -> CExpr -> CExpr
.=. (Integer -> CExpr
intE Integer
arity)
    [CodeGen ()] -> CodeGen ()
forall (t :: * -> *) (m :: * -> *) a.
(Foldable t, Monad m) =>
t (m a) -> m ()
sequence_ ([CodeGen ()] -> CodeGen ())
-> ((Integer, [CodeGen ()]) -> [CodeGen ()])
-> (Integer, [CodeGen ()])
-> CodeGen ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Integer, [CodeGen ()]) -> [CodeGen ()]
forall a b. (a, b) -> b
snd ((Integer, [CodeGen ()]) -> CodeGen ())
-> (Integer, [CodeGen ()]) -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ ((Integer, [CodeGen ()]) -> abt '[] a -> (Integer, [CodeGen ()]))
-> (Integer, [CodeGen ()])
-> [abt '[] a]
-> (Integer, [CodeGen ()])
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (\(Integer
i,[CodeGen ()]
acc) abt '[] a
e -> (Integer -> Integer
forall a. Enum a => a -> a
succ Integer
i,(abt '[] a -> Integer -> CExpr -> CodeGen ()
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> Integer -> CExpr -> CodeGen ()
assignIndex abt '[] a
e Integer
i CExpr
arrE)CodeGen () -> [CodeGen ()] -> [CodeGen ()]
forall a. a -> [a] -> [a]
:[CodeGen ()]
acc))
                            (Integer
0,[])
                            [abt '[] a]
es
    CExpr -> CodeGen ()
putExprStat (CExpr -> CodeGen ()) -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ CExpr
loc CExpr -> CExpr -> CExpr
.=. CExpr
arrE
  where assignIndex
          :: ( ABT Term abt )
          => abt '[] a
          -> Integer
          -> (CExpr -> CodeGen ())
        assignIndex :: abt '[] a -> Integer -> CExpr -> CodeGen ()
assignIndex abt '[] a
e Integer
i CExpr
loc = do
          CExpr
eE <- abt '[] a -> CodeGen CExpr
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> CodeGen CExpr
flattenWithName abt '[] a
e
          CExpr -> CodeGen ()
putExprStat (CExpr -> CodeGen ()) -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ CExpr -> CExpr
indirect ((CExpr -> CExpr
arrayData CExpr
loc) CExpr -> CExpr -> CExpr
.+. (Integer -> CExpr
intE Integer
i)) CExpr -> CExpr -> CExpr
.=. CExpr
eE

--------------
-- ArrayOps --
--------------

flattenArrayOp
  :: ( ABT Term abt
     , typs ~ UnLCs args
     , args ~ LCs typs
     )
  => ArrayOp typs a
  -> SArgs abt args
  -> (CExpr -> CodeGen ())


flattenArrayOp :: ArrayOp typs a -> SArgs abt args -> CExpr -> CodeGen ()
flattenArrayOp (Index Sing a
_)  =
  \(abt vars a
arr :* abt vars a
ind :* SArgs abt args
End) ->
    \CExpr
loc ->
      do CExpr
indE <- abt '[] a -> CodeGen CExpr
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> CodeGen CExpr
flattenWithName abt vars a
abt '[] a
ind
         CExpr
arrE <- abt '[] a -> CodeGen CExpr
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> CodeGen CExpr
flattenWithName abt vars a
abt '[] a
arr
         let valE :: CExpr
valE = CExpr -> CExpr -> CExpr
index (CExpr -> Ident -> Bool -> CExpr
CMember CExpr
arrE (String -> Ident
Ident String
"data") Bool
True) CExpr
indE
         CExpr -> CodeGen ()
putExprStat (CExpr
loc CExpr -> CExpr -> CExpr
.=. CExpr
valE)

flattenArrayOp (Size Sing a
_)   =
  \(abt vars a
arr :* SArgs abt args
End) ->
    \CExpr
loc ->
      do CExpr
arrE <- abt '[] a -> CodeGen CExpr
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> CodeGen CExpr
flattenWithName abt vars a
abt '[] a
arr
         CExpr -> CodeGen ()
putExprStat (CExpr
loc CExpr -> CExpr -> CExpr
.=. (CExpr -> Ident -> Bool -> CExpr
CMember CExpr
arrE (String -> Ident
Ident String
"size") Bool
True))

flattenArrayOp (Reduce Sing a
_) = String -> SArgs abt args -> CExpr -> CodeGen ()
forall a. HasCallStack => String -> a
error String
"TODO: flattenArrayOp"
  -- \(fun :* base :* arr :* End) ->
  -- do funE  <- flattenABT fun
  --    baseE <- flattenABT base
  --    arrE  <- flattenABT arr
  --    accI  <- genIdent' "acc"
  --    iterI <- genIdent' "iter"

  --    let sizeE = CMember arrE (Ident "size") True
  --        iterE = CVar iterI
  --        accE  = CVar accI
  --        cond  = iterE .<. sizeE
  --        inc   = CUnary CPostIncOp iterE

  --    declare (typeOf base) accI
  --    declare SInt iterI
  --    assign accI baseE
  --    forCG (iterE .=. (intE 0)) cond inc $
  --      assign accI $ CCall funE [accE]

  --    return accE

--------------------------------------------------------------------------------
--                              Bucket and Reducers                           --
--------------------------------------------------------------------------------
{- Declarations for buckets -
  since we will have some product of monoids we need unique names for each one.

  Ex:
    bucket i from 0 to 100:
      fanout(add(\_ -> 1),add(\_ -> 2))

  we will need to keep track of two ints:

  int x;
  int y;
  for (i = 0; i < 100; i++) {
    x += 1;
    y += 2;
  }

  Summary objects are nested pairs, e.g. pair(nat,pair(real,array(nat)))
-}

flattenBucket
  :: (ABT Term abt)
  => abt '[] 'HNat
  -> abt '[] 'HNat
  -> Reducer abt '[] a
  -> (CExpr -> CodeGen ())
flattenBucket :: abt '[] 'HNat
-> abt '[] 'HNat -> Reducer abt '[] a -> CExpr -> CodeGen ()
flattenBucket abt '[] 'HNat
lo abt '[] 'HNat
hi Reducer abt '[] a
red = \CExpr
loc -> do
    CStat -> CodeGen ()
putStat (CStat -> CodeGen ()) -> CStat -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ String -> CStat
opComment String
"Begin Bucket"
    CExpr
loE <- abt '[] 'HNat -> String -> CodeGen CExpr
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> String -> CodeGen CExpr
flattenWithName' abt '[] 'HNat
lo String
"lo"
    CExpr
hiE <- abt '[] 'HNat -> String -> CodeGen CExpr
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> String -> CodeGen CExpr
flattenWithName' abt '[] 'HNat
hi String
"hi"
    Ident
itId <- String -> CodeGen Ident
genIdent' String
"it"
    Sing 'HNat -> Ident -> CodeGen ()
forall (a :: Hakaru). Sing a -> Ident -> CodeGen ()
declare Sing 'HNat
SNat Ident
itId
    let itE :: CExpr
itE = Ident -> CExpr
CVar Ident
itId
    Reducer abt '[] a -> CExpr -> CodeGen ()
forall (abt :: [Hakaru] -> Hakaru -> *) (xs :: [Hakaru])
       (a :: Hakaru).
ABT Term abt =>
Reducer abt xs a -> CExpr -> CodeGen ()
initRed Reducer abt '[] a
red CExpr
loc
    Bool
isPar <- 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
    -- declare special functions for combining threads. This doesn't completely
    -- work now, because these need to capture free variables.
    Either
  CBinaryOp
  (Sing a, CExpr -> CodeGen (), CExpr -> CExpr -> CodeGen ())
-> CExpr -> CExpr -> CExpr -> CExpr -> CodeGen () -> CodeGen ()
forall (a :: Hakaru).
Either
  CBinaryOp
  (Sing a, CExpr -> CodeGen (), CExpr -> CExpr -> CodeGen ())
-> CExpr -> CExpr -> CExpr -> CExpr -> CodeGen () -> CodeGen ()
reductionCG ((Sing a, CExpr -> CodeGen (), CExpr -> CExpr -> CodeGen ())
-> Either
     CBinaryOp
     (Sing a, CExpr -> CodeGen (), CExpr -> CExpr -> CodeGen ())
forall a b. b -> Either a b
Right ( Reducer abt '[] a -> Sing a
forall (abt :: [Hakaru] -> Hakaru -> *) (xs :: [Hakaru])
       (a :: Hakaru).
Reducer abt xs a -> Sing a
typeOfReducer Reducer abt '[] a
red
                       , \CExpr
e -> CodeGen () -> CodeGen ()
forall a. CodeGen a -> CodeGen a
seqDo (  Reducer abt '[] a -> CExpr -> CodeGen ()
forall (abt :: [Hakaru] -> Hakaru -> *) (xs :: [Hakaru])
       (a :: Hakaru).
ABT Term abt =>
Reducer abt xs a -> CExpr -> CodeGen ()
initRed Reducer abt '[] a
red (CExpr -> CExpr
indirect CExpr
e)
                                     CodeGen () -> CodeGen () -> CodeGen ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> CStat -> CodeGen ()
putStat (Maybe CExpr -> CStat
CReturn Maybe CExpr
forall a. Maybe a
Nothing))
                       , \CExpr
a CExpr
b -> CodeGen () -> CodeGen ()
forall a. CodeGen a -> CodeGen a
seqDo (  Reducer abt '[] a -> CExpr -> CExpr -> CodeGen ()
forall (abt :: [Hakaru] -> Hakaru -> *) (xs :: [Hakaru])
       (a :: Hakaru).
ABT Term abt =>
Reducer abt xs a -> CExpr -> CExpr -> CodeGen ()
mulRed Reducer abt '[] a
red (CExpr -> CExpr
indirect CExpr
a) (CExpr -> CExpr
indirect CExpr
b)
                                       CodeGen () -> CodeGen () -> CodeGen ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> CStat -> CodeGen ()
putStat (Maybe CExpr -> CStat
CReturn Maybe CExpr
forall a. Maybe a
Nothing))))
                CExpr
loc
                (CExpr
itE CExpr -> CExpr -> CExpr
.=. CExpr
loE)
                (CExpr
itE CExpr -> CExpr -> CExpr
.<. CExpr
hiE)
                (CUnaryOp -> CExpr -> CExpr
CUnary CUnaryOp
CPostIncOp CExpr
itE)
                (Bool -> Reducer abt '[] a -> CExpr -> CExpr -> CodeGen ()
forall (abt :: [Hakaru] -> Hakaru -> *) (xs :: [Hakaru])
       (a :: Hakaru).
ABT Term abt =>
Bool -> Reducer abt xs a -> CExpr -> CExpr -> CodeGen ()
accumRed Bool
isPar Reducer abt '[] a
red CExpr
itE CExpr
loc)
    CStat -> CodeGen ()
putStat (CStat -> CodeGen ()) -> CStat -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ String -> CStat
opComment String
"End Bucket"

  where initRed
          :: (ABT Term abt)
          => Reducer abt xs a
          -> (CExpr -> CodeGen ())
        initRed :: Reducer abt xs a -> CExpr -> CodeGen ()
initRed Reducer abt xs a
mr = \CExpr
loc ->
          case Reducer abt xs a
mr of
            (Red_Fanout Reducer abt xs a
mr1 Reducer abt xs b
mr2) -> Reducer abt xs a -> CExpr -> CodeGen ()
forall (abt :: [Hakaru] -> Hakaru -> *) (xs :: [Hakaru])
       (a :: Hakaru).
ABT Term abt =>
Reducer abt xs a -> CExpr -> CodeGen ()
initRed Reducer abt xs a
mr1 (CExpr -> CExpr
datumFst CExpr
loc)
                                 CodeGen () -> CodeGen () -> CodeGen ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Reducer abt xs b -> CExpr -> CodeGen ()
forall (abt :: [Hakaru] -> Hakaru -> *) (xs :: [Hakaru])
       (a :: Hakaru).
ABT Term abt =>
Reducer abt xs a -> CExpr -> CodeGen ()
initRed Reducer abt xs b
mr2 (CExpr -> CExpr
datumSnd CExpr
loc)
            (Red_Split abt ('HNat : xs) HBool
_ Reducer abt xs a
mr1 Reducer abt xs b
mr2) -> Reducer abt xs a -> CExpr -> CodeGen ()
forall (abt :: [Hakaru] -> Hakaru -> *) (xs :: [Hakaru])
       (a :: Hakaru).
ABT Term abt =>
Reducer abt xs a -> CExpr -> CodeGen ()
initRed Reducer abt xs a
mr1 (CExpr -> CExpr
datumFst CExpr
loc)
                                  CodeGen () -> CodeGen () -> CodeGen ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Reducer abt xs b -> CExpr -> CodeGen ()
forall (abt :: [Hakaru] -> Hakaru -> *) (xs :: [Hakaru])
       (a :: Hakaru).
ABT Term abt =>
Reducer abt xs a -> CExpr -> CodeGen ()
initRed Reducer abt xs b
mr2 (CExpr -> CExpr
datumSnd CExpr
loc)
            (Red_Index abt xs 'HNat
s abt ('HNat : xs) 'HNat
_ Reducer abt ('HNat : xs) a
body) ->
              let (List1 Variable xs
vs,abt '[] 'HNat
s') = abt xs 'HNat -> (List1 Variable xs, abt '[] 'HNat)
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (xs :: [k]) (a :: k).
ABT syn abt =>
abt xs a -> (List1 Variable xs, abt '[] a)
caseBinds abt xs 'HNat
s
                  btyp :: Sing a
btyp     = Reducer abt ('HNat : xs) a -> Sing a
forall (abt :: [Hakaru] -> Hakaru -> *) (xs :: [Hakaru])
       (a :: Hakaru).
Reducer abt xs a -> Sing a
typeOfReducer Reducer abt ('HNat : xs) a
body in
                do [CodeGen ()] -> CodeGen ()
forall (t :: * -> *) (m :: * -> *) a.
(Foldable t, Monad m) =>
t (m a) -> m ()
sequence_ ([CodeGen ()] -> CodeGen ())
-> (List1 Variable xs -> [CodeGen ()])
-> List1 Variable xs
-> CodeGen ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (forall (i :: Hakaru). Variable i -> [CodeGen ()])
-> List1 Variable xs -> [CodeGen ()]
forall k1 k2 (f :: (k1 -> *) -> k2 -> *) m (a :: k1 -> *)
       (j :: k2).
(Foldable11 f, Monoid m) =>
(forall (i :: k1). a i -> m) -> f a j -> m
foldMap11
                     (\Variable i
v' -> case Variable i
v' of
                       (Variable _ _ typ') ->
                         [Sing i -> Ident -> CodeGen ()
forall (a :: Hakaru). Sing a -> Ident -> CodeGen ()
declare Sing i
typ' (Ident -> CodeGen ()) -> CodeGen Ident -> CodeGen ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Variable i -> CodeGen Ident
forall (a :: Hakaru). Variable a -> CodeGen Ident
createIdent Variable i
v'])
                     (List1 Variable xs -> CodeGen ())
-> List1 Variable xs -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ List1 Variable xs
vs
                   CExpr
sE <- abt '[] 'HNat -> String -> CodeGen CExpr
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> String -> CodeGen CExpr
flattenWithName' abt '[] 'HNat
s' String
"red_size"
                   CExpr -> CodeGen ()
putExprStat (CExpr -> CodeGen ()) -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ CExpr -> CExpr
arraySize CExpr
loc CExpr -> CExpr -> CExpr
.=. CExpr
sE
                   CExpr -> CExpr -> Sing a -> CodeGen ()
forall (a :: Hakaru). CExpr -> CExpr -> Sing a -> CodeGen ()
putMallocStat (CExpr -> CExpr
arrayData CExpr
loc) CExpr
sE Sing a
btyp
                   Ident
itId  <- CodeGen Ident
genIdent
                   Sing 'HNat -> Ident -> CodeGen ()
forall (a :: Hakaru). Sing a -> Ident -> CodeGen ()
declare Sing 'HNat
SNat Ident
itId
                   let itE :: CExpr
itE = Ident -> CExpr
CVar Ident
itId
                   CExpr -> CExpr -> CExpr -> CodeGen () -> CodeGen ()
forCG (CExpr
itE CExpr -> CExpr -> CExpr
.=. (Integer -> CExpr
intE Integer
0))
                         (CExpr
itE CExpr -> CExpr -> CExpr
.<. CExpr
sE)
                         (CUnaryOp -> CExpr -> CExpr
CUnary CUnaryOp
CPostIncOp CExpr
itE)
                         (Reducer abt ('HNat : xs) a -> CExpr -> CodeGen ()
forall (abt :: [Hakaru] -> Hakaru -> *) (xs :: [Hakaru])
       (a :: Hakaru).
ABT Term abt =>
Reducer abt xs a -> CExpr -> CodeGen ()
initRed Reducer abt ('HNat : xs) a
body (CExpr -> CExpr -> CExpr
index (CExpr -> CExpr
arrayData CExpr
loc) CExpr
itE))
            Reducer abt xs a
Red_Nop -> () -> CodeGen ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
            (Red_Add HSemiring a
sr abt ('HNat : xs) a
_) ->
              CExpr -> CodeGen ()
putExprStat (CExpr -> CodeGen ()) -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ CExpr
loc CExpr -> CExpr -> CExpr
.=. (Sing a -> CExpr
forall (a :: Hakaru). Sing a -> CExpr
addMonoidIdentity (Sing a -> CExpr)
-> (HSemiring a -> Sing a) -> HSemiring a -> CExpr
forall b c a. (b -> c) -> (a -> b) -> a -> c
. HSemiring a -> Sing a
forall (a :: Hakaru). HSemiring a -> Sing a
sing_HSemiring (HSemiring a -> CExpr) -> HSemiring a -> CExpr
forall a b. (a -> b) -> a -> b
$ HSemiring a
sr)

        accumRed
          :: (ABT Term abt)
          => Bool
          -> Reducer abt xs a
          -> CExpr
          -> (CExpr -> CodeGen ())
        accumRed :: Bool -> Reducer abt xs a -> CExpr -> CExpr -> CodeGen ()
accumRed Bool
isPar Reducer abt xs a
mr CExpr
itE = \CExpr
loc ->
          case Reducer abt xs a
mr of
            (Red_Index abt xs 'HNat
_ abt ('HNat : xs) 'HNat
a Reducer abt ('HNat : xs) a
body) ->
              abt ('HNat : xs) 'HNat
-> (Variable 'HNat -> abt xs 'HNat -> CodeGen ()) -> CodeGen ()
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (x :: k) (xs :: [k]) (a :: k) r.
ABT syn abt =>
abt (x : xs) a -> (Variable x -> abt xs a -> r) -> r
caseBind abt ('HNat : xs) 'HNat
a ((Variable 'HNat -> abt xs 'HNat -> CodeGen ()) -> CodeGen ())
-> (Variable 'HNat -> abt xs 'HNat -> CodeGen ()) -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ \v :: Variable 'HNat
v@(Variable Text
_ Nat
_ Sing 'HNat
typ) abt xs 'HNat
a' ->
                let (List1 Variable xs
vs,abt '[] 'HNat
a'') = abt xs 'HNat -> (List1 Variable xs, abt '[] 'HNat)
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (xs :: [k]) (a :: k).
ABT syn abt =>
abt xs a -> (List1 Variable xs, abt '[] a)
caseBinds abt xs 'HNat
a' in
                  do Ident
vId <- Variable 'HNat -> CodeGen Ident
forall (a :: Hakaru). Variable a -> CodeGen Ident
createIdent Variable 'HNat
v
                     Sing 'HNat -> Ident -> CodeGen ()
forall (a :: Hakaru). Sing a -> Ident -> CodeGen ()
declare Sing 'HNat
typ Ident
vId
                     CExpr -> CodeGen ()
putExprStat (CExpr -> CodeGen ()) -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ (Ident -> CExpr
CVar Ident
vId) CExpr -> CExpr -> CExpr
.=. CExpr
itE
                     [CodeGen ()] -> CodeGen ()
forall (t :: * -> *) (m :: * -> *) a.
(Foldable t, Monad m) =>
t (m a) -> m ()
sequence_ ([CodeGen ()] -> CodeGen ())
-> (List1 Variable xs -> [CodeGen ()])
-> List1 Variable xs
-> CodeGen ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (forall (i :: Hakaru). Variable i -> [CodeGen ()])
-> List1 Variable xs -> [CodeGen ()]
forall k1 k2 (f :: (k1 -> *) -> k2 -> *) m (a :: k1 -> *)
       (j :: k2).
(Foldable11 f, Monoid m) =>
(forall (i :: k1). a i -> m) -> f a j -> m
foldMap11
                       (\Variable i
v' -> case Variable i
v' of
                         (Variable _ _ typ') ->
                           [Sing i -> Ident -> CodeGen ()
forall (a :: Hakaru). Sing a -> Ident -> CodeGen ()
declare Sing i
typ' (Ident -> CodeGen ()) -> CodeGen Ident -> CodeGen ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Variable i -> CodeGen Ident
forall (a :: Hakaru). Variable a -> CodeGen Ident
createIdent Variable i
v'])
                       (List1 Variable xs -> CodeGen ())
-> List1 Variable xs -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ List1 Variable xs
vs
                     CExpr
aE <- abt '[] 'HNat -> String -> CodeGen CExpr
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> String -> CodeGen CExpr
flattenWithName' abt '[] 'HNat
a'' String
"index"
                     Bool -> Reducer abt ('HNat : xs) a -> CExpr -> CExpr -> CodeGen ()
forall (abt :: [Hakaru] -> Hakaru -> *) (xs :: [Hakaru])
       (a :: Hakaru).
ABT Term abt =>
Bool -> Reducer abt xs a -> CExpr -> CExpr -> CodeGen ()
accumRed Bool
isPar Reducer abt ('HNat : xs) a
body CExpr
itE (CExpr -> CExpr -> CExpr
index (CExpr -> CExpr
arrayData CExpr
loc) CExpr
aE)
            (Red_Fanout Reducer abt xs a
mr1 Reducer abt xs b
mr2) -> Bool -> Reducer abt xs a -> CExpr -> CExpr -> CodeGen ()
forall (abt :: [Hakaru] -> Hakaru -> *) (xs :: [Hakaru])
       (a :: Hakaru).
ABT Term abt =>
Bool -> Reducer abt xs a -> CExpr -> CExpr -> CodeGen ()
accumRed Bool
isPar Reducer abt xs a
mr1 CExpr
itE (CExpr -> CExpr
datumFst CExpr
loc)
                                 CodeGen () -> CodeGen () -> CodeGen ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Bool -> Reducer abt xs b -> CExpr -> CExpr -> CodeGen ()
forall (abt :: [Hakaru] -> Hakaru -> *) (xs :: [Hakaru])
       (a :: Hakaru).
ABT Term abt =>
Bool -> Reducer abt xs a -> CExpr -> CExpr -> CodeGen ()
accumRed Bool
isPar Reducer abt xs b
mr2 CExpr
itE (CExpr -> CExpr
datumSnd CExpr
loc)
            (Red_Split abt ('HNat : xs) HBool
b Reducer abt xs a
mr1 Reducer abt xs b
mr2) ->
              abt ('HNat : xs) HBool
-> (Variable 'HNat -> abt xs HBool -> CodeGen ()) -> CodeGen ()
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (x :: k) (xs :: [k]) (a :: k) r.
ABT syn abt =>
abt (x : xs) a -> (Variable x -> abt xs a -> r) -> r
caseBind abt ('HNat : xs) HBool
b ((Variable 'HNat -> abt xs HBool -> CodeGen ()) -> CodeGen ())
-> (Variable 'HNat -> abt xs HBool -> CodeGen ()) -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ \v :: Variable 'HNat
v@(Variable Text
_ Nat
_ Sing 'HNat
typ) abt xs HBool
b' ->
                let (List1 Variable xs
vs,abt '[] HBool
b'') = abt xs HBool -> (List1 Variable xs, abt '[] HBool)
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (xs :: [k]) (a :: k).
ABT syn abt =>
abt xs a -> (List1 Variable xs, abt '[] a)
caseBinds abt xs HBool
b' in
                  do Ident
vId <- Variable 'HNat -> CodeGen Ident
forall (a :: Hakaru). Variable a -> CodeGen Ident
createIdent Variable 'HNat
v
                     Sing 'HNat -> Ident -> CodeGen ()
forall (a :: Hakaru). Sing a -> Ident -> CodeGen ()
declare Sing 'HNat
typ Ident
vId
                     CExpr -> CodeGen ()
putExprStat (CExpr -> CodeGen ()) -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ (Ident -> CExpr
CVar Ident
vId) CExpr -> CExpr -> CExpr
.=. CExpr
itE
                     [CodeGen ()] -> CodeGen ()
forall (t :: * -> *) (m :: * -> *) a.
(Foldable t, Monad m) =>
t (m a) -> m ()
sequence_ ([CodeGen ()] -> CodeGen ())
-> (List1 Variable xs -> [CodeGen ()])
-> List1 Variable xs
-> CodeGen ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (forall (i :: Hakaru). Variable i -> [CodeGen ()])
-> List1 Variable xs -> [CodeGen ()]
forall k1 k2 (f :: (k1 -> *) -> k2 -> *) m (a :: k1 -> *)
       (j :: k2).
(Foldable11 f, Monoid m) =>
(forall (i :: k1). a i -> m) -> f a j -> m
foldMap11
                       (\Variable i
v' -> case Variable i
v' of
                         (Variable _ _ typ') ->
                           [Sing i -> Ident -> CodeGen ()
forall (a :: Hakaru). Sing a -> Ident -> CodeGen ()
declare Sing i
typ' (Ident -> CodeGen ()) -> CodeGen Ident -> CodeGen ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Variable i -> CodeGen Ident
forall (a :: Hakaru). Variable a -> CodeGen Ident
createIdent Variable i
v'])
                       (List1 Variable xs -> CodeGen ())
-> List1 Variable xs -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ List1 Variable xs
vs
                     CExpr
bE <- abt '[] HBool -> String -> CodeGen CExpr
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> String -> CodeGen CExpr
flattenWithName' abt '[] HBool
b'' String
"cond"
                     CExpr -> CodeGen () -> CodeGen () -> CodeGen ()
ifCG (CExpr
bE CExpr -> String -> CExpr
... String
"index" CExpr -> CExpr -> CExpr
.==. (Integer -> CExpr
intE Integer
0))
                          (Bool -> Reducer abt xs a -> CExpr -> CExpr -> CodeGen ()
forall (abt :: [Hakaru] -> Hakaru -> *) (xs :: [Hakaru])
       (a :: Hakaru).
ABT Term abt =>
Bool -> Reducer abt xs a -> CExpr -> CExpr -> CodeGen ()
accumRed Bool
isPar Reducer abt xs a
mr1 CExpr
itE (CExpr -> CExpr
datumFst CExpr
loc))
                          (Bool -> Reducer abt xs b -> CExpr -> CExpr -> CodeGen ()
forall (abt :: [Hakaru] -> Hakaru -> *) (xs :: [Hakaru])
       (a :: Hakaru).
ABT Term abt =>
Bool -> Reducer abt xs a -> CExpr -> CExpr -> CodeGen ()
accumRed Bool
isPar Reducer abt xs b
mr2 CExpr
itE (CExpr -> CExpr
datumSnd CExpr
loc))
            Reducer abt xs a
Red_Nop -> () -> CodeGen ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
            (Red_Add HSemiring a
sr abt ('HNat : xs) a
e) ->
              abt ('HNat : xs) a
-> (Variable 'HNat -> abt xs a -> CodeGen ()) -> CodeGen ()
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (x :: k) (xs :: [k]) (a :: k) r.
ABT syn abt =>
abt (x : xs) a -> (Variable x -> abt xs a -> r) -> r
caseBind abt ('HNat : xs) a
e ((Variable 'HNat -> abt xs a -> CodeGen ()) -> CodeGen ())
-> (Variable 'HNat -> abt xs a -> CodeGen ()) -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ \v :: Variable 'HNat
v@(Variable Text
_ Nat
_ Sing 'HNat
typ) abt xs a
e' ->
                let (List1 Variable xs
vs,abt '[] a
e'') = abt xs a -> (List1 Variable xs, abt '[] a)
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (xs :: [k]) (a :: k).
ABT syn abt =>
abt xs a -> (List1 Variable xs, abt '[] a)
caseBinds abt xs a
e' in
                  do Ident
vId <- Variable 'HNat -> CodeGen Ident
forall (a :: Hakaru). Variable a -> CodeGen Ident
createIdent Variable 'HNat
v
                     Sing 'HNat -> Ident -> CodeGen ()
forall (a :: Hakaru). Sing a -> Ident -> CodeGen ()
declare Sing 'HNat
typ Ident
vId
                     CExpr -> CodeGen ()
putExprStat (CExpr -> CodeGen ()) -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ (Ident -> CExpr
CVar Ident
vId) CExpr -> CExpr -> CExpr
.=. CExpr
itE
                     [CodeGen ()] -> CodeGen ()
forall (t :: * -> *) (m :: * -> *) a.
(Foldable t, Monad m) =>
t (m a) -> m ()
sequence_ ([CodeGen ()] -> CodeGen ())
-> (List1 Variable xs -> [CodeGen ()])
-> List1 Variable xs
-> CodeGen ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (forall (i :: Hakaru). Variable i -> [CodeGen ()])
-> List1 Variable xs -> [CodeGen ()]
forall k1 k2 (f :: (k1 -> *) -> k2 -> *) m (a :: k1 -> *)
       (j :: k2).
(Foldable11 f, Monoid m) =>
(forall (i :: k1). a i -> m) -> f a j -> m
foldMap11
                       (\Variable i
v' -> case Variable i
v' of
                         (Variable _ _ typ') ->
                           [Sing i -> Ident -> CodeGen ()
forall (a :: Hakaru). Sing a -> Ident -> CodeGen ()
declare Sing i
typ' (Ident -> CodeGen ()) -> CodeGen Ident -> CodeGen ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Variable i -> CodeGen Ident
forall (a :: Hakaru). Variable a -> CodeGen Ident
createIdent Variable i
v'])
                       (List1 Variable xs -> CodeGen ())
-> List1 Variable xs -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ List1 Variable xs
vs
                     CExpr
eE <- abt '[] a -> CodeGen CExpr
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> CodeGen CExpr
flattenWithName abt '[] a
e''
                     -- when isPar $  putStat . CPPStat . ompToPP $ OMP Critical
                     case HSemiring a -> Sing a
forall (a :: Hakaru). HSemiring a -> Sing a
sing_HSemiring HSemiring a
sr of
                       Sing a
SProb -> Seq CExpr -> CExpr -> CodeGen ()
logSumExpCG ([CExpr] -> Seq CExpr
forall a. [a] -> Seq a
S.fromList [CExpr
loc,CExpr
eE]) CExpr
loc
                       Sing a
_ -> CExpr -> CodeGen ()
putExprStat (CExpr -> CodeGen ()) -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ CExpr
loc CExpr -> CExpr -> CExpr
.+=. CExpr
eE

        mulRed
          :: (ABT Term abt)
          => Reducer abt xs a
          -> (CExpr -> CExpr -> CodeGen ())
        mulRed :: Reducer abt xs a -> CExpr -> CExpr -> CodeGen ()
mulRed Reducer abt xs a
mr CExpr
outp CExpr
inp =
          case Reducer abt xs a
mr of
             (Red_Index abt xs 'HNat
_ abt ('HNat : xs) 'HNat
_ Reducer abt ('HNat : xs) a
mr') ->
               do CExpr
itE <- Sing 'HNat -> CodeGen CExpr
forall (a :: Hakaru). Sing a -> CodeGen CExpr
localVar Sing 'HNat
SNat
                  CExpr -> CExpr -> CExpr -> CodeGen () -> CodeGen ()
forCG (CExpr
itE CExpr -> CExpr -> CExpr
.=. (Integer -> CExpr
intE Integer
0))
                        (CExpr
itE CExpr -> CExpr -> CExpr
.<. (Integer -> CExpr
intE Integer
0))
                        (CUnaryOp -> CExpr -> CExpr
CUnary CUnaryOp
CPostIncOp CExpr
itE)
                        (Reducer abt ('HNat : xs) a -> CExpr -> CExpr -> CodeGen ()
forall (abt :: [Hakaru] -> Hakaru -> *) (xs :: [Hakaru])
       (a :: Hakaru).
ABT Term abt =>
Reducer abt xs a -> CExpr -> CExpr -> CodeGen ()
mulRed Reducer abt ('HNat : xs) a
mr'
                                (CExpr -> CExpr -> CExpr
index (CExpr -> CExpr
arrayData CExpr
outp) CExpr
itE)
                                (CExpr -> CExpr -> CExpr
index (CExpr -> CExpr
arrayData CExpr
inp) CExpr
itE))
             (Red_Fanout Reducer abt xs a
mr1 Reducer abt xs b
mr2) -> Reducer abt xs a -> CExpr -> CExpr -> CodeGen ()
forall (abt :: [Hakaru] -> Hakaru -> *) (xs :: [Hakaru])
       (a :: Hakaru).
ABT Term abt =>
Reducer abt xs a -> CExpr -> CExpr -> CodeGen ()
mulRed Reducer abt xs a
mr1 (CExpr -> CExpr
datumFst CExpr
outp) (CExpr -> CExpr
datumFst CExpr
inp)
                                  CodeGen () -> CodeGen () -> CodeGen ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Reducer abt xs b -> CExpr -> CExpr -> CodeGen ()
forall (abt :: [Hakaru] -> Hakaru -> *) (xs :: [Hakaru])
       (a :: Hakaru).
ABT Term abt =>
Reducer abt xs a -> CExpr -> CExpr -> CodeGen ()
mulRed Reducer abt xs b
mr2 (CExpr -> CExpr
datumFst CExpr
outp) (CExpr -> CExpr
datumFst CExpr
inp)
             (Red_Split abt ('HNat : xs) HBool
_ Reducer abt xs a
mr1 Reducer abt xs b
mr2) -> Reducer abt xs a -> CExpr -> CExpr -> CodeGen ()
forall (abt :: [Hakaru] -> Hakaru -> *) (xs :: [Hakaru])
       (a :: Hakaru).
ABT Term abt =>
Reducer abt xs a -> CExpr -> CExpr -> CodeGen ()
mulRed Reducer abt xs a
mr1 (CExpr -> CExpr
datumFst CExpr
outp) (CExpr -> CExpr
datumFst CExpr
inp)
                                   CodeGen () -> CodeGen () -> CodeGen ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Reducer abt xs b -> CExpr -> CExpr -> CodeGen ()
forall (abt :: [Hakaru] -> Hakaru -> *) (xs :: [Hakaru])
       (a :: Hakaru).
ABT Term abt =>
Reducer abt xs a -> CExpr -> CExpr -> CodeGen ()
mulRed Reducer abt xs b
mr2 (CExpr -> CExpr
datumFst CExpr
outp) (CExpr -> CExpr
datumFst CExpr
inp)
             Reducer abt xs a
Red_Nop -> () -> CodeGen ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
             (Red_Add HSemiring a
_ abt ('HNat : xs) a
_) -> CExpr -> CodeGen ()
putExprStat (CExpr -> CodeGen ()) -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ CExpr
outp CExpr -> CExpr -> CExpr
.+=. CExpr
inp

addMonoidIdentity :: Sing (a :: Hakaru) -> CExpr
addMonoidIdentity :: Sing a -> CExpr
addMonoidIdentity Sing a
s =
  case Sing a
s of
    Sing a
SNat  -> Integer -> CExpr
intE Integer
0
    Sing a
SInt  -> Integer -> CExpr
intE Integer
0
    Sing a
SReal -> Float -> CExpr
floatE Float
0
    Sing a
SProb -> CExpr -> CExpr
logE (Float -> CExpr
floatE Float
0)
    SArray x -> Sing a -> CExpr
forall (a :: Hakaru). Sing a -> CExpr
addMonoidIdentity Sing a
x
    Sing a
x -> String -> CExpr
forall a. HasCallStack => String -> a
error (String -> CExpr) -> String -> CExpr
forall a b. (a -> b) -> a -> b
$ String
"addMonoidIdentity{" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Sing a -> String
forall a. Show a => a -> String
show Sing a
x String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"}"

--------------------------------------------------------------------------------
--                                 Datum and Case                             --
--------------------------------------------------------------------------------
{-

Datum are sums of products of types. This maps to a C structure. flattenDatum
will produce a literal of some datum type. This will also produce a global
struct representing that datum which will be needed for the C compiler.

-}


flattenDatum
  :: (ABT Term abt)
  => Datum (abt '[]) (HData' a)
  -> (CExpr -> CodeGen ())
flattenDatum :: Datum (abt '[]) (HData' a) -> CExpr -> CodeGen ()
flattenDatum (Datum Text
_ Sing (HData' t)
typ DatumCode (Code t) (abt '[]) (HData' t)
code) =
  \CExpr
loc ->
    do Sing (HData' a) -> CodeGen ()
forall (a :: Hakaru). Sing a -> CodeGen ()
extDeclareTypes Sing (HData' a)
Sing (HData' t)
typ
       DatumCode (Code a) (abt '[]) (HData' a) -> CExpr -> CodeGen ()
forall (abt :: [Hakaru] -> Hakaru -> *) (xss :: [[HakaruFun]])
       (c :: Hakaru).
ABT Term abt =>
DatumCode xss (abt '[]) c -> CExpr -> CodeGen ()
assignDatum DatumCode (Code a) (abt '[]) (HData' a)
DatumCode (Code t) (abt '[]) (HData' t)
code CExpr
loc

assignDatum
  :: (ABT Term abt)
  => DatumCode xss (abt '[]) c
  -> CExpr
  -> CodeGen ()
assignDatum :: DatumCode xss (abt '[]) c -> CExpr -> CodeGen ()
assignDatum DatumCode xss (abt '[]) c
code CExpr
ident =
  let ind :: Integer
ind       = DatumCode xss (abt '[]) c -> Integer
forall (xss :: [[HakaruFun]]) (b :: Hakaru -> *) (c :: Hakaru).
DatumCode xss b c -> Integer
getIndex DatumCode xss (abt '[]) c
code
      indExpr :: CExpr
indExpr = CExpr -> Ident -> Bool -> CExpr
CMember CExpr
ident (String -> Ident
Ident String
"index") Bool
True
  in  do CExpr -> CodeGen ()
putExprStat (CExpr
indExpr CExpr -> CExpr -> CExpr
.=. (Integer -> CExpr
intE Integer
ind))
         [CodeGen ()] -> CodeGen ()
forall (t :: * -> *) (m :: * -> *) a.
(Foldable t, Monad m) =>
t (m a) -> m ()
sequence_ ([CodeGen ()] -> CodeGen ()) -> [CodeGen ()] -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ DatumCode xss (abt '[]) c -> CExpr -> [CodeGen ()]
forall (abt :: [Hakaru] -> Hakaru -> *) (xs :: [[HakaruFun]])
       (c :: Hakaru).
ABT Term abt =>
DatumCode xs (abt '[]) c -> CExpr -> [CodeGen ()]
assignSum DatumCode xss (abt '[]) c
code CExpr
ident
  where getIndex :: DatumCode xss b c -> Integer
        getIndex :: DatumCode xss b c -> Integer
getIndex (Inl DatumStruct xs b c
_)    = Integer
0
        getIndex (Inr DatumCode xss b c
rest) = Integer -> Integer
forall a. Enum a => a -> a
succ (DatumCode xss b c -> Integer
forall (xss :: [[HakaruFun]]) (b :: Hakaru -> *) (c :: Hakaru).
DatumCode xss b c -> Integer
getIndex DatumCode xss b c
rest)

assignSum
  :: (ABT Term abt)
  => DatumCode xs (abt '[]) c
  -> CExpr
  -> [CodeGen ()]
assignSum :: DatumCode xs (abt '[]) c -> CExpr -> [CodeGen ()]
assignSum DatumCode xs (abt '[]) c
code CExpr
ident = ([CodeGen ()], [String]) -> [CodeGen ()]
forall a b. (a, b) -> a
fst (([CodeGen ()], [String]) -> [CodeGen ()])
-> ([CodeGen ()], [String]) -> [CodeGen ()]
forall a b. (a -> b) -> a -> b
$ State [String] [CodeGen ()] -> [String] -> ([CodeGen ()], [String])
forall s a. State s a -> s -> (a, s)
runState (DatumCode xs (abt '[]) c -> CExpr -> State [String] [CodeGen ()]
forall (abt :: [Hakaru] -> Hakaru -> *) (xs :: [[HakaruFun]])
       (c :: Hakaru).
ABT Term abt =>
DatumCode xs (abt '[]) c -> CExpr -> State [String] [CodeGen ()]
assignSum' DatumCode xs (abt '[]) c
code CExpr
ident) [String]
cNameStream

assignSum'
  :: (ABT Term abt)
  => DatumCode xs (abt '[]) c
  -> CExpr
  -> State [String] [CodeGen ()]
assignSum' :: DatumCode xs (abt '[]) c -> CExpr -> State [String] [CodeGen ()]
assignSum' (Inr DatumCode xss (abt '[]) c
rest) CExpr
topIdent =
  do [String]
names <- StateT [String] Identity [String]
forall s (m :: * -> *). MonadState s m => m s
get
     [String] -> StateT [String] Identity ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put ([String] -> [String]
forall a. [a] -> [a]
tail [String]
names)
     DatumCode xss (abt '[]) c -> CExpr -> State [String] [CodeGen ()]
forall (abt :: [Hakaru] -> Hakaru -> *) (xs :: [[HakaruFun]])
       (c :: Hakaru).
ABT Term abt =>
DatumCode xs (abt '[]) c -> CExpr -> State [String] [CodeGen ()]
assignSum' DatumCode xss (abt '[]) c
rest CExpr
topIdent
assignSum' (Inl DatumStruct xs (abt '[]) c
prod) CExpr
topIdent =
  do String
name <- [String] -> String
forall a. [a] -> a
head ([String] -> String)
-> StateT [String] Identity [String]
-> StateT [String] Identity String
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> StateT [String] Identity [String]
forall s (m :: * -> *). MonadState s m => m s
get
     [CodeGen ()] -> State [String] [CodeGen ()]
forall (m :: * -> *) a. Monad m => a -> m a
return ([CodeGen ()] -> State [String] [CodeGen ()])
-> [CodeGen ()] -> State [String] [CodeGen ()]
forall a b. (a -> b) -> a -> b
$ DatumStruct xs (abt '[]) c -> CExpr -> CExpr -> [CodeGen ()]
forall (abt :: [Hakaru] -> Hakaru -> *) (xs :: [HakaruFun])
       (c :: Hakaru).
ABT Term abt =>
DatumStruct xs (abt '[]) c -> CExpr -> CExpr -> [CodeGen ()]
assignProd DatumStruct xs (abt '[]) c
prod CExpr
topIdent (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
name)

assignProd
  :: (ABT Term abt)
  => DatumStruct xs (abt '[]) c
  -> CExpr
  -> CExpr
  -> [CodeGen ()]
assignProd :: DatumStruct xs (abt '[]) c -> CExpr -> CExpr -> [CodeGen ()]
assignProd DatumStruct xs (abt '[]) c
dstruct CExpr
topIdent CExpr
sumIdent =
  ([CodeGen ()], [String]) -> [CodeGen ()]
forall a b. (a, b) -> a
fst (([CodeGen ()], [String]) -> [CodeGen ()])
-> ([CodeGen ()], [String]) -> [CodeGen ()]
forall a b. (a -> b) -> a -> b
$ State [String] [CodeGen ()] -> [String] -> ([CodeGen ()], [String])
forall s a. State s a -> s -> (a, s)
runState (DatumStruct xs (abt '[]) c
-> CExpr -> CExpr -> State [String] [CodeGen ()]
forall (abt :: [Hakaru] -> Hakaru -> *) (xs :: [HakaruFun])
       (c :: Hakaru).
ABT Term abt =>
DatumStruct xs (abt '[]) c
-> CExpr -> CExpr -> State [String] [CodeGen ()]
assignProd' DatumStruct xs (abt '[]) c
dstruct CExpr
topIdent CExpr
sumIdent) [String]
cNameStream

assignProd'
  :: (ABT Term abt)
  => DatumStruct xs (abt '[]) c
  -> CExpr
  -> CExpr
  -> State [String] [CodeGen ()]
assignProd' :: DatumStruct xs (abt '[]) c
-> CExpr -> CExpr -> State [String] [CodeGen ()]
assignProd' DatumStruct xs (abt '[]) c
Done CExpr
_ CExpr
_ = [CodeGen ()] -> State [String] [CodeGen ()]
forall (m :: * -> *) a. Monad m => a -> m a
return []
assignProd' (Et (Konst abt '[] b
d) DatumStruct xs (abt '[]) c
rest) CExpr
topIdent (CVar Ident
sumIdent) =
  do [String]
names <- StateT [String] Identity [String]
forall s (m :: * -> *). MonadState s m => m s
get
     [String] -> StateT [String] Identity ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put ([String] -> [String]
forall a. [a] -> [a]
tail [String]
names)
     let varName :: CExpr
varName  = CExpr -> Ident -> Bool -> CExpr
CMember (CExpr -> Ident -> Bool -> CExpr
CMember (CExpr -> Ident -> Bool -> CExpr
CMember CExpr
topIdent
                                              (String -> Ident
Ident String
"sum")
                                              Bool
True)
                                     Ident
sumIdent
                                     Bool
True)
                            (String -> Ident
Ident ([String] -> String
forall a. [a] -> a
head [String]
names))
                            Bool
True
     [CodeGen ()]
rest' <- DatumStruct xs (abt '[]) c
-> CExpr -> CExpr -> State [String] [CodeGen ()]
forall (abt :: [Hakaru] -> Hakaru -> *) (xs :: [HakaruFun])
       (c :: Hakaru).
ABT Term abt =>
DatumStruct xs (abt '[]) c
-> CExpr -> CExpr -> State [String] [CodeGen ()]
assignProd' DatumStruct xs (abt '[]) c
rest CExpr
topIdent (Ident -> CExpr
CVar Ident
sumIdent)
     [CodeGen ()] -> State [String] [CodeGen ()]
forall (m :: * -> *) a. Monad m => a -> m a
return ([CodeGen ()] -> State [String] [CodeGen ()])
-> [CodeGen ()] -> State [String] [CodeGen ()]
forall a b. (a -> b) -> a -> b
$ [abt '[] b -> CExpr -> CodeGen ()
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> CExpr -> CodeGen ()
flattenABT abt '[] b
d CExpr
varName] [CodeGen ()] -> [CodeGen ()] -> [CodeGen ()]
forall a. [a] -> [a] -> [a]
++ [CodeGen ()]
rest'

assignProd' DatumStruct xs (abt '[]) c
_ CExpr
_ CExpr
_  = String -> State [String] [CodeGen ()]
forall a. HasCallStack => String -> a
error (String -> State [String] [CodeGen ()])
-> String -> State [String] [CodeGen ()]
forall a b. (a -> b) -> a -> b
$ String
"TODO: assignProd Ident"


----------
-- Case --
----------

-- currently we can only match on boolean values
flattenCase
  :: forall abt a b
  .  (ABT Term abt)
  => abt '[] a
  -> [Branch a abt b]
  -> (CExpr -> CodeGen ())

flattenCase :: abt '[] a -> [Branch a abt b] -> CExpr -> CodeGen ()
flattenCase abt '[] a
c [ Branch (PDatum Text
_ (PInl PDatumStruct xs xs (HData' t)
PDone))        abt xs b
trueB
              , Branch (PDatum Text
_ (PInr (PInl PDatumStruct xs xs (HData' t)
PDone))) abt xs b
falseB ] =
  \CExpr
loc ->
    do CExpr
cE <- abt '[] a -> CodeGen CExpr
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> CodeGen CExpr
flattenWithName abt '[] a
c
       CExpr -> CodeGen () -> CodeGen () -> CodeGen ()
ifCG ((CExpr
cE CExpr -> String -> CExpr
... String
"index") CExpr -> CExpr -> CExpr
.==. (Integer -> CExpr
intE Integer
0))
            (abt '[] b -> CExpr -> CodeGen ()
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> CExpr -> CodeGen ()
flattenABT abt xs b
abt '[] b
trueB  CExpr
loc)
            (abt '[] b -> CExpr -> CodeGen ()
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> CExpr -> CodeGen ()
flattenABT abt xs b
abt '[] b
falseB CExpr
loc)

flattenCase abt '[] a
e [ Branch (PDatum Text
_ (PInl (PEt (PKonst Pattern vars1 b
PVar)
                                            (PEt (PKonst Pattern vars1 b
PVar)
                                                 PDatumStruct xs vars2 (HData' t)
PDone)))) abt xs b
b
              ]
  = \CExpr
loc -> do
    CExpr
eE <- abt '[] a -> CodeGen CExpr
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> CodeGen CExpr
flattenWithName abt '[] a
e
    abt '[b, b] b
-> (Variable b -> abt '[b] b -> CodeGen ()) -> CodeGen ()
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (x :: k) (xs :: [k]) (a :: k) r.
ABT syn abt =>
abt (x : xs) a -> (Variable x -> abt xs a -> r) -> r
caseBind abt xs b
abt '[b, b] b
b ((Variable b -> abt '[b] b -> CodeGen ()) -> CodeGen ())
-> (Variable b -> abt '[b] b -> CodeGen ()) -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ \vfst :: Variable b
vfst@(Variable Text
_ Nat
_ Sing b
fstTyp) abt '[b] b
b' ->
      abt '[b] b -> (Variable b -> abt '[] b -> CodeGen ()) -> CodeGen ()
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (x :: k) (xs :: [k]) (a :: k) r.
ABT syn abt =>
abt (x : xs) a -> (Variable x -> abt xs a -> r) -> r
caseBind abt '[b] b
b' ((Variable b -> abt '[] b -> CodeGen ()) -> CodeGen ())
-> (Variable b -> abt '[] b -> CodeGen ()) -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ \vsnd :: Variable b
vsnd@(Variable Text
_ Nat
_ Sing b
sndTyp) abt '[] b
b'' -> do
        Ident
fstId <- Variable b -> CodeGen Ident
forall (a :: Hakaru). Variable a -> CodeGen Ident
createIdent Variable b
vfst
        Ident
sndId <- Variable b -> CodeGen Ident
forall (a :: Hakaru). Variable a -> CodeGen Ident
createIdent Variable b
vsnd
        Sing b -> Ident -> CodeGen ()
forall (a :: Hakaru). Sing a -> Ident -> CodeGen ()
declare Sing b
fstTyp Ident
fstId
        Sing b -> Ident -> CodeGen ()
forall (a :: Hakaru). Sing a -> Ident -> CodeGen ()
declare Sing b
sndTyp Ident
sndId
        CExpr -> CodeGen ()
putExprStat (CExpr -> CodeGen ()) -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ (Ident -> CExpr
CVar Ident
fstId) CExpr -> CExpr -> CExpr
.=. (CExpr -> CExpr
datumFst CExpr
eE)
        CExpr -> CodeGen ()
putExprStat (CExpr -> CodeGen ()) -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ (Ident -> CExpr
CVar Ident
sndId) CExpr -> CExpr -> CExpr
.=. (CExpr -> CExpr
datumSnd CExpr
eE)
        abt '[] b -> CExpr -> CodeGen ()
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> CExpr -> CodeGen ()
flattenABT abt '[] b
b'' CExpr
loc

flattenCase abt '[] a
e [Branch a abt b]
_ = String -> CExpr -> CodeGen ()
forall a. HasCallStack => String -> a
error (String -> CExpr -> CodeGen ()) -> String -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ String
"TODO: flattenCase{" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Sing a -> String
forall a. Show a => a -> String
show (abt '[] a -> Sing a
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> Sing a
typeOf abt '[] a
e) String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"}"


--------------------------------------------------------------------------------
--                                     PrimOp                                 --
--------------------------------------------------------------------------------

flattenPrimOp
  :: ( ABT Term abt
     , typs ~ UnLCs args
     , args ~ LCs typs)
  => PrimOp typs a
  -> SArgs abt args
  -> (CExpr -> CodeGen ())


flattenPrimOp :: PrimOp typs a -> SArgs abt args -> CExpr -> CodeGen ()
flattenPrimOp PrimOp typs a
Pi =
  \SArgs abt args
End ->
    \CExpr
loc -> let piE :: CExpr
piE = CExpr -> CExpr
log1pE ((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
"M_PI") CExpr -> CExpr -> CExpr
.-. (Integer -> CExpr
intE Integer
1)) in
      CExpr -> CodeGen ()
putExprStat (CExpr
loc CExpr -> CExpr -> CExpr
.=. CExpr
piE)

flattenPrimOp PrimOp typs a
Not =
  \(abt vars a
a :* SArgs abt args
End) ->
    \CExpr
loc ->
      do CExpr
aE <- abt '[] a -> CodeGen CExpr
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> CodeGen CExpr
flattenWithName abt vars a
abt '[] a
a
         Ident
bId <- CodeGen Ident
genIdent
         Sing a -> Ident -> CodeGen ()
forall (a :: Hakaru). Sing a -> Ident -> CodeGen ()
declare (abt '[] a -> Sing a
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> Sing a
typeOf abt vars a
abt '[] a
a) Ident
bId
         let bE :: CExpr
bE = Ident -> CExpr
CVar Ident
bId
         CExpr -> CodeGen ()
putExprStat (CExpr -> CodeGen ()) -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ CExpr -> CExpr
datumIndex CExpr
bE CExpr -> CExpr -> CExpr
.=. (CExpr -> CExpr -> CExpr -> CExpr
CCond (CExpr -> CExpr
datumIndex CExpr
aE CExpr -> CExpr -> CExpr
.==. (Integer -> CExpr
intE Integer
1))
                                               (Integer -> CExpr
intE Integer
0)
                                               (Integer -> CExpr
intE Integer
1))
         CExpr -> CodeGen ()
putExprStat (CExpr -> CodeGen ()) -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ CExpr
loc CExpr -> CExpr -> CExpr
.=. CExpr
bE

flattenPrimOp PrimOp typs a
RealPow =
  \(abt vars a
base :* abt vars a
power :* SArgs abt args
End) ->
    \CExpr
loc ->
      do CExpr
baseE <- abt '[] a -> CodeGen CExpr
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> CodeGen CExpr
flattenWithName abt vars a
abt '[] a
base
         CExpr
powerE <- abt '[] a -> CodeGen CExpr
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> CodeGen CExpr
flattenWithName abt vars a
abt '[] a
power
         let realPow :: CExpr
realPow = CExpr -> [CExpr] -> CExpr
CCall (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
"pow")
                             [ CExpr -> CExpr
expm1E CExpr
baseE CExpr -> CExpr -> CExpr
.+. (Integer -> CExpr
intE Integer
1), CExpr
powerE]
         CExpr -> CodeGen ()
putExprStat (CExpr -> CodeGen ()) -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ CExpr
loc CExpr -> CExpr -> CExpr
.=. (CExpr -> CExpr
log1pE (CExpr
realPow CExpr -> CExpr -> CExpr
.-. (Integer -> CExpr
intE Integer
1)))

flattenPrimOp (NatPow HSemiring a
baseTyp) =
  \(abt vars a
base :* abt vars a
power :* SArgs abt args
End) ->
    \CExpr
loc ->
      let sBase :: Sing a
sBase = HSemiring a -> Sing a
forall (a :: Hakaru). HSemiring a -> Sing a
sing_HSemiring HSemiring a
baseTyp in
      do Ident
baseId <- CodeGen Ident
genIdent
         Sing a -> Ident -> CodeGen ()
forall (a :: Hakaru). Sing a -> Ident -> CodeGen ()
declare Sing a
sBase Ident
baseId
         let baseE :: CExpr
baseE = Ident -> CExpr
CVar Ident
baseId
         abt '[] a -> CExpr -> CodeGen ()
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> CExpr -> CodeGen ()
flattenABT abt vars a
abt '[] a
base CExpr
baseE
         CExpr
powerE <- abt '[] a -> CodeGen CExpr
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> CodeGen CExpr
flattenWithName abt vars a
abt '[] a
power
         let powerOf :: CExpr -> CExpr -> CExpr
powerOf CExpr
x CExpr
y = CExpr -> [CExpr] -> CExpr
CCall (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
"pow") [CExpr
x,CExpr
y]
             value :: CExpr
value = case Sing a
sBase of
                       Sing a
SProb -> CExpr -> CExpr
log1pE (CExpr -> CExpr) -> CExpr -> CExpr
forall a b. (a -> b) -> a -> b
$ (CExpr -> CExpr -> CExpr
powerOf (CExpr -> CExpr
expm1E CExpr
baseE CExpr -> CExpr -> CExpr
.+. (Integer -> CExpr
intE Integer
1)) CExpr
powerE)
                                  CExpr -> CExpr -> CExpr
.-. (Integer -> CExpr
intE Integer
1)
                       Sing a
_     -> CExpr -> CExpr -> CExpr
powerOf CExpr
baseE CExpr
powerE
         CExpr -> CodeGen ()
putExprStat (CExpr -> CodeGen ()) -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ CExpr
loc CExpr -> CExpr -> CExpr
.=. CExpr
value

flattenPrimOp (NatRoot HRadical a
baseTyp) =
  \(abt vars a
base :* abt vars a
root :* SArgs abt args
End) ->
    \CExpr
loc ->
      let sBase :: Sing a
sBase = HRadical a -> Sing a
forall (a :: Hakaru). HRadical a -> Sing a
sing_HRadical HRadical a
baseTyp in
      do Ident
baseId <- CodeGen Ident
genIdent
         Sing a -> Ident -> CodeGen ()
forall (a :: Hakaru). Sing a -> Ident -> CodeGen ()
declare Sing a
sBase Ident
baseId
         let baseE :: CExpr
baseE = Ident -> CExpr
CVar Ident
baseId
         abt '[] a -> CExpr -> CodeGen ()
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> CExpr -> CodeGen ()
flattenABT abt vars a
abt '[] a
base CExpr
baseE
         CExpr
rootE <- abt '[] a -> CodeGen CExpr
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> CodeGen CExpr
flattenWithName abt vars a
abt '[] a
root
         let powerOf :: CExpr -> CExpr -> CExpr
powerOf CExpr
x CExpr
y = CExpr -> [CExpr] -> CExpr
CCall (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
"pow") [CExpr
x,CExpr
y]
             recipE :: CExpr
recipE = (Float -> CExpr
floatE Float
1) CExpr -> CExpr -> CExpr
./. CExpr
rootE
             value :: CExpr
value = case Sing a
sBase of
                       Sing a
SProb -> CExpr -> CExpr
log1pE (CExpr -> CExpr) -> CExpr -> CExpr
forall a b. (a -> b) -> a -> b
$ (CExpr -> CExpr -> CExpr
powerOf (CExpr -> CExpr
expm1E CExpr
baseE CExpr -> CExpr -> CExpr
.+. (Integer -> CExpr
intE Integer
1)) CExpr
recipE)
                                      CExpr -> CExpr -> CExpr
.-. (Integer -> CExpr
intE Integer
1)
                       Sing a
_     -> CExpr -> CExpr -> CExpr
powerOf CExpr
baseE CExpr
recipE
         CExpr -> CodeGen ()
putExprStat (CExpr -> CodeGen ()) -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ CExpr
loc CExpr -> CExpr -> CExpr
.=. CExpr
value

flattenPrimOp (Recip HFractional a
t) =
  \(abt vars a
a :* SArgs abt args
End) ->
    \CExpr
loc ->
      do CExpr
aE <- abt '[] a -> CodeGen CExpr
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> CodeGen CExpr
flattenWithName abt vars a
abt '[] a
a
         case HFractional a
t of
           HFractional a
HFractional_Real -> CExpr -> CodeGen ()
putExprStat (CExpr -> CodeGen ()) -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ CExpr
loc CExpr -> CExpr -> CExpr
.=. ((Integer -> CExpr
intE Integer
1) CExpr -> CExpr -> CExpr
./. CExpr
aE)
           HFractional a
HFractional_Prob -> CExpr -> CodeGen ()
putExprStat (CExpr -> CodeGen ()) -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ CExpr
loc CExpr -> CExpr -> CExpr
.=. (CUnaryOp -> CExpr -> CExpr
CUnary CUnaryOp
CMinOp CExpr
aE)

-- | exp : real -> prob, because of this we can just turn it into a prob without taking
--   its log, which would give us an exp in the log-domain
flattenPrimOp PrimOp typs a
Exp = \(abt vars a
a :* SArgs abt args
End) -> abt '[] a -> CExpr -> CodeGen ()
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> CExpr -> CodeGen ()
flattenABT abt vars a
abt '[] a
a

flattenPrimOp (Equal HEq a
_) =
  \(abt vars a
a :* abt vars a
b :* SArgs abt args
End) ->
    \CExpr
loc ->
      do CExpr
aE <- abt '[] a -> CodeGen CExpr
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> CodeGen CExpr
flattenWithName abt vars a
abt '[] a
a
         CExpr
bE <- abt '[] a -> CodeGen CExpr
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> CodeGen CExpr
flattenWithName abt vars a
abt '[] a
b

         -- special case for booleans
         let aE' :: CExpr
aE' = case (abt '[] a -> Sing a
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> Sing a
typeOf abt vars a
abt '[] a
a) of
                     (SData _ (SPlus SDone (SPlus SDone SVoid))) -> (CExpr -> Ident -> Bool -> CExpr
CMember CExpr
aE (String -> Ident
Ident String
"index") Bool
True)
                     Sing a
_ -> CExpr
aE
         let bE' :: CExpr
bE' = case (abt '[] a -> Sing a
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> Sing a
typeOf abt vars a
abt '[] a
b) of
                     (SData _ (SPlus SDone (SPlus SDone SVoid))) -> (CExpr -> Ident -> Bool -> CExpr
CMember CExpr
bE (String -> Ident
Ident String
"index") Bool
True)
                     Sing a
_ -> CExpr
bE

         CExpr -> CodeGen ()
putExprStat (CExpr -> CodeGen ()) -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$   (CExpr -> Ident -> Bool -> CExpr
CMember CExpr
loc (String -> Ident
Ident String
"index") Bool
True)
                     CExpr -> CExpr -> CExpr
.=. (CExpr -> CExpr -> CExpr -> CExpr
CCond (CExpr
aE' CExpr -> CExpr -> CExpr
.==. CExpr
bE') (Integer -> CExpr
intE Integer
0) (Integer -> CExpr
intE Integer
1))


flattenPrimOp (Less HOrd a
_) =
  \(abt vars a
a :* abt vars a
b :* SArgs abt args
End) ->
    \CExpr
loc ->
      do CExpr
aE <- abt '[] a -> CodeGen CExpr
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> CodeGen CExpr
flattenWithName abt vars a
abt '[] a
a
         CExpr
bE <- abt '[] a -> CodeGen CExpr
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> CodeGen CExpr
flattenWithName abt vars a
abt '[] a
b
         CExpr -> CodeGen ()
putExprStat (CExpr -> CodeGen ()) -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ (CExpr -> Ident -> Bool -> CExpr
CMember CExpr
loc (String -> Ident
Ident String
"index") Bool
True)
                     CExpr -> CExpr -> CExpr
.=. (CExpr -> CExpr -> CExpr -> CExpr
CCond (CExpr
aE CExpr -> CExpr -> CExpr
.<. CExpr
bE) (Integer -> CExpr
intE Integer
0) (Integer -> CExpr
intE Integer
1))

flattenPrimOp (Negate HRing a
_) =
 \(abt vars a
a :* SArgs abt args
End) ->
   \CExpr
loc ->
     do CExpr
aE <- abt '[] a -> CodeGen CExpr
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> CodeGen CExpr
flattenWithName abt vars a
abt '[] a
a
        CExpr -> CodeGen ()
putExprStat (CExpr -> CodeGen ()) -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ CExpr
loc CExpr -> CExpr -> CExpr
.=. (CUnaryOp -> CExpr -> CExpr
CUnary CUnaryOp
CMinOp (CExpr -> CExpr) -> CExpr -> CExpr
forall a b. (a -> b) -> a -> b
$ CExpr
aE)

flattenPrimOp PrimOp typs a
Choose = \(abt vars a
_ :* abt vars a
_ :* SArgs abt args
End) -> String -> CExpr -> CodeGen ()
forall a. HasCallStack => String -> a
error (String -> CExpr -> CodeGen ()) -> String -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ String
"TODO: flattenPrimOp: choose"

flattenPrimOp PrimOp typs a
t  = \SArgs abt args
_ -> String -> CExpr -> CodeGen ()
forall a. HasCallStack => String -> a
error (String -> CExpr -> CodeGen ()) -> String -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ String
"TODO: flattenPrimOp: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ PrimOp typs a -> String
forall a. Show a => a -> String
show PrimOp typs a
t


--------------------------------------------------------------------------------
--                           MeasureOps and Superpose                         --
--------------------------------------------------------------------------------

{-

The sections contains operations in the stochastic monad. See also
(Dirac, MBind, and Plate) found in SCon. Also see Reject found at the top level.

Remember in the C runtime. Measures are housed in a measure function, which
takes an `struct mdata` location. The MeasureOp attempts to store a value at
that location and returns 0 if it fails and 1 if it succeeds in that task.

The functions uniformCG, normalCG, and gammaCG are primitives that will generate
functions and call them (similar to logSumExpCG). The reduce code size and make
samplers a little more readable.

TODO: add inline pragmas to uniformCG, normalCG, and gammaCG
-}

uniformFun :: CFunDef
uniformFun :: CFunDef
uniformFun = [CDeclSpec] -> CDeclr -> [CDecl] -> CStat -> CFunDef
CFunDef (CTypeSpec -> CDeclSpec
CTypeSpec (CTypeSpec -> CDeclSpec) -> [CTypeSpec] -> [CDeclSpec]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [CTypeSpec]
retTyp)
                     (Maybe CPtrDeclr -> CDirectDeclr -> CDeclr
CDeclr Maybe CPtrDeclr
forall a. Maybe a
Nothing (Ident -> CDirectDeclr
CDDeclrIdent Ident
funcId))
                     [Sing 'HReal -> Ident -> CDecl
forall (a :: Hakaru). Sing a -> Ident -> CDecl
typeDeclaration Sing 'HReal
SReal Ident
loId
                     ,Sing 'HReal -> Ident -> CDecl
forall (a :: Hakaru). Sing a -> Ident -> CDecl
typeDeclaration Sing 'HReal
SReal Ident
hiId]
                     ([CCompoundBlockItem] -> CStat
CCompound ([CCompoundBlockItem] -> CStat)
-> ([[CCompoundBlockItem]] -> [CCompoundBlockItem])
-> [[CCompoundBlockItem]]
-> CStat
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [[CCompoundBlockItem]] -> [CCompoundBlockItem]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat
                     ([[CCompoundBlockItem]] -> CStat)
-> [[CCompoundBlockItem]] -> CStat
forall a b. (a -> b) -> a -> b
$ [ CDecl -> CCompoundBlockItem
CBlockDecl (CDecl -> CCompoundBlockItem) -> [CDecl] -> [CCompoundBlockItem]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [CDecl
declMD]
                       , CStat -> CCompoundBlockItem
CBlockStat (CStat -> CCompoundBlockItem) -> [CStat] -> [CCompoundBlockItem]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [CStat]
comment [CStat] -> [CStat] -> [CStat]
forall a. [a] -> [a] -> [a]
++ [CStat
assW,CStat
assS,Maybe CExpr -> CStat
CReturn (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 -> CStat) -> CExpr -> CStat
forall a b. (a -> b) -> a -> b
$ CExpr
mE]]
                     )
  where r :: CExpr
r          = [CTypeSpec] -> CExpr -> CExpr
castTo [CTypeSpec
CDouble] CExpr
randE
        rMax :: CExpr
rMax       = [CTypeSpec] -> CExpr -> CExpr
castTo [CTypeSpec
CDouble] (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
"RAND_MAX")
        retTyp :: [CTypeSpec]
retTyp     = Sing ('HMeasure 'HReal) -> [CTypeSpec]
forall (a :: Hakaru). Sing a -> [CTypeSpec]
buildType (Sing 'HReal -> Sing ('HMeasure 'HReal)
forall (a :: Hakaru). Sing a -> Sing ('HMeasure a)
SMeasure Sing 'HReal
SReal)
        (Ident
mId,CExpr
mE)   = let ident :: Ident
ident = String -> Ident
Ident String
"mdata" in (Ident
ident,Ident -> CExpr
CVar Ident
ident)
        (Ident
loId,CExpr
loE) = let ident :: Ident
ident = String -> Ident
Ident String
"lo" in (Ident
ident,Ident -> CExpr
CVar Ident
ident)
        (Ident
hiId,CExpr
hiE) = let ident :: Ident
ident = String -> Ident
Ident String
"hi" in (Ident
ident,Ident -> CExpr
CVar Ident
ident)
        value :: CExpr
value      = (CExpr
loE CExpr -> CExpr -> CExpr
.+. ((CExpr
r CExpr -> CExpr -> CExpr
./. CExpr
rMax) CExpr -> CExpr -> CExpr
.*. (CExpr
hiE CExpr -> CExpr -> CExpr
.-. CExpr
loE)))
        comment :: [CStat]
comment = (String -> CStat) -> [String] -> [CStat]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap String -> CStat
CComment
          [String
"uniform :: real -> real -> *(mdata real) -> ()"
          ,String
"------------------------------------------------"]
        declMD :: CDecl
declMD     = CTypeSpec -> Ident -> CDecl
buildDeclaration ([CTypeSpec] -> CTypeSpec
forall a. [a] -> a
head [CTypeSpec]
retTyp) Ident
mId
        assW :: CStat
assW       = 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 -> CStat) -> CExpr -> CStat
forall a b. (a -> b) -> a -> b
$ CExpr -> CExpr
mdataWeight CExpr
mE CExpr -> CExpr -> CExpr
.=. (Float -> CExpr
floatE Float
0)
        assS :: CStat
assS       = 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 -> CStat) -> CExpr -> CStat
forall a b. (a -> b) -> a -> b
$ CExpr -> CExpr
mdataSample CExpr
mE CExpr -> CExpr -> CExpr
.=. CExpr
value
        funcId :: Ident
funcId     = String -> Ident
Ident String
"uniform"


uniformCG :: CExpr -> CExpr -> (CExpr -> CodeGen ())
uniformCG :: CExpr -> CExpr -> CExpr -> CodeGen ()
uniformCG CExpr
aE CExpr
bE =
  \CExpr
loc -> do
    Ident
uId <- String -> CodeGen Ident
reserveIdent String
"uniform"
    Sing ('HMeasure 'HReal) -> CodeGen ()
forall (a :: Hakaru). Sing a -> CodeGen ()
extDeclareTypes (Sing 'HReal -> Sing ('HMeasure 'HReal)
forall (a :: Hakaru). Sing a -> Sing ('HMeasure a)
SMeasure Sing 'HReal
SReal)
    CExtDecl -> CodeGen ()
extDeclare (CExtDecl -> CodeGen ())
-> (CFunDef -> CExtDecl) -> CFunDef -> CodeGen ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CFunDef -> CExtDecl
CFunDefExt (CFunDef -> CodeGen ()) -> CFunDef -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ CFunDef
uniformFun
    CExpr -> CodeGen ()
putExprStat (CExpr -> CodeGen ()) -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ CExpr
loc CExpr -> CExpr -> CExpr
.=. CExpr -> [CExpr] -> CExpr
CCall (Ident -> CExpr
CVar Ident
uId) [CExpr
aE,CExpr
bE]


{-
  This is very cryptic, but I assure you it is only building an AST for the
  Marsaglia Polar Method
-}

normalFun :: CFunDef
normalFun :: CFunDef
normalFun = [CDeclSpec] -> CDeclr -> [CDecl] -> CStat -> CFunDef
CFunDef (CTypeSpec -> CDeclSpec
CTypeSpec (CTypeSpec -> CDeclSpec) -> [CTypeSpec] -> [CDeclSpec]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [CTypeSpec]
retTyp)
                    (Maybe CPtrDeclr -> CDirectDeclr -> CDeclr
CDeclr Maybe CPtrDeclr
forall a. Maybe a
Nothing (Ident -> CDirectDeclr
CDDeclrIdent (String -> Ident
Ident String
"normal")))
                    [Sing 'HReal -> Ident -> CDecl
forall (a :: Hakaru). Sing a -> Ident -> CDecl
typeDeclaration Sing 'HReal
SReal Ident
aId
                    ,Sing 'HProb -> Ident -> CDecl
forall (a :: Hakaru). Sing a -> Ident -> CDecl
typeDeclaration Sing 'HProb
SProb Ident
bId ]
                    ( [CCompoundBlockItem] -> CStat
CCompound ([CCompoundBlockItem] -> CStat)
-> ([[CCompoundBlockItem]] -> [CCompoundBlockItem])
-> [[CCompoundBlockItem]]
-> CStat
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [[CCompoundBlockItem]] -> [CCompoundBlockItem]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat
                    ([[CCompoundBlockItem]] -> CStat)
-> [[CCompoundBlockItem]] -> CStat
forall a b. (a -> b) -> a -> b
$ [[CDecl -> CCompoundBlockItem
CBlockDecl CDecl
declMD],[CCompoundBlockItem]
comment,[CCompoundBlockItem]
decls,[CCompoundBlockItem]
stmts])

  where r :: CExpr
r      = [CTypeSpec] -> CExpr -> CExpr
castTo [CTypeSpec
CDouble] CExpr
randE
        rMax :: CExpr
rMax   = [CTypeSpec] -> CExpr -> CExpr
castTo [CTypeSpec
CDouble] (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
"RAND_MAX")
        retTyp :: [CTypeSpec]
retTyp = Sing ('HMeasure 'HReal) -> [CTypeSpec]
forall (a :: Hakaru). Sing a -> [CTypeSpec]
buildType (Sing 'HReal -> Sing ('HMeasure 'HReal)
forall (a :: Hakaru). Sing a -> Sing ('HMeasure a)
SMeasure Sing 'HReal
SReal)
        (Ident
aId,CExpr
aE) = let ident :: Ident
ident = String -> Ident
Ident String
"a" in (Ident
ident,Ident -> CExpr
CVar Ident
ident)
        (Ident
bId,CExpr
bE) = let ident :: Ident
ident = String -> Ident
Ident String
"b" in (Ident
ident,Ident -> CExpr
CVar Ident
ident)
        (Ident
qId,CExpr
qE) = let ident :: Ident
ident = String -> Ident
Ident String
"q" in (Ident
ident,Ident -> CExpr
CVar Ident
ident)
        (Ident
uId,CExpr
uE) = let ident :: Ident
ident = String -> Ident
Ident String
"u" in (Ident
ident,Ident -> CExpr
CVar Ident
ident)
        (Ident
vId,CExpr
vE) = let ident :: Ident
ident = String -> Ident
Ident String
"v" in (Ident
ident,Ident -> CExpr
CVar Ident
ident)
        (Ident
rId,CExpr
rE) = let ident :: Ident
ident = String -> Ident
Ident String
"r" in (Ident
ident,Ident -> CExpr
CVar Ident
ident)
        (Ident
mId,CExpr
mE) = let ident :: Ident
ident = String -> Ident
Ident String
"mdata" in (Ident
ident,Ident -> CExpr
CVar Ident
ident)
        declMD :: CDecl
declMD     = CTypeSpec -> Ident -> CDecl
buildDeclaration ([CTypeSpec] -> CTypeSpec
forall a. [a] -> a
head [CTypeSpec]
retTyp) Ident
mId
        draw :: CExpr -> CStat
draw CExpr
xE = 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 -> CStat) -> CExpr -> CStat
forall a b. (a -> b) -> a -> b
$ CExpr
xE CExpr -> CExpr -> CExpr
.=. (((CExpr
r CExpr -> CExpr -> CExpr
./. CExpr
rMax) CExpr -> CExpr -> CExpr
.*. (Float -> CExpr
floatE Float
2)) CExpr -> CExpr -> CExpr
.-. (Float -> CExpr
floatE Float
1))
        body :: CStat
body = [CStat] -> CStat
seqCStat [CExpr -> CStat
draw CExpr
uE
                        ,CExpr -> CStat
draw CExpr
vE
                        ,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 -> CStat) -> CExpr -> CStat
forall a b. (a -> b) -> a -> b
$ CExpr
qE CExpr -> CExpr -> CExpr
.=. ((CExpr
uE CExpr -> CExpr -> CExpr
.*. CExpr
uE) CExpr -> CExpr -> CExpr
.+. (CExpr
vE CExpr -> CExpr -> CExpr
.*. CExpr
vE))]
        polar :: CStat
polar = CExpr -> CStat -> Bool -> CStat
CWhile (CExpr
qE CExpr -> CExpr -> CExpr
.>. (Float -> CExpr
floatE Float
1)) CStat
body Bool
True
        setR :: CStat
setR  = 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 -> CStat) -> CExpr -> CStat
forall a b. (a -> b) -> a -> b
$ CExpr
rE CExpr -> CExpr -> CExpr
.=. (CExpr -> CExpr
sqrtE (((CUnaryOp -> CExpr -> CExpr
CUnary CUnaryOp
CMinOp (Float -> CExpr
floatE Float
2)) CExpr -> CExpr -> CExpr
.*. CExpr -> CExpr
logE CExpr
qE) CExpr -> CExpr -> CExpr
./. CExpr
qE))
        finalValue :: CExpr
finalValue = CExpr
aE CExpr -> CExpr -> CExpr
.+. (CExpr
uE CExpr -> CExpr -> CExpr
.*. CExpr
rE CExpr -> CExpr -> CExpr
.*. CExpr
bE)
        comment :: [CCompoundBlockItem]
comment = (String -> CCompoundBlockItem) -> [String] -> [CCompoundBlockItem]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (CStat -> CCompoundBlockItem
CBlockStat (CStat -> CCompoundBlockItem)
-> (String -> CStat) -> String -> CCompoundBlockItem
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> CStat
CComment)
          [String
"normal :: real -> real -> *(mdata real) -> ()"
          ,String
"Marsaglia Polar Method"
          ,String
"-----------------------------------------------"]
        decls :: [CCompoundBlockItem]
decls = (CDecl -> CCompoundBlockItem
CBlockDecl (CDecl -> CCompoundBlockItem)
-> (Ident -> CDecl) -> Ident -> CCompoundBlockItem
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) (Ident -> CCompoundBlockItem) -> [Ident] -> [CCompoundBlockItem]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Ident
uId,Ident
vId,Ident
qId,Ident
rId]
        stmts :: [CCompoundBlockItem]
stmts = CStat -> CCompoundBlockItem
CBlockStat (CStat -> CCompoundBlockItem) -> [CStat] -> [CCompoundBlockItem]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [CStat
polar,CStat
setR, CStat
assW, CStat
assS,Maybe CExpr -> CStat
CReturn (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 -> CStat) -> CExpr -> CStat
forall a b. (a -> b) -> a -> b
$ CExpr
mE]
        assW :: CStat
assW = 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 -> CStat) -> CExpr -> CStat
forall a b. (a -> b) -> a -> b
$ CExpr -> CExpr
mdataWeight CExpr
mE CExpr -> CExpr -> CExpr
.=. (Float -> CExpr
floatE Float
0)
        assS :: CStat
assS = 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 -> CStat) -> CExpr -> CStat
forall a b. (a -> b) -> a -> b
$ CExpr -> CExpr
mdataSample CExpr
mE CExpr -> CExpr -> CExpr
.=. CExpr
finalValue


normalCG :: CExpr -> CExpr -> (CExpr -> CodeGen ())
normalCG :: CExpr -> CExpr -> CExpr -> CodeGen ()
normalCG CExpr
aE CExpr
bE =
  \CExpr
loc -> do
    Ident
nId <- String -> CodeGen Ident
reserveIdent String
"normal"
    Sing ('HMeasure 'HReal) -> CodeGen ()
forall (a :: Hakaru). Sing a -> CodeGen ()
extDeclareTypes (Sing 'HReal -> Sing ('HMeasure 'HReal)
forall (a :: Hakaru). Sing a -> Sing ('HMeasure a)
SMeasure Sing 'HReal
SReal)
    CExtDecl -> CodeGen ()
extDeclare (CExtDecl -> CodeGen ())
-> (CFunDef -> CExtDecl) -> CFunDef -> CodeGen ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CFunDef -> CExtDecl
CFunDefExt (CFunDef -> CodeGen ()) -> CFunDef -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ CFunDef
normalFun
    CExpr -> CodeGen ()
putExprStat (CExpr -> CodeGen ()) -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ CExpr
loc CExpr -> CExpr -> CExpr
.=. (CExpr -> [CExpr] -> CExpr
CCall (Ident -> CExpr
CVar Ident
nId) [CExpr
aE,CExpr
bE])

{-
  This method is from Marsaglia and Tsang "a simple method for generating gamma variables"
-}
gammaFun :: CFunDef
gammaFun :: CFunDef
gammaFun = [CDeclSpec] -> CDeclr -> [CDecl] -> CStat -> CFunDef
CFunDef (CTypeSpec -> CDeclSpec
CTypeSpec (CTypeSpec -> CDeclSpec) -> [CTypeSpec] -> [CDeclSpec]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [CTypeSpec]
retTyp)
                   (Maybe CPtrDeclr -> CDirectDeclr -> CDeclr
CDeclr Maybe CPtrDeclr
forall a. Maybe a
Nothing (Ident -> CDirectDeclr
CDDeclrIdent (String -> Ident
Ident String
"gamma")))
                   [Sing 'HProb -> Ident -> CDecl
forall (a :: Hakaru). Sing a -> Ident -> CDecl
typeDeclaration Sing 'HProb
SProb Ident
aId
                   ,Sing 'HProb -> Ident -> CDecl
forall (a :: Hakaru). Sing a -> Ident -> CDecl
typeDeclaration Sing 'HProb
SProb Ident
bId]
                    ( [CCompoundBlockItem] -> CStat
CCompound ([CCompoundBlockItem] -> CStat)
-> ([[CCompoundBlockItem]] -> [CCompoundBlockItem])
-> [[CCompoundBlockItem]]
-> CStat
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [[CCompoundBlockItem]] -> [CCompoundBlockItem]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat
                    ([[CCompoundBlockItem]] -> CStat)
-> [[CCompoundBlockItem]] -> CStat
forall a b. (a -> b) -> a -> b
$ [[CDecl -> CCompoundBlockItem
CBlockDecl CDecl
declMD],[CCompoundBlockItem]
comment,[CCompoundBlockItem]
decls,[CCompoundBlockItem]
stmts])
  where (Ident
aId,CExpr
aE) = let ident :: Ident
ident = String -> Ident
Ident String
"a" in (Ident
ident,Ident -> CExpr
CVar Ident
ident)
        (Ident
bId,CExpr
bE) = let ident :: Ident
ident = String -> Ident
Ident String
"b" in (Ident
ident,Ident -> CExpr
CVar Ident
ident)
        (Ident
cId,CExpr
cE) = let ident :: Ident
ident = String -> Ident
Ident String
"c" in (Ident
ident,Ident -> CExpr
CVar Ident
ident)
        (Ident
dId,CExpr
dE) = let ident :: Ident
ident = String -> Ident
Ident String
"d" in (Ident
ident,Ident -> CExpr
CVar Ident
ident)
        (Ident
xId,CExpr
xE) = let ident :: Ident
ident = String -> Ident
Ident String
"x" in (Ident
ident,Ident -> CExpr
CVar Ident
ident)
        (Ident
vId,CExpr
vE) = let ident :: Ident
ident = String -> Ident
Ident String
"v" in (Ident
ident,Ident -> CExpr
CVar Ident
ident)
        (Ident
uId,CExpr
uE) = let ident :: Ident
ident = String -> Ident
Ident String
"u" in (Ident
ident,Ident -> CExpr
CVar Ident
ident)
        (Ident
mId,CExpr
mE) = let ident :: Ident
ident = String -> Ident
Ident String
"mdata" in (Ident
ident,Ident -> CExpr
CVar Ident
ident)
        retTyp :: [CTypeSpec]
retTyp = Sing ('HMeasure 'HProb) -> [CTypeSpec]
forall (a :: Hakaru). Sing a -> [CTypeSpec]
buildType (Sing 'HProb -> Sing ('HMeasure 'HProb)
forall (a :: Hakaru). Sing a -> Sing ('HMeasure a)
SMeasure Sing 'HProb
SProb)
        declMD :: CDecl
declMD     = CTypeSpec -> Ident -> CDecl
buildDeclaration ([CTypeSpec] -> CTypeSpec
forall a. [a] -> a
head [CTypeSpec]
retTyp) Ident
mId
        comment :: [CCompoundBlockItem]
comment = (String -> CCompoundBlockItem) -> [String] -> [CCompoundBlockItem]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (CStat -> CCompoundBlockItem
CBlockStat (CStat -> CCompoundBlockItem)
-> (String -> CStat) -> String -> CCompoundBlockItem
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> CStat
CComment)
          [String
"gamma :: real -> prob -> *(mdata prob) -> ()"
          ,String
"Marsaglia and Tsang 'a simple method for generating gamma variables'"
          ,String
"--------------------------------------------------------------------"]
        decls :: [CCompoundBlockItem]
decls = (CDecl -> CCompoundBlockItem) -> [CDecl] -> [CCompoundBlockItem]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap CDecl -> CCompoundBlockItem
CBlockDecl ([CDecl] -> [CCompoundBlockItem])
-> [CDecl] -> [CCompoundBlockItem]
forall a b. (a -> b) -> a -> b
$ ((Ident -> CDecl) -> [Ident] -> [CDecl]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Sing 'HReal -> Ident -> CDecl
forall (a :: Hakaru). Sing a -> Ident -> CDecl
typeDeclaration Sing 'HReal
SReal) [Ident
dId,Ident
cId,Ident
vId])
                               [CDecl] -> [CDecl] -> [CDecl]
forall a. [a] -> [a] -> [a]
++ ((Ident -> CDecl) -> [Ident] -> [CDecl]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Sing ('HMeasure 'HReal) -> Ident -> CDecl
forall (a :: Hakaru). Sing a -> Ident -> CDecl
typeDeclaration (Sing 'HReal -> Sing ('HMeasure 'HReal)
forall (a :: Hakaru). Sing a -> Sing ('HMeasure a)
SMeasure Sing 'HReal
SReal)) [Ident
uId,Ident
xId])
        stmts :: [CCompoundBlockItem]
stmts = (CStat -> CCompoundBlockItem) -> [CStat] -> [CCompoundBlockItem]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap CStat -> CCompoundBlockItem
CBlockStat ([CStat] -> [CCompoundBlockItem])
-> [CStat] -> [CCompoundBlockItem]
forall a b. (a -> b) -> a -> b
$ [CStat
assD,CStat
assC,CStat
outerWhile]
        xS :: CExpr
xS = CExpr -> CExpr
mdataSample CExpr
xE
        uS :: CExpr
uS = CExpr -> CExpr
mdataSample CExpr
uE
        assD :: CStat
assD = 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 -> CStat) -> CExpr -> CStat
forall a b. (a -> b) -> a -> b
$ CExpr
dE CExpr -> CExpr -> CExpr
.=. (CExpr
aE CExpr -> CExpr -> CExpr
.-. ((Float -> CExpr
floatE Float
1) CExpr -> CExpr -> CExpr
./. (Float -> CExpr
floatE Float
3)))
        assC :: CStat
assC = 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 -> CStat) -> CExpr -> CStat
forall a b. (a -> b) -> a -> b
$ CExpr
cE CExpr -> CExpr -> CExpr
.=. ((Float -> CExpr
floatE Float
1) CExpr -> CExpr -> CExpr
./. (CExpr -> CExpr
sqrtE ((Float -> CExpr
floatE Float
9) CExpr -> CExpr -> CExpr
.*. CExpr
dE)))
        outerWhile :: CStat
outerWhile = CExpr -> CStat -> Bool -> CStat
CWhile (Integer -> CExpr
intE Integer
1) ([CStat] -> CStat
seqCStat [CStat
innerWhile,CStat
assV,CStat
assU,CStat
exit]) Bool
False
        innerWhile :: CStat
innerWhile = CExpr -> CStat -> Bool -> CStat
CWhile (CExpr
vE CExpr -> CExpr -> CExpr
.<=. (Float -> CExpr
floatE Float
0)) ([CStat] -> CStat
seqCStat [CStat
assX,CStat
assVIn]) Bool
True
        assX :: CStat
assX = 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 -> CStat) -> CExpr -> CStat
forall a b. (a -> b) -> a -> b
$ CExpr
xE CExpr -> CExpr -> CExpr
.=. (CExpr -> [CExpr] -> CExpr
CCall (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
"normal") [(Float -> CExpr
floatE Float
0),(Float -> CExpr
floatE Float
1)])
        assVIn :: CStat
assVIn = 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 -> CStat) -> CExpr -> CStat
forall a b. (a -> b) -> a -> b
$ CExpr
vE CExpr -> CExpr -> CExpr
.=. ((Float -> CExpr
floatE Float
1) CExpr -> CExpr -> CExpr
.+. (CExpr
cE CExpr -> CExpr -> CExpr
.*. CExpr
xS))
        assV :: CStat
assV = 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 -> CStat) -> CExpr -> CStat
forall a b. (a -> b) -> a -> b
$ CExpr
vE CExpr -> CExpr -> CExpr
.=. (CExpr
vE CExpr -> CExpr -> CExpr
.*. CExpr
vE CExpr -> CExpr -> CExpr
.*. CExpr
vE)
        assU :: CStat
assU = 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 -> CStat) -> CExpr -> CStat
forall a b. (a -> b) -> a -> b
$ CExpr
uE CExpr -> CExpr -> CExpr
.=. (CExpr -> [CExpr] -> CExpr
CCall (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
"uniform") [(Float -> CExpr
floatE Float
0),(Float -> CExpr
floatE Float
1)])
        exitC1 :: CExpr
exitC1 = CExpr
uS CExpr -> CExpr -> CExpr
.<. ((Float -> CExpr
floatE Float
1) CExpr -> CExpr -> CExpr
.-. ((Float -> CExpr
floatE Float
0.331 CExpr -> CExpr -> CExpr
.*. (CExpr
xS CExpr -> CExpr -> CExpr
.*. CExpr
xS) CExpr -> CExpr -> CExpr
.*. (CExpr
xS CExpr -> CExpr -> CExpr
.*. CExpr
xS))))
        exitC2 :: CExpr
exitC2 = (CExpr -> CExpr
logE CExpr
uS) CExpr -> CExpr -> CExpr
.<. (((Float -> CExpr
floatE Float
0.5) CExpr -> CExpr -> CExpr
.*. (CExpr
xS CExpr -> CExpr -> CExpr
.*. CExpr
xS)) CExpr -> CExpr -> CExpr
.+. (CExpr
dE CExpr -> CExpr -> CExpr
.*. ((Float -> CExpr
floatE Float
1.0) CExpr -> CExpr -> CExpr
.-. CExpr
vE CExpr -> CExpr -> CExpr
.+. (CExpr -> CExpr
logE CExpr
vE))))
        assW :: CStat
assW = 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 -> CStat) -> CExpr -> CStat
forall a b. (a -> b) -> a -> b
$ CExpr -> CExpr
mdataWeight CExpr
mE CExpr -> CExpr -> CExpr
.=. (Float -> CExpr
floatE Float
0)
        assS :: CStat
assS = 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 -> CStat) -> CExpr -> CStat
forall a b. (a -> b) -> a -> b
$ CExpr -> CExpr
mdataSample CExpr
mE CExpr -> CExpr -> CExpr
.=. (CExpr -> CExpr
logE (CExpr
dE CExpr -> CExpr -> CExpr
.*. CExpr
vE)) CExpr -> CExpr -> CExpr
.+. CExpr
bE
        exit :: CStat
exit = CExpr -> CStat -> Maybe CStat -> CStat
CIf (CExpr
exitC1 CExpr -> CExpr -> CExpr
.||. CExpr
exitC2) ([CStat] -> CStat
seqCStat [CStat
assW,CStat
assS,Maybe CExpr -> CStat
CReturn (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 -> CStat) -> CExpr -> CStat
forall a b. (a -> b) -> a -> b
$ CExpr
mE]) Maybe CStat
forall a. Maybe a
Nothing


gammaCG :: CExpr -> CExpr -> (CExpr -> CodeGen ())
gammaCG :: CExpr -> CExpr -> CExpr -> CodeGen ()
gammaCG CExpr
aE CExpr
bE =
  \CExpr
loc -> do
     Sing ('HMeasure 'HReal) -> CodeGen ()
forall (a :: Hakaru). Sing a -> CodeGen ()
extDeclareTypes (Sing 'HReal -> Sing ('HMeasure 'HReal)
forall (a :: Hakaru). Sing a -> Sing ('HMeasure a)
SMeasure Sing 'HReal
SReal)
     String -> CodeGen Ident
reserveIdent String
"uniform"
     String -> CodeGen Ident
reserveIdent String
"normal"
     Ident
gId <- String -> CodeGen Ident
reserveIdent String
"gamma"
     (CFunDef -> CodeGen ()) -> [CFunDef] -> CodeGen ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (CExtDecl -> CodeGen ()
extDeclare (CExtDecl -> CodeGen ())
-> (CFunDef -> CExtDecl) -> CFunDef -> CodeGen ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CFunDef -> CExtDecl
CFunDefExt) [CFunDef
uniformFun,CFunDef
normalFun,CFunDef
gammaFun]
     CExpr -> CodeGen ()
putExprStat (CExpr -> CodeGen ()) -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ CExpr
loc CExpr -> CExpr -> CExpr
.=. (CExpr -> [CExpr] -> CExpr
CCall (Ident -> CExpr
CVar Ident
gId) [CExpr
aE,CExpr
bE])


flattenMeasureOp
  :: forall abt typs args a .
     ( ABT Term abt
     , typs ~ UnLCs args
     , args ~ LCs typs )
  => MeasureOp typs a
  -> SArgs abt args
  -> (CExpr -> CodeGen ())


flattenMeasureOp :: MeasureOp typs a -> SArgs abt args -> CExpr -> CodeGen ()
flattenMeasureOp MeasureOp typs a
Uniform =
  \(abt vars a
a :* abt vars a
b :* SArgs abt args
End) ->
    \CExpr
loc ->
      do CExpr
aE <- abt '[] a -> CodeGen CExpr
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> CodeGen CExpr
flattenWithName abt vars a
abt '[] a
a
         CExpr
bE <- abt '[] a -> CodeGen CExpr
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> CodeGen CExpr
flattenWithName abt vars a
abt '[] a
b
         CExpr -> CExpr -> CExpr -> CodeGen ()
uniformCG CExpr
aE CExpr
bE CExpr
loc


flattenMeasureOp MeasureOp typs a
Normal  =
  \(abt vars a
a :* abt vars a
b :* SArgs abt args
End) ->
    \CExpr
loc ->
      do CExpr
aE <- abt '[] a -> CodeGen CExpr
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> CodeGen CExpr
flattenWithName abt vars a
abt '[] a
a
         CExpr
bE <- abt '[] a -> CodeGen CExpr
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> CodeGen CExpr
flattenWithName abt vars a
abt '[] a
b
         CExpr -> CExpr -> CExpr -> CodeGen ()
normalCG CExpr
aE (CExpr -> CExpr
expE CExpr
bE) CExpr
loc

flattenMeasureOp MeasureOp typs a
Poisson =
  \(abt vars a
lam :* SArgs abt args
End) ->
    \CExpr
loc ->
      do CExpr
lamE <- abt '[] a -> CodeGen CExpr
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> CodeGen CExpr
flattenWithName abt vars a
abt '[] a
lam
         Ident
lId <- String -> CodeGen Ident
genIdent' String
"l"
         Ident
kId <- String -> CodeGen Ident
genIdent' String
"k"
         Ident
pId <- String -> CodeGen Ident
genIdent' String
"p"
         Sing 'HProb -> Ident -> CodeGen ()
forall (a :: Hakaru). Sing a -> Ident -> CodeGen ()
declare Sing 'HProb
SProb Ident
lId
         Sing 'HNat -> Ident -> CodeGen ()
forall (a :: Hakaru). Sing a -> Ident -> CodeGen ()
declare Sing 'HNat
SNat  Ident
kId
         Sing 'HProb -> Ident -> CodeGen ()
forall (a :: Hakaru). Sing a -> Ident -> CodeGen ()
declare Sing 'HProb
SProb Ident
pId
         let (CExpr
lE:CExpr
kE:CExpr
pE:[]) = (Ident -> CExpr) -> [Ident] -> [CExpr]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Ident -> CExpr
CVar [Ident
lId,Ident
kId,Ident
pId]
         CExpr -> CodeGen ()
putExprStat (CExpr -> CodeGen ()) -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ CExpr
lE CExpr -> CExpr -> CExpr
.=. (CExpr -> CExpr
expE (CUnaryOp -> CExpr -> CExpr
CUnary CUnaryOp
CMinOp (CExpr -> CExpr) -> CExpr -> CExpr
forall a b. (a -> b) -> a -> b
$ CExpr -> CExpr
expE CExpr
lamE))
         CExpr -> CodeGen ()
putExprStat (CExpr -> CodeGen ()) -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ CExpr
kE CExpr -> CExpr -> CExpr
.=. (Integer -> CExpr
intE Integer
0)
         CExpr -> CodeGen ()
putExprStat (CExpr -> CodeGen ()) -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ CExpr
pE CExpr -> CExpr -> CExpr
.=. (Float -> CExpr
floatE Float
1)
         CExpr -> CodeGen () -> CodeGen ()
doWhileCG (CExpr
pE CExpr -> CExpr -> CExpr
.>. CExpr
lE) (CodeGen () -> CodeGen ()) -> CodeGen () -> CodeGen ()
forall a b. (a -> b) -> a -> b
$
           do Ident
uId <- String -> CodeGen Ident
genIdent' String
"u"
              Sing ('HMeasure 'HReal) -> Ident -> CodeGen ()
forall (a :: Hakaru). Sing a -> Ident -> CodeGen ()
declare (Sing 'HReal -> Sing ('HMeasure 'HReal)
forall (a :: Hakaru). Sing a -> Sing ('HMeasure a)
SMeasure Sing 'HReal
SReal) Ident
uId
              let uE :: CExpr
uE = Ident -> CExpr
CVar Ident
uId
              CExpr -> CExpr -> CExpr -> CodeGen ()
uniformCG (Integer -> CExpr
intE Integer
0) (Integer -> CExpr
intE Integer
1) CExpr
uE
              CExpr -> CodeGen ()
putExprStat (CExpr -> CodeGen ()) -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ CExpr
pE CExpr -> CExpr -> CExpr
.*=. (CExpr -> CExpr
mdataSample CExpr
uE)
              CExpr -> CodeGen ()
putExprStat (CExpr -> CodeGen ()) -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ CExpr
kE CExpr -> CExpr -> CExpr
.+=. (Integer -> CExpr
intE Integer
1)

         CExpr -> CodeGen ()
putExprStat (CExpr -> CodeGen ()) -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ CExpr -> CExpr
mdataWeight CExpr
loc CExpr -> CExpr -> CExpr
.=. (Integer -> CExpr
intE Integer
0)
         CExpr -> CodeGen ()
putExprStat (CExpr -> CodeGen ()) -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ CExpr -> CExpr
mdataSample CExpr
loc CExpr -> CExpr -> CExpr
.=. (CExpr
kE CExpr -> CExpr -> CExpr
.-. (Integer -> CExpr
intE Integer
1))

flattenMeasureOp MeasureOp typs a
Gamma =
  \(abt vars a
a :* abt vars a
b :* SArgs abt args
End) ->
    \CExpr
loc ->
      do CExpr
aE <- abt '[] a -> CodeGen CExpr
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> CodeGen CExpr
flattenWithName abt vars a
abt '[] a
a
         CExpr
bE <- abt '[] a -> CodeGen CExpr
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> CodeGen CExpr
flattenWithName abt vars a
abt '[] a
b
         CExpr -> CExpr -> CExpr -> CodeGen ()
gammaCG (CExpr -> CExpr
expE CExpr
aE) CExpr
bE CExpr
loc


flattenMeasureOp MeasureOp typs a
Beta =
  \(abt vars a
a :* abt vars a
b :* SArgs abt args
End) -> abt '[] ('HMeasure 'HProb) -> CExpr -> CodeGen ()
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> CExpr -> CodeGen ()
flattenABT (abt '[] 'HProb -> abt '[] 'HProb -> abt '[] ('HMeasure 'HProb)
forall (abt :: [Hakaru] -> Hakaru -> *).
ABT Term abt =>
abt '[] 'HProb -> abt '[] 'HProb -> abt '[] ('HMeasure 'HProb)
HKP.beta'' abt vars a
abt '[] 'HProb
a abt vars a
abt '[] 'HProb
b)


-- I ran into a bug here where sometime I recieved a location by reference and
-- others by value. Since measureOps assign a sample to mdata that they have a
-- reference to, we should enforce that when passing around mdata it is by
-- reference
flattenMeasureOp MeasureOp typs a
Categorical = \(abt vars a
arr :* SArgs abt args
End) ->
  \CExpr
loc ->
    do CExpr
arrE <- abt '[] a -> CodeGen CExpr
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> CodeGen CExpr
flattenWithName abt vars a
abt '[] a
arr

       Ident
itId <- String -> CodeGen Ident
genIdent' String
"it"
       Sing 'HNat -> Ident -> CodeGen ()
forall (a :: Hakaru). Sing a -> Ident -> CodeGen ()
declare Sing 'HNat
SNat Ident
itId
       let itE :: CExpr
itE = Ident -> CExpr
CVar Ident
itId

       -- Accumulator for the total probability of the input array
       Ident
wSumId <- String -> CodeGen Ident
genIdent' String
"ws"
       Sing 'HProb -> Ident -> CodeGen ()
forall (a :: Hakaru). Sing a -> Ident -> CodeGen ()
declare Sing 'HProb
SProb Ident
wSumId
       let wSumE :: CExpr
wSumE = Ident -> CExpr
CVar Ident
wSumId
       Ident -> CExpr -> CodeGen ()
assign Ident
wSumId (CExpr -> CExpr
logE (Integer -> CExpr
intE Integer
0))

       -- Accumulator for the max value in the input array
       Ident
wMaxId <- String -> CodeGen Ident
genIdent' String
"max"
       Sing 'HProb -> Ident -> CodeGen ()
forall (a :: Hakaru). Sing a -> Ident -> CodeGen ()
declare Sing 'HProb
SProb Ident
wMaxId
       let wMaxE :: CExpr
wMaxE = Ident -> CExpr
CVar Ident
wMaxId
       Ident -> CExpr -> CodeGen ()
assign Ident
wMaxId (CExpr -> CExpr
logE (Float -> CExpr
floatE Float
0))

       let currE :: CExpr
currE = CExpr -> CExpr -> CExpr
index (CExpr -> CExpr
arrayData CExpr
arrE) CExpr
itE

       CodeGen () -> CodeGen ()
forall a. CodeGen a -> CodeGen a
seqDo (CodeGen () -> CodeGen ()) -> CodeGen () -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ do
         -- Calculate the maximum value of the input array
         -- And calculate the total weight
         CExpr -> CExpr -> CExpr -> CodeGen () -> CodeGen ()
forCG (CExpr
itE CExpr -> CExpr -> CExpr
.=. (Integer -> CExpr
intE Integer
0))
               (CExpr
itE CExpr -> CExpr -> CExpr
.<. (CExpr -> CExpr
arraySize CExpr
arrE))
               (CUnaryOp -> CExpr -> CExpr
CUnary CUnaryOp
CPostIncOp CExpr
itE) (CodeGen () -> CodeGen ()) -> CodeGen () -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ do
           CExpr -> CodeGen () -> CodeGen () -> CodeGen ()
ifCG (CExpr
wMaxE CExpr -> CExpr -> CExpr
.<. CExpr
currE)
                (CExpr -> CodeGen ()
putExprStat (CExpr -> CodeGen ()) -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ CExpr
wMaxE CExpr -> CExpr -> CExpr
.=. CExpr
currE)
                (() -> CodeGen ()
forall (m :: * -> *) a. Monad m => a -> m a
return ())
           Seq CExpr -> CExpr -> CodeGen ()
logSumExpCG ([CExpr] -> Seq CExpr
forall a. [a] -> Seq a
S.fromList [CExpr
wSumE, CExpr
currE]) CExpr
wSumE
         CExpr -> CodeGen ()
putExprStat (CExpr -> CodeGen ()) -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ CExpr
wSumE CExpr -> CExpr -> CExpr
.=. (CExpr
wSumE CExpr -> CExpr -> CExpr
.-. CExpr
wMaxE)

         -- draw number from uniform(0, weightSum)
         Ident
rId <- String -> CodeGen Ident
genIdent' String
"r"
         Sing 'HReal -> Ident -> CodeGen ()
forall (a :: Hakaru). Sing a -> Ident -> CodeGen ()
declare Sing 'HReal
SReal Ident
rId
         let r :: CExpr
r    = [CTypeSpec] -> CExpr -> CExpr
castTo [CTypeSpec
CDouble] CExpr
randE
             rMax :: CExpr
rMax = [CTypeSpec] -> CExpr -> CExpr
castTo [CTypeSpec
CDouble] (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
"RAND_MAX")
             rE :: CExpr
rE = Ident -> CExpr
CVar Ident
rId
         Ident -> CExpr -> CodeGen ()
assign Ident
rId (CExpr -> CExpr
logE (CExpr
r CExpr -> CExpr -> CExpr
./. CExpr
rMax) CExpr -> CExpr -> CExpr
.+. CExpr
wSumE)

         Ident -> CExpr -> CodeGen ()
assign Ident
wSumId (CExpr -> CExpr
logE (Integer -> CExpr
intE Integer
0))
         Ident -> CExpr -> CodeGen ()
assign Ident
itId (Integer -> CExpr
intE Integer
0)
         CExpr -> CodeGen () -> CodeGen ()
whileCG (Integer -> CExpr
intE Integer
1) (CodeGen () -> CodeGen ()) -> CodeGen () -> CodeGen ()
forall a b. (a -> b) -> a -> b
$
           do CExpr -> CodeGen () -> CodeGen () -> CodeGen ()
ifCG (CExpr
rE CExpr -> CExpr -> CExpr
.<. CExpr
wSumE)
                   (do CExpr -> CodeGen ()
putExprStat (CExpr -> CodeGen ()) -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ CExpr -> CExpr
mdataWeight CExpr
loc CExpr -> CExpr -> CExpr
.=. (Integer -> CExpr
intE Integer
0)
                       CExpr -> CodeGen ()
putExprStat (CExpr -> CodeGen ()) -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ CExpr -> CExpr
mdataSample CExpr
loc CExpr -> CExpr -> CExpr
.=. (CExpr
itE CExpr -> CExpr -> CExpr
.-. (Integer -> CExpr
intE Integer
1))
                       CStat -> CodeGen ()
putStat CStat
CBreak)
                   (() -> CodeGen ()
forall (m :: * -> *) a. Monad m => a -> m a
return ())
              Seq CExpr -> CExpr -> CodeGen ()
logSumExpCG ([CExpr] -> Seq CExpr
forall a. [a] -> Seq a
S.fromList [CExpr
wSumE, CExpr
currE CExpr -> CExpr -> CExpr
.-. CExpr
wMaxE]) CExpr
wSumE
              CExpr -> CodeGen ()
putExprStat (CExpr -> CodeGen ()) -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ CUnaryOp -> CExpr -> CExpr
CUnary CUnaryOp
CPostIncOp CExpr
itE


flattenMeasureOp MeasureOp typs a
x = String -> SArgs abt args -> CExpr -> CodeGen ()
forall a. HasCallStack => String -> a
error (String -> SArgs abt args -> CExpr -> CodeGen ())
-> String -> SArgs abt args -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ String
"TODO: flattenMeasureOp: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ MeasureOp typs a -> String
forall a. Show a => a -> String
show MeasureOp typs a
x

---------------
-- Superpose --
---------------

flattenSuperpose
    :: (ABT Term abt)
    => NE.NonEmpty (abt '[] 'HProb, abt '[] ('HMeasure a))
    -> (CExpr -> CodeGen ())

-- do we need to normalize?
flattenSuperpose :: NonEmpty (abt '[] 'HProb, abt '[] ('HMeasure a))
-> CExpr -> CodeGen ()
flattenSuperpose NonEmpty (abt '[] 'HProb, abt '[] ('HMeasure a))
pairs =
  let pairs' :: [(abt '[] 'HProb, abt '[] ('HMeasure a))]
pairs' = NonEmpty (abt '[] 'HProb, abt '[] ('HMeasure a))
-> [(abt '[] 'HProb, abt '[] ('HMeasure a))]
forall a. NonEmpty a -> [a]
NE.toList NonEmpty (abt '[] 'HProb, abt '[] ('HMeasure a))
pairs in

  if [(abt '[] 'HProb, abt '[] ('HMeasure a))] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [(abt '[] 'HProb, abt '[] ('HMeasure a))]
pairs' Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1
  then \CExpr
loc -> let (abt '[] 'HProb
w,abt '[] ('HMeasure a)
m) = [(abt '[] 'HProb, abt '[] ('HMeasure a))]
-> (abt '[] 'HProb, abt '[] ('HMeasure a))
forall a. [a] -> a
head [(abt '[] 'HProb, abt '[] ('HMeasure a))]
pairs' in
         do CExpr
mE <- abt '[] ('HMeasure a) -> CodeGen CExpr
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> CodeGen CExpr
flattenWithName abt '[] ('HMeasure a)
m
            CExpr
wE <- abt '[] 'HProb -> CodeGen CExpr
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> CodeGen CExpr
flattenWithName abt '[] 'HProb
w
            CExpr -> CodeGen ()
putExprStat (CExpr -> CodeGen ()) -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ CExpr -> CExpr
mdataWeight CExpr
loc CExpr -> CExpr -> CExpr
.=. ((CExpr -> CExpr
mdataWeight CExpr
mE) CExpr -> CExpr -> CExpr
.+. CExpr
wE)
            CExpr -> CodeGen ()
putExprStat (CExpr -> CodeGen ()) -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ CExpr -> CExpr
mdataSample CExpr
loc CExpr -> CExpr -> CExpr
.=. (CExpr -> CExpr
mdataSample CExpr
mE)

  else \CExpr
loc ->
         do [CExpr]
wEs <- ((abt '[] 'HProb, abt '[] ('HMeasure a)) -> CodeGen CExpr)
-> [(abt '[] 'HProb, abt '[] ('HMeasure a))]
-> StateT CG Identity [CExpr]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (\(abt '[] 'HProb
w,abt '[] ('HMeasure a)
_) -> abt '[] 'HProb -> String -> CodeGen CExpr
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> String -> CodeGen CExpr
flattenWithName' abt '[] 'HProb
w String
"w") [(abt '[] 'HProb, abt '[] ('HMeasure a))]
pairs'

            Ident
wSumId <- String -> CodeGen Ident
genIdent' String
"ws"
            Sing 'HProb -> Ident -> CodeGen ()
forall (a :: Hakaru). Sing a -> Ident -> CodeGen ()
declare Sing 'HProb
SProb Ident
wSumId
            let wSumE :: CExpr
wSumE = Ident -> CExpr
CVar Ident
wSumId
            Seq CExpr -> CExpr -> CodeGen ()
logSumExpCG ([CExpr] -> Seq CExpr
forall a. [a] -> Seq a
S.fromList [CExpr]
wEs) CExpr
wSumE

            -- draw number from uniform(0, weightSum)
            Ident
rId <- String -> CodeGen Ident
genIdent' String
"r"
            Sing 'HReal -> Ident -> CodeGen ()
forall (a :: Hakaru). Sing a -> Ident -> CodeGen ()
declare Sing 'HReal
SReal Ident
rId
            let r :: CExpr
r    = [CTypeSpec] -> CExpr -> CExpr
castTo [CTypeSpec
CDouble] CExpr
randE
                rMax :: CExpr
rMax = [CTypeSpec] -> CExpr -> CExpr
castTo [CTypeSpec
CDouble] (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
"RAND_MAX")
                rE :: CExpr
rE = Ident -> CExpr
CVar Ident
rId
            Ident -> CExpr -> CodeGen ()
assign Ident
rId ((CExpr
r CExpr -> CExpr -> CExpr
./. CExpr
rMax) CExpr -> CExpr -> CExpr
.*. (CExpr -> CExpr
expE CExpr
wSumE))

            -- an iterator for picking a measure
            Ident
itId <- String -> CodeGen Ident
genIdent' String
"it"
            Sing 'HProb -> Ident -> CodeGen ()
forall (a :: Hakaru). Sing a -> Ident -> CodeGen ()
declare Sing 'HProb
SProb Ident
itId
            let itE :: CExpr
itE = Ident -> CExpr
CVar Ident
itId
            Ident -> CExpr -> CodeGen ()
assign Ident
itId (CExpr -> CExpr
logE (Integer -> CExpr
intE Integer
0))

            -- an output measure to assign to
            Ident
outId <- String -> CodeGen Ident
genIdent' String
"out"
            Sing ('HMeasure a) -> Ident -> CodeGen ()
forall (a :: Hakaru). Sing a -> Ident -> CodeGen ()
declare (abt '[] ('HMeasure a) -> Sing ('HMeasure a)
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> Sing a
typeOf (abt '[] ('HMeasure a) -> Sing ('HMeasure a))
-> ([(abt '[] 'HProb, abt '[] ('HMeasure a))]
    -> abt '[] ('HMeasure a))
-> [(abt '[] 'HProb, abt '[] ('HMeasure a))]
-> Sing ('HMeasure a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (abt '[] 'HProb, abt '[] ('HMeasure a)) -> abt '[] ('HMeasure a)
forall a b. (a, b) -> b
snd ((abt '[] 'HProb, abt '[] ('HMeasure a)) -> abt '[] ('HMeasure a))
-> ([(abt '[] 'HProb, abt '[] ('HMeasure a))]
    -> (abt '[] 'HProb, abt '[] ('HMeasure a)))
-> [(abt '[] 'HProb, abt '[] ('HMeasure a))]
-> abt '[] ('HMeasure a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(abt '[] 'HProb, abt '[] ('HMeasure a))]
-> (abt '[] 'HProb, abt '[] ('HMeasure a))
forall a. [a] -> a
head ([(abt '[] 'HProb, abt '[] ('HMeasure a))] -> Sing ('HMeasure a))
-> [(abt '[] 'HProb, abt '[] ('HMeasure a))] -> Sing ('HMeasure a)
forall a b. (a -> b) -> a -> b
$ [(abt '[] 'HProb, abt '[] ('HMeasure a))]
pairs') Ident
outId
            let outE :: CExpr
outE = Ident -> CExpr
CVar Ident
outId

            Ident
outLabel <- String -> CodeGen Ident
genIdent' String
"exit"

            [(CExpr, (abt '[] 'HProb, abt '[] ('HMeasure a)))]
-> ((CExpr, (abt '[] 'HProb, abt '[] ('HMeasure a))) -> CodeGen ())
-> CodeGen ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([CExpr]
-> [(abt '[] 'HProb, abt '[] ('HMeasure a))]
-> [(CExpr, (abt '[] 'HProb, abt '[] ('HMeasure a)))]
forall a b. [a] -> [b] -> [(a, b)]
zip [CExpr]
wEs [(abt '[] 'HProb, abt '[] ('HMeasure a))]
pairs')
              (((CExpr, (abt '[] 'HProb, abt '[] ('HMeasure a))) -> CodeGen ())
 -> CodeGen ())
-> ((CExpr, (abt '[] 'HProb, abt '[] ('HMeasure a))) -> CodeGen ())
-> CodeGen ()
forall a b. (a -> b) -> a -> b
$ \(CExpr
wE,(abt '[] 'HProb
_,abt '[] ('HMeasure a)
m)) ->
                  do Seq CExpr -> CExpr -> CodeGen ()
logSumExpCG ([CExpr] -> Seq CExpr
forall a. [a] -> Seq a
S.fromList [CExpr
itE,CExpr
wE]) CExpr
itE
                     CExpr -> CodeGen () -> CodeGen () -> CodeGen ()
ifCG (CExpr
rE CExpr -> CExpr -> CExpr
.<. (CExpr -> CExpr
expE CExpr
itE))
                          (abt '[] ('HMeasure a) -> CExpr -> CodeGen ()
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> CExpr -> CodeGen ()
flattenABT abt '[] ('HMeasure a)
m CExpr
outE CodeGen () -> CodeGen () -> CodeGen ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> CStat -> CodeGen ()
putStat (Ident -> CStat
CGoto Ident
outLabel))
                          (() -> CodeGen ()
forall (m :: * -> *) a. Monad m => a -> m a
return ())

            CStat -> CodeGen ()
putStat (CStat -> CodeGen ()) -> CStat -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ Ident -> CStat -> CStat
CLabel Ident
outLabel (Maybe CExpr -> CStat
CExpr Maybe CExpr
forall a. Maybe a
Nothing)
            CExpr -> CodeGen ()
putExprStat (CExpr -> CodeGen ()) -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ CExpr -> CExpr
mdataWeight CExpr
loc CExpr -> CExpr -> CExpr
.=. ((CExpr -> CExpr
mdataWeight CExpr
outE) CExpr -> CExpr -> CExpr
.+. CExpr
wSumE)
            CExpr -> CodeGen ()
putExprStat (CExpr -> CodeGen ()) -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ CExpr -> CExpr
mdataSample CExpr
loc CExpr -> CExpr -> CExpr
.=. (CExpr -> CExpr
mdataSample CExpr
outE)

--------------------------------------------------------------------------------
--                           Specialized Arithmetic                           --
--------------------------------------------------------------------------------

--------------------------------------
-- LogSumExp for NaryOp Add [SProb] --
--------------------------------------
{-

Special for addition of probabilities we have a logSumExp. This will compute the
sum of the probabilities safely. Just adding the exp(a . prob) would make us
loose any of the safety from underflow that we got from storing prob in the log
domain

-}

-- the tree traversal is a depth first search
logSumExp :: S.Seq CExpr -> CExpr
logSumExp :: Seq CExpr -> CExpr
logSumExp Seq CExpr
es = Int -> Int -> CExpr
mkCompTree Int
0 Int
1

  where lastIndex :: Int
lastIndex  = Seq CExpr -> Int
forall a. Seq a -> Int
S.length Seq CExpr
es Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1

        compIndices :: Int -> Int -> CExpr -> CExpr -> CExpr
        compIndices :: Int -> Int -> CExpr -> CExpr -> CExpr
compIndices Int
i Int
j = CExpr -> CExpr -> CExpr -> CExpr
CCond ((Seq CExpr -> Int -> CExpr
forall a. Seq a -> Int -> a
S.index Seq CExpr
es Int
i) CExpr -> CExpr -> CExpr
.>. (Seq CExpr -> Int -> CExpr
forall a. Seq a -> Int -> a
S.index Seq CExpr
es Int
j))

        mkCompTree :: Int -> Int -> CExpr
        mkCompTree :: Int -> Int -> CExpr
mkCompTree Int
i Int
j
          | Int
j Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
lastIndex = Int -> Int -> CExpr -> CExpr -> CExpr
compIndices Int
i Int
j (Int -> CExpr
logSumExp' Int
i) (Int -> CExpr
logSumExp' Int
j)
          | Bool
otherwise      = Int -> Int -> CExpr -> CExpr -> CExpr
compIndices Int
i Int
j
                               (Int -> Int -> CExpr
mkCompTree Int
i (Int -> Int
forall a. Enum a => a -> a
succ Int
j))
                               (Int -> Int -> CExpr
mkCompTree Int
j (Int -> Int
forall a. Enum a => a -> a
succ Int
j))

        diffExp :: Int -> Int -> CExpr
        diffExp :: Int -> Int -> CExpr
diffExp Int
a Int
b = CExpr -> CExpr
expm1E ((Seq CExpr -> Int -> CExpr
forall a. Seq a -> Int -> a
S.index Seq CExpr
es Int
a) CExpr -> CExpr -> CExpr
.-. (Seq CExpr -> Int -> CExpr
forall a. Seq a -> Int -> a
S.index Seq CExpr
es Int
b))

        -- given the max index, produce a logSumExp expression
        logSumExp' :: Int -> CExpr
        logSumExp' :: Int -> CExpr
logSumExp' Int
0 = Seq CExpr -> Int -> CExpr
forall a. Seq a -> Int -> a
S.index Seq CExpr
es Int
0
          CExpr -> CExpr -> CExpr
.+. (CExpr -> CExpr
log1pE (CExpr -> CExpr) -> CExpr -> CExpr
forall a b. (a -> b) -> a -> b
$ (Int -> CExpr -> CExpr) -> CExpr -> [Int] -> CExpr
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\Int
x CExpr
acc -> Int -> Int -> CExpr
diffExp Int
x Int
0 CExpr -> CExpr -> CExpr
.+. CExpr
acc)
                            (Int -> Int -> CExpr
diffExp Int
1 Int
0)
                            [Int
2..Seq CExpr -> Int
forall a. Seq a -> Int
S.length Seq CExpr
es Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
                    CExpr -> CExpr -> CExpr
.+. (Integer -> CExpr
intE (Integer -> CExpr) -> Integer -> CExpr
forall a b. (a -> b) -> a -> b
$ Int -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
lastIndex))
        logSumExp' Int
i = Seq CExpr -> Int -> CExpr
forall a. Seq a -> Int -> a
S.index Seq CExpr
es Int
i
          CExpr -> CExpr -> CExpr
.+. (CExpr -> CExpr
log1pE (CExpr -> CExpr) -> CExpr -> CExpr
forall a b. (a -> b) -> a -> b
$ (Int -> CExpr -> CExpr) -> CExpr -> [Int] -> CExpr
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\Int
x CExpr
acc -> if Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
x
                                       then CExpr
acc
                                       else Int -> Int -> CExpr
diffExp Int
x Int
i CExpr -> CExpr -> CExpr
.+. CExpr
acc)
                            (Int -> Int -> CExpr
diffExp Int
0 Int
i)
                            [Int
1..Seq CExpr -> Int
forall a. Seq a -> Int
S.length Seq CExpr
es Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
                    CExpr -> CExpr -> CExpr
.+. (Integer -> CExpr
intE (Integer -> CExpr) -> Integer -> CExpr
forall a b. (a -> b) -> a -> b
$ Int -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
lastIndex))


-- | logSumExpCG creates global functions for every n-ary logSumExp function
-- this reduces code size
logSumExpCG :: S.Seq CExpr -> (CExpr -> CodeGen ())
logSumExpCG :: Seq CExpr -> CExpr -> CodeGen ()
logSumExpCG Seq CExpr
seqE =
  let size :: Int
size   = Seq CExpr -> Int
forall a. Seq a -> Int
S.length (Seq CExpr -> Int) -> Seq CExpr -> Int
forall a b. (a -> b) -> a -> b
$ Seq CExpr
seqE
      name :: String
name   = String
"logSumExp" String -> String -> String
forall a. [a] -> [a] -> [a]
++ (Int -> String
forall a. Show a => a -> String
show Int
size)
      funcId :: Ident
funcId = String -> Ident
Ident String
name
  in \CExpr
loc -> do -- reset the names so that the function is the same for each arity
       let argIds :: [Ident]
argIds = (String -> Ident) -> [String] -> [Ident]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap String -> Ident
Ident (Int -> [String] -> [String]
forall a. Int -> [a] -> [a]
take Int
size [String]
cNameStream)
           decls :: [CDecl]
decls  = (Ident -> CDecl) -> [Ident] -> [CDecl]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Sing 'HProb -> Ident -> CDecl
forall (a :: Hakaru). Sing a -> Ident -> CDecl
typeDeclaration Sing 'HProb
SProb) [Ident]
argIds
           vars :: [CExpr]
vars   = (Ident -> CExpr) -> [Ident] -> [CExpr]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Ident -> CExpr
CVar [Ident]
argIds
       [CTypeSpec] -> Ident -> [CDecl] -> CodeGen () -> CodeGen ()
funCG [CTypeSpec
CDouble] Ident
funcId [CDecl]
decls
         (CStat -> CodeGen ()
putStat (CStat -> CodeGen ())
-> ([CExpr] -> CStat) -> [CExpr] -> CodeGen ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Maybe CExpr -> CStat
CReturn (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 -> Maybe CExpr)
-> ([CExpr] -> CExpr) -> [CExpr] -> Maybe CExpr
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Seq CExpr -> CExpr
logSumExp (Seq CExpr -> CExpr) -> ([CExpr] -> Seq CExpr) -> [CExpr] -> CExpr
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [CExpr] -> Seq CExpr
forall a. [a] -> Seq a
S.fromList ([CExpr] -> CodeGen ()) -> [CExpr] -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ [CExpr]
vars)
       CExpr -> CodeGen ()
putExprStat (CExpr -> CodeGen ()) -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ CExpr
loc CExpr -> CExpr -> CExpr
.=. (CExpr -> [CExpr] -> CExpr
CCall (Ident -> CExpr
CVar Ident
funcId) (Seq CExpr -> [CExpr]
forall (t :: * -> *) a. Foldable t => t a -> [a]
F.toList Seq CExpr
seqE))

-------------------------------------
-- LogSumExp for Summation of Prob --
-------------------------------------
{-
For summation of SProb we need a new logSumExp function that will find the max
of an array and then sum it in a loop
-}

lseSummateArrayCG
  :: ( ABT Term abt )
  => (abt '[ a ] b)
  -> CExpr
  -> (CExpr -> CodeGen ())
lseSummateArrayCG :: abt '[a] b -> CExpr -> CExpr -> CodeGen ()
lseSummateArrayCG abt '[a] b
body CExpr
arrayE =
  abt '[a] b
-> (Variable a -> abt '[] b -> CExpr -> CodeGen ())
-> CExpr
-> CodeGen ()
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (x :: k) (xs :: [k]) (a :: k) r.
ABT syn abt =>
abt (x : xs) a -> (Variable x -> abt xs a -> r) -> r
caseBind abt '[a] b
body ((Variable a -> abt '[] b -> CExpr -> CodeGen ())
 -> CExpr -> CodeGen ())
-> (Variable a -> abt '[] b -> CExpr -> CodeGen ())
-> CExpr
-> CodeGen ()
forall a b. (a -> b) -> a -> b
$ \Variable a
v abt '[] b
body' ->
    \CExpr
loc -> CodeGen () -> CodeGen ()
forall a. CodeGen a -> CodeGen a
seqDo (CodeGen () -> CodeGen ()) -> CodeGen () -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ do
      Ident
maxVId <- String -> CodeGen Ident
genIdent' String
"maxV"
      Ident
maxIId <- String -> CodeGen Ident
genIdent' String
"maxI"
      Ident
sumId <- String -> CodeGen Ident
genIdent' String
"sum"
      Ident
itId <- Variable a -> CodeGen Ident
forall (a :: Hakaru). Variable a -> CodeGen Ident
createIdent Variable a
v
      (Ident -> CodeGen ()) -> [Ident] -> CodeGen ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (Sing 'HProb -> Ident -> CodeGen ()
forall (a :: Hakaru). Sing a -> Ident -> CodeGen ()
declare Sing 'HProb
SProb) [Ident
maxVId,Ident
sumId]
      (Ident -> CodeGen ()) -> [Ident] -> CodeGen ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (Sing 'HNat -> Ident -> CodeGen ()
forall (a :: Hakaru). Sing a -> Ident -> CodeGen ()
declare Sing 'HNat
SNat)  [Ident
maxIId,Ident
itId]
      let (CExpr
maxVE:CExpr
maxIE:CExpr
sumE:CExpr
itE:[]) = (Ident -> CExpr) -> [Ident] -> [CExpr]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Ident -> CExpr
CVar [Ident
maxVId,Ident
maxIId,Ident
sumId,Ident
itId]
      CExpr -> CExpr -> CExpr -> CodeGen () -> CodeGen ()
forCG (CExpr
itE CExpr -> CExpr -> CExpr
.=. Integer -> CExpr
intE Integer
0)
            (CExpr
itE CExpr -> CExpr -> CExpr
.<. CExpr -> CExpr
arraySize CExpr
arrayE)
            (CUnaryOp -> CExpr -> CExpr
CUnary CUnaryOp
CPostIncOp CExpr
itE)
            (do Ident
tmpId <- CodeGen Ident
genIdent
                Sing 'HProb -> Ident -> CodeGen ()
forall (a :: Hakaru). Sing a -> Ident -> CodeGen ()
declare Sing 'HProb
SProb Ident
tmpId
                let tmpE :: CExpr
tmpE = Ident -> CExpr
CVar Ident
tmpId
                abt '[] b -> CExpr -> CodeGen ()
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> CExpr -> CodeGen ()
flattenABT abt '[] b
body' CExpr
tmpE
                CExpr -> CodeGen ()
putExprStat (CExpr -> CodeGen ()) -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ CExpr -> CExpr
derefIndex CExpr
itE CExpr -> CExpr -> CExpr
.=. CExpr
tmpE
                CStat -> CodeGen ()
putStat (CStat -> CodeGen ()) -> CStat -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ CExpr -> CStat -> Maybe CStat -> CStat
CIf ((CExpr
maxVE CExpr -> CExpr -> CExpr
.<. CExpr
tmpE) CExpr -> CExpr -> CExpr
.||. (CExpr
itE CExpr -> CExpr -> CExpr
.==. (Integer -> CExpr
intE Integer
0)))
                              ([CStat] -> CStat
seqCStat ([CStat] -> CStat) -> ([CExpr] -> [CStat]) -> [CExpr] -> CStat
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (CExpr -> CStat) -> [CExpr] -> [CStat]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (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] -> CStat) -> [CExpr] -> CStat
forall a b. (a -> b) -> a -> b
$
                                [ CExpr
maxVE CExpr -> CExpr -> CExpr
.=. CExpr
tmpE
                                , CExpr
maxIE CExpr -> CExpr -> CExpr
.=. CExpr
itE ])
                              Maybe CStat
forall a. Maybe a
Nothing)
      CExpr -> CodeGen ()
putExprStat (CExpr -> CodeGen ()) -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ CExpr
sumE  CExpr -> CExpr -> CExpr
.=. (Float -> CExpr
floatE Float
0) -- the sum is actually in real space
      CExpr -> CExpr -> CExpr -> CodeGen () -> CodeGen ()
forCG (CExpr
itE CExpr -> CExpr -> CExpr
.=. Integer -> CExpr
intE Integer
0)
            (CExpr
itE CExpr -> CExpr -> CExpr
.<. CExpr -> CExpr
arraySize CExpr
arrayE)
            (CUnaryOp -> CExpr -> CExpr
CUnary CUnaryOp
CPostIncOp CExpr
itE)
            (CExpr -> CodeGen () -> CodeGen () -> CodeGen ()
ifCG (CExpr
itE CExpr -> CExpr -> CExpr
.!=. CExpr
maxIE)
                  (CExpr -> CodeGen ()
putExprStat (CExpr -> CodeGen ()) -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ CExpr
sumE CExpr -> CExpr -> CExpr
.+=. (CExpr -> CExpr
expE ((CExpr -> CExpr
derefIndex CExpr
itE) CExpr -> CExpr -> CExpr
.-. (CExpr
maxVE))))
                  (() -> CodeGen ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()))

      CExpr -> CodeGen ()
putExprStat (CExpr -> CodeGen ()) -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ CExpr
loc CExpr -> CExpr -> CExpr
.=. (CExpr
maxVE CExpr -> CExpr -> CExpr
.+. (CExpr -> CExpr
log1pE CExpr
sumE))

  where derefIndex :: CExpr -> CExpr
derefIndex CExpr
xE = CExpr -> CExpr -> CExpr
index (CExpr -> CExpr
arrayData CExpr
arrayE) CExpr
xE

---------------------
-- Kahan Summation --
---------------------
-- | given a body and a size compute the kahan summation. This should work on
--   both probs and reals
kahanSummationCG
  :: ( ABT Term abt )
  => (abt '[ a ] b)
  -> CExpr
  -> CExpr
  -> (CExpr -> CodeGen ())
kahanSummationCG :: abt '[a] b -> CExpr -> CExpr -> CExpr -> CodeGen ()
kahanSummationCG abt '[a] b
body CExpr
loE CExpr
hiE =
  abt '[a] b
-> (Variable a -> abt '[] b -> CExpr -> CodeGen ())
-> CExpr
-> CodeGen ()
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (x :: k) (xs :: [k]) (a :: k) r.
ABT syn abt =>
abt (x : xs) a -> (Variable x -> abt xs a -> r) -> r
caseBind abt '[a] b
body ((Variable a -> abt '[] b -> CExpr -> CodeGen ())
 -> CExpr -> CodeGen ())
-> (Variable a -> abt '[] b -> CExpr -> CodeGen ())
-> CExpr
-> CodeGen ()
forall a b. (a -> b) -> a -> b
$ \Variable a
v abt '[] b
body' ->
    \CExpr
loc -> do
      Ident
tId <- String -> CodeGen Ident
genIdent' String
"t"
      Ident
cId <- String -> CodeGen Ident
genIdent' String
"c"
      Ident
itId <- Variable a -> CodeGen Ident
forall (a :: Hakaru). Variable a -> CodeGen Ident
createIdent Variable a
v
      Sing 'HNat -> Ident -> CodeGen ()
forall (a :: Hakaru). Sing a -> Ident -> CodeGen ()
declare Sing 'HNat
SNat Ident
itId
      (Ident -> CodeGen ()) -> [Ident] -> CodeGen ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (Sing 'HProb -> Ident -> CodeGen ()
forall (a :: Hakaru). Sing a -> Ident -> CodeGen ()
declare Sing 'HProb
SProb) [Ident
tId,Ident
cId]
      let (CExpr
tE:CExpr
cE:CExpr
itE:[]) = (Ident -> CExpr) -> [Ident] -> [CExpr]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Ident -> CExpr
CVar [Ident
tId,Ident
cId,Ident
itId]
      CExpr -> CodeGen ()
putExprStat (CExpr -> CodeGen ()) -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ CExpr
tE CExpr -> CExpr -> CExpr
.=. (Float -> CExpr
floatE Float
0)
      CExpr -> CodeGen ()
putExprStat (CExpr -> CodeGen ()) -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ CExpr
cE CExpr -> CExpr -> CExpr
.=. (Float -> CExpr
floatE Float
0)
      CExpr -> CExpr -> CExpr -> CodeGen () -> CodeGen ()
forCG (CExpr
itE CExpr -> CExpr -> CExpr
.=. CExpr
loE)
            (CExpr
itE CExpr -> CExpr -> CExpr
.<. CExpr
hiE)
            (CUnaryOp -> CExpr -> CExpr
CUnary CUnaryOp
CPostIncOp CExpr
itE)
            (do Ident
xId <- String -> CodeGen Ident
genIdent' String
"x"
                Ident
yId <- String -> CodeGen Ident
genIdent' String
"y"
                Ident
zId <- String -> CodeGen Ident
genIdent' String
"z"
                (Ident -> CodeGen ()) -> [Ident] -> CodeGen ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (Sing 'HProb -> Ident -> CodeGen ()
forall (a :: Hakaru). Sing a -> Ident -> CodeGen ()
declare Sing 'HProb
SProb) [Ident
xId,Ident
yId,Ident
zId]
                let (CExpr
xE:CExpr
yE:CExpr
zE:[]) = (Ident -> CExpr) -> [Ident] -> [CExpr]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Ident -> CExpr
CVar [Ident
xId,Ident
yId,Ident
zId]
                abt '[] b -> CExpr -> CodeGen ()
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
abt '[] a -> CExpr -> CodeGen ()
flattenABT abt '[] b
body' CExpr
xE
                CExpr -> CodeGen ()
putExprStat (CExpr -> CodeGen ()) -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ CExpr
yE CExpr -> CExpr -> CExpr
.=. (CExpr
xE CExpr -> CExpr -> CExpr
.-. CExpr
cE)
                CExpr -> CodeGen ()
putExprStat (CExpr -> CodeGen ()) -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ CExpr
zE CExpr -> CExpr -> CExpr
.=. (CExpr
tE CExpr -> CExpr -> CExpr
.+. CExpr
yE)
                CodeGen () -> CodeGen ()
whenPar (CodeGen () -> CodeGen ()) -> CodeGen () -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ CStat -> CodeGen ()
putStat (CStat -> CodeGen ()) -> (OMP -> CStat) -> OMP -> CodeGen ()
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 -> CodeGen ()) -> OMP -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ Directive -> OMP
OMP Directive
Critical
                CodeGen () -> CodeGen ()
codeBlockCG (CodeGen () -> CodeGen ()) -> CodeGen () -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ do
                  CExpr -> CodeGen ()
putExprStat (CExpr -> CodeGen ()) -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ CExpr
cE CExpr -> CExpr -> CExpr
.=.  ((CExpr
zE CExpr -> CExpr -> CExpr
.-. CExpr
tE) CExpr -> CExpr -> CExpr
.-. CExpr
yE)
                  CExpr -> CodeGen ()
putExprStat (CExpr -> CodeGen ()) -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ CExpr
tE CExpr -> CExpr -> CExpr
.=. CExpr
zE)
      CExpr -> CodeGen ()
putExprStat (CExpr -> CodeGen ()) -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ CExpr
loc CExpr -> CExpr -> CExpr
.=. CExpr
tE

--------------------------------------------------------------------------------
--                            Coercion Helpers                                --
--------------------------------------------------------------------------------

coerceToCG
  :: forall (a :: Hakaru) (b :: Hakaru)
  .  Coercion a b
  -> CExpr
  -> CodeGen CExpr
coerceToCG :: Coercion a b -> CExpr -> CodeGen CExpr
coerceToCG (CCons (Signed HRing b
HRing_Int)            Coercion b b
cs) CExpr
e = CExpr -> CodeGen CExpr
nat2int CExpr
e   CodeGen CExpr -> (CExpr -> CodeGen CExpr) -> CodeGen CExpr
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Coercion b b -> CExpr -> CodeGen CExpr
forall (a :: Hakaru) (b :: Hakaru).
Coercion a b -> CExpr -> CodeGen CExpr
coerceToCG Coercion b b
cs
coerceToCG (CCons (Signed HRing b
HRing_Real)           Coercion b b
cs) CExpr
e = CExpr -> CodeGen CExpr
prob2real CExpr
e CodeGen CExpr -> (CExpr -> CodeGen CExpr) -> CodeGen CExpr
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Coercion b b -> CExpr -> CodeGen CExpr
forall (a :: Hakaru) (b :: Hakaru).
Coercion a b -> CExpr -> CodeGen CExpr
coerceToCG Coercion b b
cs
coerceToCG (CCons (Continuous HContinuous b
HContinuous_Prob) Coercion b b
cs) CExpr
e = CExpr -> CodeGen CExpr
nat2prob CExpr
e  CodeGen CExpr -> (CExpr -> CodeGen CExpr) -> CodeGen CExpr
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Coercion b b -> CExpr -> CodeGen CExpr
forall (a :: Hakaru) (b :: Hakaru).
Coercion a b -> CExpr -> CodeGen CExpr
coerceToCG Coercion b b
cs
coerceToCG (CCons (Continuous HContinuous b
HContinuous_Real) Coercion b b
cs) CExpr
e = CExpr -> CodeGen CExpr
int2real CExpr
e  CodeGen CExpr -> (CExpr -> CodeGen CExpr) -> CodeGen CExpr
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Coercion b b -> CExpr -> CodeGen CExpr
forall (a :: Hakaru) (b :: Hakaru).
Coercion a b -> CExpr -> CodeGen CExpr
coerceToCG Coercion b b
cs
coerceToCG Coercion a b
CNil CExpr
e = CExpr -> CodeGen CExpr
forall (m :: * -> *) a. Monad m => a -> m a
return CExpr
e

coerceFromCG
  :: forall (a :: Hakaru) (b :: Hakaru)
  .  Coercion a b
  -> CExpr
  -> CodeGen CExpr
coerceFromCG :: Coercion a b -> CExpr -> CodeGen CExpr
coerceFromCG (CCons (Signed HRing b
HRing_Int)            Coercion b b
cs) CExpr
e = CExpr -> CodeGen CExpr
int2nat CExpr
e   CodeGen CExpr -> (CExpr -> CodeGen CExpr) -> CodeGen CExpr
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Coercion b b -> CExpr -> CodeGen CExpr
forall (a :: Hakaru) (b :: Hakaru).
Coercion a b -> CExpr -> CodeGen CExpr
coerceFromCG Coercion b b
cs
coerceFromCG (CCons (Signed HRing b
HRing_Real)           Coercion b b
cs) CExpr
e = CExpr -> CodeGen CExpr
real2prob CExpr
e CodeGen CExpr -> (CExpr -> CodeGen CExpr) -> CodeGen CExpr
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Coercion b b -> CExpr -> CodeGen CExpr
forall (a :: Hakaru) (b :: Hakaru).
Coercion a b -> CExpr -> CodeGen CExpr
coerceFromCG Coercion b b
cs
coerceFromCG (CCons (Continuous HContinuous b
HContinuous_Prob) Coercion b b
cs) CExpr
e = CExpr -> CodeGen CExpr
prob2nat CExpr
e  CodeGen CExpr -> (CExpr -> CodeGen CExpr) -> CodeGen CExpr
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Coercion b b -> CExpr -> CodeGen CExpr
forall (a :: Hakaru) (b :: Hakaru).
Coercion a b -> CExpr -> CodeGen CExpr
coerceFromCG Coercion b b
cs
coerceFromCG (CCons (Continuous HContinuous b
HContinuous_Real) Coercion b b
cs) CExpr
e = CExpr -> CodeGen CExpr
real2int CExpr
e  CodeGen CExpr -> (CExpr -> CodeGen CExpr) -> CodeGen CExpr
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Coercion b b -> CExpr -> CodeGen CExpr
forall (a :: Hakaru) (b :: Hakaru).
Coercion a b -> CExpr -> CodeGen CExpr
coerceFromCG Coercion b b
cs
coerceFromCG Coercion a b
CNil CExpr
e = CExpr -> CodeGen CExpr
forall (m :: * -> *) a. Monad m => a -> m a
return CExpr
e

-- safe
nat2int,nat2prob,prob2real,int2real
  :: CExpr -> CodeGen CExpr
nat2int :: CExpr -> CodeGen CExpr
nat2int   CExpr
x = CExpr -> CodeGen CExpr
forall (m :: * -> *) a. Monad m => a -> m a
return CExpr
x
nat2prob :: CExpr -> CodeGen CExpr
nat2prob  CExpr
x = do CExpr
x' <- Sing 'HProb -> String -> CodeGen CExpr
forall (a :: Hakaru). Sing a -> String -> CodeGen CExpr
localVar' Sing 'HProb
SProb String
"n2p"
                 CExpr -> CodeGen ()
putExprStat (CExpr -> CodeGen ()) -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ CExpr
x' CExpr -> CExpr -> CExpr
.=. (CExpr -> CExpr
logE CExpr
x)
                 CExpr -> CodeGen CExpr
forall (m :: * -> *) a. Monad m => a -> m a
return CExpr
x'
prob2real :: CExpr -> CodeGen CExpr
prob2real CExpr
x = do CExpr
x' <- Sing 'HProb -> String -> CodeGen CExpr
forall (a :: Hakaru). Sing a -> String -> CodeGen CExpr
localVar' Sing 'HProb
SProb String
"p2r"
                 CExpr -> CodeGen ()
putExprStat (CExpr -> CodeGen ()) -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ CExpr
x' CExpr -> CExpr -> CExpr
.=. ((CExpr -> CExpr
expm1E CExpr
x) CExpr -> CExpr -> CExpr
.+. (Float -> CExpr
floatE Float
1))
                 CExpr -> CodeGen CExpr
forall (m :: * -> *) a. Monad m => a -> m a
return CExpr
x'
int2real :: CExpr -> CodeGen CExpr
int2real  CExpr
x = CExpr -> CodeGen CExpr
forall (m :: * -> *) a. Monad m => a -> m a
return ([CTypeSpec] -> CExpr -> CExpr
castTo [CTypeSpec
CDouble] CExpr
x)

-- unsafe
{- Because of the hkc representation of reals and probs as doubles, (instead of
   rationals). we will just silently truncate values for prob2nat and real2int
-}
int2nat,prob2nat,real2prob,real2int
  :: CExpr -> CodeGen CExpr
int2nat :: CExpr -> CodeGen CExpr
int2nat CExpr
x =
  do CExpr
x' <- Sing 'HNat -> String -> CodeGen CExpr
forall (a :: Hakaru). Sing a -> String -> CodeGen CExpr
localVar' Sing 'HNat
SNat String
"i2n"
     CExpr -> CodeGen () -> CodeGen () -> CodeGen ()
ifCG (CExpr
x CExpr -> CExpr -> CExpr
.<. (Integer -> CExpr
intE Integer
0))
          (do CExpr -> CodeGen ()
putExprStat (CExpr -> CodeGen ()) -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ [CExpr] -> CExpr
printfE [ String -> CExpr
stringE String
"error: cannot coerce negative int to nat\n" ]
              CExpr -> CodeGen ()
putExprStat (CExpr -> CodeGen ()) -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ String -> [CExpr] -> CExpr
mkCallE String
"abort" [] )
          (CExpr -> CodeGen ()
putExprStat (CExpr -> CodeGen ()) -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ CExpr
x' CExpr -> CExpr -> CExpr
.=. ([CTypeSpec] -> CExpr -> CExpr
castTo [CTypeSpec
CUnsigned, CTypeSpec
CInt] CExpr
x))
     CExpr -> CodeGen CExpr
forall (m :: * -> *) a. Monad m => a -> m a
return CExpr
x'
prob2nat :: CExpr -> CodeGen CExpr
prob2nat CExpr
x = CExpr -> CodeGen CExpr
forall (m :: * -> *) a. Monad m => a -> m a
return ([CTypeSpec] -> CExpr -> CExpr
castTo [CTypeSpec
CUnsigned, CTypeSpec
CInt] CExpr
x)
real2prob :: CExpr -> CodeGen CExpr
real2prob CExpr
x =
  do CExpr
x' <- Sing 'HProb -> String -> CodeGen CExpr
forall (a :: Hakaru). Sing a -> String -> CodeGen CExpr
localVar' Sing 'HProb
SProb String
"r2p"
     CExpr -> CodeGen () -> CodeGen () -> CodeGen ()
ifCG (CExpr
x CExpr -> CExpr -> CExpr
.<. (Integer -> CExpr
intE Integer
0))
          (do CExpr -> CodeGen ()
putExprStat (CExpr -> CodeGen ()) -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ [CExpr] -> CExpr
printfE [ String -> CExpr
stringE String
"error: cannot coerce negative real to prob\n" ]
              CExpr -> CodeGen ()
putExprStat (CExpr -> CodeGen ()) -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ String -> [CExpr] -> CExpr
mkCallE String
"abort" [] )
          (CExpr -> CodeGen ()
putExprStat (CExpr -> CodeGen ()) -> CExpr -> CodeGen ()
forall a b. (a -> b) -> a -> b
$ CExpr
x' CExpr -> CExpr -> CExpr
.=. (CExpr -> CExpr
logE CExpr
x))
     CExpr -> CodeGen CExpr
forall (m :: * -> *) a. Monad m => a -> m a
return CExpr
x'
real2int :: CExpr -> CodeGen CExpr
real2int CExpr
x =  CExpr -> CodeGen CExpr
forall (m :: * -> *) a. Monad m => a -> m a
return ([CTypeSpec] -> CExpr -> CExpr
castTo [CTypeSpec
CInt] CExpr
x)

--------------------------------------------------------------------------------
--                            Parallel Helpers                                --
--------------------------------------------------------------------------------
{- SIMD (single instruction multiple data) OpenMP pragmas, should be applied to
-- the inner most loop. `hasParallelTerm` checks whether a term or any of its
-- subterms contain a parallel construct { plate, summate, product, array,
-- bucket }.
-}

hasParallelTerm :: ( ABT Term abt ) => abt '[] a -> Bool
hasParallelTerm :: abt '[] a -> Bool
hasParallelTerm abt '[] a
abt = abt '[] a -> (Variable a -> Bool) -> (Term abt a -> Bool) -> Bool
forall k (syn :: ([k] -> k -> *) -> k -> *) (abt :: [k] -> k -> *)
       (a :: k) r.
ABT syn abt =>
abt '[] a -> (Variable a -> r) -> (syn abt a -> r) -> r
caseVarSyn abt '[] a
abt (Bool -> Variable a -> Bool
forall a b. a -> b -> a
const Bool
False) Term abt a -> Bool
forall (abt :: [Hakaru] -> Hakaru -> *) (a :: Hakaru).
ABT Term abt =>
Term abt a -> Bool
hPT'
  where hPT' :: ABT Term abt => Term abt a -> Bool
        hPT' :: Term abt a -> Bool
hPT' (SCon args a
_ :$ SArgs abt args
_)          = Bool
forall a. HasCallStack => a
undefined
        hPT' (NaryOp_ NaryOp a
_ Seq (abt '[] a)
_)     = Bool
forall a. HasCallStack => a
undefined
        hPT' (Literal_ Literal a
_)      = Bool
False
        hPT' (Empty_ Sing ('HArray a)
_)        = Bool
False
        hPT' (Array_ abt '[] 'HNat
_ abt '[ 'HNat] a
_)      = Bool
True
        hPT' (ArrayLiteral_ [abt '[] a]
_) = Bool
False
        hPT' (Bucket abt '[] 'HNat
_ abt '[] 'HNat
_ Reducer abt '[] a
_)    = Bool
True
        hPT' (Datum_ Datum (abt '[]) (HData' t)
_)        = Bool
False
        hPT' (Case_ abt '[] a
_ [Branch a abt a]
_)       = Bool
forall a. HasCallStack => a
undefined
        hPT' (Superpose_ NonEmpty (abt '[] 'HProb, abt '[] ('HMeasure a))
_)    = Bool
forall a. HasCallStack => a
undefined
        hPT' (Reject_ Sing ('HMeasure a)
_)       = Bool
False