<commit_before>{-# 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  :  <EMAIL>
-- 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)
  where

import Language.Hakaru.CodeGen.CodeGenMonad
import Language.Hakaru.CodeGen.HOAS.Declaration
import Language.Hakaru.CodeGen.HOAS.Expression
import Language.Hakaru.CodeGen.HOAS.Statement

import Language.Hakaru.Syntax.AST
import Language.Hakaru.Syntax.ABT
import Language.Hakaru.Syntax.TypeOf (typeOf)
import Language.Hakaru.Syntax.Datum
import Language.Hakaru.Syntax.IClasses
import Language.Hakaru.Types.DataKind
import Language.Hakaru.Types.HClasses
import Language.Hakaru.Types.Coercion
import Language.Hakaru.Types.Sing

import Language.C.Syntax.AST
import Language.C.Data.Ident

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           Data.Functor
#endif


import Prelude hiding (log,exp,sqrt)

flattenABT :: ABT Term abt
           => abt '[] a
           -> CodeGen CExpr
flattenABT abt = caseVarSyn abt flattenVar flattenTerm


flattenVar :: Variable (a :: Hakaru) -> CodeGen CExpr
flattenVar v = varE <$> lookupIdent v

flattenTerm :: ABT Term abt => Term abt a -> CodeGen CExpr
flattenTerm (NaryOp_ t s)    = flattenNAryOp t s
flattenTerm (Literal_ x)     = flattenLit x
flattenTerm (Empty_ _)       = error "TODO: flattenTerm Empty"
flattenTerm (Datum_ d)       = flattenDatum d
flattenTerm (Case_ c bs)     = flattenCase c bs
flattenTerm (Array_ s e)     = flattenArray s e
flattenTerm (x :$ ys)        = flattenSCon x ys
flattenTerm (Reject_ _)      = error "TODO: flattenTerm Reject"
flattenTerm (Superpose_ wes) = flattenSuperpose wes

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


flattenSCon :: ( ABT Term abt )
            => SCon args a
            -> SArgs abt args
            -> CodeGen CExpr
flattenSCon Let_            =
  \(expr :* body :* End) ->
    do expr' <- flattenABT expr
       caseBind body $ \v@(Variable _ _ typ) body'->
         do ident <- createIdent v
            declare typ ident
            assign ident expr'
            flattenABT body'

-- Lambdas produce functions and then return a function pointer
flattenSCon Lam_            =
  \(body :* End) ->
    coalesceLambda body $ \vars body' ->
    let varMs = foldMap11 (\v -> [mkVarDecl v <$> createIdent v]) vars
    in  do funcId <- genIdent' "fn"
           decls  <- sequence varMs

           cg <- get
           let m       = putStat . returnS =<< flattenABT body'
               (_,cg') = runState m $ cg { statements = [] }
           put $ cg' { statements = statements cg }

           extDeclare . extFunc $ functionDef (typeOf body')
                                              funcId
                                              decls
                                              (reverse $ statements cg')
           return . varE $ funcId
  -- do at top level
  where coalesceLambda
          :: ( ABT Term abt )
          => abt '[x] a
          -> (forall (ys :: [Hakaru]) b. List1 Variable ys -> abt '[] b -> r)
          -> r
        coalesceLambda abt k =
          caseBind abt $ \v abt' ->
            caseVarSyn abt' (const (k (Cons1 v Nil1) abt')) $ \term ->
              case term of
                (Lam_ :$ body :* End) ->
                  coalesceLambda body $ \vars abt'' -> k (Cons1 v vars) abt''
                _ -> k (Cons1 v Nil1) abt'

        mkVarDecl :: Variable (a :: Hakaru) -> Ident -> CDecl
        mkVarDecl (Variable _ _ typ) = typeDeclaration typ


flattenSCon (PrimOp_ op)    = flattenPrimOp op
flattenSCon (ArrayOp_ op)   = flattenArrayOp op
flattenSCon (MeasureOp_ op) = flattenMeasureOp op
flattenSCon Dirac           = \(e :* End) -> flattenABT e

flattenSCon (Summate _ sr) = \(lo :* hi :* body :* End) ->
  do loE <- flattenABT lo
     hiE <- flattenABT hi
     caseBind body $ \v body' ->
       do iterI <- createIdent v
          declare SNat iterI
          assign iterI loE<commit_msg>fixed issue with function arguments of tyoe SArray<commit_after>{-# 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  :  <EMAIL>
-- 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)
  where

import Language.Hakaru.CodeGen.CodeGenMonad
import Language.Hakaru.CodeGen.HOAS.Declaration
import Language.Hakaru.CodeGen.HOAS.Expression
import Language.Hakaru.CodeGen.HOAS.Statement

import Language.Hakaru.Syntax.AST
import Language.Hakaru.Syntax.ABT
import Language.Hakaru.Syntax.TypeOf (typeOf)
import Language.Hakaru.Syntax.Datum
import Language.Hakaru.Syntax.IClasses
import Language.Hakaru.Types.DataKind
import Language.Hakaru.Types.HClasses
import Language.Hakaru.Types.Coercion
import Language.Hakaru.Types.Sing

import Language.C.Syntax.AST
import Language.C.Data.Ident

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           Data.Functor
#endif


import Prelude hiding (log,exp,sqrt)

flattenABT :: ABT Term abt
           => abt '[] a
           -> CodeGen CExpr
flattenABT abt = caseVarSyn abt flattenVar flattenTerm


flattenVar :: Variable (a :: Hakaru) -> CodeGen CExpr
flattenVar v = varE <$> lookupIdent v

flattenTerm :: ABT Term abt => Term abt a -> CodeGen CExpr
flattenTerm (NaryOp_ t s)    = flattenNAryOp t s
flattenTerm (Literal_ x)     = flattenLit x
flattenTerm (Empty_ _)       = error "TODO: flattenTerm Empty"
flattenTerm (Datum_ d)       = flattenDatum d
flattenTerm (Case_ c bs)     = flattenCase c bs
flattenTerm (Array_ s e)     = flattenArray s e
flattenTerm (x :$ ys)        = flattenSCon x ys
flattenTerm (Reject_ _)      = error "TODO: flattenTerm Reject"
flattenTerm (Superpose_ wes) = flattenSuperpose wes

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


flattenSCon :: ( ABT Term abt )
            => SCon args a
            -> SArgs abt args
            -> CodeGen CExpr
flattenSCon Let_            =
  \(expr :* body :* End) ->
    do expr' <- flattenABT expr
       caseBind body $ \v@(Variable _ _ typ) body'->
         do ident <- createIdent v
            declare typ ident
            assign ident expr'
            flattenABT body'

-- Lambdas produce functions and then return a function pointer
flattenSCon Lam_            =
  \(body :* End) ->
    coalesceLambda body $ \vars body' ->
    let varMs = foldMap11 (\v -> [mkVarDecl v =<< createIdent v]) vars
    in  do funcId <- genIdent' "fn"
           decls  <- sequence varMs

           cg <- get
           let m       = putStat . returnS =<< flattenABT body'
               (_,cg') = runState m $ cg { statements = [] }
           put $ cg' { statements = statements cg }

           extDeclare . extFunc $ functionDef (typeOf body')
                                              funcId
                                              decls
                                              (reverse $ statements cg')
           return . varE $ funcId
  -- do at top level
  where coalesceLambda
          :: ( ABT Term abt )
          => abt '[x] a
          -> (forall (ys :: [Hakaru]) b. List1 Variable ys -> abt '[] b -> r)
          -> r
        coalesceLambda abt k =
          caseBind abt $ \v abt' ->
            caseVarSyn abt' (const (k (Cons1 v Nil1) abt')) $ \term ->
              case term of
                (Lam_ :$ body :* End) ->
                  coalesceLambda body $ \vars abt'' -> k (Cons1 v vars) abt''
                _ -> k (Cons1 v Nil1) abt'

        mkVarDecl :: Variable (a :: Hakaru) -> Ident -> CodeGen CDecl
        mkVarDecl (Variable _ _ SInt)  = return . typeDeclaration SInt
        mkVarDecl (Variable _ _ SNat)  = return . typeDeclaration SNat
        mkVarDecl (Variable _ _ SProb) = return . typeDeclaration SProb
        mkVarDecl (Variable _ _ SReal) = return . typeDeclaration SReal
        mkVarDecl (Variable _ _ (SArray t)) = \i -> do extDeclare $ arrayStruct t
                                                       return $ arrayDeclaration t i
        mkVarDecl (Variable _ _ d@(SData _ _)) = \i -> do extDeclare $ datumStruct d
                                                          return $ datumDeclaration d i
        mkVarDecl v = error $ "flattenSCon.Lam_.mkVarDecl cannot handle vars of type " ++ show v


flattenSCon (PrimOp_ op)    = flattenPrimOp op
flattenSCon (ArrayOp_ op)   = flattenArrayOp op
flattenSCon (MeasureOp_ op) = flattenMeasureOp op
flattenSCon Dirac           = \(e :* End) -> flattenABT e

flattenSCon (Summate _ sr) = \(lo :* hi :* body :* End) ->
  do loE <- flattenABT lo
     hiE <- flattenABT hi
     caseBind body $ \v body' ->
       do iterI <- createIdent v
          declare SNat iterI
          assign iterI loE
{'fraction_non_alphanumeric': 0.1221213569039656, 'fraction_numerical': 0.004586717630195891, 'mean_word_length': 1.7184415584415584, 'pattern_counts': {'":': 0, '<': 32, '<?xml version=': 0, '>': 66, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 20, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1777076', 'n_tokens_mistral': 3187, 'n_tokens_neox': 3073, 'n_words': 919}