<commit_before>-- | The code generator cannot handle the array combinators (@map@ and
-- friends), so this module was written to transform them into the
-- equivalent do-loops.  The transformation is currently rather naive,
-- and - it's certainly worth considering when we can express such
-- transformations in-place.
module Futhark.Transform.FirstOrderTransform
  ( transformProg
  , transformBinding
  , transformBindingRecursively
  , transformLambda
  )
  where

import Control.Applicative
import Control.Monad.State
import Data.Maybe
import qualified Data.HashMap.Lazy as HM
import Data.List

import Prelude

import Futhark.Representation.Basic
import Futhark.Transform.Rename
import Futhark.MonadFreshNames
import Futhark.Tools

-- | Perform the first-order transformation on an Futhark program.
transformProg :: MonadFreshNames m => Prog -> m Prog
transformProg = intraproceduralTransformation transformFunDec

transformFunDec :: MonadFreshNames m => FunDec -> m FunDec
transformFunDec (FunDec fname rettype params body) = do
  (body',_) <-
    runBinderEmptyEnv $
    bindingIdentTypes (map paramIdent params) $
    insertBindingsM $
    transformBody body
  return $ FunDec fname rettype params body'

transformBody :: Body -> Binder Basic Body
transformBody (Body () bnds res) = insertBindingsM $ do
  mapM_ transformBindingRecursively bnds
  return $ resultBody res

-- | First transform any nested 'Body' or 'Lambda' elements, then
-- apply 'transformBinding'.
transformBindingRecursively :: Binding -> Binder Basic ()

transformBindingRecursively (Let pat () (LoopOp (DoLoop res merge form body))) = do
  body' <- bindingIdentTypes (formIdents form ++ map (paramIdent . fst) merge) $
           transformBody body
  addBinding $ Let pat () $ LoopOp $ DoLoop res merge form body'
  where formIdents (ForLoop i _) =
          [Ident i $ Basic Int]
        formIdents (WhileLoop _) =
          []

transformBindingRecursively (Let pat () (LoopOp (Kernel cs w i ispace inps returns body))) = do
  body' <- bindingIdentTypes (Ident i (Basic Int) :
                              map ((`Ident` Basic Int) . fst) ispace ++
                              map kernelInputIdent inps) $
           transformBody body
  addBinding $ Let pat () $ LoopOp $ Kernel cs w i ispace inps returns body'

transformBindingRecursively (Let pat () e) =
  transformBinding =<< liftM (Let pat ()) (mapExpM transform e)
  where transform = identityMapper { mapOnBody = transformBody
                                   , mapOnLambda = transformLambda
                                   , mapOnExtLambda = transformExtLambda
                                   }

-- | Transform a single binding _without_ recursing further.  This
-- means that if called on a 'Map' binding, the resulting loop may
-- still contain SOACs in its body.  Use 'transformBindingRecursively'
-- if this is not what you want.
transformBinding :: Binding -> Binder Basic ()

transformBinding (Let pat () (LoopOp (Map cs width fun arrs))) = do
  let i = lambdaIndex fun
      out_ts = mapType width fun
  resarr <- resultArray out_ts
  outarrs <- forM out_ts $ \t ->
             newIdent "map_outarr" $ t `setUniqueness` Unique
  let outarrs_names = map identName outarrs
      i_ident = Ident i $ Basic Int
  loopbody <- runBodyBinder $ bindingIdentTypes (i_ident:outarrs) $ do
    x <- bindLambda fun (index cs arrs (Var i))
    dests <- letwith cs outarrs_names (pexp $ Var i) $ map (PrimOp . SubExp) x
    return $ resultBody $ map Var dests
  addBinding $ Let pat () $ LoopOp $
    DoLoop outarrs_names (loopMerge outarrs (map Var resarr))
    (ForLoop i width) loopbody

transformBinding (Let pat () (LoopOp (ConcatMap cs _ fun inputs))) = do
  arrs <- forM inputs $ \input -> do
    fun' <- renameLambda fun
    let funparams = lambdaParams fun'
        (ctxparams, valparams) =<commit_msg>Produce unique arrays when first-order transforming map.<commit_after>-- | The code generator cannot handle the array combinators (@map@ and
-- friends), so this module was written to transform them into the
-- equivalent do-loops.  The transformation is currently rather naive,
-- and - it's certainly worth considering when we can express such
-- transformations in-place.
module Futhark.Transform.FirstOrderTransform
  ( transformProg
  , transformBinding
  , transformBindingRecursively
  , transformLambda
  )
  where

import Control.Applicative
import Control.Monad.State
import Data.Maybe
import qualified Data.HashMap.Lazy as HM
import Data.List

import Prelude

import Futhark.Representation.Basic
import Futhark.Transform.Rename
import Futhark.MonadFreshNames
import Futhark.Tools

-- | Perform the first-order transformation on an Futhark program.
transformProg :: MonadFreshNames m => Prog -> m Prog
transformProg = intraproceduralTransformation transformFunDec

transformFunDec :: MonadFreshNames m => FunDec -> m FunDec
transformFunDec (FunDec fname rettype params body) = do
  (body',_) <-
    runBinderEmptyEnv $
    bindingIdentTypes (map paramIdent params) $
    insertBindingsM $
    transformBody body
  return $ FunDec fname rettype params body'

transformBody :: Body -> Binder Basic Body
transformBody (Body () bnds res) = insertBindingsM $ do
  mapM_ transformBindingRecursively bnds
  return $ resultBody res

-- | First transform any nested 'Body' or 'Lambda' elements, then
-- apply 'transformBinding'.
transformBindingRecursively :: Binding -> Binder Basic ()

transformBindingRecursively (Let pat () (LoopOp (DoLoop res merge form body))) = do
  body' <- bindingIdentTypes (formIdents form ++ map (paramIdent . fst) merge) $
           transformBody body
  addBinding $ Let pat () $ LoopOp $ DoLoop res merge form body'
  where formIdents (ForLoop i _) =
          [Ident i $ Basic Int]
        formIdents (WhileLoop _) =
          []

transformBindingRecursively (Let pat () (LoopOp (Kernel cs w i ispace inps returns body))) = do
  body' <- bindingIdentTypes (Ident i (Basic Int) :
                              map ((`Ident` Basic Int) . fst) ispace ++
                              map kernelInputIdent inps) $
           transformBody body
  addBinding $ Let pat () $ LoopOp $ Kernel cs w i ispace inps returns body'

transformBindingRecursively (Let pat () e) =
  transformBinding =<< liftM (Let pat ()) (mapExpM transform e)
  where transform = identityMapper { mapOnBody = transformBody
                                   , mapOnLambda = transformLambda
                                   , mapOnExtLambda = transformExtLambda
                                   }

-- | Transform a single binding _without_ recursing further.  This
-- means that if called on a 'Map' binding, the resulting loop may
-- still contain SOACs in its body.  Use 'transformBindingRecursively'
-- if this is not what you want.
transformBinding :: Binding -> Binder Basic ()

transformBinding (Let pat () (LoopOp (Map cs width fun arrs))) = do
  let i = lambdaIndex fun
      out_ts = mapType width fun
  resarr <- resultArray out_ts
  outarrs <- forM out_ts $ \t ->
             newIdent "map_outarr" $ t `setUniqueness` Unique
  let outarrs_names = map identName outarrs
      i_ident = Ident i $ Basic Int
  loopbody <- runBodyBinder $ bindingIdentTypes (i_ident:outarrs) $ do
    x <- bindLambda fun (index cs arrs (Var i))
    dests <- letwith cs outarrs_names (pexp $ Var i) $ map (PrimOp . SubExp) x
    return $ resultBody $ map Var dests
  addBinding $ Let pat' () $ LoopOp $
    DoLoop outarrs_names (loopMerge outarrs (map Var resarr))
    (ForLoop i width) loopbody
  where pat' = basicPattern' [] $ map (`setIdentUniqueness` Unique) $
               patternValueIdents pat

transformBinding (Let pat () (LoopOp (ConcatMap cs _ fun inputs))) = do
  arrs <- forM inputs $ \input -> do
    fun' <- renameLambda fun
    let funparams = lambdaParams fun'
        (ctxparams, valparams) =
{'fraction_non_alphanumeric': 0.08089544645128466, 'fraction_numerical': 0.0, 'mean_word_length': 3.0867983367983367, 'pattern_counts': {'":': 0, '<': 27, '<?xml version=': 0, '>': 21, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 10, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5821469', 'n_tokens_mistral': 2253, 'n_tokens_neox': 2190, 'n_words': 808}