<commit_before>{-# LANGUAGE TypeFamilies, FlexibleInstances, MultiParamTypeClasses #-}
module Futhark.Representation.ExplicitMemory
       ( -- * The Lore definition
         ExplicitMemory
       , MemSummary (..)
       , MemReturn (..)
       , Returns (..)
       , ExpReturns
       , FunReturns
       , expReturns
       , bodyReturns
       , returnsToType
       , generaliseReturns
         -- * Syntax types
       , Prog
       , Body
       , Binding
       , Pattern
       , PatBindee
       , PrimOp
       , LoopOp
       , Exp
       , Lambda
       , FunDec
       , FParam
       , RetType
         -- * Module re-exports
       , module Futhark.Representation.AST.Attributes
       , module Futhark.Representation.AST.Traversals
       , module Futhark.Representation.AST.Pretty
       , module Futhark.Representation.AST.Syntax
       , AST.LambdaT(Lambda)
       , AST.BodyT(Body)
       , AST.PatternT(Pattern)
       , AST.ProgT(Prog)
       , AST.ExpT(PrimOp)
       , AST.ExpT(LoopOp)
       , AST.FunDecT(FunDec)
       )
where

import Control.Applicative
import Control.Monad.State
import qualified Data.HashSet as HS
import qualified Data.HashMap.Lazy as HM
import Data.Foldable (traverse_)
import Data.Maybe
import Data.List
import Data.Loc
import Data.Monoid
import qualified Text.PrettyPrint.Mainland as PP

import qualified Futhark.Representation.AST.Lore as Lore
import qualified Futhark.Representation.AST.Syntax as AST
import Futhark.Representation.AST.Syntax
  hiding (Prog, PrimOp, LoopOp, Exp, Body, Binding,
          Pattern, PatBindee, Lambda, FunDec, FParam, RetType)

import Futhark.TypeCheck.TypeError
import Futhark.Representation.AST.Attributes
import Futhark.Representation.AST.Traversals
import Futhark.Representation.AST.Pretty
import Futhark.Renamer
import Futhark.Binder
import Futhark.Substitute
import qualified Futhark.TypeCheck as TypeCheck
import qualified Futhark.Representation.ExplicitMemory.IndexFunction.Unsafe as IxFun

-- | A lore containing explicit memory information.
data ExplicitMemory = ExplicitMemory

type Prog = AST.Prog ExplicitMemory
type PrimOp = AST.PrimOp ExplicitMemory
type LoopOp = AST.LoopOp ExplicitMemory
type Exp = AST.Exp ExplicitMemory
type Body = AST.Body ExplicitMemory
type Binding = AST.Binding ExplicitMemory
type Pattern = AST.Pattern ExplicitMemory
type PatBindee = AST.PatBindee ExplicitMemory
type Lambda = AST.Lambda ExplicitMemory
type FunDec = AST.FunDec ExplicitMemory
type FParam = AST.FParam ExplicitMemory
type RetType = AST.RetType ExplicitMemory

instance IsRetType [FunReturns] where
  resTypeValues = map returnsToType

  basicRetType t = [ReturnsScalar t]

instance Lore.Lore ExplicitMemory where
  type LetBound ExplicitMemory = MemSummary
  type FParam   ExplicitMemory = MemSummary
  type RetType  ExplicitMemory = [FunReturns]

  representative = ExplicitMemory

  loopResultContext _ res mergevars =
    let shapeContextIdents = loopShapeContext res $ map bindeeIdent mergevars
        memContextIdents = nub $ mapMaybe memIfNecessary res
        memSizeContextIdents = nub $ mapMaybe memSizeIfNecessary memContextIdents
    in memSizeContextIdents <> memContextIdents <> shapeContextIdents
    where memIfNecessary ident =
            case find ((==ident) . bindeeIdent) mergevars of
              Just fparam | MemSummary mem _ <- bindeeLore fparam,
                            isMergeParam mem ->
                Just mem
              _ ->
                Nothing
          memSizeIfNecessary ident
            | Mem (Var sizeident) <- identType ident,
              isMergeParam sizeident =
              Just sizeident
            | otherwise =
              Nothing
          isMergeParam ident =
            any ((==ident) . bindeeIdent) mergevars

data MemSummary = MemSummary Ident IxFun.IxFun
                | Scalar
                deriving (Show)

instance Eq MemSummary where
  _ == _ = True

instance Ord MemSummary where
  _ `compare` _ = EQ

instance FreeIn MemSummary where
  freeIn (MemSummary ident ixfun) = HS.singleton ident <> freeIn ixfun
  freeIn Scalar                   = HS.empty

instance Substitute MemSummary where
  substituteNames subst (MemSummary name f) =
    MemSummary (substituteNames subst name) f
  substituteNames _ Scalar =
    Scalar

instance Rename MemSummary where
  rename (MemSummary ident f) =
    MemSummary <$> rename ident <*> pure f
  rename Scalar =
    return Scalar

instance PP.Pretty MemSummary where
  ppr Scalar = PP.text "scalar"
  ppr (MemSummary mem ixfun) =
    PP.ppr (identName mem) <> PP.text "->" <> PP.ppr ixfun

data MemReturn = ReturnsInBlock Ident
               | ReturnsNewBlock Int
               deriving (Eq, Ord, Show)

instance Rename MemReturn where
  rename (ReturnsInBlock ident) =
    ReturnsInBlock <$> rename ident
  rename (ReturnsNewBlock i) =
    return $ ReturnsNewBlock i

instance Substitute MemReturn where
  substituteNames substs (ReturnsInBlock ident) =
    ReturnsInBlock $ substituteNames substs ident
  substituteNames _ (ReturnsNewBlock i) =
    ReturnsNewBlock i

data Returns a = ReturnsArray BasicType ExtShape Uniqueness a
               | ReturnsScalar BasicType
               | ReturnsMemory SubExp
               deriving (Eq, Ord, Show)

instance FreeIn MemReturn where
  freeIn (ReturnsInBlock v) = HS.singleton v
  freeIn _                  = mempty

instance Substitute a => Substitute (Returns a) where
  substituteNames substs (ReturnsArray bt shape u x) =
    ReturnsArray
    bt
    (substituteNames substs shape)
    u
    (substituteNames substs x)
  substituteNames _ (ReturnsScalar bt) =
    ReturnsScalar bt
  substituteNames substs (ReturnsMemory size) =
    ReturnsMemory $ substituteNames substs size

instance PP.Pretty (Returns (Maybe MemReturn)) where<commit_msg>Fix FreeIn instances for MemSummary and MemReturn.<commit_after>{-# LANGUAGE TypeFamilies, FlexibleInstances, MultiParamTypeClasses #-}
module Futhark.Representation.ExplicitMemory
       ( -- * The Lore definition
         ExplicitMemory
       , MemSummary (..)
       , MemReturn (..)
       , Returns (..)
       , ExpReturns
       , FunReturns
       , expReturns
       , bodyReturns
       , returnsToType
       , generaliseReturns
         -- * Syntax types
       , Prog
       , Body
       , Binding
       , Pattern
       , PatBindee
       , PrimOp
       , LoopOp
       , Exp
       , Lambda
       , FunDec
       , FParam
       , RetType
         -- * Module re-exports
       , module Futhark.Representation.AST.Attributes
       , module Futhark.Representation.AST.Traversals
       , module Futhark.Representation.AST.Pretty
       , module Futhark.Representation.AST.Syntax
       , AST.LambdaT(Lambda)
       , AST.BodyT(Body)
       , AST.PatternT(Pattern)
       , AST.ProgT(Prog)
       , AST.ExpT(PrimOp)
       , AST.ExpT(LoopOp)
       , AST.FunDecT(FunDec)
       )
where

import Control.Applicative
import Control.Monad.State
import qualified Data.HashSet as HS
import qualified Data.HashMap.Lazy as HM
import Data.Foldable (traverse_)
import Data.Maybe
import Data.List
import Data.Loc
import Data.Monoid
import qualified Text.PrettyPrint.Mainland as PP

import qualified Futhark.Representation.AST.Lore as Lore
import qualified Futhark.Representation.AST.Syntax as AST
import Futhark.Representation.AST.Syntax
  hiding (Prog, PrimOp, LoopOp, Exp, Body, Binding,
          Pattern, PatBindee, Lambda, FunDec, FParam, RetType)

import Futhark.TypeCheck.TypeError
import Futhark.Representation.AST.Attributes
import Futhark.Representation.AST.Traversals
import Futhark.Representation.AST.Pretty
import Futhark.Renamer
import Futhark.Binder
import Futhark.Substitute
import qualified Futhark.TypeCheck as TypeCheck
import qualified Futhark.Representation.ExplicitMemory.IndexFunction.Unsafe as IxFun

-- | A lore containing explicit memory information.
data ExplicitMemory = ExplicitMemory

type Prog = AST.Prog ExplicitMemory
type PrimOp = AST.PrimOp ExplicitMemory
type LoopOp = AST.LoopOp ExplicitMemory
type Exp = AST.Exp ExplicitMemory
type Body = AST.Body ExplicitMemory
type Binding = AST.Binding ExplicitMemory
type Pattern = AST.Pattern ExplicitMemory
type PatBindee = AST.PatBindee ExplicitMemory
type Lambda = AST.Lambda ExplicitMemory
type FunDec = AST.FunDec ExplicitMemory
type FParam = AST.FParam ExplicitMemory
type RetType = AST.RetType ExplicitMemory

instance IsRetType [FunReturns] where
  resTypeValues = map returnsToType

  basicRetType t = [ReturnsScalar t]

instance Lore.Lore ExplicitMemory where
  type LetBound ExplicitMemory = MemSummary
  type FParam   ExplicitMemory = MemSummary
  type RetType  ExplicitMemory = [FunReturns]

  representative = ExplicitMemory

  loopResultContext _ res mergevars =
    let shapeContextIdents = loopShapeContext res $ map bindeeIdent mergevars
        memContextIdents = nub $ mapMaybe memIfNecessary res
        memSizeContextIdents = nub $ mapMaybe memSizeIfNecessary memContextIdents
    in memSizeContextIdents <> memContextIdents <> shapeContextIdents
    where memIfNecessary ident =
            case find ((==ident) . bindeeIdent) mergevars of
              Just fparam | MemSummary mem _ <- bindeeLore fparam,
                            isMergeParam mem ->
                Just mem
              _ ->
                Nothing
          memSizeIfNecessary ident
            | Mem (Var sizeident) <- identType ident,
              isMergeParam sizeident =
              Just sizeident
            | otherwise =
              Nothing
          isMergeParam ident =
            any ((==ident) . bindeeIdent) mergevars

data MemSummary = MemSummary Ident IxFun.IxFun
                | Scalar
                deriving (Show)

instance Eq MemSummary where
  _ == _ = True

instance Ord MemSummary where
  _ `compare` _ = EQ

instance FreeIn MemSummary where
  freeIn (MemSummary ident ixfun) = freeIn ident <> freeIn ixfun
  freeIn Scalar                   = HS.empty

instance Substitute MemSummary where
  substituteNames subst (MemSummary name f) =
    MemSummary (substituteNames subst name) f
  substituteNames _ Scalar =
    Scalar

instance Rename MemSummary where
  rename (MemSummary ident f) =
    MemSummary <$> rename ident <*> pure f
  rename Scalar =
    return Scalar

instance PP.Pretty MemSummary where
  ppr Scalar = PP.text "scalar"
  ppr (MemSummary mem ixfun) =
    PP.ppr (identName mem) <> PP.text "->" <> PP.ppr ixfun

data MemReturn = ReturnsInBlock Ident
               | ReturnsNewBlock Int
               deriving (Eq, Ord, Show)

instance Rename MemReturn where
  rename (ReturnsInBlock ident) =
    ReturnsInBlock <$> rename ident
  rename (ReturnsNewBlock i) =
    return $ ReturnsNewBlock i

instance Substitute MemReturn where
  substituteNames substs (ReturnsInBlock ident) =
    ReturnsInBlock $ substituteNames substs ident
  substituteNames _ (ReturnsNewBlock i) =
    ReturnsNewBlock i

data Returns a = ReturnsArray BasicType ExtShape Uniqueness a
               | ReturnsScalar BasicType
               | ReturnsMemory SubExp
               deriving (Eq, Ord, Show)

instance FreeIn MemReturn where
  freeIn (ReturnsInBlock v) = freeIn v
  freeIn _                  = mempty

instance Substitute a => Substitute (Returns a) where
  substituteNames substs (ReturnsArray bt shape u x) =
    ReturnsArray
    bt
    (substituteNames substs shape)
    u
    (substituteNames substs x)
  substituteNames _ (ReturnsScalar bt) =
    ReturnsScalar bt
  substituteNames substs (ReturnsMemory size) =
    ReturnsMemory $ substituteNames substs size

instance PP.Pretty (Returns (Maybe MemReturn)) where
{'fraction_non_alphanumeric': 0.061452513966480445, 'fraction_numerical': 0.0, 'mean_word_length': 2.842800528401585, 'pattern_counts': {'":': 0, '<': 23, '<?xml version=': 0, '>': 27, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5661313', 'n_tokens_mistral': 3400, 'n_tokens_neox': 3261, 'n_words': 961}