<commit_before>{-# LANGUAGE TupleSections #-}
-- |
--
-- This module implements a transformation from source to core
-- Futhark.
--
module Futhark.Internalise
  ( internaliseProg
  , internaliseValue
  )
  where

import Control.Applicative
import Control.Monad.State  hiding (mapM, sequence)
import Control.Monad.Reader hiding (mapM, sequence)
import qualified Data.HashMap.Lazy as HM
import Data.Maybe
import Data.Monoid
import Data.List
import Data.Traversable (mapM, sequence)
import Data.Loc

import Prelude hiding (mapM, sequence)

import Language.Futhark as E
import Futhark.Representation.SOACS as I hiding (bindingPattern)
import Futhark.Transform.Rename as I
import Futhark.Transform.Substitute
import Futhark.MonadFreshNames

import Futhark.Internalise.Monad
import Futhark.Internalise.AccurateSizes
import Futhark.Internalise.TypesValues
import Futhark.Internalise.Bindings
import Futhark.Internalise.Lambdas

-- | Convert a program in source Futhark to a program in the Futhark
-- core language.
internaliseProg :: MonadFreshNames m =>
                   E.Prog -> m (Either String I.Prog)
internaliseProg prog = do
  res <- do
    ftable_attempt <- buildFtable prog
    case ftable_attempt of
      Left err -> return $ Left err
      Right ftable -> do
        funs <- runInternaliseM ftable $
                mapM internaliseFun $ E.progFunctions prog
        return $ fmap I.Prog funs
  sequence $ fmap I.renameProg res

buildFtable :: MonadFreshNames m => E.Prog
            -> m (Either String FunTable)
buildFtable = fmap (HM.union builtinFtable<$>) .
              runInternaliseM mempty .
              fmap HM.fromList . mapM inspect . E.progFunctions
  where inspect (E.FunDef _ fname (TypeDecl _ (Info rettype)) params _ _) =
          bindingParams params $ \shapes values -> do
            (rettype', _) <- internaliseReturnType rettype
            let shapenames = map I.paramName shapes
            return (fname,
                    FunBinding { internalFun = (shapenames,
                                                map I.paramDeclType values,
                                                applyRetType
                                                (ExtRetType rettype')
                                                (shapes++values)
                                               )
                               , externalFun = (rettype,
                                                map E.paramType params)
                               })

        builtinFtable = HM.map addBuiltin E.builtInFunctions
        addBuiltin (t, paramts) =
          FunBinding
          ([], map (I.Prim . internalisePrimType) paramts,
           const $ Just $ ExtRetType [I.Prim $ internalisePrimType t])
          (E.Prim t, map E.Prim paramts)

internaliseFun :: E.FunDef -> InternaliseM I.FunDef
internaliseFun (E.FunDef entry fname (TypeDecl _ (Info rettype)) params body loc) =
  bindingParams params $ \shapeparams params' -> do
    (rettype', _) <- internaliseReturnType rettype
    firstbody <- internaliseBody body
    body' <- ensureResultExtShape asserting loc
             (map I.fromDecl rettype') firstbody
    return $ I.FunDef entry
      fname (ExtRetType rettype')
      (shapeparams ++ params')
      body'

internaliseIdent :: E.Ident -> InternaliseM I.VName
internaliseIdent (E.Ident name (Info tp) _) =
  case internaliseType tp of
    [I.Prim _] -> return name
    _           -> fail $ "Futhark.Internalise.internaliseIdent: asked to internalise non-prim-typed ident '"
                   ++ pretty name ++ "'."

internaliseBody :: E.Exp -> InternaliseM Body
internaliseBody e = insertBindingsM $ do
  ses <- internaliseExp "res" e
  return $ resultBody ses

internaliseBodyBindings :: E.Exp -> ([SubExp] -> InternaliseM (Body, a))
                        -> InternaliseM (Body, a)
internaliseBodyBindings e m = do
  ((Body _ bnds res,x), otherbnds) <-
    collectBindings $ m =<< internaliseExp "res" e
  (,x) <$> mkBodyM (otherbnds <> bnds) res

extraBodyBindings :: [Binding]
                  -> InternaliseM (Body, a)
                  -> InternaliseM (Body, a)
extraBodyBindings bnds m = do
  (body, x) <- m
  return (insertBindings bnds body, x)

internaliseExp :: String -> E.Exp -> InternaliseM [I.SubExp]

internaliseExp _ (E.Var var) = do
  subst <- asks $ HM.lookup (E.identName var) . envSubsts
  case subst of
    Nothing     -> (:[]) . I.Var <$> internaliseIdent var
    Just substs -> return $ map I.Var substs

internaliseExp desc (E.Index e idxs loc) = do
  idxs' <- mapM (internaliseExp1 "i") idxs
  vs <- internaliseExpToVars "indexed" e
  csidx' <- boundsChecks loc vs idxs'
  let index v = I.PrimOp $ I.Index csidx' v idxs'
  letSubExps desc (map index vs)

internaliseExp desc (E.TupLit es _) =
  concat <$> mapM (internaliseExp desc) es

internaliseExp desc (E.ArrayLit [] (Info et) _) =
  letSubExps desc $ map arrayLit $ internaliseType et
  where arrayLit et' =
          I.PrimOp $ I.ArrayLit [] $ et' `annotateArrayShape` []

internaliseExp desc (E.ArrayLit es (Info rowtype) _) = do
  aes <- mapM (internaliseExpToVars "arr_elem") es
  let es'@((e':_):_) = aes --- XXX, ugh.
  Shape rowshape <- arrayShape <$> lookupType e'
  case internaliseType rowtype of
    [et] -> letTupExp' desc $ I.PrimOp $
            I.ArrayLit (map I.Var $ concat es')
            (et `I.setArrayShape` Shape rowshape)
    ets   -> do
      let arraylit ks et =
            I.PrimOp $ I.ArrayLit (map I.Var ks)
            (et `I.setArrayShape` Shape rowshape)
      letSubExps desc (zipWith arraylit (transpose es') ets)

internaliseExp desc (E.Apply fname args _ _)
  | Just (rettype, _) <- HM.lookup fname I.builtInFunctions = do
  args' <- mapM (internaliseExp "arg" . fst) args<commit_msg>Handle empty array literals in internaliser.<commit_after>{-# LANGUAGE TupleSections #-}
-- |
--
-- This module implements a transformation from source to core
-- Futhark.
--
module Futhark.Internalise
  ( internaliseProg
  , internaliseValue
  )
  where

import Control.Applicative
import Control.Monad.State  hiding (mapM, sequence)
import Control.Monad.Reader hiding (mapM, sequence)
import qualified Data.HashMap.Lazy as HM
import Data.Maybe
import Data.Monoid
import Data.List
import Data.Traversable (mapM, sequence)
import Data.Loc

import Prelude hiding (mapM, sequence)

import Language.Futhark as E
import Futhark.Representation.SOACS as I hiding (bindingPattern)
import Futhark.Transform.Rename as I
import Futhark.Transform.Substitute
import Futhark.MonadFreshNames

import Futhark.Internalise.Monad
import Futhark.Internalise.AccurateSizes
import Futhark.Internalise.TypesValues
import Futhark.Internalise.Bindings
import Futhark.Internalise.Lambdas

-- | Convert a program in source Futhark to a program in the Futhark
-- core language.
internaliseProg :: MonadFreshNames m =>
                   E.Prog -> m (Either String I.Prog)
internaliseProg prog = do
  res <- do
    ftable_attempt <- buildFtable prog
    case ftable_attempt of
      Left err -> return $ Left err
      Right ftable -> do
        funs <- runInternaliseM ftable $
                mapM internaliseFun $ E.progFunctions prog
        return $ fmap I.Prog funs
  sequence $ fmap I.renameProg res

buildFtable :: MonadFreshNames m => E.Prog
            -> m (Either String FunTable)
buildFtable = fmap (HM.union builtinFtable<$>) .
              runInternaliseM mempty .
              fmap HM.fromList . mapM inspect . E.progFunctions
  where inspect (E.FunDef _ fname (TypeDecl _ (Info rettype)) params _ _) =
          bindingParams params $ \shapes values -> do
            (rettype', _) <- internaliseReturnType rettype
            let shapenames = map I.paramName shapes
            return (fname,
                    FunBinding { internalFun = (shapenames,
                                                map I.paramDeclType values,
                                                applyRetType
                                                (ExtRetType rettype')
                                                (shapes++values)
                                               )
                               , externalFun = (rettype,
                                                map E.paramType params)
                               })

        builtinFtable = HM.map addBuiltin E.builtInFunctions
        addBuiltin (t, paramts) =
          FunBinding
          ([], map (I.Prim . internalisePrimType) paramts,
           const $ Just $ ExtRetType [I.Prim $ internalisePrimType t])
          (E.Prim t, map E.Prim paramts)

internaliseFun :: E.FunDef -> InternaliseM I.FunDef
internaliseFun (E.FunDef entry fname (TypeDecl _ (Info rettype)) params body loc) =
  bindingParams params $ \shapeparams params' -> do
    (rettype', _) <- internaliseReturnType rettype
    firstbody <- internaliseBody body
    body' <- ensureResultExtShape asserting loc
             (map I.fromDecl rettype') firstbody
    return $ I.FunDef entry
      fname (ExtRetType rettype')
      (shapeparams ++ params')
      body'

internaliseIdent :: E.Ident -> InternaliseM I.VName
internaliseIdent (E.Ident name (Info tp) _) =
  case internaliseType tp of
    [I.Prim _] -> return name
    _           -> fail $ "Futhark.Internalise.internaliseIdent: asked to internalise non-prim-typed ident '"
                   ++ pretty name ++ "'."

internaliseBody :: E.Exp -> InternaliseM Body
internaliseBody e = insertBindingsM $ do
  ses <- internaliseExp "res" e
  return $ resultBody ses

internaliseBodyBindings :: E.Exp -> ([SubExp] -> InternaliseM (Body, a))
                        -> InternaliseM (Body, a)
internaliseBodyBindings e m = do
  ((Body _ bnds res,x), otherbnds) <-
    collectBindings $ m =<< internaliseExp "res" e
  (,x) <$> mkBodyM (otherbnds <> bnds) res

extraBodyBindings :: [Binding]
                  -> InternaliseM (Body, a)
                  -> InternaliseM (Body, a)
extraBodyBindings bnds m = do
  (body, x) <- m
  return (insertBindings bnds body, x)

internaliseExp :: String -> E.Exp -> InternaliseM [I.SubExp]

internaliseExp _ (E.Var var) = do
  subst <- asks $ HM.lookup (E.identName var) . envSubsts
  case subst of
    Nothing     -> (:[]) . I.Var <$> internaliseIdent var
    Just substs -> return $ map I.Var substs

internaliseExp desc (E.Index e idxs loc) = do
  idxs' <- mapM (internaliseExp1 "i") idxs
  vs <- internaliseExpToVars "indexed" e
  csidx' <- boundsChecks loc vs idxs'
  let index v = I.PrimOp $ I.Index csidx' v idxs'
  letSubExps desc (map index vs)

internaliseExp desc (E.TupLit es _) =
  concat <$> mapM (internaliseExp desc) es

internaliseExp desc (E.ArrayLit [] (Info et) _) =
  letSubExps desc $ map arrayLit $ internaliseType et
  where arrayLit et' =
          I.PrimOp $ I.ArrayLit [] $ et' `annotateArrayShape` []

internaliseExp desc (E.ArrayLit es (Info rowtype) _) = do
  es' <- mapM (internaliseExp "arr_elem") es
  case es' of
    [] -> do
      let rowtypes = map zeroDim $ internaliseType rowtype
          zeroDim t = t `setArrayShape`
                      Shape (replicate (I.arrayRank t) (constant (0::Int32)))
          arraylit rt = I.PrimOp $ I.ArrayLit [] rt
      letSubExps desc $ map arraylit rowtypes
    e' : _ -> do
      rowtypes <- mapM subExpType e'
      let arraylit ks rt = I.PrimOp $ I.ArrayLit ks rt
      letSubExps desc $ zipWith arraylit (transpose es') rowtypes

internaliseExp desc (E.Apply fname args _ _)
  | Just (rettype, _) <- HM.lookup fname I.builtInFunctions = do
  args' <- mapM (internaliseExp "arg" . fst) args
{'fraction_non_alphanumeric': 0.09338184338184338, 'fraction_numerical': 0.0004331254331254331, 'mean_word_length': 2.3051817921557403, 'pattern_counts': {'":': 0, '<': 54, '<?xml version=': 0, '>': 62, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 26, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6582402', 'n_tokens_mistral': 3589, 'n_tokens_neox': 3453, 'n_words': 1030}