{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE Strict #-}
{-# LANGUAGE TupleSections #-}

-- | Imperative intermediate language used as a stepping stone in code generation.
--
-- This is a generic representation parametrised on an extensible
-- arbitrary operation.
--
-- Originally inspired by the paper "Defunctionalizing Push Arrays"
-- (FHPC '14).
module Futhark.CodeGen.ImpCode
  ( Definitions (..),
    Functions (..),
    Function,
    FunctionT (..),
    Constants (..),
    ValueDesc (..),
    Signedness (..),
    ExternalValue (..),
    Param (..),
    paramName,
    SubExp (..),
    MemSize,
    DimSize,
    Space (..),
    SpaceId,
    Code (..),
    PrimValue (..),
    ExpLeaf (..),
    Exp,
    TExp,
    Volatility (..),
    Arg (..),
    var,
    vi32,
    vi64,
    index,
    ErrorMsg (..),
    ErrorMsgPart (..),
    errorMsgArgTypes,
    ArrayContents (..),
    declaredIn,
    lexicalMemoryUsage,
    calledFuncs,

    -- * Typed enumerations
    Bytes,
    Elements,
    elements,
    bytes,
    withElemType,

    -- * Re-exports from other modules.
    module Language.Futhark.Core,
    module Futhark.IR.Primitive,
    module Futhark.Analysis.PrimExp,
    module Futhark.IR.Kernels.Sizes,
    module Futhark.IR.Prop.Names,
  )
where

import Data.List (intersperse)
import qualified Data.Map as M
import qualified Data.Set as S
import Data.Traversable
import Futhark.Analysis.PrimExp
import Futhark.IR.Kernels.Sizes (Count (..))
import Futhark.IR.Pretty ()
import Futhark.IR.Primitive
import Futhark.IR.Prop.Names
import Futhark.IR.Syntax
  ( ErrorMsg (..),
    ErrorMsgPart (..),
    Space (..),
    SpaceId,
    SubExp (..),
    errorMsgArgTypes,
  )
import Futhark.Util.Pretty hiding (space)
import Language.Futhark.Core

-- | The size of a memory block.
type MemSize = SubExp

-- | The size of an array.
type DimSize = SubExp

-- | An ImpCode function parameter.
data Param
  = MemParam VName Space
  | ScalarParam VName PrimType
  deriving (Int -> Param -> ShowS
[Param] -> ShowS
Param -> String
(Int -> Param -> ShowS)
-> (Param -> String) -> ([Param] -> ShowS) -> Show Param
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Param] -> ShowS
$cshowList :: [Param] -> ShowS
show :: Param -> String
$cshow :: Param -> String
showsPrec :: Int -> Param -> ShowS
$cshowsPrec :: Int -> Param -> ShowS
Show)

-- | The name of a parameter.
paramName :: Param -> VName
paramName :: Param -> VName
paramName (MemParam VName
name Space
_) = VName
name
paramName (ScalarParam VName
name PrimType
_) = VName
name

-- | A collection of imperative functions and constants.
data Definitions a = Definitions
  { forall a. Definitions a -> Constants a
defConsts :: Constants a,
    forall a. Definitions a -> Functions a
defFuns :: Functions a
  }

-- | A collection of imperative functions.
newtype Functions a = Functions [(Name, Function a)]

instance Semigroup (Functions a) where
  Functions [(Name, Function a)]
x <> :: Functions a -> Functions a -> Functions a
<> Functions [(Name, Function a)]
y = [(Name, Function a)] -> Functions a
forall a. [(Name, Function a)] -> Functions a
Functions ([(Name, Function a)] -> Functions a)
-> [(Name, Function a)] -> Functions a
forall a b. (a -> b) -> a -> b
$ [(Name, Function a)]
x [(Name, Function a)]
-> [(Name, Function a)] -> [(Name, Function a)]
forall a. [a] -> [a] -> [a]
++ [(Name, Function a)]
y

instance Monoid (Functions a) where
  mempty :: Functions a
mempty = [(Name, Function a)] -> Functions a
forall a. [(Name, Function a)] -> Functions a
Functions []

-- | A collection of imperative constants.
data Constants a = Constants
  { -- | The constants that are made available to the functions.
    forall a. Constants a -> [Param]
constsDecl :: [Param],
    -- | Setting the value of the constants.  Note that this must not
    -- contain declarations of the names defined in 'constsDecl'.
    forall a. Constants a -> Code a
constsInit :: Code a
  }

-- | Since the core language does not care for signedness, but the
-- source language does, entry point input/output information has
-- metadata for integer types (and arrays containing these) that
-- indicate whether they are really unsigned integers.
data Signedness
  = TypeUnsigned
  | TypeDirect
  deriving (Signedness -> Signedness -> Bool
(Signedness -> Signedness -> Bool)
-> (Signedness -> Signedness -> Bool) -> Eq Signedness
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Signedness -> Signedness -> Bool
$c/= :: Signedness -> Signedness -> Bool
== :: Signedness -> Signedness -> Bool
$c== :: Signedness -> Signedness -> Bool
Eq, Int -> Signedness -> ShowS
[Signedness] -> ShowS
Signedness -> String
(Int -> Signedness -> ShowS)
-> (Signedness -> String)
-> ([Signedness] -> ShowS)
-> Show Signedness
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Signedness] -> ShowS
$cshowList :: [Signedness] -> ShowS
show :: Signedness -> String
$cshow :: Signedness -> String
showsPrec :: Int -> Signedness -> ShowS
$cshowsPrec :: Int -> Signedness -> ShowS
Show)

-- | A description of an externally meaningful value.
data ValueDesc
  = -- | An array with memory block memory space, element type,
    -- signedness of element type (if applicable), and shape.
    ArrayValue VName Space PrimType Signedness [DimSize]
  | -- | A scalar value with signedness if applicable.
    ScalarValue PrimType Signedness VName
  deriving (ValueDesc -> ValueDesc -> Bool
(ValueDesc -> ValueDesc -> Bool)
-> (ValueDesc -> ValueDesc -> Bool) -> Eq ValueDesc
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: ValueDesc -> ValueDesc -> Bool
$c/= :: ValueDesc -> ValueDesc -> Bool
== :: ValueDesc -> ValueDesc -> Bool
$c== :: ValueDesc -> ValueDesc -> Bool
Eq, Int -> ValueDesc -> ShowS
[ValueDesc] -> ShowS
ValueDesc -> String
(Int -> ValueDesc -> ShowS)
-> (ValueDesc -> String)
-> ([ValueDesc] -> ShowS)
-> Show ValueDesc
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [ValueDesc] -> ShowS
$cshowList :: [ValueDesc] -> ShowS
show :: ValueDesc -> String
$cshow :: ValueDesc -> String
showsPrec :: Int -> ValueDesc -> ShowS
$cshowsPrec :: Int -> ValueDesc -> ShowS
Show)

-- | ^ An externally visible value.  This can be an opaque value
-- (covering several physical internal values), or a single value that
-- can be used externally.
data ExternalValue
  = -- | The string is a human-readable description
    -- with no other semantics.
    OpaqueValue String [ValueDesc]
  | TransparentValue ValueDesc
  deriving (Int -> ExternalValue -> ShowS
[ExternalValue] -> ShowS
ExternalValue -> String
(Int -> ExternalValue -> ShowS)
-> (ExternalValue -> String)
-> ([ExternalValue] -> ShowS)
-> Show ExternalValue
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [ExternalValue] -> ShowS
$cshowList :: [ExternalValue] -> ShowS
show :: ExternalValue -> String
$cshow :: ExternalValue -> String
showsPrec :: Int -> ExternalValue -> ShowS
$cshowsPrec :: Int -> ExternalValue -> ShowS
Show)

-- | A imperative function, containing the body as well as its
-- low-level inputs and outputs, as well as its high-level arguments
-- and results.  The latter are only used if the function is an entry
-- point.
data FunctionT a = Function
  { forall a. FunctionT a -> Bool
functionEntry :: Bool,
    forall a. FunctionT a -> [Param]
functionOutput :: [Param],
    forall a. FunctionT a -> [Param]
functionInput :: [Param],
    forall a. FunctionT a -> Code a
functionBody :: Code a,
    forall a. FunctionT a -> [ExternalValue]
functionResult :: [ExternalValue],
    forall a. FunctionT a -> [ExternalValue]
functionArgs :: [ExternalValue]
  }
  deriving (Int -> FunctionT a -> ShowS
[FunctionT a] -> ShowS
FunctionT a -> String
(Int -> FunctionT a -> ShowS)
-> (FunctionT a -> String)
-> ([FunctionT a] -> ShowS)
-> Show (FunctionT a)
forall a. Show a => Int -> FunctionT a -> ShowS
forall a. Show a => [FunctionT a] -> ShowS
forall a. Show a => FunctionT a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [FunctionT a] -> ShowS
$cshowList :: forall a. Show a => [FunctionT a] -> ShowS
show :: FunctionT a -> String
$cshow :: forall a. Show a => FunctionT a -> String
showsPrec :: Int -> FunctionT a -> ShowS
$cshowsPrec :: forall a. Show a => Int -> FunctionT a -> ShowS
Show)

-- | Type alias for namespace control.
type Function = FunctionT

-- | The contents of a statically declared constant array.  Such
-- arrays are always unidimensional, and reshaped if necessary in the
-- code that uses them.
data ArrayContents
  = -- | Precisely these values.
    ArrayValues [PrimValue]
  | -- | This many zeroes.
    ArrayZeros Int
  deriving (Int -> ArrayContents -> ShowS
[ArrayContents] -> ShowS
ArrayContents -> String
(Int -> ArrayContents -> ShowS)
-> (ArrayContents -> String)
-> ([ArrayContents] -> ShowS)
-> Show ArrayContents
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [ArrayContents] -> ShowS
$cshowList :: [ArrayContents] -> ShowS
show :: ArrayContents -> String
$cshow :: ArrayContents -> String
showsPrec :: Int -> ArrayContents -> ShowS
$cshowsPrec :: Int -> ArrayContents -> ShowS
Show)

-- | A block of imperative code.  Parameterised by an 'Op', which
-- allows extensibility.  Concrete uses of this type will instantiate
-- the type parameter with e.g. a construct for launching GPU kernels.
data Code a
  = -- | No-op.  Crucial for the 'Monoid' instance.
    Skip
  | -- | Statement composition.  Crucial for the 'Semigroup' instance.
    Code a :>>: Code a
  | -- | A for-loop iterating the given number of times.
    -- The loop parameter starts counting from zero and will
    -- have the same (integer) type as the bound.  The bound
    -- is evaluated just once, before the loop is entered.
    For VName Exp (Code a)
  | -- | While loop.  The conditional is (of course)
    -- re-evaluated before every iteration of the loop.
    While (TExp Bool) (Code a)
  | -- | Declare a memory block variable that will point to
    -- memory in the given memory space.  Note that this is
    -- distinct from allocation.  The memory block must be the
    -- target of either an 'Allocate' or a 'SetMem' before it
    -- can be used for reading or writing.
    DeclareMem VName Space
  | -- | Declare a scalar variable with an initially undefined value.
    DeclareScalar VName Volatility PrimType
  | -- | Create an array containing the given values.  The
    -- lifetime of the array will be the entire application.
    -- This is mostly used for constant arrays, but also for
    -- some bookkeeping data, like the synchronisation
    -- counts used to implement reduction.
    DeclareArray VName Space PrimType ArrayContents
  | -- | Memory space must match the corresponding
    -- 'DeclareMem'.
    Allocate VName (Count Bytes (TExp Int64)) Space
  | -- | Indicate that some memory block will never again be
    -- referenced via the indicated variable.  However, it
    -- may still be accessed through aliases.  It is only
    -- safe to actually deallocate the memory block if this
    -- is the last reference.  There is no guarantee that
    -- all memory blocks will be freed with this statement.
    -- Backends are free to ignore it entirely.
    Free VName Space
  | -- | Destination, offset in destination, destination
    -- space, source, offset in source, offset space, number
    -- of bytes.
    Copy
      VName
      (Count Bytes (TExp Int64))
      Space
      VName
      (Count Bytes (TExp Int64))
      Space
      (Count Bytes (TExp Int64))
  | -- | @Write mem i t space vol v@ writes the value @v@ to
    -- @mem@ offset by @i@ elements of type @t@.  The
    -- 'Space' argument is the memory space of @mem@
    -- (technically redundant, but convenient).  Note that
    -- /reading/ is done with an 'Exp' ('Index').
    Write VName (Count Elements (TExp Int64)) PrimType Space Volatility Exp
  | -- | Set a scalar variable.
    SetScalar VName Exp
  | -- | Must be in same space.
    SetMem VName VName Space
  | -- | Function call.  The results are written to the
    -- provided 'VName' variables.
    Call [VName] Name [Arg]
  | -- | Conditional execution.
    If (TExp Bool) (Code a) (Code a)
  | -- | Assert that something must be true.  Should it turn
    -- out not to be true, then report a failure along with
    -- the given error message.
    Assert Exp (ErrorMsg Exp) (SrcLoc, [SrcLoc])
  | -- | Has the same semantics as the contained code, but
    -- the comment should show up in generated code for ease
    -- of inspection.
    Comment String (Code a)
  | -- | Print the given value to the screen, somehow
    -- annotated with the given string as a description.  If
    -- no type/value pair, just print the string.  This has
    -- no semantic meaning, but is used entirely for
    -- debugging.  Code generators are free to ignore this
    -- statement.
    DebugPrint String (Maybe Exp)
  | -- | Perform an extensible operation.
    Op a
  deriving (Int -> Code a -> ShowS
[Code a] -> ShowS
Code a -> String
(Int -> Code a -> ShowS)
-> (Code a -> String) -> ([Code a] -> ShowS) -> Show (Code a)
forall a. Show a => Int -> Code a -> ShowS
forall a. Show a => [Code a] -> ShowS
forall a. Show a => Code a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Code a] -> ShowS
$cshowList :: forall a. Show a => [Code a] -> ShowS
show :: Code a -> String
$cshow :: forall a. Show a => Code a -> String
showsPrec :: Int -> Code a -> ShowS
$cshowsPrec :: forall a. Show a => Int -> Code a -> ShowS
Show)

-- | The volatility of a memory access or variable.  Feel free to
-- ignore this for backends where it makes no sense (anything but C
-- and similar low-level things)
data Volatility = Volatile | Nonvolatile
  deriving (Volatility -> Volatility -> Bool
(Volatility -> Volatility -> Bool)
-> (Volatility -> Volatility -> Bool) -> Eq Volatility
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Volatility -> Volatility -> Bool
$c/= :: Volatility -> Volatility -> Bool
== :: Volatility -> Volatility -> Bool
$c== :: Volatility -> Volatility -> Bool
Eq, Eq Volatility
Eq Volatility
-> (Volatility -> Volatility -> Ordering)
-> (Volatility -> Volatility -> Bool)
-> (Volatility -> Volatility -> Bool)
-> (Volatility -> Volatility -> Bool)
-> (Volatility -> Volatility -> Bool)
-> (Volatility -> Volatility -> Volatility)
-> (Volatility -> Volatility -> Volatility)
-> Ord Volatility
Volatility -> Volatility -> Bool
Volatility -> Volatility -> Ordering
Volatility -> Volatility -> Volatility
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: Volatility -> Volatility -> Volatility
$cmin :: Volatility -> Volatility -> Volatility
max :: Volatility -> Volatility -> Volatility
$cmax :: Volatility -> Volatility -> Volatility
>= :: Volatility -> Volatility -> Bool
$c>= :: Volatility -> Volatility -> Bool
> :: Volatility -> Volatility -> Bool
$c> :: Volatility -> Volatility -> Bool
<= :: Volatility -> Volatility -> Bool
$c<= :: Volatility -> Volatility -> Bool
< :: Volatility -> Volatility -> Bool
$c< :: Volatility -> Volatility -> Bool
compare :: Volatility -> Volatility -> Ordering
$ccompare :: Volatility -> Volatility -> Ordering
Ord, Int -> Volatility -> ShowS
[Volatility] -> ShowS
Volatility -> String
(Int -> Volatility -> ShowS)
-> (Volatility -> String)
-> ([Volatility] -> ShowS)
-> Show Volatility
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Volatility] -> ShowS
$cshowList :: [Volatility] -> ShowS
show :: Volatility -> String
$cshow :: Volatility -> String
showsPrec :: Int -> Volatility -> ShowS
$cshowsPrec :: Int -> Volatility -> ShowS
Show)

instance Semigroup (Code a) where
  Code a
Skip <> :: Code a -> Code a -> Code a
<> Code a
y = Code a
y
  Code a
x <> Code a
Skip = Code a
x
  Code a
x <> Code a
y = Code a
x Code a -> Code a -> Code a
forall a. Code a -> Code a -> Code a
:>>: Code a
y

instance Monoid (Code a) where
  mempty :: Code a
mempty = Code a
forall a. Code a
Skip

-- | Find those memory blocks that are used only lexically.  That is,
-- are not used as the source or target of a 'SetMem', or are the
-- result of the function, nor passed as arguments to other functions.
-- This is interesting because such memory blocks do not need
-- reference counting, but can be managed in a purely stack-like
-- fashion.
--
-- We do not look inside any 'Op's.  We assume that no 'Op' is going
-- to 'SetMem' a memory block declared outside it.
lexicalMemoryUsage :: Function a -> M.Map VName Space
lexicalMemoryUsage :: forall a. Function a -> Map VName Space
lexicalMemoryUsage Function a
func =
  (VName -> Space -> Bool) -> Map VName Space -> Map VName Space
forall k a. (k -> a -> Bool) -> Map k a -> Map k a
M.filterWithKey (Bool -> Space -> Bool
forall a b. a -> b -> a
const (Bool -> Space -> Bool)
-> (VName -> Bool) -> VName -> Space -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Bool -> Bool
not (Bool -> Bool) -> (VName -> Bool) -> VName -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VName -> Names -> Bool
`nameIn` Names
nonlexical)) (Map VName Space -> Map VName Space)
-> Map VName Space -> Map VName Space
forall a b. (a -> b) -> a -> b
$
    Code a -> Map VName Space
forall {a}. Code a -> Map VName Space
declared (Code a -> Map VName Space) -> Code a -> Map VName Space
forall a b. (a -> b) -> a -> b
$ Function a -> Code a
forall a. FunctionT a -> Code a
functionBody Function a
func
  where
    nonlexical :: Names
nonlexical =
      Code a -> Names
forall {a}. Code a -> Names
set (Function a -> Code a
forall a. FunctionT a -> Code a
functionBody Function a
func)
        Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> [VName] -> Names
namesFromList ((Param -> VName) -> [Param] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param -> VName
paramName (Function a -> [Param]
forall a. FunctionT a -> [Param]
functionOutput Function a
func))

    go :: (Code a -> a) -> Code a -> a
go Code a -> a
f (Code a
x :>>: Code a
y) = Code a -> a
f Code a
x a -> a -> a
forall a. Semigroup a => a -> a -> a
<> Code a -> a
f Code a
y
    go Code a -> a
f (If TExp Bool
_ Code a
x Code a
y) = Code a -> a
f Code a
x a -> a -> a
forall a. Semigroup a => a -> a -> a
<> Code a -> a
f Code a
y
    go Code a -> a
f (For VName
_ Exp
_ Code a
x) = Code a -> a
f Code a
x
    go Code a -> a
f (While TExp Bool
_ Code a
x) = Code a -> a
f Code a
x
    go Code a -> a
f (Comment String
_ Code a
x) = Code a -> a
f Code a
x
    go Code a -> a
_ Code a
_ = a
forall a. Monoid a => a
mempty

    declared :: Code a -> Map VName Space
declared (DeclareMem VName
mem Space
space) =
      VName -> Space -> Map VName Space
forall k a. k -> a -> Map k a
M.singleton VName
mem Space
space
    declared Code a
x = (Code a -> Map VName Space) -> Code a -> Map VName Space
forall {a} {a}. Monoid a => (Code a -> a) -> Code a -> a
go Code a -> Map VName Space
declared Code a
x

    set :: Code a -> Names
set (SetMem VName
x VName
y Space
_) = [VName] -> Names
namesFromList [VName
x, VName
y]
    set (Call [VName]
_ Name
_ [Arg]
args) = (Arg -> Names) -> [Arg] -> Names
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Arg -> Names
onArg [Arg]
args
      where
        onArg :: Arg -> Names
onArg ExpArg {} = Names
forall a. Monoid a => a
mempty
        onArg (MemArg VName
x) = VName -> Names
oneName VName
x
    set Code a
x = (Code a -> Names) -> Code a -> Names
forall {a} {a}. Monoid a => (Code a -> a) -> Code a -> a
go Code a -> Names
set Code a
x

-- | The set of functions that are called by this code.  Assumes there
-- are no function calls in 'Op's.
calledFuncs :: Code a -> S.Set Name
calledFuncs :: forall a. Code a -> Set Name
calledFuncs (Code a
x :>>: Code a
y) = Code a -> Set Name
forall a. Code a -> Set Name
calledFuncs Code a
x Set Name -> Set Name -> Set Name
forall a. Semigroup a => a -> a -> a
<> Code a -> Set Name
forall a. Code a -> Set Name
calledFuncs Code a
y
calledFuncs (If TExp Bool
_ Code a
x Code a
y) = Code a -> Set Name
forall a. Code a -> Set Name
calledFuncs Code a
x Set Name -> Set Name -> Set Name
forall a. Semigroup a => a -> a -> a
<> Code a -> Set Name
forall a. Code a -> Set Name
calledFuncs Code a
y
calledFuncs (For VName
_ Exp
_ Code a
x) = Code a -> Set Name
forall a. Code a -> Set Name
calledFuncs Code a
x
calledFuncs (While TExp Bool
_ Code a
x) = Code a -> Set Name
forall a. Code a -> Set Name
calledFuncs Code a
x
calledFuncs (Comment String
_ Code a
x) = Code a -> Set Name
forall a. Code a -> Set Name
calledFuncs Code a
x
calledFuncs (Call [VName]
_ Name
f [Arg]
_) = Name -> Set Name
forall a. a -> Set a
S.singleton Name
f
calledFuncs Code a
_ = Set Name
forall a. Monoid a => a
mempty

-- | The leaves of an 'Exp'.
data ExpLeaf
  = -- | A scalar variable.  The type is stored in the
    -- 'LeafExp' constructor itself.
    ScalarVar VName
  | -- | Reading a value from memory.  The arguments have
    -- the same meaning as with 'Write'.
    Index VName (Count Elements (TExp Int64)) PrimType Space Volatility
  deriving (ExpLeaf -> ExpLeaf -> Bool
(ExpLeaf -> ExpLeaf -> Bool)
-> (ExpLeaf -> ExpLeaf -> Bool) -> Eq ExpLeaf
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: ExpLeaf -> ExpLeaf -> Bool
$c/= :: ExpLeaf -> ExpLeaf -> Bool
== :: ExpLeaf -> ExpLeaf -> Bool
$c== :: ExpLeaf -> ExpLeaf -> Bool
Eq, Int -> ExpLeaf -> ShowS
[ExpLeaf] -> ShowS
ExpLeaf -> String
(Int -> ExpLeaf -> ShowS)
-> (ExpLeaf -> String) -> ([ExpLeaf] -> ShowS) -> Show ExpLeaf
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [ExpLeaf] -> ShowS
$cshowList :: [ExpLeaf] -> ShowS
show :: ExpLeaf -> String
$cshow :: ExpLeaf -> String
showsPrec :: Int -> ExpLeaf -> ShowS
$cshowsPrec :: Int -> ExpLeaf -> ShowS
Show)

-- | A side-effect free expression whose execution will produce a
-- single primitive value.
type Exp = PrimExp ExpLeaf

-- | Like 'Exp', but with a required/known type.
type TExp t = TPrimExp t ExpLeaf

-- | A function call argument.
data Arg
  = ExpArg Exp
  | MemArg VName
  deriving (Int -> Arg -> ShowS
[Arg] -> ShowS
Arg -> String
(Int -> Arg -> ShowS)
-> (Arg -> String) -> ([Arg] -> ShowS) -> Show Arg
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Arg] -> ShowS
$cshowList :: [Arg] -> ShowS
show :: Arg -> String
$cshow :: Arg -> String
showsPrec :: Int -> Arg -> ShowS
$cshowsPrec :: Int -> Arg -> ShowS
Show)

-- | Phantom type for a count of elements.
data Elements

-- | Phantom type for a count of bytes.
data Bytes

-- | This expression counts elements.
elements :: a -> Count Elements a
elements :: forall a. a -> Count Elements a
elements = a -> Count Elements a
forall u e. e -> Count u e
Count

-- | This expression counts bytes.
bytes :: a -> Count Bytes a
bytes :: forall a. a -> Count Bytes a
bytes = a -> Count Bytes a
forall u e. e -> Count u e
Count

-- | Convert a count of elements into a count of bytes, given the
-- per-element size.
withElemType :: Count Elements (TExp Int64) -> PrimType -> Count Bytes (TExp Int64)
withElemType :: Count Elements (TPrimExp Int64 ExpLeaf)
-> PrimType -> Count Bytes (TPrimExp Int64 ExpLeaf)
withElemType (Count TPrimExp Int64 ExpLeaf
e) PrimType
t = TPrimExp Int64 ExpLeaf -> Count Bytes (TPrimExp Int64 ExpLeaf)
forall a. a -> Count Bytes a
bytes (TPrimExp Int64 ExpLeaf -> Count Bytes (TPrimExp Int64 ExpLeaf))
-> TPrimExp Int64 ExpLeaf -> Count Bytes (TPrimExp Int64 ExpLeaf)
forall a b. (a -> b) -> a -> b
$ TPrimExp Int64 ExpLeaf -> TPrimExp Int64 ExpLeaf
forall t v. IntExp t => TPrimExp t v -> TPrimExp Int64 v
sExt64 TPrimExp Int64 ExpLeaf
e TPrimExp Int64 ExpLeaf
-> TPrimExp Int64 ExpLeaf -> TPrimExp Int64 ExpLeaf
forall a. Num a => a -> a -> a
* PrimType -> TPrimExp Int64 ExpLeaf
forall a. Num a => PrimType -> a
primByteSize PrimType
t

-- | Turn a 'VName' into a 'Imp.ScalarVar'.
var :: VName -> PrimType -> Exp
var :: VName -> PrimType -> Exp
var = ExpLeaf -> PrimType -> Exp
forall v. v -> PrimType -> PrimExp v
LeafExp (ExpLeaf -> PrimType -> Exp)
-> (VName -> ExpLeaf) -> VName -> PrimType -> Exp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> ExpLeaf
ScalarVar

-- | Turn a 'VName' into a v'Int32' 'Imp.ScalarVar'.
vi32 :: VName -> TExp Int32
vi32 :: VName -> TExp Int32
vi32 = Exp -> TExp Int32
forall t v. PrimExp v -> TPrimExp t v
TPrimExp (Exp -> TExp Int32) -> (VName -> Exp) -> VName -> TExp Int32
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VName -> PrimType -> Exp) -> PrimType -> VName -> Exp
forall a b c. (a -> b -> c) -> b -> a -> c
flip VName -> PrimType -> Exp
var (IntType -> PrimType
IntType IntType
Int32)

-- | Turn a 'VName' into a v'Int64' 'Imp.ScalarVar'.
vi64 :: VName -> TExp Int64
vi64 :: VName -> TPrimExp Int64 ExpLeaf
vi64 = Exp -> TPrimExp Int64 ExpLeaf
forall t v. PrimExp v -> TPrimExp t v
TPrimExp (Exp -> TPrimExp Int64 ExpLeaf)
-> (VName -> Exp) -> VName -> TPrimExp Int64 ExpLeaf
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VName -> PrimType -> Exp) -> PrimType -> VName -> Exp
forall a b c. (a -> b -> c) -> b -> a -> c
flip VName -> PrimType -> Exp
var (IntType -> PrimType
IntType IntType
Int64)

-- | Concise wrapper for using 'Index'.
index :: VName -> Count Elements (TExp Int64) -> PrimType -> Space -> Volatility -> Exp
index :: VName
-> Count Elements (TPrimExp Int64 ExpLeaf)
-> PrimType
-> Space
-> Volatility
-> Exp
index VName
arr Count Elements (TPrimExp Int64 ExpLeaf)
i PrimType
t Space
s Volatility
vol = ExpLeaf -> PrimType -> Exp
forall v. v -> PrimType -> PrimExp v
LeafExp (VName
-> Count Elements (TPrimExp Int64 ExpLeaf)
-> PrimType
-> Space
-> Volatility
-> ExpLeaf
Index VName
arr Count Elements (TPrimExp Int64 ExpLeaf)
i PrimType
t Space
s Volatility
vol) PrimType
t

-- Prettyprinting definitions.

instance Pretty op => Pretty (Definitions op) where
  ppr :: Definitions op -> Doc
ppr (Definitions Constants op
consts Functions op
funs) =
    Constants op -> Doc
forall a. Pretty a => a -> Doc
ppr Constants op
consts Doc -> Doc -> Doc
</> Functions op -> Doc
forall a. Pretty a => a -> Doc
ppr Functions op
funs

instance Pretty op => Pretty (Functions op) where
  ppr :: Functions op -> Doc
ppr (Functions [(Name, Function op)]
funs) = [Doc] -> Doc
stack ([Doc] -> Doc) -> [Doc] -> Doc
forall a b. (a -> b) -> a -> b
$ Doc -> [Doc] -> [Doc]
forall a. a -> [a] -> [a]
intersperse Doc
forall a. Monoid a => a
mempty ([Doc] -> [Doc]) -> [Doc] -> [Doc]
forall a b. (a -> b) -> a -> b
$ ((Name, Function op) -> Doc) -> [(Name, Function op)] -> [Doc]
forall a b. (a -> b) -> [a] -> [b]
map (Name, Function op) -> Doc
forall {a} {a}. (Pretty a, Pretty a) => (a, a) -> Doc
ppFun [(Name, Function op)]
funs
    where
      ppFun :: (a, a) -> Doc
ppFun (a
name, a
fun) =
        String -> Doc
text String
"Function " Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> a -> Doc
forall a. Pretty a => a -> Doc
ppr a
name Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Doc
colon Doc -> Doc -> Doc
</> Int -> Doc -> Doc
indent Int
2 (a -> Doc
forall a. Pretty a => a -> Doc
ppr a
fun)

instance Pretty op => Pretty (Constants op) where
  ppr :: Constants op -> Doc
ppr (Constants [Param]
decls Code op
code) =
    String -> Doc
text String
"Constants:" Doc -> Doc -> Doc
</> Int -> Doc -> Doc
indent Int
2 ([Doc] -> Doc
stack ([Doc] -> Doc) -> [Doc] -> Doc
forall a b. (a -> b) -> a -> b
$ (Param -> Doc) -> [Param] -> [Doc]
forall a b. (a -> b) -> [a] -> [b]
map Param -> Doc
forall a. Pretty a => a -> Doc
ppr [Param]
decls)
      Doc -> Doc -> Doc
</> Doc
forall a. Monoid a => a
mempty
      Doc -> Doc -> Doc
</> String -> Doc
text String
"Initialisation:"
      Doc -> Doc -> Doc
</> Int -> Doc -> Doc
indent Int
2 (Code op -> Doc
forall a. Pretty a => a -> Doc
ppr Code op
code)

instance Pretty op => Pretty (FunctionT op) where
  ppr :: FunctionT op -> Doc
ppr (Function Bool
_ [Param]
outs [Param]
ins Code op
body [ExternalValue]
results [ExternalValue]
args) =
    String -> Doc
text String
"Inputs:" Doc -> Doc -> Doc
</> [Param] -> Doc
forall a. Pretty a => [a] -> Doc
block [Param]
ins
      Doc -> Doc -> Doc
</> String -> Doc
text String
"Outputs:"
      Doc -> Doc -> Doc
</> [Param] -> Doc
forall a. Pretty a => [a] -> Doc
block [Param]
outs
      Doc -> Doc -> Doc
</> String -> Doc
text String
"Arguments:"
      Doc -> Doc -> Doc
</> [ExternalValue] -> Doc
forall a. Pretty a => [a] -> Doc
block [ExternalValue]
args
      Doc -> Doc -> Doc
</> String -> Doc
text String
"Result:"
      Doc -> Doc -> Doc
</> [ExternalValue] -> Doc
forall a. Pretty a => [a] -> Doc
block [ExternalValue]
results
      Doc -> Doc -> Doc
</> String -> Doc
text String
"Body:"
      Doc -> Doc -> Doc
</> Int -> Doc -> Doc
indent Int
2 (Code op -> Doc
forall a. Pretty a => a -> Doc
ppr Code op
body)
    where
      block :: Pretty a => [a] -> Doc
      block :: forall a. Pretty a => [a] -> Doc
block = Int -> Doc -> Doc
indent Int
2 (Doc -> Doc) -> ([a] -> Doc) -> [a] -> Doc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Doc] -> Doc
stack ([Doc] -> Doc) -> ([a] -> [Doc]) -> [a] -> Doc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> Doc) -> [a] -> [Doc]
forall a b. (a -> b) -> [a] -> [b]
map a -> Doc
forall a. Pretty a => a -> Doc
ppr

instance Pretty Param where
  ppr :: Param -> Doc
ppr (ScalarParam VName
name PrimType
ptype) = PrimType -> Doc
forall a. Pretty a => a -> Doc
ppr PrimType
ptype Doc -> Doc -> Doc
<+> VName -> Doc
forall a. Pretty a => a -> Doc
ppr VName
name
  ppr (MemParam VName
name Space
space) = String -> Doc
text String
"mem" Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Space -> Doc
forall a. Pretty a => a -> Doc
ppr Space
space Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> String -> Doc
text String
" " Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> VName -> Doc
forall a. Pretty a => a -> Doc
ppr VName
name

instance Pretty ValueDesc where
  ppr :: ValueDesc -> Doc
ppr (ScalarValue PrimType
t Signedness
ept VName
name) =
    PrimType -> Doc
forall a. Pretty a => a -> Doc
ppr PrimType
t Doc -> Doc -> Doc
<+> VName -> Doc
forall a. Pretty a => a -> Doc
ppr VName
name Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Doc
ept'
    where
      ept' :: Doc
ept' = case Signedness
ept of
        Signedness
TypeUnsigned -> String -> Doc
text String
" (unsigned)"
        Signedness
TypeDirect -> Doc
forall a. Monoid a => a
mempty
  ppr (ArrayValue VName
mem Space
space PrimType
et Signedness
ept [DimSize]
shape) =
    (DimSize -> Doc -> Doc) -> Doc -> [DimSize] -> Doc
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr DimSize -> Doc -> Doc
forall {a}. Pretty a => a -> Doc -> Doc
f (PrimType -> Doc
forall a. Pretty a => a -> Doc
ppr PrimType
et) [DimSize]
shape Doc -> Doc -> Doc
<+> String -> Doc
text String
"at" Doc -> Doc -> Doc
<+> VName -> Doc
forall a. Pretty a => a -> Doc
ppr VName
mem Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Space -> Doc
forall a. Pretty a => a -> Doc
ppr Space
space Doc -> Doc -> Doc
<+> Doc
ept'
    where
      f :: a -> Doc -> Doc
f a
e Doc
s = Doc -> Doc
brackets (Doc -> Doc) -> Doc -> Doc
forall a b. (a -> b) -> a -> b
$ Doc
s Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Doc
comma Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> a -> Doc
forall a. Pretty a => a -> Doc
ppr a
e
      ept' :: Doc
ept' = case Signedness
ept of
        Signedness
TypeUnsigned -> String -> Doc
text String
" (unsigned)"
        Signedness
TypeDirect -> Doc
forall a. Monoid a => a
mempty

instance Pretty ExternalValue where
  ppr :: ExternalValue -> Doc
ppr (TransparentValue ValueDesc
v) = ValueDesc -> Doc
forall a. Pretty a => a -> Doc
ppr ValueDesc
v
  ppr (OpaqueValue String
desc [ValueDesc]
vs) =
    String -> Doc
text String
"opaque" Doc -> Doc -> Doc
<+> String -> Doc
text String
desc
      Doc -> Doc -> Doc
<+> String -> String -> Doc -> Doc
nestedBlock String
"{" String
"}" ([Doc] -> Doc
stack ([Doc] -> Doc) -> [Doc] -> Doc
forall a b. (a -> b) -> a -> b
$ (ValueDesc -> Doc) -> [ValueDesc] -> [Doc]
forall a b. (a -> b) -> [a] -> [b]
map ValueDesc -> Doc
forall a. Pretty a => a -> Doc
ppr [ValueDesc]
vs)

instance Pretty ArrayContents where
  ppr :: ArrayContents -> Doc
ppr (ArrayValues [PrimValue]
vs) = Doc -> Doc
braces ([Doc] -> Doc
commasep ([Doc] -> Doc) -> [Doc] -> Doc
forall a b. (a -> b) -> a -> b
$ (PrimValue -> Doc) -> [PrimValue] -> [Doc]
forall a b. (a -> b) -> [a] -> [b]
map PrimValue -> Doc
forall a. Pretty a => a -> Doc
ppr [PrimValue]
vs)
  ppr (ArrayZeros Int
n) = Doc -> Doc
braces (String -> Doc
text String
"0") Doc -> Doc -> Doc
<+> String -> Doc
text String
"*" Doc -> Doc -> Doc
<+> Int -> Doc
forall a. Pretty a => a -> Doc
ppr Int
n

instance Pretty op => Pretty (Code op) where
  ppr :: Code op -> Doc
ppr (Op op
op) = op -> Doc
forall a. Pretty a => a -> Doc
ppr op
op
  ppr Code op
Skip = String -> Doc
text String
"skip"
  ppr (Code op
c1 :>>: Code op
c2) = Code op -> Doc
forall a. Pretty a => a -> Doc
ppr Code op
c1 Doc -> Doc -> Doc
</> Code op -> Doc
forall a. Pretty a => a -> Doc
ppr Code op
c2
  ppr (For VName
i Exp
limit Code op
body) =
    String -> Doc
text String
"for" Doc -> Doc -> Doc
<+> VName -> Doc
forall a. Pretty a => a -> Doc
ppr VName
i Doc -> Doc -> Doc
<+> Doc
langle Doc -> Doc -> Doc
<+> Exp -> Doc
forall a. Pretty a => a -> Doc
ppr Exp
limit Doc -> Doc -> Doc
<+> String -> Doc
text String
"{"
      Doc -> Doc -> Doc
</> Int -> Doc -> Doc
indent Int
2 (Code op -> Doc
forall a. Pretty a => a -> Doc
ppr Code op
body)
      Doc -> Doc -> Doc
</> String -> Doc
text String
"}"
  ppr (While TExp Bool
cond Code op
body) =
    String -> Doc
text String
"while" Doc -> Doc -> Doc
<+> TExp Bool -> Doc
forall a. Pretty a => a -> Doc
ppr TExp Bool
cond Doc -> Doc -> Doc
<+> String -> Doc
text String
"{"
      Doc -> Doc -> Doc
</> Int -> Doc -> Doc
indent Int
2 (Code op -> Doc
forall a. Pretty a => a -> Doc
ppr Code op
body)
      Doc -> Doc -> Doc
</> String -> Doc
text String
"}"
  ppr (DeclareMem VName
name Space
space) =
    String -> Doc
text String
"var" Doc -> Doc -> Doc
<+> VName -> Doc
forall a. Pretty a => a -> Doc
ppr VName
name Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> String -> Doc
text String
": mem" Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Space -> Doc
forall a. Pretty a => a -> Doc
ppr Space
space
  ppr (DeclareScalar VName
name Volatility
vol PrimType
t) =
    String -> Doc
text String
"var" Doc -> Doc -> Doc
<+> VName -> Doc
forall a. Pretty a => a -> Doc
ppr VName
name Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> String -> Doc
text String
":" Doc -> Doc -> Doc
<+> Doc
vol' Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> PrimType -> Doc
forall a. Pretty a => a -> Doc
ppr PrimType
t
    where
      vol' :: Doc
vol' = case Volatility
vol of
        Volatility
Volatile -> String -> Doc
text String
"volatile "
        Volatility
Nonvolatile -> Doc
forall a. Monoid a => a
mempty
  ppr (DeclareArray VName
name Space
space PrimType
t ArrayContents
vs) =
    String -> Doc
text String
"array" Doc -> Doc -> Doc
<+> VName -> Doc
forall a. Pretty a => a -> Doc
ppr VName
name Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> String -> Doc
text String
"@" Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Space -> Doc
forall a. Pretty a => a -> Doc
ppr Space
space Doc -> Doc -> Doc
<+> String -> Doc
text String
":" Doc -> Doc -> Doc
<+> PrimType -> Doc
forall a. Pretty a => a -> Doc
ppr PrimType
t
      Doc -> Doc -> Doc
<+> Doc
equals
      Doc -> Doc -> Doc
<+> ArrayContents -> Doc
forall a. Pretty a => a -> Doc
ppr ArrayContents
vs
  ppr (Allocate VName
name Count Bytes (TPrimExp Int64 ExpLeaf)
e Space
space) =
    VName -> Doc
forall a. Pretty a => a -> Doc
ppr VName
name Doc -> Doc -> Doc
<+> String -> Doc
text String
"<-" Doc -> Doc -> Doc
<+> String -> Doc
text String
"malloc" Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Doc -> Doc
parens (Count Bytes (TPrimExp Int64 ExpLeaf) -> Doc
forall a. Pretty a => a -> Doc
ppr Count Bytes (TPrimExp Int64 ExpLeaf)
e) Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Space -> Doc
forall a. Pretty a => a -> Doc
ppr Space
space
  ppr (Free VName
name Space
space) =
    String -> Doc
text String
"free" Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Doc -> Doc
parens (VName -> Doc
forall a. Pretty a => a -> Doc
ppr VName
name) Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Space -> Doc
forall a. Pretty a => a -> Doc
ppr Space
space
  ppr (Write VName
name Count Elements (TPrimExp Int64 ExpLeaf)
i PrimType
bt Space
space Volatility
vol Exp
val) =
    VName -> Doc
forall a. Pretty a => a -> Doc
ppr VName
name Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Doc
langle Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Doc
vol' Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> PrimType -> Doc
forall a. Pretty a => a -> Doc
ppr PrimType
bt Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Space -> Doc
forall a. Pretty a => a -> Doc
ppr Space
space Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Doc
rangle Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Doc -> Doc
brackets (Count Elements (TPrimExp Int64 ExpLeaf) -> Doc
forall a. Pretty a => a -> Doc
ppr Count Elements (TPrimExp Int64 ExpLeaf)
i)
      Doc -> Doc -> Doc
<+> String -> Doc
text String
"<-"
      Doc -> Doc -> Doc
<+> Exp -> Doc
forall a. Pretty a => a -> Doc
ppr Exp
val
    where
      vol' :: Doc
vol' = case Volatility
vol of
        Volatility
Volatile -> String -> Doc
text String
"volatile "
        Volatility
Nonvolatile -> Doc
forall a. Monoid a => a
mempty
  ppr (SetScalar VName
name Exp
val) =
    VName -> Doc
forall a. Pretty a => a -> Doc
ppr VName
name Doc -> Doc -> Doc
<+> String -> Doc
text String
"<-" Doc -> Doc -> Doc
<+> Exp -> Doc
forall a. Pretty a => a -> Doc
ppr Exp
val
  ppr (SetMem VName
dest VName
from Space
space) =
    VName -> Doc
forall a. Pretty a => a -> Doc
ppr VName
dest Doc -> Doc -> Doc
<+> String -> Doc
text String
"<-" Doc -> Doc -> Doc
<+> VName -> Doc
forall a. Pretty a => a -> Doc
ppr VName
from Doc -> Doc -> Doc
<+> String -> Doc
text String
"@" Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Space -> Doc
forall a. Pretty a => a -> Doc
ppr Space
space
  ppr (Assert Exp
e ErrorMsg Exp
msg (SrcLoc, [SrcLoc])
_) =
    String -> Doc
text String
"assert" Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Doc -> Doc
parens ([Doc] -> Doc
commasep [ErrorMsg Exp -> Doc
forall a. Pretty a => a -> Doc
ppr ErrorMsg Exp
msg, Exp -> Doc
forall a. Pretty a => a -> Doc
ppr Exp
e])
  ppr (Copy VName
dest Count Bytes (TPrimExp Int64 ExpLeaf)
destoffset Space
destspace VName
src Count Bytes (TPrimExp Int64 ExpLeaf)
srcoffset Space
srcspace Count Bytes (TPrimExp Int64 ExpLeaf)
size) =
    String -> Doc
text String
"memcpy"
      Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Doc -> Doc
parens
        ( VName -> Count Bytes (TPrimExp Int64 ExpLeaf) -> Doc
forall {a} {a}. (Pretty a, Pretty a) => a -> a -> Doc
ppMemLoc VName
dest Count Bytes (TPrimExp Int64 ExpLeaf)
destoffset Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Space -> Doc
forall a. Pretty a => a -> Doc
ppr Space
destspace Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Doc
comma
            Doc -> Doc -> Doc
</> VName -> Count Bytes (TPrimExp Int64 ExpLeaf) -> Doc
forall {a} {a}. (Pretty a, Pretty a) => a -> a -> Doc
ppMemLoc VName
src Count Bytes (TPrimExp Int64 ExpLeaf)
srcoffset Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Space -> Doc
forall a. Pretty a => a -> Doc
ppr Space
srcspace Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Doc
comma
            Doc -> Doc -> Doc
</> Count Bytes (TPrimExp Int64 ExpLeaf) -> Doc
forall a. Pretty a => a -> Doc
ppr Count Bytes (TPrimExp Int64 ExpLeaf)
size
        )
    where
      ppMemLoc :: a -> a -> Doc
ppMemLoc a
base a
offset =
        a -> Doc
forall a. Pretty a => a -> Doc
ppr a
base Doc -> Doc -> Doc
<+> String -> Doc
text String
"+" Doc -> Doc -> Doc
<+> a -> Doc
forall a. Pretty a => a -> Doc
ppr a
offset
  ppr (If TExp Bool
cond Code op
tbranch Code op
fbranch) =
    String -> Doc
text String
"if" Doc -> Doc -> Doc
<+> TExp Bool -> Doc
forall a. Pretty a => a -> Doc
ppr TExp Bool
cond Doc -> Doc -> Doc
<+> String -> Doc
text String
"then {"
      Doc -> Doc -> Doc
</> Int -> Doc -> Doc
indent Int
2 (Code op -> Doc
forall a. Pretty a => a -> Doc
ppr Code op
tbranch)
      Doc -> Doc -> Doc
</> String -> Doc
text String
"} else {"
      Doc -> Doc -> Doc
</> Int -> Doc -> Doc
indent Int
2 (Code op -> Doc
forall a. Pretty a => a -> Doc
ppr Code op
fbranch)
      Doc -> Doc -> Doc
</> String -> Doc
text String
"}"
  ppr (Call [VName]
dests Name
fname [Arg]
args) =
    [Doc] -> Doc
commasep ((VName -> Doc) -> [VName] -> [Doc]
forall a b. (a -> b) -> [a] -> [b]
map VName -> Doc
forall a. Pretty a => a -> Doc
ppr [VName]
dests) Doc -> Doc -> Doc
<+> String -> Doc
text String
"<-"
      Doc -> Doc -> Doc
<+> Name -> Doc
forall a. Pretty a => a -> Doc
ppr Name
fname Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Doc -> Doc
parens ([Doc] -> Doc
commasep ([Doc] -> Doc) -> [Doc] -> Doc
forall a b. (a -> b) -> a -> b
$ (Arg -> Doc) -> [Arg] -> [Doc]
forall a b. (a -> b) -> [a] -> [b]
map Arg -> Doc
forall a. Pretty a => a -> Doc
ppr [Arg]
args)
  ppr (Comment String
s Code op
code) =
    String -> Doc
text String
"--" Doc -> Doc -> Doc
<+> String -> Doc
text String
s Doc -> Doc -> Doc
</> Code op -> Doc
forall a. Pretty a => a -> Doc
ppr Code op
code
  ppr (DebugPrint String
desc (Just Exp
e)) =
    String -> Doc
text String
"debug" Doc -> Doc -> Doc
<+> Doc -> Doc
parens ([Doc] -> Doc
commasep [String -> Doc
text (ShowS
forall a. Show a => a -> String
show String
desc), Exp -> Doc
forall a. Pretty a => a -> Doc
ppr Exp
e])
  ppr (DebugPrint String
desc Maybe Exp
Nothing) =
    String -> Doc
text String
"debug" Doc -> Doc -> Doc
<+> Doc -> Doc
parens (String -> Doc
text (ShowS
forall a. Show a => a -> String
show String
desc))

instance Pretty Arg where
  ppr :: Arg -> Doc
ppr (MemArg VName
m) = VName -> Doc
forall a. Pretty a => a -> Doc
ppr VName
m
  ppr (ExpArg Exp
e) = Exp -> Doc
forall a. Pretty a => a -> Doc
ppr Exp
e

instance Pretty ExpLeaf where
  ppr :: ExpLeaf -> Doc
ppr (ScalarVar VName
v) =
    VName -> Doc
forall a. Pretty a => a -> Doc
ppr VName
v
  ppr (Index VName
v Count Elements (TPrimExp Int64 ExpLeaf)
is PrimType
bt Space
space Volatility
vol) =
    VName -> Doc
forall a. Pretty a => a -> Doc
ppr VName
v Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Doc
langle Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Doc
vol' Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> PrimType -> Doc
forall a. Pretty a => a -> Doc
ppr PrimType
bt Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Space -> Doc
forall a. Pretty a => a -> Doc
ppr Space
space Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Doc
rangle Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Doc -> Doc
brackets (Count Elements (TPrimExp Int64 ExpLeaf) -> Doc
forall a. Pretty a => a -> Doc
ppr Count Elements (TPrimExp Int64 ExpLeaf)
is)
    where
      vol' :: Doc
vol' = case Volatility
vol of
        Volatility
Volatile -> String -> Doc
text String
"volatile "
        Volatility
Nonvolatile -> Doc
forall a. Monoid a => a
mempty

instance Functor Functions where
  fmap :: forall a b. (a -> b) -> Functions a -> Functions b
fmap = (a -> b) -> Functions a -> Functions b
forall (t :: * -> *) a b. Traversable t => (a -> b) -> t a -> t b
fmapDefault

instance Foldable Functions where
  foldMap :: forall m a. Monoid m => (a -> m) -> Functions a -> m
foldMap = (a -> m) -> Functions a -> m
forall (t :: * -> *) m a.
(Traversable t, Monoid m) =>
(a -> m) -> t a -> m
foldMapDefault

instance Traversable Functions where
  traverse :: forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Functions a -> f (Functions b)
traverse a -> f b
f (Functions [(Name, Function a)]
funs) =
    [(Name, Function b)] -> Functions b
forall a. [(Name, Function a)] -> Functions a
Functions ([(Name, Function b)] -> Functions b)
-> f [(Name, Function b)] -> f (Functions b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((Name, Function a) -> f (Name, Function b))
-> [(Name, Function a)] -> f [(Name, Function b)]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (Name, Function a) -> f (Name, Function b)
forall {t :: * -> *} {t}. Traversable t => (t, t a) -> f (t, t b)
f' [(Name, Function a)]
funs
    where
      f' :: (t, t a) -> f (t, t b)
f' (t
name, t a
fun) = (t
name,) (t b -> (t, t b)) -> f (t b) -> f (t, t b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (a -> f b) -> t a -> f (t b)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse a -> f b
f t a
fun

instance Functor FunctionT where
  fmap :: forall a b. (a -> b) -> FunctionT a -> FunctionT b
fmap = (a -> b) -> FunctionT a -> FunctionT b
forall (t :: * -> *) a b. Traversable t => (a -> b) -> t a -> t b
fmapDefault

instance Foldable FunctionT where
  foldMap :: forall m a. Monoid m => (a -> m) -> FunctionT a -> m
foldMap = (a -> m) -> FunctionT a -> m
forall (t :: * -> *) m a.
(Traversable t, Monoid m) =>
(a -> m) -> t a -> m
foldMapDefault

instance Traversable FunctionT where
  traverse :: forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> FunctionT a -> f (FunctionT b)
traverse a -> f b
f (Function Bool
entry [Param]
outs [Param]
ins Code a
body [ExternalValue]
results [ExternalValue]
args) =
    Bool
-> [Param]
-> [Param]
-> Code b
-> [ExternalValue]
-> [ExternalValue]
-> FunctionT b
forall a.
Bool
-> [Param]
-> [Param]
-> Code a
-> [ExternalValue]
-> [ExternalValue]
-> FunctionT a
Function Bool
entry [Param]
outs [Param]
ins (Code b -> [ExternalValue] -> [ExternalValue] -> FunctionT b)
-> f (Code b)
-> f ([ExternalValue] -> [ExternalValue] -> FunctionT b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (a -> f b) -> Code a -> f (Code b)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse a -> f b
f Code a
body f ([ExternalValue] -> [ExternalValue] -> FunctionT b)
-> f [ExternalValue] -> f ([ExternalValue] -> FunctionT b)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> [ExternalValue] -> f [ExternalValue]
forall (f :: * -> *) a. Applicative f => a -> f a
pure [ExternalValue]
results f ([ExternalValue] -> FunctionT b)
-> f [ExternalValue] -> f (FunctionT b)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> [ExternalValue] -> f [ExternalValue]
forall (f :: * -> *) a. Applicative f => a -> f a
pure [ExternalValue]
args

instance Functor Code where
  fmap :: forall a b. (a -> b) -> Code a -> Code b
fmap = (a -> b) -> Code a -> Code b
forall (t :: * -> *) a b. Traversable t => (a -> b) -> t a -> t b
fmapDefault

instance Foldable Code where
  foldMap :: forall m a. Monoid m => (a -> m) -> Code a -> m
foldMap = (a -> m) -> Code a -> m
forall (t :: * -> *) m a.
(Traversable t, Monoid m) =>
(a -> m) -> t a -> m
foldMapDefault

instance Traversable Code where
  traverse :: forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Code a -> f (Code b)
traverse a -> f b
f (Code a
x :>>: Code a
y) =
    Code b -> Code b -> Code b
forall a. Code a -> Code a -> Code a
(:>>:) (Code b -> Code b -> Code b) -> f (Code b) -> f (Code b -> Code b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (a -> f b) -> Code a -> f (Code b)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse a -> f b
f Code a
x f (Code b -> Code b) -> f (Code b) -> f (Code b)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (a -> f b) -> Code a -> f (Code b)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse a -> f b
f Code a
y
  traverse a -> f b
f (For VName
i Exp
bound Code a
code) =
    VName -> Exp -> Code b -> Code b
forall a. VName -> Exp -> Code a -> Code a
For VName
i Exp
bound (Code b -> Code b) -> f (Code b) -> f (Code b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (a -> f b) -> Code a -> f (Code b)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse a -> f b
f Code a
code
  traverse a -> f b
f (While TExp Bool
cond Code a
code) =
    TExp Bool -> Code b -> Code b
forall a. TExp Bool -> Code a -> Code a
While TExp Bool
cond (Code b -> Code b) -> f (Code b) -> f (Code b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (a -> f b) -> Code a -> f (Code b)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse a -> f b
f Code a
code
  traverse a -> f b
f (If TExp Bool
cond Code a
x Code a
y) =
    TExp Bool -> Code b -> Code b -> Code b
forall a. TExp Bool -> Code a -> Code a -> Code a
If TExp Bool
cond (Code b -> Code b -> Code b) -> f (Code b) -> f (Code b -> Code b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (a -> f b) -> Code a -> f (Code b)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse a -> f b
f Code a
x f (Code b -> Code b) -> f (Code b) -> f (Code b)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (a -> f b) -> Code a -> f (Code b)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse a -> f b
f Code a
y
  traverse a -> f b
f (Op a
kernel) =
    b -> Code b
forall a. a -> Code a
Op (b -> Code b) -> f b -> f (Code b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> f b
f a
kernel
  traverse a -> f b
_ Code a
Skip =
    Code b -> f (Code b)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Code b
forall a. Code a
Skip
  traverse a -> f b
_ (DeclareMem VName
name Space
space) =
    Code b -> f (Code b)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Code b -> f (Code b)) -> Code b -> f (Code b)
forall a b. (a -> b) -> a -> b
$ VName -> Space -> Code b
forall a. VName -> Space -> Code a
DeclareMem VName
name Space
space
  traverse a -> f b
_ (DeclareScalar VName
name Volatility
vol PrimType
bt) =
    Code b -> f (Code b)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Code b -> f (Code b)) -> Code b -> f (Code b)
forall a b. (a -> b) -> a -> b
$ VName -> Volatility -> PrimType -> Code b
forall a. VName -> Volatility -> PrimType -> Code a
DeclareScalar VName
name Volatility
vol PrimType
bt
  traverse a -> f b
_ (DeclareArray VName
name Space
space PrimType
t ArrayContents
vs) =
    Code b -> f (Code b)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Code b -> f (Code b)) -> Code b -> f (Code b)
forall a b. (a -> b) -> a -> b
$ VName -> Space -> PrimType -> ArrayContents -> Code b
forall a. VName -> Space -> PrimType -> ArrayContents -> Code a
DeclareArray VName
name Space
space PrimType
t ArrayContents
vs
  traverse a -> f b
_ (Allocate VName
name Count Bytes (TPrimExp Int64 ExpLeaf)
size Space
s) =
    Code b -> f (Code b)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Code b -> f (Code b)) -> Code b -> f (Code b)
forall a b. (a -> b) -> a -> b
$ VName -> Count Bytes (TPrimExp Int64 ExpLeaf) -> Space -> Code b
forall a.
VName -> Count Bytes (TPrimExp Int64 ExpLeaf) -> Space -> Code a
Allocate VName
name Count Bytes (TPrimExp Int64 ExpLeaf)
size Space
s
  traverse a -> f b
_ (Free VName
name Space
space) =
    Code b -> f (Code b)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Code b -> f (Code b)) -> Code b -> f (Code b)
forall a b. (a -> b) -> a -> b
$ VName -> Space -> Code b
forall a. VName -> Space -> Code a
Free VName
name Space
space
  traverse a -> f b
_ (Copy VName
dest Count Bytes (TPrimExp Int64 ExpLeaf)
destoffset Space
destspace VName
src Count Bytes (TPrimExp Int64 ExpLeaf)
srcoffset Space
srcspace Count Bytes (TPrimExp Int64 ExpLeaf)
size) =
    Code b -> f (Code b)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Code b -> f (Code b)) -> Code b -> f (Code b)
forall a b. (a -> b) -> a -> b
$ VName
-> Count Bytes (TPrimExp Int64 ExpLeaf)
-> Space
-> VName
-> Count Bytes (TPrimExp Int64 ExpLeaf)
-> Space
-> Count Bytes (TPrimExp Int64 ExpLeaf)
-> Code b
forall a.
VName
-> Count Bytes (TPrimExp Int64 ExpLeaf)
-> Space
-> VName
-> Count Bytes (TPrimExp Int64 ExpLeaf)
-> Space
-> Count Bytes (TPrimExp Int64 ExpLeaf)
-> Code a
Copy VName
dest Count Bytes (TPrimExp Int64 ExpLeaf)
destoffset Space
destspace VName
src Count Bytes (TPrimExp Int64 ExpLeaf)
srcoffset Space
srcspace Count Bytes (TPrimExp Int64 ExpLeaf)
size
  traverse a -> f b
_ (Write VName
name Count Elements (TPrimExp Int64 ExpLeaf)
i PrimType
bt Space
val Volatility
space Exp
vol) =
    Code b -> f (Code b)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Code b -> f (Code b)) -> Code b -> f (Code b)
forall a b. (a -> b) -> a -> b
$ VName
-> Count Elements (TPrimExp Int64 ExpLeaf)
-> PrimType
-> Space
-> Volatility
-> Exp
-> Code b
forall a.
VName
-> Count Elements (TPrimExp Int64 ExpLeaf)
-> PrimType
-> Space
-> Volatility
-> Exp
-> Code a
Write VName
name Count Elements (TPrimExp Int64 ExpLeaf)
i PrimType
bt Space
val Volatility
space Exp
vol
  traverse a -> f b
_ (SetScalar VName
name Exp
val) =
    Code b -> f (Code b)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Code b -> f (Code b)) -> Code b -> f (Code b)
forall a b. (a -> b) -> a -> b
$ VName -> Exp -> Code b
forall a. VName -> Exp -> Code a
SetScalar VName
name Exp
val
  traverse a -> f b
_ (SetMem VName
dest VName
from Space
space) =
    Code b -> f (Code b)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Code b -> f (Code b)) -> Code b -> f (Code b)
forall a b. (a -> b) -> a -> b
$ VName -> VName -> Space -> Code b
forall a. VName -> VName -> Space -> Code a
SetMem VName
dest VName
from Space
space
  traverse a -> f b
_ (Assert Exp
e ErrorMsg Exp
msg (SrcLoc, [SrcLoc])
loc) =
    Code b -> f (Code b)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Code b -> f (Code b)) -> Code b -> f (Code b)
forall a b. (a -> b) -> a -> b
$ Exp -> ErrorMsg Exp -> (SrcLoc, [SrcLoc]) -> Code b
forall a. Exp -> ErrorMsg Exp -> (SrcLoc, [SrcLoc]) -> Code a
Assert Exp
e ErrorMsg Exp
msg (SrcLoc, [SrcLoc])
loc
  traverse a -> f b
_ (Call [VName]
dests Name
fname [Arg]
args) =
    Code b -> f (Code b)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Code b -> f (Code b)) -> Code b -> f (Code b)
forall a b. (a -> b) -> a -> b
$ [VName] -> Name -> [Arg] -> Code b
forall a. [VName] -> Name -> [Arg] -> Code a
Call [VName]
dests Name
fname [Arg]
args
  traverse a -> f b
f (Comment String
s Code a
code) =
    String -> Code b -> Code b
forall a. String -> Code a -> Code a
Comment String
s (Code b -> Code b) -> f (Code b) -> f (Code b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (a -> f b) -> Code a -> f (Code b)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse a -> f b
f Code a
code
  traverse a -> f b
_ (DebugPrint String
s Maybe Exp
v) =
    Code b -> f (Code b)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Code b -> f (Code b)) -> Code b -> f (Code b)
forall a b. (a -> b) -> a -> b
$ String -> Maybe Exp -> Code b
forall a. String -> Maybe Exp -> Code a
DebugPrint String
s Maybe Exp
v

-- | The names declared with 'DeclareMem', 'DeclareScalar', and
-- 'DeclareArray' in the given code.
declaredIn :: Code a -> Names
declaredIn :: forall {a}. Code a -> Names
declaredIn (DeclareMem VName
name Space
_) = VName -> Names
oneName VName
name
declaredIn (DeclareScalar VName
name Volatility
_ PrimType
_) = VName -> Names
oneName VName
name
declaredIn (DeclareArray VName
name Space
_ PrimType
_ ArrayContents
_) = VName -> Names
oneName VName
name
declaredIn (If TExp Bool
_ Code a
t Code a
f) = Code a -> Names
forall {a}. Code a -> Names
declaredIn Code a
t Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> Code a -> Names
forall {a}. Code a -> Names
declaredIn Code a
f
declaredIn (Code a
x :>>: Code a
y) = Code a -> Names
forall {a}. Code a -> Names
declaredIn Code a
x Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> Code a -> Names
forall {a}. Code a -> Names
declaredIn Code a
y
declaredIn (For VName
i Exp
_ Code a
body) = VName -> Names
oneName VName
i Names -> Names -> Names
forall a. Semigroup a => a -> a -> a
<> Code a -> Names
forall {a}. Code a -> Names
declaredIn Code a
body
declaredIn (While TExp Bool
_ Code a
body) = Code a -> Names
forall {a}. Code a -> Names
declaredIn Code a
body
declaredIn (Comment String
_ Code a
body) = Code a -> Names
forall {a}. Code a -> Names
declaredIn Code a
body
declaredIn Code a
_ = Names
forall a. Monoid a => a
mempty

instance FreeIn a => FreeIn (Functions a) where
  freeIn' :: Functions a -> FV
freeIn' (Functions [(Name, Function a)]
fs) = ((Name, Function a) -> FV) -> [(Name, Function a)] -> FV
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (Function a -> FV
forall {a}. FreeIn a => FunctionT a -> FV
onFun (Function a -> FV)
-> ((Name, Function a) -> Function a) -> (Name, Function a) -> FV
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Name, Function a) -> Function a
forall a b. (a, b) -> b
snd) [(Name, Function a)]
fs
    where
      onFun :: FunctionT a -> FV
onFun FunctionT a
f =
        Names -> FV -> FV
fvBind Names
pnames (FV -> FV) -> FV -> FV
forall a b. (a -> b) -> a -> b
$
          Code a -> FV
forall a. FreeIn a => a -> FV
freeIn' (FunctionT a -> Code a
forall a. FunctionT a -> Code a
functionBody FunctionT a
f) FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> [ExternalValue] -> FV
forall a. FreeIn a => a -> FV
freeIn' (FunctionT a -> [ExternalValue]
forall a. FunctionT a -> [ExternalValue]
functionResult FunctionT a
f [ExternalValue] -> [ExternalValue] -> [ExternalValue]
forall a. Semigroup a => a -> a -> a
<> FunctionT a -> [ExternalValue]
forall a. FunctionT a -> [ExternalValue]
functionArgs FunctionT a
f)
        where
          pnames :: Names
pnames =
            [VName] -> Names
namesFromList ([VName] -> Names) -> [VName] -> Names
forall a b. (a -> b) -> a -> b
$ (Param -> VName) -> [Param] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param -> VName
paramName ([Param] -> [VName]) -> [Param] -> [VName]
forall a b. (a -> b) -> a -> b
$ FunctionT a -> [Param]
forall a. FunctionT a -> [Param]
functionInput FunctionT a
f [Param] -> [Param] -> [Param]
forall a. Semigroup a => a -> a -> a
<> FunctionT a -> [Param]
forall a. FunctionT a -> [Param]
functionOutput FunctionT a
f

instance FreeIn ValueDesc where
  freeIn' :: ValueDesc -> FV
freeIn' (ArrayValue VName
mem Space
_ PrimType
_ Signedness
_ [DimSize]
dims) = VName -> FV
forall a. FreeIn a => a -> FV
freeIn' VName
mem FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> [DimSize] -> FV
forall a. FreeIn a => a -> FV
freeIn' [DimSize]
dims
  freeIn' ScalarValue {} = FV
forall a. Monoid a => a
mempty

instance FreeIn ExternalValue where
  freeIn' :: ExternalValue -> FV
freeIn' (TransparentValue ValueDesc
vd) = ValueDesc -> FV
forall a. FreeIn a => a -> FV
freeIn' ValueDesc
vd
  freeIn' (OpaqueValue String
_ [ValueDesc]
vds) = (ValueDesc -> FV) -> [ValueDesc] -> FV
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap ValueDesc -> FV
forall a. FreeIn a => a -> FV
freeIn' [ValueDesc]
vds

instance FreeIn a => FreeIn (Code a) where
  freeIn' :: Code a -> FV
freeIn' (Code a
x :>>: Code a
y) =
    Names -> FV -> FV
fvBind (Code a -> Names
forall {a}. Code a -> Names
declaredIn Code a
x) (FV -> FV) -> FV -> FV
forall a b. (a -> b) -> a -> b
$ Code a -> FV
forall a. FreeIn a => a -> FV
freeIn' Code a
x FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> Code a -> FV
forall a. FreeIn a => a -> FV
freeIn' Code a
y
  freeIn' Code a
Skip =
    FV
forall a. Monoid a => a
mempty
  freeIn' (For VName
i Exp
bound Code a
body) =
    Names -> FV -> FV
fvBind (VName -> Names
oneName VName
i) (FV -> FV) -> FV -> FV
forall a b. (a -> b) -> a -> b
$ Exp -> FV
forall a. FreeIn a => a -> FV
freeIn' Exp
bound FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> Code a -> FV
forall a. FreeIn a => a -> FV
freeIn' Code a
body
  freeIn' (While TExp Bool
cond Code a
body) =
    TExp Bool -> FV
forall a. FreeIn a => a -> FV
freeIn' TExp Bool
cond FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> Code a -> FV
forall a. FreeIn a => a -> FV
freeIn' Code a
body
  freeIn' (DeclareMem VName
_ Space
space) =
    Space -> FV
forall a. FreeIn a => a -> FV
freeIn' Space
space
  freeIn' DeclareScalar {} =
    FV
forall a. Monoid a => a
mempty
  freeIn' DeclareArray {} =
    FV
forall a. Monoid a => a
mempty
  freeIn' (Allocate VName
name Count Bytes (TPrimExp Int64 ExpLeaf)
size Space
space) =
    VName -> FV
forall a. FreeIn a => a -> FV
freeIn' VName
name FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> Count Bytes (TPrimExp Int64 ExpLeaf) -> FV
forall a. FreeIn a => a -> FV
freeIn' Count Bytes (TPrimExp Int64 ExpLeaf)
size FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> Space -> FV
forall a. FreeIn a => a -> FV
freeIn' Space
space
  freeIn' (Free VName
name Space
_) =
    VName -> FV
forall a. FreeIn a => a -> FV
freeIn' VName
name
  freeIn' (Copy VName
dest Count Bytes (TPrimExp Int64 ExpLeaf)
x Space
_ VName
src Count Bytes (TPrimExp Int64 ExpLeaf)
y Space
_ Count Bytes (TPrimExp Int64 ExpLeaf)
n) =
    VName -> FV
forall a. FreeIn a => a -> FV
freeIn' VName
dest FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> Count Bytes (TPrimExp Int64 ExpLeaf) -> FV
forall a. FreeIn a => a -> FV
freeIn' Count Bytes (TPrimExp Int64 ExpLeaf)
x FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> VName -> FV
forall a. FreeIn a => a -> FV
freeIn' VName
src FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> Count Bytes (TPrimExp Int64 ExpLeaf) -> FV
forall a. FreeIn a => a -> FV
freeIn' Count Bytes (TPrimExp Int64 ExpLeaf)
y FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> Count Bytes (TPrimExp Int64 ExpLeaf) -> FV
forall a. FreeIn a => a -> FV
freeIn' Count Bytes (TPrimExp Int64 ExpLeaf)
n
  freeIn' (SetMem VName
x VName
y Space
_) =
    VName -> FV
forall a. FreeIn a => a -> FV
freeIn' VName
x FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> VName -> FV
forall a. FreeIn a => a -> FV
freeIn' VName
y
  freeIn' (Write VName
v Count Elements (TPrimExp Int64 ExpLeaf)
i PrimType
_ Space
_ Volatility
_ Exp
e) =
    VName -> FV
forall a. FreeIn a => a -> FV
freeIn' VName
v FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> Count Elements (TPrimExp Int64 ExpLeaf) -> FV
forall a. FreeIn a => a -> FV
freeIn' Count Elements (TPrimExp Int64 ExpLeaf)
i FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> Exp -> FV
forall a. FreeIn a => a -> FV
freeIn' Exp
e
  freeIn' (SetScalar VName
x Exp
y) =
    VName -> FV
forall a. FreeIn a => a -> FV
freeIn' VName
x FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> Exp -> FV
forall a. FreeIn a => a -> FV
freeIn' Exp
y
  freeIn' (Call [VName]
dests Name
_ [Arg]
args) =
    [VName] -> FV
forall a. FreeIn a => a -> FV
freeIn' [VName]
dests FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> [Arg] -> FV
forall a. FreeIn a => a -> FV
freeIn' [Arg]
args
  freeIn' (If TExp Bool
cond Code a
t Code a
f) =
    TExp Bool -> FV
forall a. FreeIn a => a -> FV
freeIn' TExp Bool
cond FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> Code a -> FV
forall a. FreeIn a => a -> FV
freeIn' Code a
t FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> Code a -> FV
forall a. FreeIn a => a -> FV
freeIn' Code a
f
  freeIn' (Assert Exp
e ErrorMsg Exp
msg (SrcLoc, [SrcLoc])
_) =
    Exp -> FV
forall a. FreeIn a => a -> FV
freeIn' Exp
e FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> (Exp -> FV) -> ErrorMsg Exp -> FV
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Exp -> FV
forall a. FreeIn a => a -> FV
freeIn' ErrorMsg Exp
msg
  freeIn' (Op a
op) =
    a -> FV
forall a. FreeIn a => a -> FV
freeIn' a
op
  freeIn' (Comment String
_ Code a
code) =
    Code a -> FV
forall a. FreeIn a => a -> FV
freeIn' Code a
code
  freeIn' (DebugPrint String
_ Maybe Exp
v) =
    FV -> (Exp -> FV) -> Maybe Exp -> FV
forall b a. b -> (a -> b) -> Maybe a -> b
maybe FV
forall a. Monoid a => a
mempty Exp -> FV
forall a. FreeIn a => a -> FV
freeIn' Maybe Exp
v

instance FreeIn ExpLeaf where
  freeIn' :: ExpLeaf -> FV
freeIn' (Index VName
v Count Elements (TPrimExp Int64 ExpLeaf)
e PrimType
_ Space
_ Volatility
_) = VName -> FV
forall a. FreeIn a => a -> FV
freeIn' VName
v FV -> FV -> FV
forall a. Semigroup a => a -> a -> a
<> Count Elements (TPrimExp Int64 ExpLeaf) -> FV
forall a. FreeIn a => a -> FV
freeIn' Count Elements (TPrimExp Int64 ExpLeaf)
e
  freeIn' (ScalarVar VName
v) = VName -> FV
forall a. FreeIn a => a -> FV
freeIn' VName
v

instance FreeIn Arg where
  freeIn' :: Arg -> FV
freeIn' (MemArg VName
m) = VName -> FV
forall a. FreeIn a => a -> FV
freeIn' VName
m
  freeIn' (ExpArg Exp
e) = Exp -> FV
forall a. FreeIn a => a -> FV
freeIn' Exp
e