{-# LANGUAGE TypeFamilies #-}

-- | = Constructing Futhark ASTs
--
-- This module re-exports and defines a bunch of building blocks for
-- constructing fragments of Futhark ASTs.  More importantly, it also
-- contains a basic introduction on how to use them.
--
-- The "Futhark.IR.Syntax" module contains the core
-- AST definition.  One important invariant is that all bound names in
-- a Futhark program must be /globally/ unique.  In principle, you
-- could use the facilities from "Futhark.MonadFreshNames" (or your
-- own bespoke source of unique names) to manually construct
-- expressions, statements, and entire ASTs.  In practice, this would
-- be very tedious.  Instead, we have defined a collection of building
-- blocks (centered around the 'MonadBuilder' type class) that permits
-- a more abstract way of generating code.
--
-- Constructing ASTs with these building blocks requires you to ensure
-- that all free variables are in scope.  See
-- "Futhark.IR.Prop.Scope".
--
-- == 'MonadBuilder'
--
-- A monad that implements 'MonadBuilder' tracks the statements added
-- so far, the current names in scope, and allows you to add
-- additional statements with 'addStm'.  Any monad that implements
-- 'MonadBuilder' also implements the t'Rep' type family, which
-- indicates which rep it works with.  Inside a 'MonadBuilder' we can
-- use 'collectStms' to gather up the 'Stms' added with 'addStm' in
-- some nested computation.
--
-- The 'BuilderT' monad (and its convenient 'Builder' version) provides
-- the simplest implementation of 'MonadBuilder'.
--
-- == Higher-level building blocks
--
-- On top of the raw facilities provided by 'MonadBuilder', we have
-- more convenient facilities.  For example, 'letSubExp' lets us
-- conveniently create a 'Stm' for an 'Exp' that produces a /single/
-- value, and returns the (fresh) name for the resulting variable:
--
-- @
-- z <- letExp "z" $ BasicOp $ BinOp (Add Int32) (Var x) (Var y)
-- @
--
-- == Monadic expression builders
--
-- This module also contains "monadic expression" functions that let
-- us build nested expressions in a "direct" style, rather than using
-- 'letExp' and friends to bind every sub-part first.  See functions
-- such as 'eIf' and 'eBody' for example.  See also
-- "Futhark.Analysis.PrimExp" and the 'ToExp' type class.
--
-- == Examples
--
-- The "Futhark.Transform.FirstOrderTransform" module is a
-- (relatively) simple example of how to use these components.  As are
-- some of the high-level building blocks in this very module.
module Futhark.Construct
  ( -- * Basic building blocks
    module Futhark.Builder,
    letSubExp,
    letExp,
    letTupExp,
    letTupExp',
    letInPlace,

    -- * Monadic expression builders
    eSubExp,
    eVar,
    eParam,
    eMatch',
    eMatch,
    eIf,
    eIf',
    eBinOp,
    eUnOp,
    eCmpOp,
    eConvOp,
    eSignum,
    eCopy,
    eBody,
    eLambda,
    eBlank,
    eAll,
    eAny,
    eDimInBounds,
    eShapeInBounds,
    eOutOfBounds,
    eIndex,
    eLast,

    -- * Other building blocks
    asIntZ,
    asIntS,
    resultBody,
    resultBodyM,
    insertStmsM,
    buildBody,
    buildBody_,
    mapResult,
    foldBinOp,
    binOpLambda,
    cmpOpLambda,
    mkLambda,
    sliceDim,
    fullSlice,
    fullSliceNum,
    isIdentitySlice,
    sliceAt,
    iota64,

    -- * Result types
    instantiateShapes,
    instantiateShapes',
    removeExistentials,

    -- * Convenience
    simpleMkLetNames,
    ToExp (..),
    toSubExp,
  )
where

import Control.Monad
import Control.Monad.Identity
import Control.Monad.State
import Data.List qualified as L
import Data.Map.Strict qualified as M
import Futhark.Builder
import Futhark.IR
import Futhark.Util (maybeNth)

-- | @letSubExp desc e@ binds the expression @e@, which must produce a
-- single value.  Returns a t'SubExp' corresponding to the resulting
-- value.  For expressions that produce multiple values, see
-- 'letTupExp'.
letSubExp ::
  (MonadBuilder m) =>
  Name ->
  Exp (Rep m) ->
  m SubExp
letSubExp :: forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
_ (BasicOp (SubExp SubExp
se)) = SubExp -> m SubExp
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SubExp
se
letSubExp Name
desc Exp (Rep m)
e = VName -> SubExp
Var (VName -> SubExp) -> m VName -> m SubExp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Name -> Exp (Rep m) -> m VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m VName
letExp Name
desc Exp (Rep m)
e

-- | Like 'letSubExp', but returns a name rather than a t'SubExp'.
letExp ::
  (MonadBuilder m) =>
  Name ->
  Exp (Rep m) ->
  m VName
letExp :: forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m VName
letExp Name
_ (BasicOp (SubExp (Var VName
v))) =
  VName -> m VName
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure VName
v
letExp Name
desc Exp (Rep m)
e = do
  n <- [ExtType] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([ExtType] -> Int) -> m [ExtType] -> m Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp (Rep m) -> m [ExtType]
forall rep (m :: * -> *).
(HasScope rep m, TypedOp (OpC rep)) =>
Exp rep -> m [ExtType]
expExtType Exp (Rep m)
e
  vs <- replicateM n $ newVName desc
  letBindNames vs e
  case vs of
    [VName
v] -> VName -> m VName
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure VName
v
    [VName]
_ -> [Char] -> m VName
forall a. HasCallStack => [Char] -> a
error ([Char] -> m VName) -> [Char] -> m VName
forall a b. (a -> b) -> a -> b
$ [Char]
"letExp: tuple-typed expression given:\n" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Exp (Rep m) -> [Char]
forall a. Pretty a => a -> [Char]
prettyString Exp (Rep m)
e

-- | Like 'letExp', but the 'VName' and 'Slice' denote an array that
-- is 'Update'd with the result of the expression.  The name of the
-- updated array is returned.
letInPlace ::
  (MonadBuilder m) =>
  Name ->
  VName ->
  Slice SubExp ->
  Exp (Rep m) ->
  m VName
letInPlace :: forall (m :: * -> *).
MonadBuilder m =>
Name -> VName -> Slice SubExp -> Exp (Rep m) -> m VName
letInPlace Name
desc VName
src Slice SubExp
slice Exp (Rep m)
e = do
  tmp <- Name -> Exp (Rep m) -> m SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp (Name
desc Name -> Name -> Name
forall a. Semigroup a => a -> a -> a
<> Name
"_tmp") Exp (Rep m)
e
  letExp desc $ BasicOp $ Update Unsafe src slice tmp

-- | Like 'letExp', but the expression may return multiple values.
letTupExp ::
  (MonadBuilder m) =>
  Name ->
  Exp (Rep m) ->
  m [VName]
letTupExp :: forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m [VName]
letTupExp Name
_ (BasicOp (SubExp (Var VName
v))) =
  [VName] -> m [VName]
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [VName
v]
letTupExp Name
name Exp (Rep m)
e = do
  e_t <- Exp (Rep m) -> m [ExtType]
forall rep (m :: * -> *).
(HasScope rep m, TypedOp (OpC rep)) =>
Exp rep -> m [ExtType]
expExtType Exp (Rep m)
e
  names <- replicateM (length e_t) $ newVName name
  letBindNames names e
  pure names

-- | Like 'letTupExp', but returns t'SubExp's instead of 'VName's.
letTupExp' ::
  (MonadBuilder m) =>
  Name ->
  Exp (Rep m) ->
  m [SubExp]
letTupExp' :: forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m [SubExp]
letTupExp' Name
_ (BasicOp (SubExp SubExp
se)) = [SubExp] -> m [SubExp]
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [SubExp
se]
letTupExp' Name
name Exp (Rep m)
ses = (VName -> SubExp) -> [VName] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
Var ([VName] -> [SubExp]) -> m [VName] -> m [SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Name -> Exp (Rep m) -> m [VName]
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m [VName]
letTupExp Name
name Exp (Rep m)
ses

-- | Turn a subexpression into a monad expression.  Does not actually
-- lead to any code generation.  This is supposed to be used alongside
-- the other monadic expression functions, such as 'eIf'.
eSubExp ::
  (MonadBuilder m) =>
  SubExp ->
  m (Exp (Rep m))
eSubExp :: forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp = Exp (Rep m) -> m (Exp (Rep m))
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp (Rep m) -> m (Exp (Rep m)))
-> (SubExp -> Exp (Rep m)) -> SubExp -> m (Exp (Rep m))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BasicOp -> Exp (Rep m)
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp (Rep m))
-> (SubExp -> BasicOp) -> SubExp -> Exp (Rep m)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SubExp -> BasicOp
SubExp

-- | Turn a variable into a monad expression, through 'eSubExp'.
eVar ::
  (MonadBuilder m) =>
  VName ->
  m (Exp (Rep m))
eVar :: forall (m :: * -> *). MonadBuilder m => VName -> m (Exp (Rep m))
eVar = SubExp -> m (Exp (Rep m))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp (SubExp -> m (Exp (Rep m)))
-> (VName -> SubExp) -> VName -> m (Exp (Rep m))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> SubExp
Var

-- | Treat a parameter as a monadic expression.
eParam ::
  (MonadBuilder m) =>
  Param t ->
  m (Exp (Rep m))
eParam :: forall (m :: * -> *) t.
MonadBuilder m =>
Param t -> m (Exp (Rep m))
eParam = SubExp -> m (Exp (Rep m))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp (SubExp -> m (Exp (Rep m)))
-> (Param t -> SubExp) -> Param t -> m (Exp (Rep m))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> SubExp
Var (VName -> SubExp) -> (Param t -> VName) -> Param t -> SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Param t -> VName
forall dec. Param dec -> VName
paramName

removeRedundantScrutinees :: [SubExp] -> [Case b] -> ([SubExp], [Case b])
removeRedundantScrutinees :: forall b. [SubExp] -> [Case b] -> ([SubExp], [Case b])
removeRedundantScrutinees [SubExp]
ses [Case b]
cases =
  let ([SubExp]
ses', [[Maybe PrimValue]]
vs) =
        [(SubExp, [Maybe PrimValue])] -> ([SubExp], [[Maybe PrimValue]])
forall a b. [(a, b)] -> ([a], [b])
unzip ([(SubExp, [Maybe PrimValue])] -> ([SubExp], [[Maybe PrimValue]]))
-> [(SubExp, [Maybe PrimValue])] -> ([SubExp], [[Maybe PrimValue]])
forall a b. (a -> b) -> a -> b
$ ((SubExp, [Maybe PrimValue]) -> Bool)
-> [(SubExp, [Maybe PrimValue])] -> [(SubExp, [Maybe PrimValue])]
forall a. (a -> Bool) -> [a] -> [a]
filter (SubExp, [Maybe PrimValue]) -> Bool
forall {a}. (a, [Maybe PrimValue]) -> Bool
interesting ([(SubExp, [Maybe PrimValue])] -> [(SubExp, [Maybe PrimValue])])
-> [(SubExp, [Maybe PrimValue])] -> [(SubExp, [Maybe PrimValue])]
forall a b. (a -> b) -> a -> b
$ [SubExp] -> [[Maybe PrimValue]] -> [(SubExp, [Maybe PrimValue])]
forall a b. [a] -> [b] -> [(a, b)]
zip [SubExp]
ses ([[Maybe PrimValue]] -> [(SubExp, [Maybe PrimValue])])
-> [[Maybe PrimValue]] -> [(SubExp, [Maybe PrimValue])]
forall a b. (a -> b) -> a -> b
$ [[Maybe PrimValue]] -> [[Maybe PrimValue]]
forall a. [[a]] -> [[a]]
L.transpose ((Case b -> [Maybe PrimValue]) -> [Case b] -> [[Maybe PrimValue]]
forall a b. (a -> b) -> [a] -> [b]
map Case b -> [Maybe PrimValue]
forall body. Case body -> [Maybe PrimValue]
casePat [Case b]
cases)
   in ([SubExp]
ses', ([Maybe PrimValue] -> b -> Case b)
-> [[Maybe PrimValue]] -> [b] -> [Case b]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith [Maybe PrimValue] -> b -> Case b
forall body. [Maybe PrimValue] -> body -> Case body
Case ([[Maybe PrimValue]] -> [[Maybe PrimValue]]
forall a. [[a]] -> [[a]]
L.transpose [[Maybe PrimValue]]
vs) ([b] -> [Case b]) -> [b] -> [Case b]
forall a b. (a -> b) -> a -> b
$ (Case b -> b) -> [Case b] -> [b]
forall a b. (a -> b) -> [a] -> [b]
map Case b -> b
forall body. Case body -> body
caseBody [Case b]
cases)
  where
    interesting :: (a, [Maybe PrimValue]) -> Bool
interesting = (Maybe PrimValue -> Bool) -> [Maybe PrimValue] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (Maybe PrimValue -> Maybe PrimValue -> Bool
forall a. Eq a => a -> a -> Bool
/= Maybe PrimValue
forall a. Maybe a
Nothing) ([Maybe PrimValue] -> Bool)
-> ((a, [Maybe PrimValue]) -> [Maybe PrimValue])
-> (a, [Maybe PrimValue])
-> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a, [Maybe PrimValue]) -> [Maybe PrimValue]
forall a b. (a, b) -> b
snd

-- | As 'eMatch', but an 'MatchSort' can be given.
eMatch' ::
  (MonadBuilder m, BranchType (Rep m) ~ ExtType) =>
  [SubExp] ->
  [Case (m (Body (Rep m)))] ->
  m (Body (Rep m)) ->
  MatchSort ->
  m (Exp (Rep m))
eMatch' :: forall (m :: * -> *).
(MonadBuilder m, BranchType (Rep m) ~ ExtType) =>
[SubExp]
-> [Case (m (Body (Rep m)))]
-> m (Body (Rep m))
-> MatchSort
-> m (Exp (Rep m))
eMatch' [SubExp]
ses [Case (m (Body (Rep m)))]
cases_m m (Body (Rep m))
defbody_m MatchSort
sort = do
  cases <- (Case (m (Body (Rep m))) -> m (Case (Body (Rep m))))
-> [Case (m (Body (Rep m)))] -> m [Case (Body (Rep m))]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM ((m (Body (Rep m)) -> m (Body (Rep m)))
-> Case (m (Body (Rep m))) -> m (Case (Body (Rep m)))
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Case a -> f (Case b)
traverse m (Body (Rep m)) -> m (Body (Rep m))
forall (m :: * -> *).
MonadBuilder m =>
m (Body (Rep m)) -> m (Body (Rep m))
insertStmsM) [Case (m (Body (Rep m)))]
cases_m
  defbody <- insertStmsM defbody_m
  ts <-
    L.foldl' generaliseExtTypes
      <$> bodyExtType defbody
      <*> mapM (bodyExtType . caseBody) cases
  cases' <- mapM (traverse $ addContextForBranch ts) cases
  defbody' <- addContextForBranch ts defbody
  let ts' = Int -> ExtType -> [ExtType]
forall a. Int -> a -> [a]
replicate (Set Int -> Int
forall a. Set a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([ExtType] -> Set Int
forall o. [TypeBase (ShapeBase (Ext SubExp)) o] -> Set Int
shapeContext [ExtType]
ts)) (PrimType -> ExtType
forall shape o. PrimType -> TypeBase shape o
Prim PrimType
int64) [ExtType] -> [ExtType] -> [ExtType]
forall a. [a] -> [a] -> [a]
++ [ExtType]
ts
      (ses', cases'') = removeRedundantScrutinees ses cases'
  pure $ Match ses' cases'' defbody' $ MatchDec ts' sort
  where
    addContextForBranch :: [TypeBase (ShapeBase (Ext SubExp)) o]
-> GBody (Rep m) SubExpRes -> m (GBody (Rep m) SubExpRes)
addContextForBranch [TypeBase (ShapeBase (Ext SubExp)) o]
ts (Body BodyDec (Rep m)
_ Stms (Rep m)
stms [SubExpRes]
val_res) = do
      body_ts <- ExtendedScope (Rep m) m [Type] -> Scope (Rep m) -> m [Type]
forall rep (m :: * -> *) a.
ExtendedScope rep m a -> Scope rep -> m a
extendedScope ((SubExpRes -> ExtendedScope (Rep m) m Type)
-> [SubExpRes] -> ExtendedScope (Rep m) m [Type]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> [a] -> f [b]
traverse SubExpRes -> ExtendedScope (Rep m) m Type
forall t (m :: * -> *). HasScope t m => SubExpRes -> m Type
subExpResType [SubExpRes]
val_res) Scope (Rep m)
stmsscope
      let ctx_res =
            ((Int, SubExp) -> SubExp) -> [(Int, SubExp)] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map (Int, SubExp) -> SubExp
forall a b. (a, b) -> b
snd ([(Int, SubExp)] -> [SubExp]) -> [(Int, SubExp)] -> [SubExp]
forall a b. (a -> b) -> a -> b
$ ((Int, SubExp) -> Int) -> [(Int, SubExp)] -> [(Int, SubExp)]
forall b a. Ord b => (a -> b) -> [a] -> [a]
L.sortOn (Int, SubExp) -> Int
forall a b. (a, b) -> a
fst ([(Int, SubExp)] -> [(Int, SubExp)])
-> [(Int, SubExp)] -> [(Int, SubExp)]
forall a b. (a -> b) -> a -> b
$ Map Int SubExp -> [(Int, SubExp)]
forall k a. Map k a -> [(k, a)]
M.toList (Map Int SubExp -> [(Int, SubExp)])
-> Map Int SubExp -> [(Int, SubExp)]
forall a b. (a -> b) -> a -> b
$ [TypeBase (ShapeBase (Ext SubExp)) o] -> [Type] -> Map Int SubExp
forall o o1.
[TypeBase (ShapeBase (Ext SubExp)) o]
-> [TypeBase Shape o1] -> Map Int SubExp
shapeExtMapping [TypeBase (ShapeBase (Ext SubExp)) o]
ts [Type]
body_ts
      mkBodyM stms $ subExpsRes ctx_res ++ val_res
      where
        stmsscope :: Scope (Rep m)
stmsscope = Stms (Rep m) -> Scope (Rep m)
forall rep a. Scoped rep a => a -> Scope rep
scopeOf Stms (Rep m)
stms

-- | Construct a 'Match' expression.  The main convenience here is
-- that the existential context of the return type is automatically
-- deduced, and the necessary elements added to the branches.
eMatch ::
  (MonadBuilder m, BranchType (Rep m) ~ ExtType) =>
  [SubExp] ->
  [Case (m (Body (Rep m)))] ->
  m (Body (Rep m)) ->
  m (Exp (Rep m))
eMatch :: forall (m :: * -> *).
(MonadBuilder m, BranchType (Rep m) ~ ExtType) =>
[SubExp]
-> [Case (m (Body (Rep m)))] -> m (Body (Rep m)) -> m (Exp (Rep m))
eMatch [SubExp]
ses [Case (m (Body (Rep m)))]
cases_m m (Body (Rep m))
defbody_m = [SubExp]
-> [Case (m (Body (Rep m)))]
-> m (Body (Rep m))
-> MatchSort
-> m (Exp (Rep m))
forall (m :: * -> *).
(MonadBuilder m, BranchType (Rep m) ~ ExtType) =>
[SubExp]
-> [Case (m (Body (Rep m)))]
-> m (Body (Rep m))
-> MatchSort
-> m (Exp (Rep m))
eMatch' [SubExp]
ses [Case (m (Body (Rep m)))]
cases_m m (Body (Rep m))
defbody_m MatchSort
MatchNormal

-- | Construct a 'Match' modelling an if-expression from a monadic
-- condition and monadic branches.  'eBody' might be convenient for
-- constructing the branches.
eIf ::
  (MonadBuilder m, BranchType (Rep m) ~ ExtType) =>
  m (Exp (Rep m)) ->
  m (Body (Rep m)) ->
  m (Body (Rep m)) ->
  m (Exp (Rep m))
eIf :: forall (m :: * -> *).
(MonadBuilder m, BranchType (Rep m) ~ ExtType) =>
m (Exp (Rep m))
-> m (Body (Rep m)) -> m (Body (Rep m)) -> m (Exp (Rep m))
eIf m (Exp (Rep m))
ce m (Body (Rep m))
te m (Body (Rep m))
fe = m (Exp (Rep m))
-> m (Body (Rep m))
-> m (Body (Rep m))
-> MatchSort
-> m (Exp (Rep m))
forall (m :: * -> *).
(MonadBuilder m, BranchType (Rep m) ~ ExtType) =>
m (Exp (Rep m))
-> m (Body (Rep m))
-> m (Body (Rep m))
-> MatchSort
-> m (Exp (Rep m))
eIf' m (Exp (Rep m))
ce m (Body (Rep m))
te m (Body (Rep m))
fe MatchSort
MatchNormal

-- | As 'eIf', but an 'MatchSort' can be given.
eIf' ::
  (MonadBuilder m, BranchType (Rep m) ~ ExtType) =>
  m (Exp (Rep m)) ->
  m (Body (Rep m)) ->
  m (Body (Rep m)) ->
  MatchSort ->
  m (Exp (Rep m))
eIf' :: forall (m :: * -> *).
(MonadBuilder m, BranchType (Rep m) ~ ExtType) =>
m (Exp (Rep m))
-> m (Body (Rep m))
-> m (Body (Rep m))
-> MatchSort
-> m (Exp (Rep m))
eIf' m (Exp (Rep m))
ce m (Body (Rep m))
te m (Body (Rep m))
fe MatchSort
if_sort = do
  ce' <- Name -> Exp (Rep m) -> m SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"cond" (Exp (Rep m) -> m SubExp) -> m (Exp (Rep m)) -> m SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< m (Exp (Rep m))
ce
  eMatch' [ce'] [Case [Just $ BoolValue True] te] fe if_sort

-- The type of a body.  Watch out: this only works for the degenerate
-- case where the body does not already return its context.
bodyExtType :: (HasScope rep m, Monad m) => Body rep -> m [ExtType]
bodyExtType :: forall rep (m :: * -> *).
(HasScope rep m, Monad m) =>
Body rep -> m [ExtType]
bodyExtType (Body BodyDec rep
_ Stms rep
stms [SubExpRes]
res) =
  [VName] -> [ExtType] -> [ExtType]
existentialiseExtTypes (Map VName (NameInfo rep) -> [VName]
forall k a. Map k a -> [k]
M.keys Map VName (NameInfo rep)
stmsscope) ([ExtType] -> [ExtType])
-> ([Type] -> [ExtType]) -> [Type] -> [ExtType]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Type] -> [ExtType]
forall o.
[TypeBase Shape o] -> [TypeBase (ShapeBase (Ext SubExp)) o]
staticShapes
    ([Type] -> [ExtType]) -> m [Type] -> m [ExtType]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ExtendedScope rep m [Type] -> Map VName (NameInfo rep) -> m [Type]
forall rep (m :: * -> *) a.
ExtendedScope rep m a -> Scope rep -> m a
extendedScope ((SubExpRes -> ExtendedScope rep m Type)
-> [SubExpRes] -> ExtendedScope rep m [Type]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> [a] -> f [b]
traverse SubExpRes -> ExtendedScope rep m Type
forall t (m :: * -> *). HasScope t m => SubExpRes -> m Type
subExpResType [SubExpRes]
res) Map VName (NameInfo rep)
stmsscope
  where
    stmsscope :: Map VName (NameInfo rep)
stmsscope = Stms rep -> Map VName (NameInfo rep)
forall rep a. Scoped rep a => a -> Scope rep
scopeOf Stms rep
stms

-- | Construct a v'BinOp' expression with the given operator.
eBinOp ::
  (MonadBuilder m) =>
  BinOp ->
  m (Exp (Rep m)) ->
  m (Exp (Rep m)) ->
  m (Exp (Rep m))
eBinOp :: forall (m :: * -> *).
MonadBuilder m =>
BinOp -> m (Exp (Rep m)) -> m (Exp (Rep m)) -> m (Exp (Rep m))
eBinOp BinOp
op m (Exp (Rep m))
x m (Exp (Rep m))
y = do
  x' <- Name -> Exp (Rep m) -> m SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"x" (Exp (Rep m) -> m SubExp) -> m (Exp (Rep m)) -> m SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< m (Exp (Rep m))
x
  y' <- letSubExp "y" =<< y
  pure $ BasicOp $ BinOp op x' y'

-- | Construct a v'UnOp' expression with the given operator.
eUnOp ::
  (MonadBuilder m) =>
  UnOp ->
  m (Exp (Rep m)) ->
  m (Exp (Rep m))
eUnOp :: forall (m :: * -> *).
MonadBuilder m =>
UnOp -> m (Exp (Rep m)) -> m (Exp (Rep m))
eUnOp UnOp
op m (Exp (Rep m))
x = BasicOp -> Exp (Rep m)
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp (Rep m))
-> (SubExp -> BasicOp) -> SubExp -> Exp (Rep m)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. UnOp -> SubExp -> BasicOp
UnOp UnOp
op (SubExp -> Exp (Rep m)) -> m SubExp -> m (Exp (Rep m))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Name -> Exp (Rep m) -> m SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"x" (Exp (Rep m) -> m SubExp) -> m (Exp (Rep m)) -> m SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< m (Exp (Rep m))
x)

-- | Construct a v'CmpOp' expression with the given comparison.
eCmpOp ::
  (MonadBuilder m) =>
  CmpOp ->
  m (Exp (Rep m)) ->
  m (Exp (Rep m)) ->
  m (Exp (Rep m))
eCmpOp :: forall (m :: * -> *).
MonadBuilder m =>
CmpOp -> m (Exp (Rep m)) -> m (Exp (Rep m)) -> m (Exp (Rep m))
eCmpOp CmpOp
op m (Exp (Rep m))
x m (Exp (Rep m))
y = do
  x' <- Name -> Exp (Rep m) -> m SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"x" (Exp (Rep m) -> m SubExp) -> m (Exp (Rep m)) -> m SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< m (Exp (Rep m))
x
  y' <- letSubExp "y" =<< y
  pure $ BasicOp $ CmpOp op x' y'

-- | Construct a v'ConvOp' expression with the given conversion.
eConvOp ::
  (MonadBuilder m) =>
  ConvOp ->
  m (Exp (Rep m)) ->
  m (Exp (Rep m))
eConvOp :: forall (m :: * -> *).
MonadBuilder m =>
ConvOp -> m (Exp (Rep m)) -> m (Exp (Rep m))
eConvOp ConvOp
op m (Exp (Rep m))
x = do
  x' <- Name -> Exp (Rep m) -> m SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"x" (Exp (Rep m) -> m SubExp) -> m (Exp (Rep m)) -> m SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< m (Exp (Rep m))
x
  pure $ BasicOp $ ConvOp op x'

-- | Construct a 'SSignum' expression.  Fails if the provided
-- expression is not of integer type.
eSignum ::
  (MonadBuilder m) =>
  m (Exp (Rep m)) ->
  m (Exp (Rep m))
eSignum :: forall (m :: * -> *).
MonadBuilder m =>
m (Exp (Rep m)) -> m (Exp (Rep m))
eSignum m (Exp (Rep m))
em = do
  e <- m (Exp (Rep m))
em
  e' <- letSubExp "signum_arg" e
  t <- subExpType e'
  case t of
    Prim (IntType IntType
int_t) ->
      Exp (Rep m) -> m (Exp (Rep m))
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp (Rep m) -> m (Exp (Rep m))) -> Exp (Rep m) -> m (Exp (Rep m))
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep m)
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp (Rep m)) -> BasicOp -> Exp (Rep m)
forall a b. (a -> b) -> a -> b
$ UnOp -> SubExp -> BasicOp
UnOp (IntType -> UnOp
SSignum IntType
int_t) SubExp
e'
    Type
_ ->
      [Char] -> m (Exp (Rep m))
forall a. HasCallStack => [Char] -> a
error ([Char] -> m (Exp (Rep m))) -> [Char] -> m (Exp (Rep m))
forall a b. (a -> b) -> a -> b
$ [Char]
"eSignum: operand " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Exp (Rep m) -> [Char]
forall a. Pretty a => a -> [Char]
prettyString Exp (Rep m)
e [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
" has invalid type."

-- | Copy a value.
eCopy ::
  (MonadBuilder m) =>
  m (Exp (Rep m)) ->
  m (Exp (Rep m))
eCopy :: forall (m :: * -> *).
MonadBuilder m =>
m (Exp (Rep m)) -> m (Exp (Rep m))
eCopy m (Exp (Rep m))
e = BasicOp -> Exp (Rep m)
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp (Rep m))
-> (SubExp -> BasicOp) -> SubExp -> Exp (Rep m)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Shape -> SubExp -> BasicOp
Replicate Shape
forall a. Monoid a => a
mempty (SubExp -> Exp (Rep m)) -> m SubExp -> m (Exp (Rep m))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Name -> Exp (Rep m) -> m SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"copy_arg" (Exp (Rep m) -> m SubExp) -> m (Exp (Rep m)) -> m SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< m (Exp (Rep m))
e)

-- | Construct a body from expressions.  If multiple expressions are
-- provided, their results will be concatenated in order and returned
-- as the result.
--
-- /Beware/: this will not produce correct code if the type of the
-- body would be existential.  That is, the type of the results being
-- returned should be invariant to the body.
eBody ::
  (MonadBuilder m) =>
  [m (Exp (Rep m))] ->
  m (Body (Rep m))
eBody :: forall (m :: * -> *).
MonadBuilder m =>
[m (Exp (Rep m))] -> m (Body (Rep m))
eBody [m (Exp (Rep m))]
es = m [SubExpRes] -> m (GBody (Rep m) SubExpRes)
forall (m :: * -> *) res.
(MonadBuilder m, IsResult res) =>
m [res] -> m (GBody (Rep m) res)
buildBody_ (m [SubExpRes] -> m (GBody (Rep m) SubExpRes))
-> m [SubExpRes] -> m (GBody (Rep m) SubExpRes)
forall a b. (a -> b) -> a -> b
$ do
  es' <- [m (Exp (Rep m))] -> m [Exp (Rep m)]
forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
forall (m :: * -> *) a. Monad m => [m a] -> m [a]
sequence [m (Exp (Rep m))]
es
  xs <- mapM (letTupExp "x") es'
  pure $ varsRes $ concat xs

-- | Bind each lambda parameter to the result of an expression, then
-- bind the body of the lambda.  The expressions must produce only a
-- single value each.
eLambda ::
  (MonadBuilder m) =>
  Lambda (Rep m) ->
  [m (Exp (Rep m))] ->
  m [SubExpRes]
eLambda :: forall (m :: * -> *).
MonadBuilder m =>
Lambda (Rep m) -> [m (Exp (Rep m))] -> m [SubExpRes]
eLambda Lambda (Rep m)
lam [m (Exp (Rep m))]
args = do
  (Param (LParamInfo (Rep m)) -> m (Exp (Rep m)) -> m ())
-> [Param (LParamInfo (Rep m))] -> [m (Exp (Rep m))] -> m ()
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m ()
zipWithM_ Param (LParamInfo (Rep m)) -> m (Exp (Rep m)) -> m ()
forall {m :: * -> *} {dec}.
MonadBuilder m =>
Param dec -> m (Exp (Rep m)) -> m ()
bindParam (Lambda (Rep m) -> [Param (LParamInfo (Rep m))]
forall rep t. GLambda rep t -> [LParam rep]
lambdaParams Lambda (Rep m)
lam) [m (Exp (Rep m))]
args
  Body (Rep m) -> m [SubExpRes]
forall (m :: * -> *).
MonadBuilder m =>
Body (Rep m) -> m [SubExpRes]
bodyBind (Body (Rep m) -> m [SubExpRes]) -> Body (Rep m) -> m [SubExpRes]
forall a b. (a -> b) -> a -> b
$ Lambda (Rep m) -> Body (Rep m)
forall rep t. GLambda rep t -> Body rep
lambdaBody Lambda (Rep m)
lam
  where
    bindParam :: Param dec -> m (Exp (Rep m)) -> m ()
bindParam Param dec
param m (Exp (Rep m))
arg = [VName] -> Exp (Rep m) -> m ()
forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [Param dec -> VName
forall dec. Param dec -> VName
paramName Param dec
param] (Exp (Rep m) -> m ()) -> m (Exp (Rep m)) -> m ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< m (Exp (Rep m))
arg

-- | @eInBoundsForDim w i@ produces @0 <= i < w@.
eDimInBounds :: (MonadBuilder m) => m (Exp (Rep m)) -> m (Exp (Rep m)) -> m (Exp (Rep m))
eDimInBounds :: forall (m :: * -> *).
MonadBuilder m =>
m (Exp (Rep m)) -> m (Exp (Rep m)) -> m (Exp (Rep m))
eDimInBounds m (Exp (Rep m))
w m (Exp (Rep m))
i =
  BinOp -> m (Exp (Rep m)) -> m (Exp (Rep m)) -> m (Exp (Rep m))
forall (m :: * -> *).
MonadBuilder m =>
BinOp -> m (Exp (Rep m)) -> m (Exp (Rep m)) -> m (Exp (Rep m))
eBinOp
    BinOp
LogAnd
    (CmpOp -> m (Exp (Rep m)) -> m (Exp (Rep m)) -> m (Exp (Rep m))
forall (m :: * -> *).
MonadBuilder m =>
CmpOp -> m (Exp (Rep m)) -> m (Exp (Rep m)) -> m (Exp (Rep m))
eCmpOp (IntType -> CmpOp
CmpSle IntType
Int64) (SubExp -> m (Exp (Rep m))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp (IntType -> Integer -> SubExp
intConst IntType
Int64 Integer
0)) m (Exp (Rep m))
i)
    (CmpOp -> m (Exp (Rep m)) -> m (Exp (Rep m)) -> m (Exp (Rep m))
forall (m :: * -> *).
MonadBuilder m =>
CmpOp -> m (Exp (Rep m)) -> m (Exp (Rep m)) -> m (Exp (Rep m))
eCmpOp (IntType -> CmpOp
CmpSlt IntType
Int64) m (Exp (Rep m))
i m (Exp (Rep m))
w)

-- | Check if the given indexes are in-bounds for the given shape. The shape may have extra dimensions.
eShapeInBounds :: (MonadBuilder m) => Shape -> [m (Exp (Rep m))] -> m (Exp (Rep m))
eShapeInBounds :: forall (m :: * -> *).
MonadBuilder m =>
Shape -> [m (Exp (Rep m))] -> m (Exp (Rep m))
eShapeInBounds (Shape [SubExp]
ds) [m (Exp (Rep m))]
is =
  [SubExp] -> m (Exp (Rep m))
forall (m :: * -> *). MonadBuilder m => [SubExp] -> m (Exp (Rep m))
eAll
    ([SubExp] -> m (Exp (Rep m))) -> m [SubExp] -> m (Exp (Rep m))
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< (Exp (Rep m) -> m SubExp) -> [Exp (Rep m)] -> m [SubExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (Name -> Exp (Rep m) -> m SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"dim_in_bounds")
    ([Exp (Rep m)] -> m [SubExp]) -> m [Exp (Rep m)] -> m [SubExp]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< (m (Exp (Rep m)) -> m (Exp (Rep m)) -> m (Exp (Rep m)))
-> [m (Exp (Rep m))] -> [m (Exp (Rep m))] -> m [Exp (Rep m)]
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM m (Exp (Rep m)) -> m (Exp (Rep m)) -> m (Exp (Rep m))
forall (m :: * -> *).
MonadBuilder m =>
m (Exp (Rep m)) -> m (Exp (Rep m)) -> m (Exp (Rep m))
eDimInBounds ((SubExp -> m (Exp (Rep m))) -> [SubExp] -> [m (Exp (Rep m))]
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> m (Exp (Rep m))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp (Int -> [SubExp] -> [SubExp]
forall a. Int -> [a] -> [a]
take ([m (Exp (Rep m))] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [m (Exp (Rep m))]
is) [SubExp]
ds)) [m (Exp (Rep m))]
is

-- | Are these indexes out-of-bounds for the array?
eOutOfBounds ::
  (MonadBuilder m) =>
  VName ->
  [m (Exp (Rep m))] ->
  m (Exp (Rep m))
eOutOfBounds :: forall (m :: * -> *).
MonadBuilder m =>
VName -> [m (Exp (Rep m))] -> m (Exp (Rep m))
eOutOfBounds VName
arr [m (Exp (Rep m))]
is = do
  arr_t <- VName -> m Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
arr
  let ws = Type -> [SubExp]
forall d m. TypeBase (ShapeBase d) m -> [d]
arrayDims Type
arr_t
  is' <- mapM (letSubExp "write_i") =<< sequence is
  let checkDim SubExp
w SubExp
i = do
        less_than_zero <-
          Name -> Exp (Rep m) -> m SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"less_than_zero" (Exp (Rep m) -> m SubExp) -> Exp (Rep m) -> m SubExp
forall a b. (a -> b) -> a -> b
$
            BasicOp -> Exp (Rep m)
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp (Rep m)) -> BasicOp -> Exp (Rep m)
forall a b. (a -> b) -> a -> b
$
              CmpOp -> SubExp -> SubExp -> BasicOp
CmpOp (IntType -> CmpOp
CmpSlt IntType
Int64) SubExp
i (Int64 -> SubExp
forall v. IsValue v => v -> SubExp
constant (Int64
0 :: Int64))
        greater_than_size <-
          letSubExp "greater_than_size" $
            BasicOp $
              CmpOp (CmpSle Int64) w i
        letSubExp "outside_bounds_dim" $
          BasicOp $
            BinOp LogOr less_than_zero greater_than_size
  foldBinOp LogOr (constant False) =<< zipWithM checkDim ws is'

-- | The array element at this index.  Returns array unmodified if
-- indexes are null (does not even need to be an array in that case).
eIndex :: (MonadBuilder m) => VName -> [m (Exp (Rep m))] -> m (Exp (Rep m))
eIndex :: forall (m :: * -> *).
MonadBuilder m =>
VName -> [m (Exp (Rep m))] -> m (Exp (Rep m))
eIndex VName
arr [] = SubExp -> m (Exp (Rep m))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp (SubExp -> m (Exp (Rep m))) -> SubExp -> m (Exp (Rep m))
forall a b. (a -> b) -> a -> b
$ VName -> SubExp
Var VName
arr
eIndex VName
arr [m (Exp (Rep m))]
is = do
  is' <- (m (Exp (Rep m)) -> m SubExp) -> [m (Exp (Rep m))] -> m [SubExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (Name -> Exp (Rep m) -> m SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"i" (Exp (Rep m) -> m SubExp) -> m (Exp (Rep m)) -> m SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<<) [m (Exp (Rep m))]
is
  arr_t <- lookupType arr
  pure $ BasicOp $ Index arr $ fullSlice arr_t $ map DimFix is'

-- | The last element of the given array.
eLast :: (MonadBuilder m) => VName -> m (Exp (Rep m))
eLast :: forall (m :: * -> *). MonadBuilder m => VName -> m (Exp (Rep m))
eLast VName
arr = do
  n <- Int -> Type -> SubExp
forall o. Int -> TypeBase Shape o -> SubExp
arraySize Int
0 (Type -> SubExp) -> m Type -> m SubExp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> m Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
arr
  nm1 <-
    letSubExp "nm1" . BasicOp $
      BinOp (Sub Int64 OverflowUndef) n (intConst Int64 1)
  eIndex arr [eSubExp nm1]

-- | Construct an unspecified value of the given type.
eBlank :: (MonadBuilder m) => Type -> m (Exp (Rep m))
eBlank :: forall (m :: * -> *). MonadBuilder m => Type -> m (Exp (Rep m))
eBlank (Prim PrimType
t) = Exp (Rep m) -> m (Exp (Rep m))
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp (Rep m) -> m (Exp (Rep m))) -> Exp (Rep m) -> m (Exp (Rep m))
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep m)
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp (Rep m)) -> BasicOp -> Exp (Rep m)
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
SubExp (SubExp -> BasicOp) -> SubExp -> BasicOp
forall a b. (a -> b) -> a -> b
$ PrimValue -> SubExp
Constant (PrimValue -> SubExp) -> PrimValue -> SubExp
forall a b. (a -> b) -> a -> b
$ PrimType -> PrimValue
blankPrimValue PrimType
t
eBlank (Array PrimType
t Shape
shape NoMode
_) = Exp (Rep m) -> m (Exp (Rep m))
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp (Rep m) -> m (Exp (Rep m))) -> Exp (Rep m) -> m (Exp (Rep m))
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep m)
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp (Rep m)) -> BasicOp -> Exp (Rep m)
forall a b. (a -> b) -> a -> b
$ PrimType -> [SubExp] -> BasicOp
Scratch PrimType
t ([SubExp] -> BasicOp) -> [SubExp] -> BasicOp
forall a b. (a -> b) -> a -> b
$ Shape -> [SubExp]
forall d. ShapeBase d -> [d]
shapeDims Shape
shape
eBlank Acc {} = [Char] -> m (Exp (Rep m))
forall a. HasCallStack => [Char] -> a
error [Char]
"eBlank: cannot create blank accumulator"
eBlank Mem {} = [Char] -> m (Exp (Rep m))
forall a. HasCallStack => [Char] -> a
error [Char]
"eBlank: cannot create blank memory"

-- | Sign-extend to the given integer type.
asIntS :: (MonadBuilder m) => IntType -> SubExp -> m SubExp
asIntS :: forall (m :: * -> *).
MonadBuilder m =>
IntType -> SubExp -> m SubExp
asIntS = (IntType -> IntType -> ConvOp) -> IntType -> SubExp -> m SubExp
forall (m :: * -> *).
MonadBuilder m =>
(IntType -> IntType -> ConvOp) -> IntType -> SubExp -> m SubExp
asInt IntType -> IntType -> ConvOp
SExt

-- | Zero-extend to the given integer type.
asIntZ :: (MonadBuilder m) => IntType -> SubExp -> m SubExp
asIntZ :: forall (m :: * -> *).
MonadBuilder m =>
IntType -> SubExp -> m SubExp
asIntZ = (IntType -> IntType -> ConvOp) -> IntType -> SubExp -> m SubExp
forall (m :: * -> *).
MonadBuilder m =>
(IntType -> IntType -> ConvOp) -> IntType -> SubExp -> m SubExp
asInt IntType -> IntType -> ConvOp
ZExt

asInt ::
  (MonadBuilder m) =>
  (IntType -> IntType -> ConvOp) ->
  IntType ->
  SubExp ->
  m SubExp
asInt :: forall (m :: * -> *).
MonadBuilder m =>
(IntType -> IntType -> ConvOp) -> IntType -> SubExp -> m SubExp
asInt IntType -> IntType -> ConvOp
ext IntType
to_it SubExp
e = do
  e_t <- SubExp -> m Type
forall t (m :: * -> *). HasScope t m => SubExp -> m Type
subExpType SubExp
e
  case e_t of
    Prim (IntType IntType
from_it)
      | IntType
to_it IntType -> IntType -> Bool
forall a. Eq a => a -> a -> Bool
== IntType
from_it -> SubExp -> m SubExp
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SubExp
e
      | Bool
otherwise -> Name -> Exp (Rep m) -> m SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
s (Exp (Rep m) -> m SubExp) -> Exp (Rep m) -> m SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep m)
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp (Rep m)) -> BasicOp -> Exp (Rep m)
forall a b. (a -> b) -> a -> b
$ ConvOp -> SubExp -> BasicOp
ConvOp (IntType -> IntType -> ConvOp
ext IntType
from_it IntType
to_it) SubExp
e
    Type
_ -> [Char] -> m SubExp
forall a. HasCallStack => [Char] -> a
error [Char]
"asInt: wrong type"
  where
    s :: Name
s = case SubExp
e of
      Var VName
v -> VName -> Name
baseName VName
v
      SubExp
_ -> Text -> Name
nameFromText (Text -> Name) -> Text -> Name
forall a b. (a -> b) -> a -> b
$ Text
"to_" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> IntType -> Text
forall a. Pretty a => a -> Text
prettyText IntType
to_it

-- | Apply a binary operator to several subexpressions.  A left-fold.
foldBinOp ::
  (MonadBuilder m) =>
  BinOp ->
  SubExp ->
  [SubExp] ->
  m (Exp (Rep m))
foldBinOp :: forall (m :: * -> *).
MonadBuilder m =>
BinOp -> SubExp -> [SubExp] -> m (Exp (Rep m))
foldBinOp BinOp
_ SubExp
ne [] =
  Exp (Rep m) -> m (Exp (Rep m))
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp (Rep m) -> m (Exp (Rep m))) -> Exp (Rep m) -> m (Exp (Rep m))
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep m)
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp (Rep m)) -> BasicOp -> Exp (Rep m)
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
SubExp SubExp
ne
foldBinOp BinOp
bop SubExp
ne (SubExp
e : [SubExp]
es) =
  BinOp -> m (Exp (Rep m)) -> m (Exp (Rep m)) -> m (Exp (Rep m))
forall (m :: * -> *).
MonadBuilder m =>
BinOp -> m (Exp (Rep m)) -> m (Exp (Rep m)) -> m (Exp (Rep m))
eBinOp BinOp
bop (Exp (Rep m) -> m (Exp (Rep m))
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp (Rep m) -> m (Exp (Rep m))) -> Exp (Rep m) -> m (Exp (Rep m))
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep m)
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp (Rep m)) -> BasicOp -> Exp (Rep m)
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
SubExp SubExp
e) (BinOp -> SubExp -> [SubExp] -> m (Exp (Rep m))
forall (m :: * -> *).
MonadBuilder m =>
BinOp -> SubExp -> [SubExp] -> m (Exp (Rep m))
foldBinOp BinOp
bop SubExp
ne [SubExp]
es)

-- | True if all operands are true.
eAll :: (MonadBuilder m) => [SubExp] -> m (Exp (Rep m))
eAll :: forall (m :: * -> *). MonadBuilder m => [SubExp] -> m (Exp (Rep m))
eAll [] = Exp (Rep m) -> m (Exp (Rep m))
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp (Rep m) -> m (Exp (Rep m))) -> Exp (Rep m) -> m (Exp (Rep m))
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep m)
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp (Rep m)) -> BasicOp -> Exp (Rep m)
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
SubExp (SubExp -> BasicOp) -> SubExp -> BasicOp
forall a b. (a -> b) -> a -> b
$ Bool -> SubExp
forall v. IsValue v => v -> SubExp
constant Bool
True
eAll [SubExp
x] = SubExp -> m (Exp (Rep m))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
x
eAll (SubExp
x : [SubExp]
xs) = BinOp -> SubExp -> [SubExp] -> m (Exp (Rep m))
forall (m :: * -> *).
MonadBuilder m =>
BinOp -> SubExp -> [SubExp] -> m (Exp (Rep m))
foldBinOp BinOp
LogAnd SubExp
x [SubExp]
xs

-- | True if any operand is true.
eAny :: (MonadBuilder m) => [SubExp] -> m (Exp (Rep m))
eAny :: forall (m :: * -> *). MonadBuilder m => [SubExp] -> m (Exp (Rep m))
eAny [] = Exp (Rep m) -> m (Exp (Rep m))
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp (Rep m) -> m (Exp (Rep m))) -> Exp (Rep m) -> m (Exp (Rep m))
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep m)
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp (Rep m)) -> BasicOp -> Exp (Rep m)
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
SubExp (SubExp -> BasicOp) -> SubExp -> BasicOp
forall a b. (a -> b) -> a -> b
$ Bool -> SubExp
forall v. IsValue v => v -> SubExp
constant Bool
False
eAny [SubExp
x] = SubExp -> m (Exp (Rep m))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
x
eAny (SubExp
x : [SubExp]
xs) = BinOp -> SubExp -> [SubExp] -> m (Exp (Rep m))
forall (m :: * -> *).
MonadBuilder m =>
BinOp -> SubExp -> [SubExp] -> m (Exp (Rep m))
foldBinOp BinOp
LogOr SubExp
x [SubExp]
xs

-- | Create a two-parameter lambda whose body applies the given binary
-- operation to its arguments.  It is assumed that both argument and
-- result types are the same.  (This assumption should be fixed at
-- some point.)
binOpLambda ::
  (MonadBuilder m, Buildable (Rep m)) =>
  BinOp ->
  PrimType ->
  m (Lambda (Rep m))
binOpLambda :: forall (m :: * -> *).
(MonadBuilder m, Buildable (Rep m)) =>
BinOp -> PrimType -> m (Lambda (Rep m))
binOpLambda BinOp
bop PrimType
t = (SubExp -> SubExp -> BasicOp)
-> PrimType -> PrimType -> m (GLambda (Rep m) Type)
forall (m :: * -> *).
(MonadBuilder m, Buildable (Rep m)) =>
(SubExp -> SubExp -> BasicOp)
-> PrimType -> PrimType -> m (Lambda (Rep m))
binLambda (BinOp -> SubExp -> SubExp -> BasicOp
BinOp BinOp
bop) PrimType
t PrimType
t

-- | As 'binOpLambda', but for t'CmpOp's.
cmpOpLambda ::
  (MonadBuilder m, Buildable (Rep m)) =>
  CmpOp ->
  m (Lambda (Rep m))
cmpOpLambda :: forall (m :: * -> *).
(MonadBuilder m, Buildable (Rep m)) =>
CmpOp -> m (Lambda (Rep m))
cmpOpLambda CmpOp
cop = (SubExp -> SubExp -> BasicOp)
-> PrimType -> PrimType -> m (GLambda (Rep m) Type)
forall (m :: * -> *).
(MonadBuilder m, Buildable (Rep m)) =>
(SubExp -> SubExp -> BasicOp)
-> PrimType -> PrimType -> m (Lambda (Rep m))
binLambda (CmpOp -> SubExp -> SubExp -> BasicOp
CmpOp CmpOp
cop) (CmpOp -> PrimType
cmpOpType CmpOp
cop) PrimType
Bool

binLambda ::
  (MonadBuilder m, Buildable (Rep m)) =>
  (SubExp -> SubExp -> BasicOp) ->
  PrimType ->
  PrimType ->
  m (Lambda (Rep m))
binLambda :: forall (m :: * -> *).
(MonadBuilder m, Buildable (Rep m)) =>
(SubExp -> SubExp -> BasicOp)
-> PrimType -> PrimType -> m (Lambda (Rep m))
binLambda SubExp -> SubExp -> BasicOp
bop PrimType
arg_t PrimType
ret_t = do
  x <- Name -> m VName
forall (m :: * -> *). MonadFreshNames m => Name -> m VName
newVName Name
"x"
  y <- newVName "y"
  body <-
    buildBody_ . fmap (pure . subExpRes) $
      letSubExp "binlam_res" $
        BasicOp $
          bop (Var x) (Var y)
  pure
    Lambda
      { lambdaParams =
          [ Param mempty x (Prim arg_t),
            Param mempty y (Prim arg_t)
          ],
        lambdaReturnType = [Prim ret_t],
        lambdaBody = body
      }

-- | Easily construct a t'Lambda' within a 'MonadBuilder'.  See also
-- 'runLambdaBuilder'.
mkLambda ::
  (MonadBuilder m) =>
  [LParam (Rep m)] ->
  m Result ->
  m (Lambda (Rep m))
mkLambda :: forall (m :: * -> *).
MonadBuilder m =>
[LParam (Rep m)] -> m [SubExpRes] -> m (Lambda (Rep m))
mkLambda [LParam (Rep m)]
params m [SubExpRes]
m = do
  (body, ret) <- m ([SubExpRes], [Type]) -> m (Body (Rep m), [Type])
forall (m :: * -> *) res a.
(MonadBuilder m, IsResult res) =>
m ([res], a) -> m (GBody (Rep m) res, a)
buildBody (m ([SubExpRes], [Type]) -> m (Body (Rep m), [Type]))
-> (m ([SubExpRes], [Type]) -> m ([SubExpRes], [Type]))
-> m ([SubExpRes], [Type])
-> m (Body (Rep m), [Type])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Scope (Rep m) -> m ([SubExpRes], [Type]) -> m ([SubExpRes], [Type])
forall a. Scope (Rep m) -> m a -> m a
forall rep (m :: * -> *) a.
LocalScope rep m =>
Scope rep -> m a -> m a
localScope ([LParam (Rep m)] -> Scope (Rep m)
forall rep dec. (LParamInfo rep ~ dec) => [Param dec] -> Scope rep
scopeOfLParams [LParam (Rep m)]
params) (m ([SubExpRes], [Type]) -> m (Body (Rep m), [Type]))
-> m ([SubExpRes], [Type]) -> m (Body (Rep m), [Type])
forall a b. (a -> b) -> a -> b
$ do
    res <- m [SubExpRes]
m
    ret <- mapM subExpResType res
    pure (res, ret)
  pure $ Lambda params ret body

-- | Slice a full dimension of the given size.
sliceDim :: SubExp -> DimIndex SubExp
sliceDim :: SubExp -> DimIndex SubExp
sliceDim SubExp
d = SubExp -> SubExp -> SubExp -> DimIndex SubExp
forall d. d -> d -> d -> DimIndex d
DimSlice (Int64 -> SubExp
forall v. IsValue v => v -> SubExp
constant (Int64
0 :: Int64)) SubExp
d (Int64 -> SubExp
forall v. IsValue v => v -> SubExp
constant (Int64
1 :: Int64))

-- | @fullSlice t slice@ returns @slice@, but with 'DimSlice's of
-- entire dimensions appended to the full dimensionality of @t@.  This
-- function is used to turn incomplete indexing complete, as required
-- by 'Index'.
fullSlice :: Type -> [DimIndex SubExp] -> Slice SubExp
fullSlice :: Type -> [DimIndex SubExp] -> Slice SubExp
fullSlice Type
t [DimIndex SubExp]
slice =
  [DimIndex SubExp] -> Slice SubExp
forall d. [DimIndex d] -> Slice d
Slice ([DimIndex SubExp] -> Slice SubExp)
-> [DimIndex SubExp] -> Slice SubExp
forall a b. (a -> b) -> a -> b
$ [DimIndex SubExp]
slice [DimIndex SubExp] -> [DimIndex SubExp] -> [DimIndex SubExp]
forall a. [a] -> [a] -> [a]
++ (SubExp -> DimIndex SubExp) -> [SubExp] -> [DimIndex SubExp]
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> DimIndex SubExp
sliceDim (Int -> [SubExp] -> [SubExp]
forall a. Int -> [a] -> [a]
drop ([DimIndex SubExp] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [DimIndex SubExp]
slice) ([SubExp] -> [SubExp]) -> [SubExp] -> [SubExp]
forall a b. (a -> b) -> a -> b
$ Type -> [SubExp]
forall d m. TypeBase (ShapeBase d) m -> [d]
arrayDims Type
t)

-- | @ sliceAt t n slice@ returns @slice@ but with 'DimSlice's of the
-- outer @n@ dimensions prepended, and as many appended as to make it
-- a full slice.  This is a generalisation of 'fullSlice'.
sliceAt :: Type -> Int -> [DimIndex SubExp] -> Slice SubExp
sliceAt :: Type -> Int -> [DimIndex SubExp] -> Slice SubExp
sliceAt Type
t Int
n [DimIndex SubExp]
slice =
  Type -> [DimIndex SubExp] -> Slice SubExp
fullSlice Type
t ([DimIndex SubExp] -> Slice SubExp)
-> [DimIndex SubExp] -> Slice SubExp
forall a b. (a -> b) -> a -> b
$ (SubExp -> DimIndex SubExp) -> [SubExp] -> [DimIndex SubExp]
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> DimIndex SubExp
sliceDim (Int -> [SubExp] -> [SubExp]
forall a. Int -> [a] -> [a]
take Int
n ([SubExp] -> [SubExp]) -> [SubExp] -> [SubExp]
forall a b. (a -> b) -> a -> b
$ Type -> [SubExp]
forall d m. TypeBase (ShapeBase d) m -> [d]
arrayDims Type
t) [DimIndex SubExp] -> [DimIndex SubExp] -> [DimIndex SubExp]
forall a. [a] -> [a] -> [a]
++ [DimIndex SubExp]
slice

-- | Produce a straightforward `Int64` `Iota` of the given length with offset 0
-- and stride 1.
iota64 :: SubExp -> Exp rep
iota64 :: forall rep. SubExp -> Exp rep
iota64 SubExp
n = BasicOp -> Exp rep
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp rep) -> BasicOp -> Exp rep
forall a b. (a -> b) -> a -> b
$ SubExp -> SubExp -> SubExp -> IntType -> BasicOp
Iota SubExp
n (IntType -> Integer -> SubExp
intConst IntType
Int64 Integer
0) (IntType -> Integer -> SubExp
intConst IntType
Int64 Integer
1) IntType
Int64

-- | Like 'fullSlice', but the dimensions are simply numeric.
fullSliceNum :: (Num d) => [d] -> [DimIndex d] -> Slice d
fullSliceNum :: forall d. Num d => [d] -> [DimIndex d] -> Slice d
fullSliceNum [d]
dims [DimIndex d]
slice =
  [DimIndex d] -> Slice d
forall d. [DimIndex d] -> Slice d
Slice ([DimIndex d] -> Slice d) -> [DimIndex d] -> Slice d
forall a b. (a -> b) -> a -> b
$ [DimIndex d]
slice [DimIndex d] -> [DimIndex d] -> [DimIndex d]
forall a. [a] -> [a] -> [a]
++ (d -> DimIndex d) -> [d] -> [DimIndex d]
forall a b. (a -> b) -> [a] -> [b]
map (\d
d -> d -> d -> d -> DimIndex d
forall d. d -> d -> d -> DimIndex d
DimSlice d
0 d
d d
1) (Int -> [d] -> [d]
forall a. Int -> [a] -> [a]
drop ([DimIndex d] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [DimIndex d]
slice) [d]
dims)

-- | Does indexing an array of this shape with this slice produce the array
-- itself? This is a slice that selects every element (which may be fixing unit
-- dimensions) and has zero offset and unit stride.
isIdentitySlice :: Shape -> Slice SubExp -> Bool
isIdentitySlice :: Shape -> Slice SubExp -> Bool
isIdentitySlice Shape
shape Slice SubExp
slice =
  [SubExp] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (Shape -> [SubExp]
forall d. ShapeBase d -> [d]
shapeDims Shape
shape) Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [DimIndex SubExp] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (Slice SubExp -> [DimIndex SubExp]
forall d. Slice d -> [DimIndex d]
unSlice Slice SubExp
slice)
    Bool -> Bool -> Bool
&& [Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
and ((SubExp -> DimIndex SubExp -> Bool)
-> [SubExp] -> [DimIndex SubExp] -> [Bool]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith SubExp -> DimIndex SubExp -> Bool
allOfIt (Shape -> [SubExp]
forall d. ShapeBase d -> [d]
shapeDims Shape
shape) (Slice SubExp -> [DimIndex SubExp]
forall d. Slice d -> [DimIndex d]
unSlice Slice SubExp
slice))
  where
    allOfIt :: SubExp -> DimIndex SubExp -> Bool
allOfIt (Constant PrimValue
v) DimFix {} = PrimValue -> Bool
oneIsh PrimValue
v
    allOfIt SubExp
d (DimSlice (Constant PrimValue
off) SubExp
n (Constant PrimValue
s)) =
      SubExp
d SubExp -> SubExp -> Bool
forall a. Eq a => a -> a -> Bool
== SubExp
n Bool -> Bool -> Bool
&& PrimValue -> Bool
zeroIsh PrimValue
off Bool -> Bool -> Bool
&& PrimValue -> Bool
oneIsh PrimValue
s
    allOfIt SubExp
_ DimIndex SubExp
_ = Bool
False

-- | Conveniently construct a body that contains no bindings.
resultBody :: (Buildable rep) => [SubExp] -> Body rep
resultBody :: forall rep. Buildable rep => [SubExp] -> Body rep
resultBody = Stms rep -> [SubExpRes] -> GBody rep SubExpRes
forall res. IsResult res => Stms rep -> [res] -> GBody rep res
forall rep res.
(Buildable rep, IsResult res) =>
Stms rep -> [res] -> GBody rep res
mkBody Stms rep
forall a. Monoid a => a
mempty ([SubExpRes] -> GBody rep SubExpRes)
-> ([SubExp] -> [SubExpRes]) -> [SubExp] -> GBody rep SubExpRes
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [SubExp] -> [SubExpRes]
subExpsRes

-- | Conveniently construct a body that contains no bindings - but
-- this time, monadically!
resultBodyM :: (MonadBuilder m) => [SubExp] -> m (Body (Rep m))
resultBodyM :: forall (m :: * -> *).
MonadBuilder m =>
[SubExp] -> m (Body (Rep m))
resultBodyM = Stms (Rep m) -> [SubExpRes] -> m (Body (Rep m))
forall res.
IsResult res =>
Stms (Rep m) -> [res] -> m (GBody (Rep m) res)
forall (m :: * -> *) res.
(MonadBuilder m, IsResult res) =>
Stms (Rep m) -> [res] -> m (GBody (Rep m) res)
mkBodyM Stms (Rep m)
forall a. Monoid a => a
mempty ([SubExpRes] -> m (Body (Rep m)))
-> ([SubExp] -> [SubExpRes]) -> [SubExp] -> m (Body (Rep m))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [SubExp] -> [SubExpRes]
subExpsRes

-- | Evaluate the action, producing a body, then wrap it in all the
-- bindings it created using 'addStm'.
insertStmsM ::
  (MonadBuilder m) =>
  m (Body (Rep m)) ->
  m (Body (Rep m))
insertStmsM :: forall (m :: * -> *).
MonadBuilder m =>
m (Body (Rep m)) -> m (Body (Rep m))
insertStmsM m (Body (Rep m))
m = do
  (Body _ stms res, otherstms) <- m (Body (Rep m)) -> m (Body (Rep m), Stms (Rep m))
forall a. m a -> m (a, Stms (Rep m))
forall (m :: * -> *) a.
MonadBuilder m =>
m a -> m (a, Stms (Rep m))
collectStms m (Body (Rep m))
m
  mkBodyM (otherstms <> stms) res

-- | Evaluate an action that produces a 'Result' and an auxiliary
-- value, then return the body constructed from the 'Result' and any
-- statements added during the action, along the auxiliary value.
buildBody ::
  (MonadBuilder m, IsResult res) =>
  m ([res], a) ->
  m (GBody (Rep m) res, a)
buildBody :: forall (m :: * -> *) res a.
(MonadBuilder m, IsResult res) =>
m ([res], a) -> m (GBody (Rep m) res, a)
buildBody m ([res], a)
m = do
  ((res, v), stms) <- m ([res], a) -> m (([res], a), Stms (Rep m))
forall a. m a -> m (a, Stms (Rep m))
forall (m :: * -> *) a.
MonadBuilder m =>
m a -> m (a, Stms (Rep m))
collectStms m ([res], a)
m
  body <- mkBodyM stms res
  pure (body, v)

-- | As 'buildBody', but there is no auxiliary value.
buildBody_ ::
  (MonadBuilder m, IsResult res) =>
  m [res] ->
  m (GBody (Rep m) res)
buildBody_ :: forall (m :: * -> *) res.
(MonadBuilder m, IsResult res) =>
m [res] -> m (GBody (Rep m) res)
buildBody_ m [res]
m = (GBody (Rep m) res, ()) -> GBody (Rep m) res
forall a b. (a, b) -> a
fst ((GBody (Rep m) res, ()) -> GBody (Rep m) res)
-> m (GBody (Rep m) res, ()) -> m (GBody (Rep m) res)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> m ([res], ()) -> m (GBody (Rep m) res, ())
forall (m :: * -> *) res a.
(MonadBuilder m, IsResult res) =>
m ([res], a) -> m (GBody (Rep m) res, a)
buildBody ((,()) ([res] -> ([res], ())) -> m [res] -> m ([res], ())
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> m [res]
m)

-- | Change that result where evaluation of the body would stop.  Also
-- change type annotations at branches.
mapResult ::
  (Buildable rep) =>
  (Result -> Body rep) ->
  Body rep ->
  Body rep
mapResult :: forall rep.
Buildable rep =>
([SubExpRes] -> Body rep) -> Body rep -> Body rep
mapResult [SubExpRes] -> Body rep
f (Body BodyDec rep
_ Stms rep
stms [SubExpRes]
res) =
  let Body BodyDec rep
_ Stms rep
stms2 [SubExpRes]
newres = [SubExpRes] -> Body rep
f [SubExpRes]
res
   in Stms rep -> [SubExpRes] -> Body rep
forall res. IsResult res => Stms rep -> [res] -> GBody rep res
forall rep res.
(Buildable rep, IsResult res) =>
Stms rep -> [res] -> GBody rep res
mkBody (Stms rep
stms Stms rep -> Stms rep -> Stms rep
forall a. Semigroup a => a -> a -> a
<> Stms rep
stms2) [SubExpRes]
newres

-- | Instantiate all existential parts dimensions of the given
-- type, using a monadic action to create the necessary t'SubExp's.
-- You should call this function within some monad that allows you to
-- collect the actions performed (say, 'State').
instantiateShapes ::
  (Monad m) =>
  (Int -> m SubExp) ->
  [TypeBase ExtShape o] ->
  m [TypeBase Shape o]
instantiateShapes :: forall (m :: * -> *) o.
Monad m =>
(Int -> m SubExp)
-> [TypeBase (ShapeBase (Ext SubExp)) o] -> m [TypeBase Shape o]
instantiateShapes Int -> m SubExp
f [TypeBase (ShapeBase (Ext SubExp)) o]
ts = StateT (Map Int SubExp) m [TypeBase Shape o]
-> Map Int SubExp -> m [TypeBase Shape o]
forall (m :: * -> *) s a. Monad m => StateT s m a -> s -> m a
evalStateT ((TypeBase (ShapeBase (Ext SubExp)) o
 -> StateT (Map Int SubExp) m (TypeBase Shape o))
-> [TypeBase (ShapeBase (Ext SubExp)) o]
-> StateT (Map Int SubExp) m [TypeBase Shape o]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM TypeBase (ShapeBase (Ext SubExp)) o
-> StateT (Map Int SubExp) m (TypeBase Shape o)
instantiate [TypeBase (ShapeBase (Ext SubExp)) o]
ts) Map Int SubExp
forall k a. Map k a
M.empty
  where
    instantiate :: TypeBase (ShapeBase (Ext SubExp)) o
-> StateT (Map Int SubExp) m (TypeBase Shape o)
instantiate TypeBase (ShapeBase (Ext SubExp)) o
t = do
      shape <- (Ext SubExp -> StateT (Map Int SubExp) m SubExp)
-> [Ext SubExp] -> StateT (Map Int SubExp) m [SubExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM Ext SubExp -> StateT (Map Int SubExp) m SubExp
instantiate' ([Ext SubExp] -> StateT (Map Int SubExp) m [SubExp])
-> [Ext SubExp] -> StateT (Map Int SubExp) m [SubExp]
forall a b. (a -> b) -> a -> b
$ ShapeBase (Ext SubExp) -> [Ext SubExp]
forall d. ShapeBase d -> [d]
shapeDims (ShapeBase (Ext SubExp) -> [Ext SubExp])
-> ShapeBase (Ext SubExp) -> [Ext SubExp]
forall a b. (a -> b) -> a -> b
$ TypeBase (ShapeBase (Ext SubExp)) o -> ShapeBase (Ext SubExp)
forall shape o. ArrayShape shape => TypeBase shape o -> shape
arrayShape TypeBase (ShapeBase (Ext SubExp)) o
t
      pure $ t `setArrayShape` Shape shape
    instantiate' :: Ext SubExp -> StateT (Map Int SubExp) m SubExp
instantiate' (Ext Int
x) = do
      m <- StateT (Map Int SubExp) m (Map Int SubExp)
forall s (m :: * -> *). MonadState s m => m s
get
      case M.lookup x m of
        Just SubExp
se -> SubExp -> StateT (Map Int SubExp) m SubExp
forall a. a -> StateT (Map Int SubExp) m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SubExp
se
        Maybe SubExp
Nothing -> do
          se <- m SubExp -> StateT (Map Int SubExp) m SubExp
forall (m :: * -> *) a.
Monad m =>
m a -> StateT (Map Int SubExp) m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (m SubExp -> StateT (Map Int SubExp) m SubExp)
-> m SubExp -> StateT (Map Int SubExp) m SubExp
forall a b. (a -> b) -> a -> b
$ Int -> m SubExp
f Int
x
          put $ M.insert x se m
          pure se
    instantiate' (Free SubExp
se) = SubExp -> StateT (Map Int SubExp) m SubExp
forall a. a -> StateT (Map Int SubExp) m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SubExp
se

-- | Like 'instantiateShapes', but obtains names from the provided
-- list.  If an 'Ext' is out of bounds of this list, the function
-- fails with 'error'.
instantiateShapes' :: [VName] -> [TypeBase ExtShape o] -> [TypeBase Shape o]
instantiateShapes' :: forall o.
[VName]
-> [TypeBase (ShapeBase (Ext SubExp)) o] -> [TypeBase Shape o]
instantiateShapes' [VName]
names [TypeBase (ShapeBase (Ext SubExp)) o]
ts =
  -- Carefully ensure that the order of idents we produce corresponds
  -- to their existential index.
  Identity [TypeBase Shape o] -> [TypeBase Shape o]
forall a. Identity a -> a
runIdentity (Identity [TypeBase Shape o] -> [TypeBase Shape o])
-> Identity [TypeBase Shape o] -> [TypeBase Shape o]
forall a b. (a -> b) -> a -> b
$ (Int -> Identity SubExp)
-> [TypeBase (ShapeBase (Ext SubExp)) o]
-> Identity [TypeBase Shape o]
forall (m :: * -> *) o.
Monad m =>
(Int -> m SubExp)
-> [TypeBase (ShapeBase (Ext SubExp)) o] -> m [TypeBase Shape o]
instantiateShapes Int -> Identity SubExp
instantiate [TypeBase (ShapeBase (Ext SubExp)) o]
ts
  where
    instantiate :: Int -> Identity SubExp
instantiate Int
x =
      case Int -> [VName] -> Maybe VName
forall int a. Integral int => int -> [a] -> Maybe a
maybeNth Int
x [VName]
names of
        Maybe VName
Nothing -> [Char] -> Identity SubExp
forall a. HasCallStack => [Char] -> a
error ([Char] -> Identity SubExp) -> [Char] -> Identity SubExp
forall a b. (a -> b) -> a -> b
$ [Char]
"instantiateShapes': " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [VName] -> [Char]
forall a. Pretty a => a -> [Char]
prettyString [VName]
names [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
", " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show Int
x
        Just VName
name -> SubExp -> Identity SubExp
forall a. a -> Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SubExp -> Identity SubExp) -> SubExp -> Identity SubExp
forall a b. (a -> b) -> a -> b
$ VName -> SubExp
Var VName
name

-- | Remove existentials by imposing sizes from another type where
-- needed.
removeExistentials :: ExtType -> Type -> Type
removeExistentials :: ExtType -> Type -> Type
removeExistentials ExtType
t1 Type
t2 =
  ExtType
t1
    ExtType -> [SubExp] -> Type
forall oldshape o.
TypeBase oldshape o -> [SubExp] -> TypeBase Shape o
`setArrayDims` (Ext SubExp -> SubExp -> SubExp)
-> [Ext SubExp] -> [SubExp] -> [SubExp]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith
      Ext SubExp -> SubExp -> SubExp
forall {p}. Ext p -> p -> p
nonExistential
      (ShapeBase (Ext SubExp) -> [Ext SubExp]
forall d. ShapeBase d -> [d]
shapeDims (ShapeBase (Ext SubExp) -> [Ext SubExp])
-> ShapeBase (Ext SubExp) -> [Ext SubExp]
forall a b. (a -> b) -> a -> b
$ ExtType -> ShapeBase (Ext SubExp)
forall shape o. ArrayShape shape => TypeBase shape o -> shape
arrayShape ExtType
t1)
      (Type -> [SubExp]
forall d m. TypeBase (ShapeBase d) m -> [d]
arrayDims Type
t2)
  where
    nonExistential :: Ext p -> p -> p
nonExistential (Ext Int
_) p
dim = p
dim
    nonExistential (Free p
dim) p
_ = p
dim

-- | Can be used as the definition of 'mkLetNames' for a 'Buildable'
-- instance for simple representations.
simpleMkLetNames ::
  ( ExpDec rep ~ (),
    LetDec rep ~ Type,
    MonadFreshNames m,
    TypedOp (OpC rep),
    HasScope rep m
  ) =>
  [VName] ->
  Exp rep ->
  m (Stm rep)
simpleMkLetNames :: forall rep (m :: * -> *).
(ExpDec rep ~ (), LetDec rep ~ Type, MonadFreshNames m,
 TypedOp (OpC rep), HasScope rep m) =>
[VName] -> Exp rep -> m (Stm rep)
simpleMkLetNames [VName]
names Exp rep
e = do
  et <- Exp rep -> m [ExtType]
forall rep (m :: * -> *).
(HasScope rep m, TypedOp (OpC rep)) =>
Exp rep -> m [ExtType]
expExtType Exp rep
e
  let ts = [VName] -> [ExtType] -> [Type]
forall o.
[VName]
-> [TypeBase (ShapeBase (Ext SubExp)) o] -> [TypeBase Shape o]
instantiateShapes' [VName]
names [ExtType]
et
  pure $ Let (Pat $ zipWith PatElem names ts) (defAux ()) e

-- | Instances of this class can be converted to Futhark expressions
-- within a 'MonadBuilder'.
class ToExp a where
  toExp :: (MonadBuilder m) => a -> m (Exp (Rep m))

instance ToExp SubExp where
  toExp :: forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
toExp = Exp (Rep m) -> m (Exp (Rep m))
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp (Rep m) -> m (Exp (Rep m)))
-> (SubExp -> Exp (Rep m)) -> SubExp -> m (Exp (Rep m))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BasicOp -> Exp (Rep m)
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp (Rep m))
-> (SubExp -> BasicOp) -> SubExp -> Exp (Rep m)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SubExp -> BasicOp
SubExp

instance ToExp VName where
  toExp :: forall (m :: * -> *). MonadBuilder m => VName -> m (Exp (Rep m))
toExp = Exp (Rep m) -> m (Exp (Rep m))
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp (Rep m) -> m (Exp (Rep m)))
-> (VName -> Exp (Rep m)) -> VName -> m (Exp (Rep m))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BasicOp -> Exp (Rep m)
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp (Rep m))
-> (VName -> BasicOp) -> VName -> Exp (Rep m)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SubExp -> BasicOp
SubExp (SubExp -> BasicOp) -> (VName -> SubExp) -> VName -> BasicOp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> SubExp
Var

-- | A convenient composition of 'letSubExp' and 'toExp'.
toSubExp :: (MonadBuilder m, ToExp a) => Name -> a -> m SubExp
toSubExp :: forall (m :: * -> *) a.
(MonadBuilder m, ToExp a) =>
Name -> a -> m SubExp
toSubExp Name
s a
e = Name -> Exp (Rep m) -> m SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
s (Exp (Rep m) -> m SubExp) -> m (Exp (Rep m)) -> m SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< a -> m (Exp (Rep m))
forall a (m :: * -> *).
(ToExp a, MonadBuilder m) =>
a -> m (Exp (Rep m))
forall (m :: * -> *). MonadBuilder m => a -> m (Exp (Rep m))
toExp a
e