{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE Strict #-}
{-# LANGUAGE TypeFamilies #-}

-- | Conversion of a monomorphic, first-order, defunctorised source
-- program to a core Futhark program.
module Futhark.Internalise.Exps (transformProg) where

import Control.Monad
import Control.Monad.Reader
import Data.Bifunctor
import Data.Foldable (toList)
import Data.List (elemIndex, find, intercalate, intersperse, transpose)
import Data.List.NonEmpty (NonEmpty (..))
import Data.List.NonEmpty qualified as NE
import Data.Map.Strict qualified as M
import Data.Set qualified as S
import Data.Text qualified as T
import Futhark.Error (compilerLimitation)
import Futhark.IR.SOACS as I hiding (stmPat)
import Futhark.Internalise.AccurateSizes
import Futhark.Internalise.Bindings
import Futhark.Internalise.Entry
import Futhark.Internalise.Monad as I
import Futhark.Internalise.TypesValues
import Futhark.Transform.Rename as I
import Futhark.Util (lookupWithIndex, splitAt3)
import Futhark.Util.Pretty (align, docText, pretty)
import Language.Futhark as E hiding (TypeArg)
import Language.Futhark.TypeChecker.Types qualified as E

-- | Convert a program in source Futhark to a program in the Futhark
-- core language.
transformProg ::
  (MonadFreshNames m) =>
  Bool ->
  Bool ->
  VisibleTypes ->
  [E.ValBind] ->
  m (I.Prog SOACS)
transformProg :: forall (m :: * -> *).
MonadFreshNames m =>
Bool -> Bool -> VisibleTypes -> [ValBind] -> m (Prog SOACS)
transformProg Bool
always_safe Bool
strip_provenance VisibleTypes
types [ValBind]
vbinds = do
  (opaques, consts, funs) <-
    Bool
-> Bool
-> InternaliseM ()
-> m (OpaqueTypes, Stms SOACS, [FunDef SOACS])
forall (m :: * -> *).
MonadFreshNames m =>
Bool
-> Bool
-> InternaliseM ()
-> m (OpaqueTypes, Stms SOACS, [FunDef SOACS])
runInternaliseM Bool
always_safe Bool
strip_provenance (InternaliseM () -> m (OpaqueTypes, Stms SOACS, [FunDef SOACS]))
-> InternaliseM () -> m (OpaqueTypes, Stms SOACS, [FunDef SOACS])
forall a b. (a -> b) -> a -> b
$
      VisibleTypes -> [ValBind] -> InternaliseM ()
internaliseValBinds VisibleTypes
types [ValBind]
vbinds
  I.renameProg $ I.Prog opaques consts funs

internaliseValBinds :: VisibleTypes -> [E.ValBind] -> InternaliseM ()
internaliseValBinds :: VisibleTypes -> [ValBind] -> InternaliseM ()
internaliseValBinds VisibleTypes
types [ValBind]
vbinds = do
  -- Internalise every function's parameters and register its calling
  -- information *before* internalising any body, so that references between
  -- functions resolve regardless of definition order.
  headers <- [ValBind]
-> (ValBind -> InternaliseM (ValBind, FunParams))
-> InternaliseM [(ValBind, FunParams)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [ValBind]
vbinds ((ValBind -> InternaliseM (ValBind, FunParams))
 -> InternaliseM [(ValBind, FunParams)])
-> (ValBind -> InternaliseM (ValBind, FunParams))
-> InternaliseM [(ValBind, FunParams)]
forall a b. (a -> b) -> a -> b
$ \ValBind
vb -> do
    let E.ValBind Maybe (Info EntryPoint)
_ VName
fname SrcLoc
_ Maybe (TypeExp Exp VName)
_ (Info ResRetType
rettype) [TypeParamBase VName]
tparams [PatBase Info VName ParamType]
params Exp
_ Maybe DocComment
_ [AttrInfo VName]
_ SrcLoc
_ = ValBind
vb
    fps <- [TypeParamBase VName]
-> [PatBase Info VName ParamType] -> InternaliseM FunParams
internaliseFParams [TypeParamBase VName]
tparams [PatBase Info VName ParamType]
params
    let (_, _, _, info) = funHeader (funShapeParams fps) (funValueParams fps) rettype
    unless (null (funValueParams fps)) $ addFunInfo fname info
    pure (vb, fps)
  forM_ headers $ \(ValBind
vb, FunParams
fps) ->
    FunParams -> InternaliseM () -> InternaliseM ()
forall a. FunParams -> InternaliseM a -> InternaliseM a
bindFParams FunParams
fps (InternaliseM () -> InternaliseM ())
-> InternaliseM () -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$ VisibleTypes -> ValBind -> FunParams -> InternaliseM ()
internaliseValBindBody VisibleTypes
types ValBind
vb FunParams
fps

-- | A function's internalised parameters, return type, full internalised
-- return type, and calling information, all computed from its signature alone
-- (no body required).
funHeader ::
  [I.FParam I.SOACS] ->
  [[Tree (I.FParam I.SOACS)]] ->
  E.ResRetType ->
  ([Tree (I.FParam I.SOACS)], [I.ExtType], [(I.ExtType, RetAls)], FunInfo)
funHeader :: [FParam SOACS]
-> [[Tree (FParam SOACS)]]
-> ResRetType
-> ([Tree (FParam SOACS)], [TypeBase ExtShape NoMode],
    [(TypeBase ExtShape NoMode, RetAls)], FunInfo)
funHeader [FParam SOACS]
shapeparams [[Tree (FParam SOACS)]]
params' ResRetType
rettype =
  ([Tree (Param DeclType)]
[Tree (FParam SOACS)]
all_params, [TypeBase ExtShape NoMode]
rettype', [(TypeBase ExtShape NoMode, RetAls)]
fun_rettype, ([VName], [DeclType], [Param DeclType],
 [(SubExp, TypeBase Shape NoMode)]
 -> Maybe [(TypeBase ExtShape NoMode, RetAls)])
FunInfo
info)
  where
    shapenames :: [VName]
shapenames = (Param DeclType -> VName) -> [Param DeclType] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param DeclType -> VName
forall dec. Param dec -> VName
I.paramName [Param DeclType]
[FParam SOACS]
shapeparams
    all_params :: [Tree (Param DeclType)]
all_params = (Param DeclType -> Tree (Param DeclType))
-> [Param DeclType] -> [Tree (Param DeclType)]
forall a b. (a -> b) -> [a] -> [b]
map Param DeclType -> Tree (Param DeclType)
forall a. a -> Free [] a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [Param DeclType]
[FParam SOACS]
shapeparams [Tree (Param DeclType)]
-> [Tree (Param DeclType)] -> [Tree (Param DeclType)]
forall a. [a] -> [a] -> [a]
++ [[Tree (Param DeclType)]] -> [Tree (Param DeclType)]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[Tree (Param DeclType)]]
[[Tree (FParam SOACS)]]
params'
    zeroExts :: [TypeBase ExtShape o] -> [TypeBase ExtShape o]
zeroExts [TypeBase ExtShape o]
ts = [TypeBase ExtShape o]
-> [TypeBase ExtShape o] -> [TypeBase ExtShape o]
forall o.
[TypeBase ExtShape o]
-> [TypeBase ExtShape o] -> [TypeBase ExtShape o]
generaliseExtTypes [TypeBase ExtShape o]
ts [TypeBase ExtShape o]
ts
    ([TypeBase ExtShape NoMode]
rettype', [RetAls]
retals) =
      ([TypeBase ExtShape NoMode] -> [TypeBase ExtShape NoMode])
-> ([TypeBase ExtShape NoMode], [RetAls])
-> ([TypeBase ExtShape NoMode], [RetAls])
forall a b c. (a -> b) -> (a, c) -> (b, c)
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first [TypeBase ExtShape NoMode] -> [TypeBase ExtShape NoMode]
forall {o}. [TypeBase ExtShape o] -> [TypeBase ExtShape o]
zeroExts (([TypeBase ExtShape NoMode], [RetAls])
 -> ([TypeBase ExtShape NoMode], [RetAls]))
-> ([(TypeBase ExtShape NoMode, RetAls)]
    -> ([TypeBase ExtShape NoMode], [RetAls]))
-> [(TypeBase ExtShape NoMode, RetAls)]
-> ([TypeBase ExtShape NoMode], [RetAls])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(TypeBase ExtShape NoMode, RetAls)]
-> ([TypeBase ExtShape NoMode], [RetAls])
forall a b. [(a, b)] -> ([a], [b])
unzip ([(TypeBase ExtShape NoMode, RetAls)]
 -> ([TypeBase ExtShape NoMode], [RetAls]))
-> [(TypeBase ExtShape NoMode, RetAls)]
-> ([TypeBase ExtShape NoMode], [RetAls])
forall a b. (a -> b) -> a -> b
$
        [Tree DeclType]
-> ResRetType -> [(TypeBase ExtShape NoMode, RetAls)]
internaliseReturnType ((Tree (Param DeclType) -> Tree DeclType)
-> [Tree (Param DeclType)] -> [Tree DeclType]
forall a b. (a -> b) -> [a] -> [b]
map ((Param DeclType -> DeclType)
-> Tree (Param DeclType) -> Tree DeclType
forall a b. (a -> b) -> Free [] a -> Free [] b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Param DeclType -> DeclType
forall dec. DeclTyped dec => Param dec -> DeclType
paramDeclType) [Tree (Param DeclType)]
all_params) ResRetType
rettype
    num_ctx :: Int
num_ctx = Set Int -> Int
forall a. Set a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([TypeBase ExtShape NoMode] -> Set Int
forall o. [TypeBase ExtShape o] -> Set Int
shapeContext [TypeBase ExtShape NoMode]
rettype')
    fun_rettype :: [(TypeBase ExtShape NoMode, RetAls)]
fun_rettype =
      Int
-> (TypeBase ExtShape NoMode, RetAls)
-> [(TypeBase ExtShape NoMode, RetAls)]
forall a. Int -> a -> [a]
replicate Int
num_ctx (PrimType -> TypeBase ExtShape NoMode
forall shape o. PrimType -> TypeBase shape o
I.Prim PrimType
int64, RetAls
forall a. Monoid a => a
mempty)
        [(TypeBase ExtShape NoMode, RetAls)]
-> [(TypeBase ExtShape NoMode, RetAls)]
-> [(TypeBase ExtShape NoMode, RetAls)]
forall a. [a] -> [a] -> [a]
++ [TypeBase ExtShape NoMode]
-> [RetAls] -> [(TypeBase ExtShape NoMode, RetAls)]
forall a b. [a] -> [b] -> [(a, b)]
zip [TypeBase ExtShape NoMode]
rettype' ((RetAls -> RetAls) -> [RetAls] -> [RetAls]
forall a b. (a -> b) -> [a] -> [b]
map (Int -> RetAls -> RetAls
shiftRetAls Int
num_ctx) [RetAls]
retals)
    info :: ([VName], [DeclType], [Param DeclType],
 [(SubExp, TypeBase Shape NoMode)]
 -> Maybe [(TypeBase ExtShape NoMode, RetAls)])
info =
      ( [VName]
shapenames,
        (Param DeclType -> DeclType) -> [Param DeclType] -> [DeclType]
forall a b. (a -> b) -> [a] -> [b]
map Param DeclType -> DeclType
forall t. DeclTyped t => t -> DeclType
declTypeOf ([Param DeclType] -> [DeclType]) -> [Param DeclType] -> [DeclType]
forall a b. (a -> b) -> a -> b
$ ([Tree (Param DeclType)] -> [Param DeclType])
-> [[Tree (Param DeclType)]] -> [Param DeclType]
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap ((Tree (Param DeclType) -> [Param DeclType])
-> [Tree (Param DeclType)] -> [Param DeclType]
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Tree (Param DeclType) -> [Param DeclType]
forall a. Free [] a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList) [[Tree (Param DeclType)]]
[[Tree (FParam SOACS)]]
params',
        (Tree (Param DeclType) -> [Param DeclType])
-> [Tree (Param DeclType)] -> [Param DeclType]
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Tree (Param DeclType) -> [Param DeclType]
forall a. Free [] a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList [Tree (Param DeclType)]
all_params,
        ([TypeBase ExtShape NoMode]
 -> [(TypeBase ExtShape NoMode, RetAls)])
-> Maybe [TypeBase ExtShape NoMode]
-> Maybe [(TypeBase ExtShape NoMode, RetAls)]
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ([TypeBase ExtShape NoMode]
-> [RetAls] -> [(TypeBase ExtShape NoMode, RetAls)]
forall a b. [a] -> [b] -> [(a, b)]
`zip` ((TypeBase ExtShape NoMode, RetAls) -> RetAls)
-> [(TypeBase ExtShape NoMode, RetAls)] -> [RetAls]
forall a b. (a -> b) -> [a] -> [b]
map (TypeBase ExtShape NoMode, RetAls) -> RetAls
forall a b. (a, b) -> b
snd [(TypeBase ExtShape NoMode, RetAls)]
fun_rettype)
          (Maybe [TypeBase ExtShape NoMode]
 -> Maybe [(TypeBase ExtShape NoMode, RetAls)])
-> ([(SubExp, TypeBase Shape NoMode)]
    -> Maybe [TypeBase ExtShape NoMode])
-> [(SubExp, TypeBase Shape NoMode)]
-> Maybe [(TypeBase ExtShape NoMode, RetAls)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [TypeBase ExtShape NoMode]
-> [Param DeclType]
-> [(SubExp, TypeBase Shape NoMode)]
-> Maybe [TypeBase ExtShape NoMode]
forall dec.
Typed dec =>
[TypeBase ExtShape NoMode]
-> [Param dec]
-> [(SubExp, TypeBase Shape NoMode)]
-> Maybe [TypeBase ExtShape NoMode]
forall rt dec.
(IsRetType rt, Typed dec) =>
[rt]
-> [Param dec] -> [(SubExp, TypeBase Shape NoMode)] -> Maybe [rt]
applyRetType (((TypeBase ExtShape NoMode, RetAls) -> TypeBase ExtShape NoMode)
-> [(TypeBase ExtShape NoMode, RetAls)]
-> [TypeBase ExtShape NoMode]
forall a b. (a -> b) -> [a] -> [b]
map (TypeBase ExtShape NoMode, RetAls) -> TypeBase ExtShape NoMode
forall a b. (a, b) -> a
fst [(TypeBase ExtShape NoMode, RetAls)]
fun_rettype) ((Tree (Param DeclType) -> [Param DeclType])
-> [Tree (Param DeclType)] -> [Param DeclType]
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Tree (Param DeclType) -> [Param DeclType]
forall a. Free [] a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList [Tree (Param DeclType)]
all_params)
      )

internaliseFunName :: VName -> Name
internaliseFunName :: VName -> Name
internaliseFunName = Text -> Name
nameFromText (Text -> Name) -> (VName -> Text) -> VName -> Name
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> Text
forall a. Pretty a => a -> Text
prettyText

shiftRetAls :: Int -> RetAls -> RetAls
shiftRetAls :: Int -> RetAls -> RetAls
shiftRetAls Int
d (RetAls [Int]
pals [Int]
rals) = [Int] -> [Int] -> RetAls
RetAls [Int]
pals ([Int] -> RetAls) -> [Int] -> RetAls
forall a b. (a -> b) -> a -> b
$ (Int -> Int) -> [Int] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
d) [Int]
rals

-- | Internalise a function's body, given its already-bound parameters (see
-- 'internaliseFParams' and 'bindFParams'). Its calling information has already
-- been registered by 'internaliseValBinds'.
internaliseValBindBody :: VisibleTypes -> E.ValBind -> FunParams -> InternaliseM ()
internaliseValBindBody :: VisibleTypes -> ValBind -> FunParams -> InternaliseM ()
internaliseValBindBody VisibleTypes
types fb :: ValBind
fb@(E.ValBind Maybe (Info EntryPoint)
entry VName
fname SrcLoc
_ Maybe (TypeExp Exp VName)
_ (Info ResRetType
rettype) [TypeParamBase VName]
_ [PatBase Info VName ParamType]
_ Exp
body Maybe DocComment
_ [AttrInfo VName]
attrs SrcLoc
_) FunParams
fps = do
  let ([Tree (FParam SOACS)]
all_params, [TypeBase ExtShape NoMode]
rettype', [(TypeBase ExtShape NoMode, RetAls)]
fun_rettype, FunInfo
_) =
        [FParam SOACS]
-> [[Tree (FParam SOACS)]]
-> ResRetType
-> ([Tree (FParam SOACS)], [TypeBase ExtShape NoMode],
    [(TypeBase ExtShape NoMode, RetAls)], FunInfo)
funHeader (FunParams -> [FParam SOACS]
funShapeParams FunParams
fps) (FunParams -> [[Tree (FParam SOACS)]]
funValueParams FunParams
fps) ResRetType
rettype
      params' :: [[Tree (FParam SOACS)]]
params' = FunParams -> [[Tree (FParam SOACS)]]
funValueParams FunParams
fps
      msg :: ErrorMsg SubExp
msg =
        [ErrorMsgPart SubExp] -> ErrorMsg SubExp
forall a. [ErrorMsgPart a] -> ErrorMsg a
errorMsg
          [ ErrorMsgPart SubExp
"Internal runtime error.\n",
            ErrorMsgPart SubExp
"Return value of ",
            Text -> ErrorMsgPart SubExp
forall a. Text -> ErrorMsgPart a
ErrorString (VName -> Text
forall a. Pretty a => a -> Text
prettyText VName
fname),
            ErrorMsgPart SubExp
" does not match type shape.\n",
            ErrorMsgPart SubExp
"This is a bug in the Futhark compiler. Please report this:\n",
            ErrorMsgPart SubExp
"  https://github.com/diku-dk/futhark/issues"
          ]

  body' <- InternaliseM [SubExpRes]
-> InternaliseM (GBody (Rep InternaliseM) SubExpRes)
forall (m :: * -> *) res.
(MonadBuilder m, IsResult res) =>
m [res] -> m (GBody (Rep m) res)
buildBody_ (InternaliseM [SubExpRes]
 -> InternaliseM (GBody (Rep InternaliseM) SubExpRes))
-> InternaliseM [SubExpRes]
-> InternaliseM (GBody (Rep InternaliseM) SubExpRes)
forall a b. (a -> b) -> a -> b
$ do
    body_res <- Name -> Exp -> InternaliseM [SubExp]
internaliseExp (VName -> Name
baseName VName
fname Name -> Name -> Name
forall a. Semigroup a => a -> a -> a
<> Name
"_res") Exp
body

    when (null params') $
      bindExtSizes (E.AppRes (E.toStruct $ E.retType rettype) (E.retDims rettype)) body_res

    ensureResultExtShape msg rettype' $ subExpsRes body_res

  attrs' <- internaliseAttrs attrs

  let fd =
        Maybe EntryPoint
-> Attrs
-> Name
-> [(RetType SOACS, RetAls)]
-> [FParam SOACS]
-> Body SOACS
-> FunDef SOACS
forall rep.
Maybe EntryPoint
-> Attrs
-> Name
-> [(RetType rep, RetAls)]
-> [FParam rep]
-> Body rep
-> FunDef rep
I.FunDef
          Maybe EntryPoint
forall a. Maybe a
Nothing
          Attrs
attrs'
          (VName -> Name
internaliseFunName VName
fname)
          [(TypeBase ExtShape NoMode, RetAls)]
[(RetType SOACS, RetAls)]
fun_rettype
          ((Tree (Param DeclType) -> [Param DeclType])
-> [Tree (Param DeclType)] -> [Param DeclType]
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Tree (Param DeclType) -> [Param DeclType]
forall a. Free [] a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList [Tree (Param DeclType)]
[Tree (FParam SOACS)]
all_params)
          Body SOACS
body'

  if null params'
    then addConstant fname fd
    else addFunDef fd

  case entry of
    Just (Info EntryPoint
entry') -> VisibleTypes -> EntryPoint -> ValBind -> InternaliseM ()
generateEntryPoint VisibleTypes
types EntryPoint
entry' ValBind
fb
    Maybe (Info EntryPoint)
Nothing -> () -> InternaliseM ()
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

generateEntryPoint :: VisibleTypes -> E.EntryPoint -> E.ValBind -> InternaliseM ()
generateEntryPoint :: VisibleTypes -> EntryPoint -> ValBind -> InternaliseM ()
generateEntryPoint VisibleTypes
types (E.EntryPoint [EntryParam]
e_params EntryType
e_rettype Maybe Text
doc) ValBind
vb = do
  let (E.ValBind Maybe (Info EntryPoint)
_ VName
ofname SrcLoc
_ Maybe (TypeExp Exp VName)
_ (Info ResRetType
rettype) [TypeParamBase VName]
tparams [PatBase Info VName ParamType]
params Exp
_ Maybe DocComment
_ [AttrInfo VName]
attrs SrcLoc
_) = ValBind
vb
  [TypeParamBase VName]
-> [PatBase Info VName ParamType]
-> ([FParam SOACS] -> [[Tree (FParam SOACS)]] -> InternaliseM ())
-> InternaliseM ()
forall a.
[TypeParamBase VName]
-> [PatBase Info VName ParamType]
-> ([FParam SOACS] -> [[Tree (FParam SOACS)]] -> InternaliseM a)
-> InternaliseM a
bindingFParams [TypeParamBase VName]
tparams [PatBase Info VName ParamType]
params (([FParam SOACS] -> [[Tree (FParam SOACS)]] -> InternaliseM ())
 -> InternaliseM ())
-> ([FParam SOACS] -> [[Tree (FParam SOACS)]] -> InternaliseM ())
-> InternaliseM ()
forall a b. (a -> b) -> a -> b
$ \[FParam SOACS]
shapeparams [[Tree (FParam SOACS)]]
params' -> do
    let all_params :: [Tree (Param DeclType)]
all_params = (Param DeclType -> Tree (Param DeclType))
-> [Param DeclType] -> [Tree (Param DeclType)]
forall a b. (a -> b) -> [a] -> [b]
map Param DeclType -> Tree (Param DeclType)
forall a. a -> Free [] a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [Param DeclType]
[FParam SOACS]
shapeparams [Tree (Param DeclType)]
-> [Tree (Param DeclType)] -> [Tree (Param DeclType)]
forall a. [a] -> [a] -> [a]
++ [[Tree (Param DeclType)]] -> [Tree (Param DeclType)]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[Tree (Param DeclType)]]
[[Tree (FParam SOACS)]]
params'
        ([[TypeBase ExtShape Freshness]]
entry_rettype, [[RetAls]]
retals) =
          [([TypeBase ExtShape Freshness], [RetAls])]
-> ([[TypeBase ExtShape Freshness]], [[RetAls]])
forall a b. [(a, b)] -> ([a], [b])
unzip ([([TypeBase ExtShape Freshness], [RetAls])]
 -> ([[TypeBase ExtShape Freshness]], [[RetAls]]))
-> [([TypeBase ExtShape Freshness], [RetAls])]
-> ([[TypeBase ExtShape Freshness]], [[RetAls]])
forall a b. (a -> b) -> a -> b
$ ([(TypeBase ExtShape Freshness, RetAls)]
 -> ([TypeBase ExtShape Freshness], [RetAls]))
-> [[(TypeBase ExtShape Freshness, RetAls)]]
-> [([TypeBase ExtShape Freshness], [RetAls])]
forall a b. (a -> b) -> [a] -> [b]
map [(TypeBase ExtShape Freshness, RetAls)]
-> ([TypeBase ExtShape Freshness], [RetAls])
forall a b. [(a, b)] -> ([a], [b])
unzip ([[(TypeBase ExtShape Freshness, RetAls)]]
 -> [([TypeBase ExtShape Freshness], [RetAls])])
-> [[(TypeBase ExtShape Freshness, RetAls)]]
-> [([TypeBase ExtShape Freshness], [RetAls])]
forall a b. (a -> b) -> a -> b
$ [Tree DeclType]
-> ResRetType -> [[(TypeBase ExtShape Freshness, RetAls)]]
internaliseEntryReturnType ((Tree (Param DeclType) -> Tree DeclType)
-> [Tree (Param DeclType)] -> [Tree DeclType]
forall a b. (a -> b) -> [a] -> [b]
map ((Param DeclType -> DeclType)
-> Tree (Param DeclType) -> Tree DeclType
forall a b. (a -> b) -> Free [] a -> Free [] b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Param DeclType -> DeclType
forall dec. DeclTyped dec => Param dec -> DeclType
paramDeclType) [Tree (Param DeclType)]
all_params) ResRetType
rettype
        (EntryPoint
entry', OpaqueTypes
opaques) =
          VisibleTypes
-> Name
-> Maybe Text
-> [(EntryParam, [Param DeclType])]
-> (EntryType, [[TypeBase Rank Freshness]])
-> (EntryPoint, OpaqueTypes)
entryPoint
            VisibleTypes
types
            (VName -> Name
baseName VName
ofname)
            Maybe Text
doc
            ([EntryParam]
-> [[Param DeclType]] -> [(EntryParam, [Param DeclType])]
forall a b. [a] -> [b] -> [(a, b)]
zip [EntryParam]
e_params ([[Param DeclType]] -> [(EntryParam, [Param DeclType])])
-> [[Param DeclType]] -> [(EntryParam, [Param DeclType])]
forall a b. (a -> b) -> a -> b
$ ([Tree (Param DeclType)] -> [Param DeclType])
-> [[Tree (Param DeclType)]] -> [[Param DeclType]]
forall a b. (a -> b) -> [a] -> [b]
map ((Tree (Param DeclType) -> [Param DeclType])
-> [Tree (Param DeclType)] -> [Param DeclType]
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Tree (Param DeclType) -> [Param DeclType]
forall a. Free [] a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList) [[Tree (Param DeclType)]]
[[Tree (FParam SOACS)]]
params')
            (EntryType
e_rettype, ([TypeBase ExtShape Freshness] -> [TypeBase Rank Freshness])
-> [[TypeBase ExtShape Freshness]] -> [[TypeBase Rank Freshness]]
forall a b. (a -> b) -> [a] -> [b]
map ((TypeBase ExtShape Freshness -> TypeBase Rank Freshness)
-> [TypeBase ExtShape Freshness] -> [TypeBase Rank Freshness]
forall a b. (a -> b) -> [a] -> [b]
map TypeBase ExtShape Freshness -> TypeBase Rank Freshness
forall shape o.
ArrayShape shape =>
TypeBase shape o -> TypeBase Rank o
I.rankShaped) [[TypeBase ExtShape Freshness]]
entry_rettype)
        args :: [SubExp]
args = (Param DeclType -> SubExp) -> [Param DeclType] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map (VName -> SubExp
I.Var (VName -> SubExp)
-> (Param DeclType -> VName) -> Param DeclType -> SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Param DeclType -> VName
forall dec. Param dec -> VName
I.paramName) ([Param DeclType] -> [SubExp]) -> [Param DeclType] -> [SubExp]
forall a b. (a -> b) -> a -> b
$ ([Tree (Param DeclType)] -> [Param DeclType])
-> [[Tree (Param DeclType)]] -> [Param DeclType]
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap ((Tree (Param DeclType) -> [Param DeclType])
-> [Tree (Param DeclType)] -> [Param DeclType]
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Tree (Param DeclType) -> [Param DeclType]
forall a. Free [] a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList) [[Tree (Param DeclType)]]
[[Tree (FParam SOACS)]]
params'

    OpaqueTypes -> InternaliseM ()
addOpaques OpaqueTypes
opaques

    (entry_body, ctx_ts) <- InternaliseM ([SubExpRes], [(TypeBase ExtShape NoMode, RetAls)])
-> InternaliseM
     (GBody (Rep InternaliseM) SubExpRes,
      [(TypeBase ExtShape NoMode, RetAls)])
forall (m :: * -> *) res a.
(MonadBuilder m, IsResult res) =>
m ([res], a) -> m (GBody (Rep m) res, a)
buildBody (InternaliseM ([SubExpRes], [(TypeBase ExtShape NoMode, RetAls)])
 -> InternaliseM
      (GBody (Rep InternaliseM) SubExpRes,
       [(TypeBase ExtShape NoMode, RetAls)]))
-> InternaliseM ([SubExpRes], [(TypeBase ExtShape NoMode, RetAls)])
-> InternaliseM
     (GBody (Rep InternaliseM) SubExpRes,
      [(TypeBase ExtShape NoMode, RetAls)])
forall a b. (a -> b) -> a -> b
$ do
      -- Special case the (rare) situation where the entry point is
      -- not a function.
      maybe_const <- VName -> InternaliseM (Maybe [SubExp])
lookupConst VName
ofname
      vals <- case maybe_const of
        Just [SubExp]
ses ->
          [SubExp] -> InternaliseM [SubExp]
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [SubExp]
ses
        Maybe [SubExp]
Nothing ->
          Name -> QualName VName -> [SubExp] -> InternaliseM [SubExp]
funcall Name
"entry_result" (VName -> QualName VName
forall v. v -> QualName v
E.qualName VName
ofname) [SubExp]
args
      ctx <-
        extractShapeContext (zeroExts $ concat entry_rettype)
          <$> mapM (fmap I.arrayDims . subExpType) vals
      pure (subExpsRes $ ctx ++ vals, map (const (I.Prim int64, mempty)) ctx)

    attrs' <- internaliseAttrs attrs
    let num_ctx = [(TypeBase ExtShape NoMode, RetAls)] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [(TypeBase ExtShape NoMode, RetAls)]
ctx_ts
    addFunDef $
      I.FunDef
        (Just entry')
        attrs'
        ("entry_" <> baseName ofname)
        ( ctx_ts
            ++ zip
              (map I.fromDecl (zeroExts (concat entry_rettype)))
              (map (shiftRetAls num_ctx) $ concat retals)
        )
        (shapeparams ++ foldMap (foldMap toList) params')
        entry_body
  where
    zeroExts :: [TypeBase ExtShape o] -> [TypeBase ExtShape o]
zeroExts [TypeBase ExtShape o]
ts = [TypeBase ExtShape o]
-> [TypeBase ExtShape o] -> [TypeBase ExtShape o]
forall o.
[TypeBase ExtShape o]
-> [TypeBase ExtShape o] -> [TypeBase ExtShape o]
generaliseExtTypes [TypeBase ExtShape o]
ts [TypeBase ExtShape o]
ts

internaliseBody :: Name -> E.Exp -> InternaliseM (Body SOACS)
internaliseBody :: Name -> Exp -> InternaliseM (Body SOACS)
internaliseBody Name
desc Exp
e =
  InternaliseM [SubExpRes]
-> InternaliseM (GBody (Rep InternaliseM) SubExpRes)
forall (m :: * -> *) res.
(MonadBuilder m, IsResult res) =>
m [res] -> m (GBody (Rep m) res)
buildBody_ (InternaliseM [SubExpRes]
 -> InternaliseM (GBody (Rep InternaliseM) SubExpRes))
-> InternaliseM [SubExpRes]
-> InternaliseM (GBody (Rep InternaliseM) SubExpRes)
forall a b. (a -> b) -> a -> b
$ [SubExp] -> [SubExpRes]
subExpsRes ([SubExp] -> [SubExpRes])
-> InternaliseM [SubExp] -> InternaliseM [SubExpRes]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Name -> Exp -> InternaliseM [SubExp]
internaliseExp (Name
desc Name -> Name -> Name
forall a. Semigroup a => a -> a -> a
<> Name
"_res") Exp
e

bodyFromStms ::
  InternaliseM (Result, a) ->
  InternaliseM (Body SOACS, a)
bodyFromStms :: forall a.
InternaliseM ([SubExpRes], a) -> InternaliseM (Body SOACS, a)
bodyFromStms InternaliseM ([SubExpRes], a)
m = do
  ((res, a), stms) <- InternaliseM ([SubExpRes], a)
-> InternaliseM (([SubExpRes], a), Stms (Rep InternaliseM))
forall a.
InternaliseM a -> InternaliseM (a, Stms (Rep InternaliseM))
forall (m :: * -> *) a.
MonadBuilder m =>
m a -> m (a, Stms (Rep m))
collectStms InternaliseM ([SubExpRes], a)
m
  (,a) <$> mkBodyM stms res

-- | Only returns those pattern names that are not used in the pattern
-- itself (the "non-existential" part, you could say).
letValExp :: Name -> I.Exp SOACS -> InternaliseM [VName]
letValExp :: Name -> Exp SOACS -> InternaliseM [VName]
letValExp Name
name Exp SOACS
e = do
  e_t <- Exp SOACS -> InternaliseM [TypeBase ExtShape NoMode]
forall rep (m :: * -> *).
(HasScope rep m, TypedOp (OpC rep)) =>
Exp rep -> m [TypeBase ExtShape NoMode]
expExtType Exp SOACS
e
  names <- replicateM (length e_t) $ newVName name
  letBindNames names e
  let ctx = [TypeBase ExtShape NoMode] -> Set Int
forall o. [TypeBase ExtShape o] -> Set Int
shapeContext [TypeBase ExtShape NoMode]
e_t
  pure $ map fst $ filter ((`S.notMember` ctx) . snd) $ zip names [0 ..]

letValExp' :: Name -> I.Exp SOACS -> InternaliseM [SubExp]
letValExp' :: Name -> Exp SOACS -> InternaliseM [SubExp]
letValExp' Name
_ (BasicOp (SubExp SubExp
se)) = [SubExp] -> InternaliseM [SubExp]
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [SubExp
se]
letValExp' Name
name Exp SOACS
ses = (VName -> SubExp) -> [VName] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
I.Var ([VName] -> [SubExp])
-> InternaliseM [VName] -> InternaliseM [SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Name -> Exp SOACS -> InternaliseM [VName]
letValExp Name
name Exp SOACS
ses

internaliseAppExp :: Name -> E.AppRes -> E.AppExp -> InternaliseM [I.SubExp]
internaliseAppExp :: Name -> AppRes -> AppExp -> InternaliseM [SubExp]
internaliseAppExp Name
desc AppRes
_ (E.Index Exp
e SliceBase Info VName
idxs SrcLoc
_) = do
  vs <- Name -> Exp -> InternaliseM [VName]
internaliseExpToVars Name
"indexed" Exp
e
  dims <- case vs of
    [] -> [SubExp] -> InternaliseM [SubExp]
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [] -- Will this happen?
    VName
v : [VName]
_ -> TypeBase Shape NoMode -> [SubExp]
forall d m. TypeBase (ShapeBase d) m -> [d]
I.arrayDims (TypeBase Shape NoMode -> [SubExp])
-> InternaliseM (TypeBase Shape NoMode) -> InternaliseM [SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> InternaliseM (TypeBase Shape NoMode)
forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoMode)
lookupType VName
v
  (idxs', cs) <- internaliseSlice dims idxs
  let index VName
v = do
        v_t <- VName -> InternaliseM (TypeBase Shape NoMode)
forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoMode)
lookupType VName
v
        pure $ I.BasicOp $ I.Index v $ fullSlice v_t idxs'
  certifying cs $ mapM (letSubExp desc <=< index) vs
internaliseAppExp Name
desc AppRes
_ (E.Range Exp
start Maybe Exp
maybe_second Inclusiveness Exp
end SrcLoc
_) = do
  start' <- Name -> Exp -> InternaliseM SubExp
internaliseExp1 Name
"range_start" Exp
start
  end' <- internaliseExp1 "range_end" $ case end of
    DownToExclusive Exp
e -> Exp
e
    ToInclusive Exp
e -> Exp
e
    UpToExclusive Exp
e -> Exp
e
  maybe_second' <-
    traverse (internaliseExp1 "range_second") maybe_second

  -- Construct an error message in case the range is invalid.
  let conv = case Exp -> StructType
E.typeOf Exp
start of
        E.Scalar (E.Prim (E.Unsigned IntType
_)) -> IntType -> SubExp -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
IntType -> SubExp -> m SubExp
asIntZ IntType
Int64
        StructType
_ -> IntType -> SubExp -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
IntType -> SubExp -> m SubExp
asIntS IntType
Int64
  start'_i64 <- conv start'
  end'_i64 <- conv end'
  maybe_second'_i64 <- traverse conv maybe_second'
  let errmsg =
        [ErrorMsgPart SubExp] -> ErrorMsg SubExp
forall a. [ErrorMsgPart a] -> ErrorMsg a
errorMsg ([ErrorMsgPart SubExp] -> ErrorMsg SubExp)
-> [ErrorMsgPart SubExp] -> ErrorMsg SubExp
forall a b. (a -> b) -> a -> b
$
          [ErrorMsgPart SubExp
"Range "]
            [ErrorMsgPart SubExp]
-> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
forall a. [a] -> [a] -> [a]
++ [PrimType -> SubExp -> ErrorMsgPart SubExp
forall a. PrimType -> a -> ErrorMsgPart a
ErrorVal PrimType
int64 SubExp
start'_i64]
            [ErrorMsgPart SubExp]
-> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
forall a. [a] -> [a] -> [a]
++ ( case Maybe SubExp
maybe_second'_i64 of
                   Maybe SubExp
Nothing -> []
                   Just SubExp
second_i64 -> [ErrorMsgPart SubExp
"..", PrimType -> SubExp -> ErrorMsgPart SubExp
forall a. PrimType -> a -> ErrorMsgPart a
ErrorVal PrimType
int64 SubExp
second_i64]
               )
            [ErrorMsgPart SubExp]
-> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
forall a. [a] -> [a] -> [a]
++ ( case Inclusiveness Exp
end of
                   DownToExclusive {} -> [ErrorMsgPart SubExp
"..>"]
                   ToInclusive {} -> [ErrorMsgPart SubExp
"..."]
                   UpToExclusive {} -> [ErrorMsgPart SubExp
"..<"]
               )
            [ErrorMsgPart SubExp]
-> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
forall a. [a] -> [a] -> [a]
++ [PrimType -> SubExp -> ErrorMsgPart SubExp
forall a. PrimType -> a -> ErrorMsgPart a
ErrorVal PrimType
int64 SubExp
end'_i64, ErrorMsgPart SubExp
" is invalid."]

  (it, lt_op) <-
    case E.typeOf start of
      E.Scalar (E.Prim (E.Signed IntType
it)) -> (IntType, CmpOp) -> InternaliseM (IntType, CmpOp)
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (IntType
it, IntType -> CmpOp
CmpSlt IntType
it)
      E.Scalar (E.Prim (E.Unsigned IntType
it)) -> (IntType, CmpOp) -> InternaliseM (IntType, CmpOp)
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (IntType
it, IntType -> CmpOp
CmpUlt IntType
it)
      StructType
start_t -> [Char] -> InternaliseM (IntType, CmpOp)
forall a. HasCallStack => [Char] -> a
error ([Char] -> InternaliseM (IntType, CmpOp))
-> [Char] -> InternaliseM (IntType, CmpOp)
forall a b. (a -> b) -> a -> b
$ [Char]
"Start value in range has type " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ StructType -> [Char]
forall a. Pretty a => a -> [Char]
prettyString StructType
start_t

  let one = IntType -> Integer -> SubExp
intConst IntType
it Integer
1
      negone = IntType -> Integer -> SubExp
intConst IntType
it (-Integer
1)
      default_step = case Inclusiveness Exp
end of
        DownToExclusive {} -> SubExp
negone
        ToInclusive {} -> SubExp
one
        UpToExclusive {} -> SubExp
one

  (step, step_zero) <- case maybe_second' of
    Just SubExp
second' -> do
      subtracted_step <-
        Name -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"subtracted_step" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
          BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
forall a b. (a -> b) -> a -> b
$
            BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
I.Sub IntType
it Overflow
I.OverflowWrap) SubExp
second' SubExp
start'
      step_zero <- letSubExp "step_zero" $ I.BasicOp $ I.CmpOp (I.CmpEq $ IntType it) start' second'
      pure (subtracted_step, step_zero)
    Maybe SubExp
Nothing ->
      (SubExp, SubExp) -> InternaliseM (SubExp, SubExp)
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SubExp
default_step, Bool -> SubExp
forall v. IsValue v => v -> SubExp
constant Bool
False)

  step_sign <- letSubExp "s_sign" $ BasicOp $ I.UnOp (I.SSignum it) step
  step_sign_i64 <- asIntS Int64 step_sign

  bounds_invalid_downwards <-
    letSubExp "bounds_invalid_downwards" $
      I.BasicOp $
        I.CmpOp lt_op start' end'
  bounds_invalid_upwards <-
    letSubExp "bounds_invalid_upwards" $
      I.BasicOp $
        I.CmpOp lt_op end' start'

  (distance, step_wrong_dir, bounds_invalid) <- case end of
    DownToExclusive {} -> do
      step_wrong_dir <-
        Name -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"step_wrong_dir" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
          BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
forall a b. (a -> b) -> a -> b
$
            CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
I.CmpEq (PrimType -> CmpOp) -> PrimType -> CmpOp
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
IntType IntType
it) SubExp
step_sign SubExp
one
      distance <-
        letSubExp "distance" $
          I.BasicOp $
            I.BinOp (Sub it I.OverflowWrap) start' end'
      distance_i64 <- asIntS Int64 distance
      pure (distance_i64, step_wrong_dir, bounds_invalid_downwards)
    UpToExclusive {} -> do
      step_wrong_dir <-
        Name -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"step_wrong_dir" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
          BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
forall a b. (a -> b) -> a -> b
$
            CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
I.CmpEq (PrimType -> CmpOp) -> PrimType -> CmpOp
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
IntType IntType
it) SubExp
step_sign SubExp
negone
      distance <- letSubExp "distance" $ I.BasicOp $ I.BinOp (Sub it I.OverflowWrap) end' start'
      distance_i64 <- asIntS Int64 distance
      pure (distance_i64, step_wrong_dir, bounds_invalid_upwards)
    ToInclusive {} -> do
      downwards <-
        Name -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"downwards" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
          BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
forall a b. (a -> b) -> a -> b
$
            CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
I.CmpEq (PrimType -> CmpOp) -> PrimType -> CmpOp
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
IntType IntType
it) SubExp
step_sign SubExp
negone
      distance_downwards_exclusive <-
        letSubExp "distance_downwards_exclusive" $
          I.BasicOp $
            I.BinOp (Sub it I.OverflowWrap) start' end'
      distance_upwards_exclusive <-
        letSubExp "distance_upwards_exclusive" $
          I.BasicOp $
            I.BinOp (Sub it I.OverflowWrap) end' start'

      bounds_invalid <-
        letSubExp "bounds_invalid"
          =<< eIf
            (eSubExp downwards)
            (resultBodyM [bounds_invalid_downwards])
            (resultBodyM [bounds_invalid_upwards])
      distance_exclusive <-
        letSubExp "distance_exclusive"
          =<< eIf
            (eSubExp downwards)
            (resultBodyM [distance_downwards_exclusive])
            (resultBodyM [distance_upwards_exclusive])
      distance_exclusive_i64 <- asIntS Int64 distance_exclusive
      distance <-
        letSubExp "distance" $
          I.BasicOp $
            I.BinOp
              (Add Int64 I.OverflowWrap)
              distance_exclusive_i64
              (intConst Int64 1)
      pure (distance, constant False, bounds_invalid)

  step_invalid <-
    letSubExp "step_invalid" $
      I.BasicOp $
        I.BinOp I.LogOr step_wrong_dir step_zero

  invalid <-
    letSubExp "range_invalid" $
      I.BasicOp $
        I.BinOp I.LogOr step_invalid bounds_invalid
  valid <- letSubExp "valid" $ I.BasicOp $ I.UnOp (I.Neg I.Bool) invalid
  cs <- assert "range_valid_c" valid errmsg

  step_i64 <- asIntS Int64 step
  pos_step <-
    letSubExp "pos_step" $
      I.BasicOp $
        I.BinOp (Mul Int64 I.OverflowWrap) step_i64 step_sign_i64

  num_elems <-
    certifying cs $
      letSubExp "num_elems" $
        I.BasicOp $
          I.BinOp (SCeilDiv Int64 I.Unsafe) distance pos_step

  se <- letSubExp desc (I.BasicOp $ I.Iota num_elems start' step it)
  pure [se]
internaliseAppExp Name
desc (E.AppRes StructType
et [VName]
ext) e :: AppExp
e@E.Apply {} =
  case AppExp -> (Function, [(Exp, Maybe VName)])
findFuncall AppExp
e of
    (FunctionHole SrcLoc
loc, [(Exp, Maybe VName)]
_args) -> do
      -- The function we are supposed to call doesn't exist, but we
      -- have to synthesize some fake values of the right type.  The
      -- easy way to do this is to just ignore the arguments and
      -- create a hole whose type is the type of the entire
      -- application.  One caveat is that we need to replace any
      -- existential sizes, too (with zeroes, because they don't
      -- matter).
      let subst :: [(VName, Subst StructRetType)]
subst = (VName -> (VName, Subst StructRetType))
-> [VName] -> [(VName, Subst StructRetType)]
forall a b. (a -> b) -> [a] -> [b]
map (,Exp -> Subst StructRetType
forall t. Exp -> Subst t
E.ExpSubst (Integer -> SrcLoc -> Exp
E.sizeFromInteger Integer
0 SrcLoc
forall a. Monoid a => a
mempty)) [VName]
ext
          et' :: StructType
et' = TypeSubs -> StructType -> StructType
forall a. Substitutable a => TypeSubs -> a -> a
E.applySubst (VName
-> [(VName, Subst StructRetType)] -> Maybe (Subst StructRetType)
forall a b. Eq a => a -> [(a, b)] -> Maybe b
`lookup` [(VName, Subst StructRetType)]
subst) StructType
et
      Name -> Exp -> InternaliseM [SubExp]
internaliseExp Name
desc (Info StructType -> SrcLoc -> Exp
forall (f :: * -> *) vn. f StructType -> SrcLoc -> ExpBase f vn
E.Hole (StructType -> Info StructType
forall a. a -> Info a
Info StructType
et') SrcLoc
loc)
    (FunctionName QualName VName
qfname, [(Exp, Maybe VName)]
args) -> do
      -- Argument evaluation is outermost-in so that any existential sizes
      -- created by function applications can be brought into scope.
      let fname :: Name
fname = VName -> Name
baseName (VName -> Name) -> VName -> Name
forall a b. (a -> b) -> a -> b
$ QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
qfname
          arg_desc :: Name
arg_desc = Name
fname Name -> Name -> Name
forall a. Semigroup a => a -> a -> a
<> Name
"_arg"

      -- Some functions are magical (overloaded) and we handle that here.
      case () of
        ()
          -- Short-circuiting operators are magical.
          | VName -> Bool
isIntrinsic (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
qfname),
            VName -> Name
baseName (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
qfname) Name -> Name -> Bool
forall a. Eq a => a -> a -> Bool
== Name
"&&",
            [(Exp
x, Maybe VName
_), (Exp
y, Maybe VName
_)] <- [(Exp, Maybe VName)]
args ->
              Name -> Exp -> InternaliseM [SubExp]
internaliseExp Name
desc (Exp -> InternaliseM [SubExp]) -> Exp -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$
                AppExp -> Info AppRes -> Exp
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
E.AppExp
                  (Exp -> Exp -> Exp -> SrcLoc -> AppExp
forall (f :: * -> *) vn.
ExpBase f vn
-> ExpBase f vn -> ExpBase f vn -> SrcLoc -> AppExpBase f vn
E.If Exp
x Exp
y (PrimValue -> SrcLoc -> Exp
forall (f :: * -> *) vn. PrimValue -> SrcLoc -> ExpBase f vn
E.Literal (Bool -> PrimValue
E.BoolValue Bool
False) SrcLoc
forall a. Monoid a => a
mempty) SrcLoc
forall a. Monoid a => a
mempty)
                  (AppRes -> Info AppRes
forall a. a -> Info a
Info (AppRes -> Info AppRes) -> AppRes -> Info AppRes
forall a b. (a -> b) -> a -> b
$ StructType -> [VName] -> AppRes
AppRes (ScalarTypeBase Exp NoMode -> StructType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
E.Scalar (ScalarTypeBase Exp NoMode -> StructType)
-> ScalarTypeBase Exp NoMode -> StructType
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase Exp NoMode
forall dim o. PrimType -> ScalarTypeBase dim o
E.Prim PrimType
E.Bool) [])
          | VName -> Bool
isIntrinsic (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
qfname),
            VName -> Name
baseName (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
qfname) Name -> Name -> Bool
forall a. Eq a => a -> a -> Bool
== Name
"||",
            [(Exp
x, Maybe VName
_), (Exp
y, Maybe VName
_)] <- [(Exp, Maybe VName)]
args ->
              Name -> Exp -> InternaliseM [SubExp]
internaliseExp Name
desc (Exp -> InternaliseM [SubExp]) -> Exp -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$
                AppExp -> Info AppRes -> Exp
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
E.AppExp
                  (Exp -> Exp -> Exp -> SrcLoc -> AppExp
forall (f :: * -> *) vn.
ExpBase f vn
-> ExpBase f vn -> ExpBase f vn -> SrcLoc -> AppExpBase f vn
E.If Exp
x (PrimValue -> SrcLoc -> Exp
forall (f :: * -> *) vn. PrimValue -> SrcLoc -> ExpBase f vn
E.Literal (Bool -> PrimValue
E.BoolValue Bool
True) SrcLoc
forall a. Monoid a => a
mempty) Exp
y SrcLoc
forall a. Monoid a => a
mempty)
                  (AppRes -> Info AppRes
forall a. a -> Info a
Info (AppRes -> Info AppRes) -> AppRes -> Info AppRes
forall a b. (a -> b) -> a -> b
$ StructType -> [VName] -> AppRes
AppRes (ScalarTypeBase Exp NoMode -> StructType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
E.Scalar (ScalarTypeBase Exp NoMode -> StructType)
-> ScalarTypeBase Exp NoMode -> StructType
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase Exp NoMode
forall dim o. PrimType -> ScalarTypeBase dim o
E.Prim PrimType
E.Bool) [])
          -- Overloaded and intrinsic functions never take array
          -- arguments (except equality, but those cannot be
          -- existential), so we can safely ignore the existential
          -- dimensions.
          | Just [(StructType, [SubExp])] -> InternaliseM [SubExp]
internalise <- QualName VName
-> Name
-> Maybe ([(StructType, [SubExp])] -> InternaliseM [SubExp])
isOverloadedFunction QualName VName
qfname Name
desc -> do
              let prepareArg :: (Exp, b) -> InternaliseM (StructType, [SubExp])
prepareArg (Exp
arg, b
_) =
                    (StructType -> StructType
forall dim o. TypeBase dim o -> TypeBase dim NoMode
E.toStruct (Exp -> StructType
E.typeOf Exp
arg),) ([SubExp] -> (StructType, [SubExp]))
-> InternaliseM [SubExp] -> InternaliseM (StructType, [SubExp])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Name -> Exp -> InternaliseM [SubExp]
internaliseExp Name
"arg" Exp
arg
              [(StructType, [SubExp])] -> InternaliseM [SubExp]
internalise ([(StructType, [SubExp])] -> InternaliseM [SubExp])
-> InternaliseM [(StructType, [SubExp])] -> InternaliseM [SubExp]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< ((Exp, Maybe VName) -> InternaliseM (StructType, [SubExp]))
-> [(Exp, Maybe VName)] -> InternaliseM [(StructType, [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 (Exp, Maybe VName) -> InternaliseM (StructType, [SubExp])
forall {b}. (Exp, b) -> InternaliseM (StructType, [SubExp])
prepareArg [(Exp, Maybe VName)]
args
          | Just Name -> InternaliseM [SubExp]
internalise <- QualName VName -> [Exp] -> Maybe (Name -> InternaliseM [SubExp])
isIntrinsicFunction QualName VName
qfname (((Exp, Maybe VName) -> Exp) -> [(Exp, Maybe VName)] -> [Exp]
forall a b. (a -> b) -> [a] -> [b]
map (Exp, Maybe VName) -> Exp
forall a b. (a, b) -> a
fst [(Exp, Maybe VName)]
args) ->
              Name -> InternaliseM [SubExp]
internalise Name
desc
          | VName -> Bool
isIntrinsic (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
qfname),
            Just (PrimType
rettype, [PrimType]
_) <- Name
-> Map Name (PrimType, [PrimType]) -> Maybe (PrimType, [PrimType])
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name
fname Map Name (PrimType, [PrimType])
I.builtInFunctions -> do
              let tag :: [a] -> [(a, Diet)]
tag [a]
ses = [(a
se, Diet
I.Observe) | a
se <- [a]
ses]
              args' <- [[SubExp]] -> [[SubExp]]
forall a. [a] -> [a]
reverse ([[SubExp]] -> [[SubExp]])
-> InternaliseM [[SubExp]] -> InternaliseM [[SubExp]]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((Exp, Maybe VName) -> InternaliseM [SubExp])
-> [(Exp, Maybe VName)] -> InternaliseM [[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, Maybe VName) -> InternaliseM [SubExp]
internaliseArg Name
arg_desc) ([(Exp, Maybe VName)] -> [(Exp, Maybe VName)]
forall a. [a] -> [a]
reverse [(Exp, Maybe VName)]
args)
              let args'' = ([SubExp] -> [(SubExp, Diet)]) -> [[SubExp]] -> [(SubExp, Diet)]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap [SubExp] -> [(SubExp, Diet)]
forall {a}. [a] -> [(a, Diet)]
tag [[SubExp]]
args'
              letValExp' desc $ I.Apply fname args'' [(I.Prim rettype, mempty)] Safe
          | Bool
otherwise -> do
              args' <-
                [[SubExp]] -> [SubExp]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[SubExp]] -> [SubExp])
-> ([[SubExp]] -> [[SubExp]]) -> [[SubExp]] -> [SubExp]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [[SubExp]] -> [[SubExp]]
forall a. [a] -> [a]
reverse
                  ([[SubExp]] -> [SubExp])
-> InternaliseM [[SubExp]] -> InternaliseM [SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((Exp, Maybe VName) -> InternaliseM [SubExp])
-> [(Exp, Maybe VName)] -> InternaliseM [[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, Maybe VName) -> InternaliseM [SubExp]
internaliseArg Name
arg_desc) ([(Exp, Maybe VName)] -> [(Exp, Maybe VName)]
forall a. [a] -> [a]
reverse [(Exp, Maybe VName)]
args)
              funcall desc qfname args'
internaliseAppExp Name
desc AppRes
_ (E.LetPat [SizeBinder VName]
sizes PatBase Info VName StructType
pat Exp
e Exp
body SrcLoc
_) =
  Name
-> [SizeBinder VName]
-> PatBase Info VName StructType
-> Exp
-> InternaliseM [SubExp]
-> InternaliseM [SubExp]
forall a.
Name
-> [SizeBinder VName]
-> PatBase Info VName StructType
-> Exp
-> InternaliseM a
-> InternaliseM a
internalisePat Name
desc [SizeBinder VName]
sizes PatBase Info VName StructType
pat Exp
e (InternaliseM [SubExp] -> InternaliseM [SubExp])
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ Name -> Exp -> InternaliseM [SubExp]
internaliseExp Name
desc Exp
body
internaliseAppExp Name
_ AppRes
_ (E.LetFun (VName
ofname, SrcLoc
_) ([TypeParamBase VName], [PatBase Info VName ParamType],
 Maybe (TypeExp Exp VName), Info ResRetType, Exp)
_ Exp
_ SrcLoc
_) =
  [Char] -> InternaliseM [SubExp]
forall a. HasCallStack => [Char] -> a
error ([Char] -> InternaliseM [SubExp])
-> [Char] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ [Char]
"Unexpected LetFun " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ VName -> [Char]
forall a. Pretty a => a -> [Char]
prettyString VName
ofname
internaliseAppExp Name
desc AppRes
_ (E.Loop [VName]
sparams PatBase Info VName ParamType
mergepat LoopInitBase Info VName
loopinit LoopFormBase Info VName
form Exp
loopbody SrcLoc
_) = do
  ses <- Name -> Exp -> InternaliseM [SubExp]
internaliseExp Name
"loop_init" (Exp -> InternaliseM [SubExp]) -> Exp -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ LoopInitBase Info VName -> Exp
loopInitExp LoopInitBase Info VName
loopinit
  ((loopbody', (form', shapepat, mergepat', mergeinit')), initstms) <-
    collectStms $ handleForm ses form

  addStms initstms
  mergeinit_ts' <- mapM subExpType mergeinit'

  ctxinit <- argShapes (map I.paramName shapepat) mergepat' mergeinit_ts'

  -- Ensure that the initial loop values match the shapes of the loop
  -- parameters.  XXX: Ideally they should already match (by the
  -- source language type rules), but some of our transformations
  -- (esp. defunctionalisation) strips out some size information.  For
  -- a type-correct source program, these reshapes should simplify
  -- away.
  let args = [SubExp]
ctxinit [SubExp] -> [SubExp] -> [SubExp]
forall a. [a] -> [a] -> [a]
++ [SubExp]
mergeinit'
  args' <-
    ensureArgShapes
      "initial loop values have right shape"
      (map I.paramName shapepat)
      (map paramType $ shapepat ++ mergepat')
      args

  let dropCond = case LoopFormBase Info VName
form of
        E.While {} -> Int -> [VName] -> [VName]
forall a. Int -> [a] -> [a]
drop Int
1
        LoopFormBase Info VName
_ -> [VName] -> [VName]
forall a. a -> a
id

  -- As above, ensure that the result has the right shape.
  let merge = [Param DeclType] -> [SubExp] -> [(Param DeclType, SubExp)]
forall a b. [a] -> [b] -> [(a, b)]
zip ([Param DeclType]
shapepat [Param DeclType] -> [Param DeclType] -> [Param DeclType]
forall a. [a] -> [a] -> [a]
++ [Param DeclType]
mergepat') [SubExp]
args'
      merge_ts = ((Param DeclType, SubExp) -> TypeBase Shape NoMode)
-> [(Param DeclType, SubExp)] -> [TypeBase Shape NoMode]
forall a b. (a -> b) -> [a] -> [b]
map (Param DeclType -> TypeBase Shape NoMode
forall dec. Typed dec => Param dec -> TypeBase Shape NoMode
I.paramType (Param DeclType -> TypeBase Shape NoMode)
-> ((Param DeclType, SubExp) -> Param DeclType)
-> (Param DeclType, SubExp)
-> TypeBase Shape NoMode
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Param DeclType, SubExp) -> Param DeclType
forall a b. (a, b) -> a
fst) [(Param DeclType, SubExp)]
merge
  loopbody'' <-
    localScope (scopeOfFParams (map fst merge) <> scopeOfLoopForm form') . buildBody_ $
      fmap subExpsRes
        . ensureArgShapes
          "shape of loop result does not match shapes in loop parameter"
          (map (I.paramName . fst) merge)
          merge_ts
        . map resSubExp
        =<< bodyBind loopbody'

  attrs <- asks envAttrs
  map I.Var . dropCond
    <$> attributing
      attrs
      (letValExp desc (I.Loop merge form' loopbody''))
  where
    sparams' :: [TypeParamBase VName]
sparams' = (VName -> TypeParamBase VName) -> [VName] -> [TypeParamBase VName]
forall a b. (a -> b) -> [a] -> [b]
map (VName -> SrcLoc -> TypeParamBase VName
forall vn. vn -> SrcLoc -> TypeParamBase vn
`TypeParamDim` SrcLoc
forall a. Monoid a => a
mempty) [VName]
sparams

    -- Attributes that apply to loops.
    loopAttrs :: Attrs
loopAttrs = Attr -> Attrs
oneAttr Attr
"unroll"
    -- Remove those attributes from the attribute set that apply to
    -- the loop itself.
    noLoopAttrs :: InternaliseEnv -> InternaliseEnv
noLoopAttrs InternaliseEnv
env = InternaliseEnv
env {envAttrs = envAttrs env `withoutAttrs` loopAttrs}

    loopBody :: InternaliseM [SubExp]
loopBody = (InternaliseEnv -> InternaliseEnv)
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a.
(InternaliseEnv -> InternaliseEnv)
-> InternaliseM a -> InternaliseM a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local InternaliseEnv -> InternaliseEnv
noLoopAttrs (InternaliseM [SubExp] -> InternaliseM [SubExp])
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ Name -> Exp -> InternaliseM [SubExp]
internaliseExp Name
"loopres" Exp
loopbody

    forLoop :: [Param DeclType]
-> [Param DeclType]
-> [SubExp]
-> VName
-> [(Param (TypeBase Shape NoMode), VName)]
-> LoopForm
-> InternaliseM
     (Body SOACS,
      (LoopForm, [Param DeclType], [Param DeclType], [SubExp]))
forLoop [Param DeclType]
mergepat' [Param DeclType]
shapepat [SubExp]
mergeinit VName
i [(Param (TypeBase Shape NoMode), VName)]
loopvars LoopForm
form' =
      InternaliseM
  ([SubExpRes],
   (LoopForm, [Param DeclType], [Param DeclType], [SubExp]))
-> InternaliseM
     (Body SOACS,
      (LoopForm, [Param DeclType], [Param DeclType], [SubExp]))
forall a.
InternaliseM ([SubExpRes], a) -> InternaliseM (Body SOACS, a)
bodyFromStms (InternaliseM
   ([SubExpRes],
    (LoopForm, [Param DeclType], [Param DeclType], [SubExp]))
 -> InternaliseM
      (Body SOACS,
       (LoopForm, [Param DeclType], [Param DeclType], [SubExp])))
-> (InternaliseM
      ([SubExpRes],
       (LoopForm, [Param DeclType], [Param DeclType], [SubExp]))
    -> InternaliseM
         ([SubExpRes],
          (LoopForm, [Param DeclType], [Param DeclType], [SubExp])))
-> InternaliseM
     ([SubExpRes],
      (LoopForm, [Param DeclType], [Param DeclType], [SubExp]))
-> InternaliseM
     (Body SOACS,
      (LoopForm, [Param DeclType], [Param DeclType], [SubExp]))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Scope SOACS
-> InternaliseM
     ([SubExpRes],
      (LoopForm, [Param DeclType], [Param DeclType], [SubExp]))
-> InternaliseM
     ([SubExpRes],
      (LoopForm, [Param DeclType], [Param DeclType], [SubExp]))
forall a. Scope SOACS -> InternaliseM a -> InternaliseM a
forall rep (m :: * -> *) a.
LocalScope rep m =>
Scope rep -> m a -> m a
localScope (LoopForm -> Scope SOACS
forall rep. LoopForm -> Scope rep
scopeOfLoopForm LoopForm
form') (InternaliseM
   ([SubExpRes],
    (LoopForm, [Param DeclType], [Param DeclType], [SubExp]))
 -> InternaliseM
      (Body SOACS,
       (LoopForm, [Param DeclType], [Param DeclType], [SubExp])))
-> InternaliseM
     ([SubExpRes],
      (LoopForm, [Param DeclType], [Param DeclType], [SubExp]))
-> InternaliseM
     (Body SOACS,
      (LoopForm, [Param DeclType], [Param DeclType], [SubExp]))
forall a b. (a -> b) -> a -> b
$ do
        [(Param (TypeBase Shape NoMode), VName)]
-> ((Param (TypeBase Shape NoMode), VName) -> InternaliseM ())
-> InternaliseM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [(Param (TypeBase Shape NoMode), VName)]
loopvars (((Param (TypeBase Shape NoMode), VName) -> InternaliseM ())
 -> InternaliseM ())
-> ((Param (TypeBase Shape NoMode), VName) -> InternaliseM ())
-> InternaliseM ()
forall a b. (a -> b) -> a -> b
$ \(Param (TypeBase Shape NoMode)
p, VName
arr) ->
          [VName] -> Exp (Rep InternaliseM) -> InternaliseM ()
forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [Param (TypeBase Shape NoMode) -> VName
forall dec. Param dec -> VName
I.paramName Param (TypeBase Shape NoMode)
p] (Exp SOACS -> InternaliseM ())
-> InternaliseM (Exp SOACS) -> InternaliseM ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< VName
-> [InternaliseM (Exp (Rep InternaliseM))]
-> InternaliseM (Exp (Rep InternaliseM))
forall (m :: * -> *).
MonadBuilder m =>
VName -> [m (Exp (Rep m))] -> m (Exp (Rep m))
eIndex VName
arr [SubExp -> InternaliseM (Exp (Rep InternaliseM))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp (VName -> SubExp
I.Var VName
i)]
        ses <- InternaliseM [SubExp]
loopBody
        sets <- mapM subExpType ses
        shapeargs <- argShapes (map I.paramName shapepat) mergepat' sets
        pure
          ( subExpsRes $ shapeargs ++ ses,
            ( form',
              shapepat,
              mergepat',
              mergeinit
            )
          )

    handleForm :: [SubExp]
-> LoopFormBase Info VName
-> InternaliseM
     (Body SOACS,
      (LoopForm, [Param DeclType], [Param DeclType], [SubExp]))
handleForm [SubExp]
mergeinit (E.ForIn PatBase Info VName StructType
x Exp
arr) = do
      arr' <- Name -> Exp -> InternaliseM [VName]
internaliseExpToVars Name
"for_in_arr" Exp
arr
      arr_ts <- mapM lookupType arr'
      let w = Int -> [TypeBase Shape NoMode] -> SubExp
forall o. Int -> [TypeBase Shape o] -> SubExp
arraysSize Int
0 [TypeBase Shape NoMode]
arr_ts

      i <- newVName "i"

      ts <- mapM subExpType mergeinit
      bindingLoopParams sparams' mergepat ts $ \[FParam SOACS]
shapepat [FParam SOACS]
mergepat' ->
        [PatBase Info VName ParamType]
-> [TypeBase Shape NoMode]
-> ([LParam SOACS]
    -> InternaliseM
         (Body SOACS,
          (LoopForm, [Param DeclType], [Param DeclType], [SubExp])))
-> InternaliseM
     (Body SOACS,
      (LoopForm, [Param DeclType], [Param DeclType], [SubExp]))
forall a.
[PatBase Info VName ParamType]
-> [TypeBase Shape NoMode]
-> ([LParam SOACS] -> InternaliseM a)
-> InternaliseM a
bindingLambdaParams [Diet -> StructType -> ParamType
forall o. Diet -> TypeBase Exp o -> ParamType
toParam Diet
E.Observe (StructType -> ParamType)
-> PatBase Info VName StructType -> PatBase Info VName ParamType
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> PatBase Info VName StructType
x] ((TypeBase Shape NoMode -> TypeBase Shape NoMode)
-> [TypeBase Shape NoMode] -> [TypeBase Shape NoMode]
forall a b. (a -> b) -> [a] -> [b]
map TypeBase Shape NoMode -> TypeBase Shape NoMode
forall d m. TypeBase (ShapeBase d) m -> TypeBase (ShapeBase d) m
rowType [TypeBase Shape NoMode]
arr_ts) (([LParam SOACS]
  -> InternaliseM
       (Body SOACS,
        (LoopForm, [Param DeclType], [Param DeclType], [SubExp])))
 -> InternaliseM
      (Body SOACS,
       (LoopForm, [Param DeclType], [Param DeclType], [SubExp])))
-> ([LParam SOACS]
    -> InternaliseM
         (Body SOACS,
          (LoopForm, [Param DeclType], [Param DeclType], [SubExp])))
-> InternaliseM
     (Body SOACS,
      (LoopForm, [Param DeclType], [Param DeclType], [SubExp]))
forall a b. (a -> b) -> a -> b
$ \[LParam SOACS]
x_params -> do
          let loopvars :: [(Param (TypeBase Shape NoMode), VName)]
loopvars = [Param (TypeBase Shape NoMode)]
-> [VName] -> [(Param (TypeBase Shape NoMode), VName)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Param (TypeBase Shape NoMode)]
[LParam SOACS]
x_params [VName]
arr'
          [Param DeclType]
-> [Param DeclType]
-> [SubExp]
-> VName
-> [(Param (TypeBase Shape NoMode), VName)]
-> LoopForm
-> InternaliseM
     (Body SOACS,
      (LoopForm, [Param DeclType], [Param DeclType], [SubExp]))
forLoop [Param DeclType]
[FParam SOACS]
mergepat' [Param DeclType]
[FParam SOACS]
shapepat [SubExp]
mergeinit VName
i [(Param (TypeBase Shape NoMode), VName)]
loopvars (LoopForm
 -> InternaliseM
      (Body SOACS,
       (LoopForm, [Param DeclType], [Param DeclType], [SubExp])))
-> LoopForm
-> InternaliseM
     (Body SOACS,
      (LoopForm, [Param DeclType], [Param DeclType], [SubExp]))
forall a b. (a -> b) -> a -> b
$ VName -> IntType -> SubExp -> LoopForm
I.ForLoop VName
i IntType
Int64 SubExp
w
    handleForm [SubExp]
mergeinit (E.For IdentBase Info VName StructType
i Exp
num_iterations) = do
      num_iterations' <- Name -> Exp -> InternaliseM SubExp
internaliseExp1 Name
"upper_bound" Exp
num_iterations
      num_iterations_t <- I.subExpType num_iterations'
      it <- case num_iterations_t of
        I.Prim (IntType IntType
it) -> IntType -> InternaliseM IntType
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure IntType
it
        TypeBase Shape NoMode
_ -> [Char] -> InternaliseM IntType
forall a. HasCallStack => [Char] -> a
error [Char]
"internaliseExp Loop: invalid type"

      ts <- mapM subExpType mergeinit
      bindingLoopParams sparams' mergepat ts $ \[FParam SOACS]
shapepat [FParam SOACS]
mergepat' ->
        [Param DeclType]
-> [Param DeclType]
-> [SubExp]
-> VName
-> [(Param (TypeBase Shape NoMode), VName)]
-> LoopForm
-> InternaliseM
     (Body SOACS,
      (LoopForm, [Param DeclType], [Param DeclType], [SubExp]))
forLoop [Param DeclType]
[FParam SOACS]
mergepat' [Param DeclType]
[FParam SOACS]
shapepat [SubExp]
mergeinit (IdentBase Info VName StructType -> VName
forall {k} (f :: k -> *) vn (t :: k). IdentBase f vn t -> vn
E.identName IdentBase Info VName StructType
i) [] (LoopForm
 -> InternaliseM
      (Body SOACS,
       (LoopForm, [Param DeclType], [Param DeclType], [SubExp])))
-> LoopForm
-> InternaliseM
     (Body SOACS,
      (LoopForm, [Param DeclType], [Param DeclType], [SubExp]))
forall a b. (a -> b) -> a -> b
$
          VName -> IntType -> SubExp -> LoopForm
I.ForLoop (IdentBase Info VName StructType -> VName
forall {k} (f :: k -> *) vn (t :: k). IdentBase f vn t -> vn
E.identName IdentBase Info VName StructType
i) IntType
it SubExp
num_iterations'
    handleForm [SubExp]
mergeinit (E.While Exp
cond) = do
      ts <- (SubExp -> InternaliseM (TypeBase Shape NoMode))
-> [SubExp] -> InternaliseM [TypeBase Shape NoMode]
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 SubExp -> InternaliseM (TypeBase Shape NoMode)
forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoMode)
subExpType [SubExp]
mergeinit
      bindingLoopParams sparams' mergepat ts $ \[FParam SOACS]
shapepat [FParam SOACS]
mergepat' -> do
        mergeinit_ts <- (SubExp -> InternaliseM (TypeBase Shape NoMode))
-> [SubExp] -> InternaliseM [TypeBase Shape NoMode]
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 SubExp -> InternaliseM (TypeBase Shape NoMode)
forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoMode)
subExpType [SubExp]
mergeinit
        -- We need to insert 'cond' twice - once for the initial
        -- condition (do we enter the loop at all?), and once with the
        -- result values of the loop (do we continue into the next
        -- iteration?).  This is safe, as the type rules for the
        -- external language guarantees that 'cond' does not consume
        -- anything.
        shapeinit <- argShapes (map I.paramName shapepat) mergepat' mergeinit_ts

        (loop_initial_cond, init_loop_cond_stms) <- collectStms $ do
          forM_ (zip shapepat shapeinit) $ \(Param DeclType
p, SubExp
se) ->
            [VName] -> Exp (Rep InternaliseM) -> InternaliseM ()
forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [Param DeclType -> VName
forall dec. Param dec -> VName
I.paramName Param DeclType
p] (Exp (Rep InternaliseM) -> InternaliseM ())
-> Exp (Rep InternaliseM) -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
SubExp SubExp
se
          forM_ (zip mergepat' mergeinit) $ \(Param DeclType
p, SubExp
se) ->
            Bool -> InternaliseM () -> InternaliseM ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (SubExp
se SubExp -> SubExp -> Bool
forall a. Eq a => a -> a -> Bool
== VName -> SubExp
I.Var (Param DeclType -> VName
forall dec. Param dec -> VName
I.paramName Param DeclType
p)) (InternaliseM () -> InternaliseM ())
-> InternaliseM () -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$
              [VName] -> Exp (Rep InternaliseM) -> InternaliseM ()
forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [Param DeclType -> VName
forall dec. Param dec -> VName
I.paramName Param DeclType
p] (Exp (Rep InternaliseM) -> InternaliseM ())
-> Exp (Rep InternaliseM) -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$
                case SubExp
se of
                  I.Var VName
v
                    | Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ TypeBase Shape NoMode -> Bool
forall shape o. TypeBase shape o -> Bool
primType (TypeBase Shape NoMode -> Bool) -> TypeBase Shape NoMode -> Bool
forall a b. (a -> b) -> a -> b
$ Param DeclType -> TypeBase Shape NoMode
forall dec. Typed dec => Param dec -> TypeBase Shape NoMode
paramType Param DeclType
p ->
                        [SubExp] -> VName -> Exp SOACS
forall rep. [SubExp] -> VName -> Exp rep
shapeCoerce (TypeBase Shape NoMode -> [SubExp]
forall d m. TypeBase (ShapeBase d) m -> [d]
I.arrayDims (TypeBase Shape NoMode -> [SubExp])
-> TypeBase Shape NoMode -> [SubExp]
forall a b. (a -> b) -> a -> b
$ Param DeclType -> TypeBase Shape NoMode
forall dec. Typed dec => Param dec -> TypeBase Shape NoMode
paramType Param DeclType
p) VName
v
                  SubExp
_ -> BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
SubExp SubExp
se

          -- As the condition expression is inserted twice, we have to
          -- avoid shadowing (#1935).
          (cond_stms, cond') <-
            uncurry (flip renameStmsWith)
              =<< collectStms (internaliseExp1 "loop_cond" cond)
          addStms cond_stms
          pure cond'

        addStms init_loop_cond_stms

        bodyFromStms $ do
          ses <- loopBody
          sets <- mapM subExpType ses
          loop_while <- newParam "loop_while" $ I.Prim I.Bool
          shapeargs <- argShapes (map I.paramName shapepat) mergepat' sets

          -- Careful not to clobber anything.
          loop_end_cond_body <- renameBody <=< buildBody_ $ do
            -- Shape arguments may refer to other shape parameters, which
            -- must not be clobbered before they are read.
            let shapepat_names = (Param DeclType -> VName) -> [Param DeclType] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param DeclType -> VName
forall dec. Param dec -> VName
I.paramName [Param DeclType]
[FParam SOACS]
shapepat
            shapeargs' <- forM (zip shapepat shapeargs) $ \case
              (Param DeclType
p, I.Var VName
v)
                | VName
v VName -> VName -> Bool
forall a. Eq a => a -> a -> Bool
/= Param DeclType -> VName
forall dec. Param dec -> VName
I.paramName Param DeclType
p,
                  VName
v VName -> [VName] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [VName]
shapepat_names -> do
                    v' <- Name -> InternaliseM VName
forall (m :: * -> *). MonadFreshNames m => Name -> m VName
newVName (Name -> InternaliseM VName) -> Name -> InternaliseM VName
forall a b. (a -> b) -> a -> b
$ VName -> Name
baseName VName
v Name -> Name -> Name
forall a. Semigroup a => a -> a -> a
<> Name
"_tmp"
                    letBindNames [v'] $ I.BasicOp $ I.SubExp $ I.Var v
                    pure $ I.Var v'
              (Param DeclType
_, SubExp
se) -> SubExp -> InternaliseM SubExp
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SubExp
se
            forM_ (zip shapepat shapeargs') $ \(Param DeclType
p, SubExp
se) ->
              Bool -> InternaliseM () -> InternaliseM ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (SubExp
se SubExp -> SubExp -> Bool
forall a. Eq a => a -> a -> Bool
== VName -> SubExp
I.Var (Param DeclType -> VName
forall dec. Param dec -> VName
I.paramName Param DeclType
p)) (InternaliseM () -> InternaliseM ())
-> InternaliseM () -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$
                [VName] -> Exp (Rep InternaliseM) -> InternaliseM ()
forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [Param DeclType -> VName
forall dec. Param dec -> VName
I.paramName Param DeclType
p] (Exp (Rep InternaliseM) -> InternaliseM ())
-> Exp (Rep InternaliseM) -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$
                  BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
forall a b. (a -> b) -> a -> b
$
                    SubExp -> BasicOp
SubExp SubExp
se
            -- Is may be that one of the names in mergepat' are also in ses (and
            -- not in the same position), in which case we must be careful to
            -- avoid clobbering.
            let mergepat_names = (Param DeclType -> VName) -> [Param DeclType] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param DeclType -> VName
forall dec. Param dec -> VName
I.paramName [Param DeclType]
[FParam SOACS]
mergepat'
            ses' <- forM (zip mergepat' ses) $ \case
              (Param DeclType
p, I.Var VName
v)
                | VName
v VName -> [VName] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [VName]
mergepat_names -> do
                    v' <- Name -> InternaliseM VName
forall (m :: * -> *). MonadFreshNames m => Name -> m VName
newVName (Name -> InternaliseM VName) -> Name -> InternaliseM VName
forall a b. (a -> b) -> a -> b
$ VName -> Name
baseName VName
v Name -> Name -> Name
forall a. Semigroup a => a -> a -> a
<> Name
"_tmp"
                    letBindNames [v'] $
                      if primType $ paramType p
                        then I.BasicOp (I.SubExp $ I.Var v)
                        -- Need administrative coerce due to the renaming.
                        else shapeCoerce (I.arrayDims $ paramType p) v
                    pure $ I.Var v'
              (Param DeclType
_, SubExp
se) -> SubExp -> InternaliseM SubExp
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SubExp
se
            forM_ (zip mergepat' ses') $ \(Param DeclType
p, SubExp
se) ->
              [VName] -> Exp (Rep InternaliseM) -> InternaliseM ()
forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [Param DeclType -> VName
forall dec. Param dec -> VName
I.paramName Param DeclType
p] (Exp (Rep InternaliseM) -> InternaliseM ())
-> Exp (Rep InternaliseM) -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$
                case SubExp
se of
                  I.Var VName
v
                    | Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ TypeBase Shape NoMode -> Bool
forall shape o. TypeBase shape o -> Bool
primType (TypeBase Shape NoMode -> Bool) -> TypeBase Shape NoMode -> Bool
forall a b. (a -> b) -> a -> b
$ Param DeclType -> TypeBase Shape NoMode
forall dec. Typed dec => Param dec -> TypeBase Shape NoMode
paramType Param DeclType
p ->
                        [SubExp] -> VName -> Exp SOACS
forall rep. [SubExp] -> VName -> Exp rep
shapeCoerce (TypeBase Shape NoMode -> [SubExp]
forall d m. TypeBase (ShapeBase d) m -> [d]
I.arrayDims (TypeBase Shape NoMode -> [SubExp])
-> TypeBase Shape NoMode -> [SubExp]
forall a b. (a -> b) -> a -> b
$ Param DeclType -> TypeBase Shape NoMode
forall dec. Typed dec => Param dec -> TypeBase Shape NoMode
paramType Param DeclType
p) VName
v
                  SubExp
_ -> BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
SubExp SubExp
se
            subExpsRes <$> internaliseExp "loop_cond" cond
          loop_end_cond <- bodyBind loop_end_cond_body

          pure
            ( subExpsRes shapeargs ++ loop_end_cond ++ subExpsRes ses,
              ( I.WhileLoop $ I.paramName loop_while,
                shapepat,
                loop_while : mergepat',
                loop_initial_cond : mergeinit
              )
            )
internaliseAppExp Name
desc AppRes
_ (E.Match Exp
e NonEmpty (CaseBase Info VName)
orig_cs SrcLoc
_) = do
  ses <- Name -> Exp -> InternaliseM [SubExp]
internaliseExp (Name
desc Name -> Name -> Name
forall a. Semigroup a => a -> a -> a
<> Name
"_scrutinee") Exp
e
  cs <- mapM (onCase ses) orig_cs
  case NE.uncons cs of
    (I.Case [Maybe PrimValue]
_ InternaliseM (Body SOACS)
body, Maybe (NonEmpty (Case (InternaliseM (Body SOACS))))
Nothing) ->
      ([SubExpRes] -> [SubExp])
-> InternaliseM [SubExpRes] -> InternaliseM [SubExp]
forall a b. (a -> b) -> InternaliseM a -> InternaliseM b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((SubExpRes -> SubExp) -> [SubExpRes] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map SubExpRes -> SubExp
resSubExp) (InternaliseM [SubExpRes] -> InternaliseM [SubExp])
-> InternaliseM [SubExpRes] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ GBody (Rep InternaliseM) SubExpRes -> InternaliseM [SubExpRes]
Body SOACS -> InternaliseM [SubExpRes]
forall (m :: * -> *).
MonadBuilder m =>
Body (Rep m) -> m [SubExpRes]
bodyBind (Body SOACS -> InternaliseM [SubExpRes])
-> InternaliseM (Body SOACS) -> InternaliseM [SubExpRes]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< InternaliseM (Body SOACS)
body
    (Case (InternaliseM (Body SOACS)),
 Maybe (NonEmpty (Case (InternaliseM (Body SOACS)))))
_ -> do
      Name -> Exp SOACS -> InternaliseM [SubExp]
letValExp' Name
desc (Exp SOACS -> InternaliseM [SubExp])
-> InternaliseM (Exp SOACS) -> InternaliseM [SubExp]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< [SubExp]
-> [Case (InternaliseM (GBody (Rep InternaliseM) SubExpRes))]
-> InternaliseM (GBody (Rep InternaliseM) SubExpRes)
-> InternaliseM (Exp (Rep InternaliseM))
forall (m :: * -> *).
(MonadBuilder m, BranchType (Rep m) ~ TypeBase ExtShape NoMode) =>
[SubExp]
-> [Case (m (Body (Rep m)))] -> m (Body (Rep m)) -> m (Exp (Rep m))
eMatch [SubExp]
ses (NonEmpty (Case (InternaliseM (Body SOACS)))
-> [Case (InternaliseM (Body SOACS))]
forall a. NonEmpty a -> [a]
NE.init NonEmpty (Case (InternaliseM (Body SOACS)))
cs) (Case (InternaliseM (GBody (Rep InternaliseM) SubExpRes))
-> InternaliseM (GBody (Rep InternaliseM) SubExpRes)
forall body. Case body -> body
I.caseBody (Case (InternaliseM (GBody (Rep InternaliseM) SubExpRes))
 -> InternaliseM (GBody (Rep InternaliseM) SubExpRes))
-> Case (InternaliseM (GBody (Rep InternaliseM) SubExpRes))
-> InternaliseM (GBody (Rep InternaliseM) SubExpRes)
forall a b. (a -> b) -> a -> b
$ NonEmpty (Case (InternaliseM (GBody (Rep InternaliseM) SubExpRes)))
-> Case (InternaliseM (GBody (Rep InternaliseM) SubExpRes))
forall a. NonEmpty a -> a
NE.last NonEmpty (Case (InternaliseM (GBody (Rep InternaliseM) SubExpRes)))
NonEmpty (Case (InternaliseM (Body SOACS)))
cs)
  where
    onCase :: [SubExp]
-> CaseBase Info VName
-> InternaliseM (Case (InternaliseM (Body SOACS)))
onCase [SubExp]
ses (E.CasePat PatBase Info VName StructType
p Exp
case_e SrcLoc
_) = do
      (cmps, pertinent) <- PatBase Info VName StructType
-> [SubExp] -> InternaliseM ([Maybe PrimValue], [SubExp])
generateCond PatBase Info VName StructType
p [SubExp]
ses
      pure . I.Case cmps $
        internalisePat' [] p pertinent $
          internaliseBody "case" case_e
internaliseAppExp Name
desc AppRes
_ (E.If Exp
ce Exp
te Exp
fe SrcLoc
_) =
  Name -> Exp SOACS -> InternaliseM [SubExp]
letValExp' Name
desc
    (Exp SOACS -> InternaliseM [SubExp])
-> InternaliseM (Exp SOACS) -> InternaliseM [SubExp]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< InternaliseM (Exp (Rep InternaliseM))
-> InternaliseM (GBody (Rep InternaliseM) SubExpRes)
-> InternaliseM (GBody (Rep InternaliseM) SubExpRes)
-> InternaliseM (Exp (Rep InternaliseM))
forall (m :: * -> *).
(MonadBuilder m, BranchType (Rep m) ~ TypeBase ExtShape NoMode) =>
m (Exp (Rep m))
-> m (Body (Rep m)) -> m (Body (Rep m)) -> m (Exp (Rep m))
eIf
      (BasicOp -> Exp SOACS
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp SOACS)
-> (SubExp -> BasicOp) -> SubExp -> Exp SOACS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SubExp -> BasicOp
SubExp (SubExp -> Exp SOACS)
-> InternaliseM SubExp -> InternaliseM (Exp SOACS)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Name -> Exp -> InternaliseM SubExp
internaliseExp1 Name
"cond" Exp
ce)
      (Name -> Exp -> InternaliseM (Body SOACS)
internaliseBody (Name
desc Name -> Name -> Name
forall a. Semigroup a => a -> a -> a
<> Name
"_t") Exp
te)
      (Name -> Exp -> InternaliseM (Body SOACS)
internaliseBody (Name
desc Name -> Name -> Name
forall a. Semigroup a => a -> a -> a
<> Name
"_f") Exp
fe)
internaliseAppExp Name
_ AppRes
_ e :: AppExp
e@E.BinOp {} =
  [Char] -> InternaliseM [SubExp]
forall a. HasCallStack => [Char] -> a
error ([Char] -> InternaliseM [SubExp])
-> [Char] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ [Char]
"internaliseAppExp: Unexpected BinOp " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ AppExp -> [Char]
forall a. Pretty a => a -> [Char]
prettyString AppExp
e
internaliseAppExp Name
_ AppRes
_ e :: AppExp
e@(E.LetWith {}) =
  [Char] -> InternaliseM [SubExp]
forall a. HasCallStack => [Char] -> a
error ([Char] -> InternaliseM [SubExp])
-> [Char] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ [Char]
"internaliseAppExp: Unexpected LetWith at " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ SrcLoc -> [Char]
forall a. Located a => a -> [Char]
locStr (AppExp -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf AppExp
e)

internaliseExp :: Name -> E.Exp -> InternaliseM [I.SubExp]
internaliseExp :: Name -> Exp -> InternaliseM [SubExp]
internaliseExp Name
desc (E.Parens Exp
e SrcLoc
_) =
  Name -> Exp -> InternaliseM [SubExp]
internaliseExp Name
desc Exp
e
internaliseExp Name
desc (E.Hole (Info StructType
t) SrcLoc
loc) = do
  let msg :: Text
msg = Doc (ZonkAny 0) -> Text
forall a. Doc a -> Text
docText (Doc (ZonkAny 0) -> Text) -> Doc (ZonkAny 0) -> Text
forall a b. (a -> b) -> a -> b
$ Doc (ZonkAny 0)
"Reached hole of type: " Doc (ZonkAny 0) -> Doc (ZonkAny 0) -> Doc (ZonkAny 0)
forall a. Semigroup a => a -> a -> a
<> Doc (ZonkAny 0) -> Doc (ZonkAny 0)
forall ann. Doc ann -> Doc ann
align (StructType -> Doc (ZonkAny 0)
forall a ann. Pretty a => a -> Doc ann
forall ann. StructType -> Doc ann
pretty StructType
t)
      ts :: [TypeBase ExtShape NoMode]
ts = (Tree (TypeBase ExtShape NoMode) -> [TypeBase ExtShape NoMode])
-> [Tree (TypeBase ExtShape NoMode)] -> [TypeBase ExtShape NoMode]
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Tree (TypeBase ExtShape NoMode) -> [TypeBase ExtShape NoMode]
forall a. Free [] a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList ([Tree (TypeBase ExtShape NoMode)] -> [TypeBase ExtShape NoMode])
-> [Tree (TypeBase ExtShape NoMode)] -> [TypeBase ExtShape NoMode]
forall a b. (a -> b) -> a -> b
$ StructType -> [Tree (TypeBase ExtShape NoMode)]
internaliseType (StructType -> StructType
forall dim o. TypeBase dim o -> TypeBase dim NoMode
E.toStruct StructType
t)
  c <- Name -> SubExp -> ErrorMsg SubExp -> InternaliseM Certs
assert Name
"hole_c" (Bool -> SubExp
forall v. IsValue v => v -> SubExp
constant Bool
False) (ErrorMsg SubExp -> InternaliseM Certs)
-> ErrorMsg SubExp -> InternaliseM Certs
forall a b. (a -> b) -> a -> b
$ [ErrorMsgPart SubExp] -> ErrorMsg SubExp
forall a. [ErrorMsgPart a] -> ErrorMsg a
errorMsg [Text -> ErrorMsgPart SubExp
forall a. Text -> ErrorMsgPart a
ErrorString Text
msg]
  case mapM hasStaticShape ts of
    Maybe [TypeBase Shape NoMode]
Nothing ->
      [Char] -> InternaliseM [SubExp]
forall a. HasCallStack => [Char] -> a
error ([Char] -> InternaliseM [SubExp])
-> [Char] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ [Char]
"Hole at " [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> SrcLoc -> [Char]
forall a. Located a => a -> [Char]
locStr SrcLoc
loc [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> [Char]
" has existential type:\n" [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> [TypeBase ExtShape NoMode] -> [Char]
forall a. Show a => a -> [Char]
show [TypeBase ExtShape NoMode]
ts
    Just [TypeBase Shape NoMode]
ts' ->
      -- Make sure we always generate a binding, even for primitives.
      Certs -> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a. Certs -> InternaliseM a -> InternaliseM a
forall (m :: * -> *) a. MonadBuilder m => Certs -> m a -> m a
certifying Certs
c (InternaliseM [SubExp] -> InternaliseM [SubExp])
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ (TypeBase Shape NoMode -> InternaliseM SubExp)
-> [TypeBase Shape NoMode] -> InternaliseM [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 ((VName -> SubExp) -> InternaliseM VName -> InternaliseM SubExp
forall a b. (a -> b) -> InternaliseM a -> InternaliseM b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap VName -> SubExp
I.Var (InternaliseM VName -> InternaliseM SubExp)
-> (Exp SOACS -> InternaliseM VName)
-> Exp SOACS
-> InternaliseM SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Name -> Exp (Rep InternaliseM) -> InternaliseM VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m VName
letExp Name
desc (Exp SOACS -> InternaliseM SubExp)
-> (TypeBase Shape NoMode -> InternaliseM (Exp SOACS))
-> TypeBase Shape NoMode
-> InternaliseM SubExp
forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< TypeBase Shape NoMode -> InternaliseM (Exp (Rep InternaliseM))
TypeBase Shape NoMode -> InternaliseM (Exp SOACS)
forall (m :: * -> *).
MonadBuilder m =>
TypeBase Shape NoMode -> m (Exp (Rep m))
eBlank (TypeBase Shape NoMode -> InternaliseM (Exp SOACS))
-> (TypeBase Shape NoMode -> TypeBase Shape NoMode)
-> TypeBase Shape NoMode
-> InternaliseM (Exp SOACS)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeBase Shape NoMode -> TypeBase Shape NoMode
forall shape o. TypeBase shape o -> TypeBase shape NoMode
I.fromDecl) [TypeBase Shape NoMode]
ts'
internaliseExp Name
desc (E.QualParens (QualName VName, SrcLoc)
_ Exp
e SrcLoc
_) =
  Name -> Exp -> InternaliseM [SubExp]
internaliseExp Name
desc Exp
e
internaliseExp Name
desc (E.StringLit [Word8]
vs SrcLoc
loc) =
  SrcLoc -> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a b. Located a => a -> InternaliseM b -> InternaliseM b
locating SrcLoc
loc (InternaliseM [SubExp] -> InternaliseM [SubExp])
-> (Exp SOACS -> InternaliseM [SubExp])
-> Exp SOACS
-> InternaliseM [SubExp]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (SubExp -> [SubExp])
-> InternaliseM SubExp -> InternaliseM [SubExp]
forall a b. (a -> b) -> InternaliseM a -> InternaliseM b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap SubExp -> [SubExp]
forall a. a -> [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (InternaliseM SubExp -> InternaliseM [SubExp])
-> (Exp SOACS -> InternaliseM SubExp)
-> Exp SOACS
-> InternaliseM [SubExp]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Name -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
desc (Exp SOACS -> InternaliseM [SubExp])
-> Exp SOACS -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$
    BasicOp -> Exp SOACS
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp SOACS) -> BasicOp -> Exp SOACS
forall a b. (a -> b) -> a -> b
$
      [SubExp] -> TypeBase Shape NoMode -> BasicOp
I.ArrayLit ((Word8 -> SubExp) -> [Word8] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map Word8 -> SubExp
forall v. IsValue v => v -> SubExp
constant [Word8]
vs) (TypeBase Shape NoMode -> BasicOp)
-> TypeBase Shape NoMode -> BasicOp
forall a b. (a -> b) -> a -> b
$
        PrimType -> TypeBase Shape NoMode
forall shape o. PrimType -> TypeBase shape o
I.Prim PrimType
int8
internaliseExp Name
_ (E.Var (E.QualName [VName]
_ VName
name) Info StructType
_ SrcLoc
_) = do
  subst <- VName -> InternaliseM (Maybe [SubExp])
lookupSubst VName
name
  case subst of
    Just [SubExp]
substs -> [SubExp] -> InternaliseM [SubExp]
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [SubExp]
substs
    Maybe [SubExp]
Nothing -> [SubExp] -> InternaliseM [SubExp]
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [VName -> SubExp
I.Var VName
name]
internaliseExp Name
desc (E.AppExp AppExp
e (Info AppRes
appres)) = do
  ses <- AppExp -> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a b. Located a => a -> InternaliseM b -> InternaliseM b
locating AppExp
e (InternaliseM [SubExp] -> InternaliseM [SubExp])
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ Name -> AppRes -> AppExp -> InternaliseM [SubExp]
internaliseAppExp Name
desc AppRes
appres AppExp
e
  bindExtSizes appres ses
  pure ses
internaliseExp Name
_ (E.TupLit [] SrcLoc
_) =
  [SubExp] -> InternaliseM [SubExp]
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [PrimValue -> SubExp
forall v. IsValue v => v -> SubExp
constant PrimValue
UnitValue]
internaliseExp Name
_ (E.RecordLit [] SrcLoc
_) =
  [SubExp] -> InternaliseM [SubExp]
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [PrimValue -> SubExp
forall v. IsValue v => v -> SubExp
constant PrimValue
UnitValue]
internaliseExp Name
desc (E.TupLit [Exp]
es SrcLoc
_) = [[SubExp]] -> [SubExp]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[SubExp]] -> [SubExp])
-> InternaliseM [[SubExp]] -> InternaliseM [SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Exp -> InternaliseM [SubExp]) -> [Exp] -> InternaliseM [[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 -> InternaliseM [SubExp]
internaliseExp Name
desc) [Exp]
es
internaliseExp Name
desc (E.RecordLit [FieldBase Info VName]
orig_fields SrcLoc
_) =
  ((Name, [SubExp]) -> [SubExp]) -> [(Name, [SubExp])] -> [SubExp]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (Name, [SubExp]) -> [SubExp]
forall a b. (a, b) -> b
snd ([(Name, [SubExp])] -> [SubExp])
-> ([Map Name [SubExp]] -> [(Name, [SubExp])])
-> [Map Name [SubExp]]
-> [SubExp]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map Name [SubExp] -> [(Name, [SubExp])]
forall a. Map Name a -> [(Name, a)]
sortFields (Map Name [SubExp] -> [(Name, [SubExp])])
-> ([Map Name [SubExp]] -> Map Name [SubExp])
-> [Map Name [SubExp]]
-> [(Name, [SubExp])]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Map Name [SubExp]] -> Map Name [SubExp]
forall (f :: * -> *) k a.
(Foldable f, Ord k) =>
f (Map k a) -> Map k a
M.unions ([Map Name [SubExp]] -> [SubExp])
-> InternaliseM [Map Name [SubExp]] -> InternaliseM [SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (FieldBase Info VName -> InternaliseM (Map Name [SubExp]))
-> [FieldBase Info VName] -> InternaliseM [Map Name [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 FieldBase Info VName -> InternaliseM (Map Name [SubExp])
internaliseField [FieldBase Info VName]
orig_fields
  where
    internaliseField :: FieldBase Info VName -> InternaliseM (Map Name [SubExp])
internaliseField (E.RecordFieldExplicit (L Loc
_ Name
name) Exp
e SrcLoc
_) =
      Name -> [SubExp] -> Map Name [SubExp]
forall k a. k -> a -> Map k a
M.singleton Name
name ([SubExp] -> Map Name [SubExp])
-> InternaliseM [SubExp] -> InternaliseM (Map Name [SubExp])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Name -> Exp -> InternaliseM [SubExp]
internaliseExp Name
desc Exp
e
    internaliseField (E.RecordFieldImplicit (L Loc
_ VName
name) Info StructType
t SrcLoc
loc) =
      FieldBase Info VName -> InternaliseM (Map Name [SubExp])
internaliseField (FieldBase Info VName -> InternaliseM (Map Name [SubExp]))
-> FieldBase Info VName -> InternaliseM (Map Name [SubExp])
forall a b. (a -> b) -> a -> b
$
        L Name -> Exp -> SrcLoc -> FieldBase Info VName
forall (f :: * -> *) vn.
L Name -> ExpBase f vn -> SrcLoc -> FieldBase f vn
E.RecordFieldExplicit
          (Loc -> Name -> L Name
forall a. Loc -> a -> L a
L Loc
forall a. IsLocation a => a
noLoc (VName -> Name
baseName VName
name))
          (QualName VName -> Info StructType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
QualName vn -> f StructType -> SrcLoc -> ExpBase f vn
E.Var (VName -> QualName VName
forall v. v -> QualName v
E.qualName VName
name) Info StructType
t SrcLoc
loc)
          SrcLoc
loc
internaliseExp Name
desc (E.ArrayVal [PrimValue]
vs PrimType
t SrcLoc
loc) =
  SrcLoc -> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a b. Located a => a -> InternaliseM b -> InternaliseM b
locating SrcLoc
loc (InternaliseM [SubExp] -> InternaliseM [SubExp])
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$
    (SubExp -> [SubExp])
-> InternaliseM SubExp -> InternaliseM [SubExp]
forall a b. (a -> b) -> InternaliseM a -> InternaliseM b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap SubExp -> [SubExp]
forall a. a -> [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (InternaliseM SubExp -> InternaliseM [SubExp])
-> (BasicOp -> InternaliseM SubExp)
-> BasicOp
-> InternaliseM [SubExp]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Name -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
desc (Exp SOACS -> InternaliseM SubExp)
-> (BasicOp -> Exp SOACS) -> BasicOp -> InternaliseM SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BasicOp -> Exp SOACS
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> InternaliseM [SubExp])
-> BasicOp -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$
      [PrimValue] -> PrimType -> BasicOp
I.ArrayVal ((PrimValue -> PrimValue) -> [PrimValue] -> [PrimValue]
forall a b. (a -> b) -> [a] -> [b]
map PrimValue -> PrimValue
internalisePrimValue [PrimValue]
vs) (PrimType -> PrimType
internalisePrimType PrimType
t)
internaliseExp Name
desc (E.ArrayLit [Exp]
es (Info StructType
arr_t) SrcLoc
loc)
  -- If this is a multidimensional array literal of primitives, we
  -- treat it specially by flattening it out followed by a reshape.
  -- This cuts down on the amount of statements that are produced, and
  -- thus allows us to efficiently handle huge array literals - a
  -- corner case, but an important one.
  | Just (([Int]
eshape, [Exp]
e') : [([Int], [Exp])]
es') <- (Exp -> Maybe ([Int], [Exp])) -> [Exp] -> Maybe [([Int], [Exp])]
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 Exp -> Maybe ([Int], [Exp])
isArrayLiteral [Exp]
es,
    Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ [Int] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Int]
eshape,
    (([Int], [Exp]) -> Bool) -> [([Int], [Exp])] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (([Int]
eshape [Int] -> [Int] -> Bool
forall a. Eq a => a -> a -> Bool
==) ([Int] -> Bool)
-> (([Int], [Exp]) -> [Int]) -> ([Int], [Exp]) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([Int], [Exp]) -> [Int]
forall a b. (a, b) -> a
fst) [([Int], [Exp])]
es',
    Just StructType
basetype <- Int -> StructType -> Maybe StructType
forall dim o. Int -> TypeBase dim o -> Maybe (TypeBase dim o)
E.peelArray ([Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
eshape) StructType
arr_t = do
      let flat_lit :: Exp
flat_lit = [Exp] -> Info StructType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
[ExpBase f vn] -> f StructType -> SrcLoc -> ExpBase f vn
E.ArrayLit ([Exp]
e' [Exp] -> [Exp] -> [Exp]
forall a. [a] -> [a] -> [a]
++ (([Int], [Exp]) -> [Exp]) -> [([Int], [Exp])] -> [Exp]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap ([Int], [Exp]) -> [Exp]
forall a b. (a, b) -> b
snd [([Int], [Exp])]
es') (StructType -> Info StructType
forall a. a -> Info a
Info StructType
basetype) SrcLoc
loc
          new_shape :: [Int]
new_shape = [Exp] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Exp]
es Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: [Int]
eshape
      flat_arrs <- Name -> Exp -> InternaliseM [VName]
internaliseExpToVars Name
"flat_literal" Exp
flat_lit
      forM flat_arrs $ \VName
flat_arr -> do
        flat_arr_t <- VName -> InternaliseM (TypeBase Shape NoMode)
forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoMode)
lookupType VName
flat_arr
        let new_shape' =
              Shape -> Int -> Shape -> Shape
reshapeOuter
                ([SubExp] -> Shape
forall d. [d] -> ShapeBase d
I.Shape ([SubExp] -> Shape) -> [SubExp] -> Shape
forall a b. (a -> b) -> a -> b
$ (Int -> SubExp) -> [Int] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map (IntType -> Integer -> SubExp
intConst IntType
Int64 (Integer -> SubExp) -> (Int -> Integer) -> Int -> SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Integer
forall a. Integral a => a -> Integer
toInteger) [Int]
new_shape)
                Int
1
                (Shape -> Shape) -> Shape -> Shape
forall a b. (a -> b) -> a -> b
$ TypeBase Shape NoMode -> Shape
forall shape o. ArrayShape shape => TypeBase shape o -> shape
I.arrayShape TypeBase Shape NoMode
flat_arr_t
        letSubExp desc $ I.BasicOp $ I.Reshape flat_arr (reshapeAll (I.arrayShape flat_arr_t) new_shape')
  | Bool
otherwise = do
      es' <- (Exp -> InternaliseM [SubExp]) -> [Exp] -> InternaliseM [[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 -> InternaliseM [SubExp]
internaliseExp Name
"arr_elem") [Exp]
es
      let arr_t_ext = (Tree (TypeBase ExtShape NoMode) -> [TypeBase ExtShape NoMode])
-> [Tree (TypeBase ExtShape NoMode)] -> [TypeBase ExtShape NoMode]
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Tree (TypeBase ExtShape NoMode) -> [TypeBase ExtShape NoMode]
forall a. Free [] a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList ([Tree (TypeBase ExtShape NoMode)] -> [TypeBase ExtShape NoMode])
-> [Tree (TypeBase ExtShape NoMode)] -> [TypeBase ExtShape NoMode]
forall a b. (a -> b) -> a -> b
$ StructType -> [Tree (TypeBase ExtShape NoMode)]
internaliseType (StructType -> [Tree (TypeBase ExtShape NoMode)])
-> StructType -> [Tree (TypeBase ExtShape NoMode)]
forall a b. (a -> b) -> a -> b
$ StructType -> StructType
forall dim o. TypeBase dim o -> TypeBase dim NoMode
E.toStruct StructType
arr_t

      rowtypes <-
        case mapM (fmap rowType . hasStaticShape . I.fromDecl) arr_t_ext of
          Just [TypeBase Shape NoMode]
ts -> [TypeBase Shape NoMode] -> InternaliseM [TypeBase Shape NoMode]
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [TypeBase Shape NoMode]
ts
          Maybe [TypeBase Shape NoMode]
Nothing ->
            -- XXX: the monomorphiser may create single-element array
            -- literals with an unknown row type.  In those cases we
            -- need to look at the types of the actual elements.
            -- Fixing this in the monomorphiser is a lot more tricky
            -- than just working around it here.
            case [[SubExp]]
es' of
              [] -> [Char] -> InternaliseM [TypeBase Shape NoMode]
forall a. HasCallStack => [Char] -> a
error ([Char] -> InternaliseM [TypeBase Shape NoMode])
-> [Char] -> InternaliseM [TypeBase Shape NoMode]
forall a b. (a -> b) -> a -> b
$ [Char]
"internaliseExp ArrayLit: existential type: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ StructType -> [Char]
forall a. Pretty a => a -> [Char]
prettyString StructType
arr_t
              [SubExp]
e' : [[SubExp]]
_ -> (SubExp -> InternaliseM (TypeBase Shape NoMode))
-> [SubExp] -> InternaliseM [TypeBase Shape NoMode]
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 SubExp -> InternaliseM (TypeBase Shape NoMode)
forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoMode)
subExpType [SubExp]
e'

      let arraylit [SubExp]
ks TypeBase Shape NoMode
rt = do
            ks' <-
              (SubExp -> InternaliseM SubExp)
-> [SubExp] -> InternaliseM [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
                ( ErrorMsg SubExp
-> TypeBase Shape NoMode -> Name -> SubExp -> InternaliseM SubExp
ensureShape
                    ErrorMsg SubExp
"shape of element differs from shape of first element"
                    TypeBase Shape NoMode
rt
                    Name
"elem_reshaped"
                )
                [SubExp]
ks
            pure $ I.BasicOp $ I.ArrayLit ks' rt

      mapM (letSubExp desc)
        =<< if null es'
          then mapM (arraylit []) rowtypes
          else zipWithM arraylit (transpose es') rowtypes
  where
    isArrayLiteral :: E.Exp -> Maybe ([Int], [E.Exp])
    isArrayLiteral :: Exp -> Maybe ([Int], [Exp])
isArrayLiteral (E.ArrayLit [Exp]
inner_es Info StructType
_ SrcLoc
_) = do
      (eshape, e) : inner_es' <- (Exp -> Maybe ([Int], [Exp])) -> [Exp] -> Maybe [([Int], [Exp])]
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 Exp -> Maybe ([Int], [Exp])
isArrayLiteral [Exp]
inner_es
      guard $ all ((eshape ==) . fst) inner_es'
      pure (length inner_es : eshape, e ++ concatMap snd inner_es')
    isArrayLiteral Exp
e =
      ([Int], [Exp]) -> Maybe ([Int], [Exp])
forall a. a -> Maybe a
Just ([], [Exp
e])
internaliseExp Name
desc (E.Ascript Exp
e TypeExp Exp VName
_ SrcLoc
_) =
  Name -> Exp -> InternaliseM [SubExp]
internaliseExp Name
desc Exp
e
internaliseExp Name
desc (E.Coerce Exp
e TypeExp Exp VName
_ (Info StructType
et) SrcLoc
_) = do
  ses <- Name -> Exp -> InternaliseM [SubExp]
internaliseExp Name
desc Exp
e
  let ts = StructType -> [TypeBase ExtShape NoMode]
internaliseCoerceType (StructType -> StructType
forall dim o. TypeBase dim o -> TypeBase dim NoMode
E.toStruct StructType
et)
  dt' <- typeExpForError $ toStruct et
  forM (zip ses ts) $ \(SubExp
e', TypeBase ExtShape NoMode
t') -> do
    dims <- TypeBase Shape NoMode -> [SubExp]
forall d m. TypeBase (ShapeBase d) m -> [d]
arrayDims (TypeBase Shape NoMode -> [SubExp])
-> InternaliseM (TypeBase Shape NoMode) -> InternaliseM [SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SubExp -> InternaliseM (TypeBase Shape NoMode)
forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoMode)
subExpType SubExp
e'
    let parts =
          [ErrorMsgPart SubExp
"Value of (desugared) shape ["]
            [ErrorMsgPart SubExp]
-> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
forall a. [a] -> [a] -> [a]
++ ErrorMsgPart SubExp
-> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
forall a. a -> [a] -> [a]
intersperse ErrorMsgPart SubExp
"][" ((SubExp -> ErrorMsgPart SubExp)
-> [SubExp] -> [ErrorMsgPart SubExp]
forall a b. (a -> b) -> [a] -> [b]
map (PrimType -> SubExp -> ErrorMsgPart SubExp
forall a. PrimType -> a -> ErrorMsgPart a
ErrorVal PrimType
int64) [SubExp]
dims)
            [ErrorMsgPart SubExp]
-> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
forall a. [a] -> [a] -> [a]
++ [ErrorMsgPart SubExp
"] cannot match shape of type \""]
            [ErrorMsgPart SubExp]
-> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
forall a. [a] -> [a] -> [a]
++ [ErrorMsgPart SubExp]
dt'
            [ErrorMsgPart SubExp]
-> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
forall a. [a] -> [a] -> [a]
++ [ErrorMsgPart SubExp
"\"."]
    ensureExtShape (errorMsg parts) t' desc e'
internaliseExp Name
desc (E.Negate Exp
e SrcLoc
loc) = SrcLoc -> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a b. Located a => a -> InternaliseM b -> InternaliseM b
locating SrcLoc
loc (InternaliseM [SubExp] -> InternaliseM [SubExp])
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ do
  e' <- Name -> Exp -> InternaliseM SubExp
internaliseExp1 Name
"negate_arg" Exp
e
  et <- subExpType e'
  case et of
    I.Prim PrimType
pt ->
      Name -> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m [SubExp]
letTupExp' Name
desc (Exp (Rep InternaliseM) -> InternaliseM [SubExp])
-> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
forall a b. (a -> b) -> a -> b
$ UnOp -> SubExp -> BasicOp
I.UnOp (PrimType -> UnOp
I.Neg PrimType
pt) SubExp
e'
    TypeBase Shape NoMode
_ -> [Char] -> InternaliseM [SubExp]
forall a. HasCallStack => [Char] -> a
error [Char]
"Futhark.Internalise.internaliseExp: non-primitive type in Negate"
internaliseExp Name
desc (E.Not Exp
e SrcLoc
loc) = SrcLoc -> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a b. Located a => a -> InternaliseM b -> InternaliseM b
locating SrcLoc
loc (InternaliseM [SubExp] -> InternaliseM [SubExp])
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ do
  e' <- Name -> Exp -> InternaliseM SubExp
internaliseExp1 Name
"not_arg" Exp
e
  et <- subExpType e'
  case et of
    I.Prim (I.IntType IntType
t) ->
      Name -> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m [SubExp]
letTupExp' Name
desc (Exp (Rep InternaliseM) -> InternaliseM [SubExp])
-> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
forall a b. (a -> b) -> a -> b
$ UnOp -> SubExp -> BasicOp
I.UnOp (IntType -> UnOp
I.Complement IntType
t) SubExp
e'
    I.Prim PrimType
I.Bool ->
      Name -> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m [SubExp]
letTupExp' Name
desc (Exp (Rep InternaliseM) -> InternaliseM [SubExp])
-> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
forall a b. (a -> b) -> a -> b
$ UnOp -> SubExp -> BasicOp
I.UnOp (PrimType -> UnOp
I.Neg PrimType
I.Bool) SubExp
e'
    TypeBase Shape NoMode
_ ->
      [Char] -> InternaliseM [SubExp]
forall a. HasCallStack => [Char] -> a
error [Char]
"Futhark.Internalise.internaliseExp: non-int/bool type in Not"
internaliseExp Name
desc (E.Update Exp
src [UpdateStep Info VName]
steps Exp
ve Info StructType
_ SrcLoc
loc) = SrcLoc -> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a b. Located a => a -> InternaliseM b -> InternaliseM b
locating SrcLoc
loc (InternaliseM [SubExp] -> InternaliseM [SubExp])
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ do
  src_vals <- Name -> Exp -> InternaliseM [VName]
internaliseExpToVars Name
desc Exp
src
  ve_vals <- internaliseExp "update_val" ve
  lowerPath (E.typeOf src) src_vals steps ve_vals
  where
    lowerPath ::
      E.StructType ->
      [I.VName] ->
      [E.UpdateStep Info VName] ->
      [I.SubExp] ->
      InternaliseM [I.SubExp]
    lowerPath :: StructType
-> [VName]
-> [UpdateStep Info VName]
-> [SubExp]
-> InternaliseM [SubExp]
lowerPath StructType
base_t [VName]
base_vs [UpdateStep Info VName]
path_steps [SubExp]
new_vals
      | (UpdateStep Info VName -> Bool) -> [UpdateStep Info VName] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all UpdateStep Info VName -> Bool
forall {f :: * -> *} {vn}. UpdateStep f vn -> Bool
noSlice [UpdateStep Info VName]
path_steps =
          StructType
-> [VName]
-> [UpdateStep Info VName]
-> [SubExp]
-> InternaliseM [SubExp]
lowerPathFields StructType
base_t [VName]
base_vs [UpdateStep Info VName]
path_steps [SubExp]
new_vals
      where
        noSlice :: UpdateStep f vn -> Bool
noSlice E.UpdateStepSlice {} = Bool
False
        noSlice UpdateStep f vn
_ = Bool
True
    lowerPath StructType
base_t [VName]
base_vs [UpdateStep Info VName]
path_steps [SubExp]
new_vals = do
      let focused0 :: [(Int, Maybe (Slice SubExp))]
focused0 = [Int] -> [Maybe (Slice SubExp)] -> [(Int, Maybe (Slice SubExp))]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 ..] ((VName -> Maybe (Slice SubExp))
-> [VName] -> [Maybe (Slice SubExp)]
forall a b. (a -> b) -> [a] -> [b]
map (Maybe (Slice SubExp) -> VName -> Maybe (Slice SubExp)
forall a b. a -> b -> a
const Maybe (Slice SubExp)
forall a. Maybe a
Nothing) [VName]
base_vs)
      updates <- [VName]
-> Certs
-> StructType
-> [(Int, Maybe (Slice SubExp))]
-> [UpdateStep Info VName]
-> [SubExp]
-> InternaliseM (Map Int SubExp)
go [VName]
base_vs Certs
forall a. Monoid a => a
mempty StructType
base_t [(Int, Maybe (Slice SubExp))]
focused0 [UpdateStep Info VName]
path_steps [SubExp]
new_vals
      pure $
        [ M.findWithDefault (I.Var v) i updates
        | (i, v) <- zip [0 ..] base_vs
        ]

    lowerPathFields ::
      E.StructType ->
      [I.VName] ->
      [E.UpdateStep Info VName] ->
      [I.SubExp] ->
      InternaliseM [I.SubExp]
    lowerPathFields :: StructType
-> [VName]
-> [UpdateStep Info VName]
-> [SubExp]
-> InternaliseM [SubExp]
lowerPathFields StructType
_ [VName]
_ [] [SubExp]
new_vals =
      [SubExp] -> InternaliseM [SubExp]
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [SubExp]
new_vals
    lowerPathFields StructType
base_t [VName]
base_vs (E.UpdateStepField Name
f : [UpdateStep Info VName]
rest) [SubExp]
new_vals = do
      (before, field_vals, after, field_t) <- StructType
-> Name
-> [VName]
-> InternaliseM ([VName], [VName], [VName], StructType)
forall a.
StructType
-> Name -> [a] -> InternaliseM ([a], [a], [a], StructType)
splitField StructType
base_t Name
f [VName]
base_vs
      updated_field_vals <- lowerPathFields field_t field_vals rest new_vals
      pure $ map I.Var before ++ updated_field_vals ++ map I.Var after
    lowerPathFields StructType
_ [VName]
_ [UpdateStep Info VName]
_ [SubExp]
_ =
      [Char] -> InternaliseM [SubExp]
forall a. HasCallStack => [Char] -> a
error [Char]
"lowerPathFields: unexpected slice step"

    go ::
      [I.VName] ->
      Certs ->
      E.StructType ->
      [(Int, Maybe (I.Slice I.SubExp))] ->
      [E.UpdateStep Info VName] ->
      [I.SubExp] ->
      InternaliseM (M.Map Int I.SubExp)
    go :: [VName]
-> Certs
-> StructType
-> [(Int, Maybe (Slice SubExp))]
-> [UpdateStep Info VName]
-> [SubExp]
-> InternaliseM (Map Int SubExp)
go [VName]
root_vs Certs
certs StructType
_ [(Int, Maybe (Slice SubExp))]
focused [] [SubExp]
new_vals = do
      Bool -> InternaliseM () -> InternaliseM ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when ([(Int, Maybe (Slice SubExp))] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [(Int, Maybe (Slice SubExp))]
focused Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= [SubExp] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [SubExp]
new_vals) (InternaliseM () -> InternaliseM ())
-> InternaliseM () -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$
        [Char] -> InternaliseM ()
forall a. HasCallStack => [Char] -> a
error [Char]
"internaliseExp Update: update arity mismatch"

      let errormsg :: ErrorMsg a
errormsg = [ErrorMsgPart a] -> ErrorMsg a
forall a. [ErrorMsgPart a] -> ErrorMsg a
errorMsg [Text -> ErrorMsgPart a
forall a. Text -> ErrorMsgPart a
ErrorString Text
"Shape of slice does not match shape of value."]

          writeBack :: (Int, Maybe (Slice SubExp)) -> SubExp -> InternaliseM SubExp
writeBack (Int
i, Maybe (Slice SubExp)
m_slice) SubExp
ve' = do
            let sname :: VName
sname = [VName]
root_vs [VName] -> Int -> VName
forall a. HasCallStack => [a] -> Int -> a
!! Int
i
            case Maybe (Slice SubExp)
m_slice of
              Maybe (Slice SubExp)
Nothing -> SubExp -> InternaliseM SubExp
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SubExp
ve'
              Just Slice SubExp
sl -> do
                sname_t <- VName -> InternaliseM (TypeBase Shape NoMode)
forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoMode)
lookupType VName
sname
                let rowtype = TypeBase Shape NoMode
sname_t TypeBase Shape NoMode -> [SubExp] -> TypeBase Shape NoMode
forall oldshape o.
TypeBase oldshape o -> [SubExp] -> TypeBase Shape o
`setArrayDims` Slice SubExp -> [SubExp]
forall d. Slice d -> [d]
sliceDims Slice SubExp
sl
                ve'' <- ensureShape errormsg rowtype "update_val_correct_shape" ve'
                certifying certs $ I.Var <$> letInPlace desc sname sl (BasicOp $ SubExp ve'')

      ([(Int, SubExp)] -> Map Int SubExp)
-> InternaliseM [(Int, SubExp)] -> InternaliseM (Map Int SubExp)
forall a b. (a -> b) -> InternaliseM a -> InternaliseM b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [(Int, SubExp)] -> Map Int SubExp
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList (InternaliseM [(Int, SubExp)] -> InternaliseM (Map Int SubExp))
-> ((((Int, Maybe (Slice SubExp)), SubExp)
     -> InternaliseM (Int, SubExp))
    -> InternaliseM [(Int, SubExp)])
-> (((Int, Maybe (Slice SubExp)), SubExp)
    -> InternaliseM (Int, SubExp))
-> InternaliseM (Map Int SubExp)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [((Int, Maybe (Slice SubExp)), SubExp)]
-> (((Int, Maybe (Slice SubExp)), SubExp)
    -> InternaliseM (Int, SubExp))
-> InternaliseM [(Int, SubExp)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM ([(Int, Maybe (Slice SubExp))]
-> [SubExp] -> [((Int, Maybe (Slice SubExp)), SubExp)]
forall a b. [a] -> [b] -> [(a, b)]
zip [(Int, Maybe (Slice SubExp))]
focused [SubExp]
new_vals) ((((Int, Maybe (Slice SubExp)), SubExp)
  -> InternaliseM (Int, SubExp))
 -> InternaliseM (Map Int SubExp))
-> (((Int, Maybe (Slice SubExp)), SubExp)
    -> InternaliseM (Int, SubExp))
-> InternaliseM (Map Int SubExp)
forall a b. (a -> b) -> a -> b
$ \((Int, Maybe (Slice SubExp))
f, SubExp
ve') -> do
        res <- (Int, Maybe (Slice SubExp)) -> SubExp -> InternaliseM SubExp
writeBack (Int, Maybe (Slice SubExp))
f SubExp
ve'
        pure (fst f, res)
    go [VName]
root_vs Certs
certs StructType
base_t [(Int, Maybe (Slice SubExp))]
focused (E.UpdateStepField Name
f : [UpdateStep Info VName]
rest) [SubExp]
new_vals = do
      (_before, field_vals, _after, field_t) <- StructType
-> Name
-> [(Int, Maybe (Slice SubExp))]
-> InternaliseM
     ([(Int, Maybe (Slice SubExp))], [(Int, Maybe (Slice SubExp))],
      [(Int, Maybe (Slice SubExp))], StructType)
forall a.
StructType
-> Name -> [a] -> InternaliseM ([a], [a], [a], StructType)
splitField StructType
base_t Name
f [(Int, Maybe (Slice SubExp))]
focused
      go root_vs certs field_t field_vals rest new_vals
    go [VName]
root_vs Certs
certs StructType
base_t [(Int, Maybe (Slice SubExp))]
focused (E.UpdateStepSlice SliceBase Info VName
slice : [UpdateStep Info VName]
rest) [SubExp]
new_vals = do
      current_dims <- StructType -> InternaliseM [SubExp]
dimsOfType StructType
base_t

      (idxs', cs) <- internaliseSlice current_dims slice

      let current_slice =
            [DimIndex SubExp] -> Slice SubExp
forall d. [DimIndex d] -> Slice d
I.Slice ([DimIndex SubExp] -> Slice SubExp)
-> [DimIndex SubExp] -> Slice SubExp
forall a b. (a -> b) -> a -> b
$
              [DimIndex SubExp]
idxs'
                [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
d -> SubExp -> SubExp -> SubExp -> DimIndex SubExp
forall d. d -> d -> d -> DimIndex d
I.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)))
                  (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]
idxs') [SubExp]
current_dims)
      focused' <-
        forM focused $ \(Int
i, Maybe (Slice SubExp)
m_slice) -> do
          slice' <- case Maybe (Slice SubExp)
m_slice of
            Maybe (Slice SubExp)
Nothing -> do
              root_dims <- TypeBase Shape NoMode -> [SubExp]
forall d m. TypeBase (ShapeBase d) m -> [d]
I.arrayDims (TypeBase Shape NoMode -> [SubExp])
-> InternaliseM (TypeBase Shape NoMode) -> InternaliseM [SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> InternaliseM (TypeBase Shape NoMode)
forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoMode)
lookupType ([VName]
root_vs [VName] -> Int -> VName
forall a. HasCallStack => [a] -> Int -> a
!! Int
i)
              pure . I.Slice $
                idxs'
                  ++ map
                    (\SubExp
d -> SubExp -> SubExp -> SubExp -> DimIndex SubExp
forall d. d -> d -> d -> DimIndex d
I.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)))
                    (drop (length idxs') root_dims)
            Just Slice SubExp
sl -> Slice (TPrimExp Int64 VName) -> InternaliseM (Slice SubExp)
forall (m :: * -> *).
MonadBuilder m =>
Slice (TPrimExp Int64 VName) -> m (Slice SubExp)
subExpSlice (Slice (TPrimExp Int64 VName) -> InternaliseM (Slice SubExp))
-> Slice (TPrimExp Int64 VName) -> InternaliseM (Slice SubExp)
forall a b. (a -> b) -> a -> b
$ Slice (TPrimExp Int64 VName)
-> Slice (TPrimExp Int64 VName) -> Slice (TPrimExp Int64 VName)
forall d. Num d => Slice d -> Slice d -> Slice d
sliceSlice (Slice SubExp -> Slice (TPrimExp Int64 VName)
primExpSlice Slice SubExp
sl) (Slice SubExp -> Slice (TPrimExp Int64 VName)
primExpSlice Slice SubExp
current_slice)
          pure (i, Just slice')

      go root_vs (certs <> cs) (indexType base_t slice) focused' rest new_vals

    dimsOfType :: E.StructType -> InternaliseM [SubExp]
    dimsOfType :: StructType -> InternaliseM [SubExp]
dimsOfType (E.Array NoMode
_ Shape Exp
shape ScalarTypeBase Exp NoMode
_) = (Exp -> InternaliseM SubExp) -> [Exp] -> InternaliseM [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 Exp -> InternaliseM SubExp
dimOfSize ([Exp] -> InternaliseM [SubExp]) -> [Exp] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ Shape Exp -> [Exp]
forall dim. Shape dim -> [dim]
E.shapeDims Shape Exp
shape
    dimsOfType StructType
_ = [SubExp] -> InternaliseM [SubExp]
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure []

    dimOfSize :: E.Size -> InternaliseM SubExp
    dimOfSize :: Exp -> InternaliseM SubExp
dimOfSize Exp
d
      | Just Int
_ <- Exp -> Maybe Int
E.isAnySize Exp
d =
          [Char] -> InternaliseM SubExp
forall a. HasCallStack => [Char] -> a
error [Char]
"lowerPath: unexpected anonymous size in slice type"
    dimOfSize (E.IntLit Integer
n Info StructType
_ SrcLoc
_) = SubExp -> InternaliseM SubExp
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SubExp -> InternaliseM SubExp) -> SubExp -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ IntType -> Integer -> SubExp
intConst IntType
Int64 Integer
n
    dimOfSize (E.Var (E.QualName [VName]
_ VName
v) Info StructType
_ SrcLoc
_) = SubExp -> InternaliseM SubExp
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SubExp -> InternaliseM SubExp) -> SubExp -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ VName -> SubExp
I.Var VName
v
    dimOfSize Exp
d =
      [Char] -> InternaliseM SubExp
forall a. HasCallStack => [Char] -> a
error ([Char] -> InternaliseM SubExp) -> [Char] -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ [Char]
"lowerPath: unexpected size expression " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Exp -> [Char]
forall a. Pretty a => a -> [Char]
prettyString Exp
d

    splitField ::
      E.StructType ->
      Name ->
      [a] ->
      InternaliseM ([a], [a], [a], E.StructType)
    splitField :: forall a.
StructType
-> Name -> [a] -> InternaliseM ([a], [a], [a], StructType)
splitField (E.Scalar (E.Record Map Name StructType
m)) Name
f [a]
vals
      | Just StructType
field_t <- Name -> Map Name StructType -> Maybe StructType
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name
f Map Name StructType
m =
          let i :: Int
i =
                [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ([Int] -> Int)
-> ([(Name, StructType)] -> [Int]) -> [(Name, StructType)] -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Name, StructType) -> Int) -> [(Name, StructType)] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (StructType -> Int
forall o. TypeBase Exp o -> Int
internalisedTypeSize (StructType -> Int)
-> ((Name, StructType) -> StructType) -> (Name, StructType) -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Name, StructType) -> StructType
forall a b. (a, b) -> b
snd) ([(Name, StructType)] -> Int) -> [(Name, StructType)] -> Int
forall a b. (a -> b) -> a -> b
$
                  ((Name, StructType) -> Bool)
-> [(Name, StructType)] -> [(Name, StructType)]
forall a. (a -> Bool) -> [a] -> [a]
takeWhile ((Name -> Name -> Bool
forall a. Eq a => a -> a -> Bool
/= Name
f) (Name -> Bool)
-> ((Name, StructType) -> Name) -> (Name, StructType) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Name, StructType) -> Name
forall a b. (a, b) -> a
fst) ([(Name, StructType)] -> [(Name, StructType)])
-> [(Name, StructType)] -> [(Name, StructType)]
forall a b. (a -> b) -> a -> b
$
                    Map Name StructType -> [(Name, StructType)]
forall a. Map Name a -> [(Name, a)]
sortFields Map Name StructType
m
              k :: Int
k = StructType -> Int
forall o. TypeBase Exp o -> Int
internalisedTypeSize StructType
field_t
              ([a]
before, [a]
field_vals, [a]
after) = Int -> Int -> [a] -> ([a], [a], [a])
forall a. Int -> Int -> [a] -> ([a], [a], [a])
splitAt3 Int
i Int
k [a]
vals
           in ([a], [a], [a], StructType)
-> InternaliseM ([a], [a], [a], StructType)
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([a]
before, [a]
field_vals, [a]
after, StructType
field_t)
    splitField StructType
t Name
f [a]
_ =
      [Char] -> InternaliseM ([a], [a], [a], StructType)
forall a. HasCallStack => [Char] -> a
error ([Char] -> InternaliseM ([a], [a], [a], StructType))
-> [Char] -> InternaliseM ([a], [a], [a], StructType)
forall a b. (a -> b) -> a -> b
$
        [Char]
"lowerPath: missing field "
          [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Name -> [Char]
forall a. Pretty a => a -> [Char]
prettyString Name
f
          [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
" in type "
          [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ StructType -> [Char]
forall a. Pretty a => a -> [Char]
prettyString StructType
t

    indexType :: E.StructType -> [E.DimIndex] -> E.StructType
    indexType :: StructType -> SliceBase Info VName -> StructType
indexType (E.Array NoMode
o (E.Shape [Exp]
dims) ScalarTypeBase Exp NoMode
et) SliceBase Info VName
idxs =
      case [Exp]
dims' of
        [] -> ScalarTypeBase Exp NoMode -> StructType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
E.Scalar ScalarTypeBase Exp NoMode
et
        [Exp]
ds -> NoMode -> Shape Exp -> ScalarTypeBase Exp NoMode -> StructType
forall dim o.
o -> Shape dim -> ScalarTypeBase dim NoMode -> TypeBase dim o
E.Array NoMode
o ([Exp] -> Shape Exp
forall dim. [dim] -> Shape dim
E.Shape [Exp]
ds) ScalarTypeBase Exp NoMode
et
      where
        dims' :: [Exp]
dims' = [Exp]
keptPrefix [Exp] -> [Exp] -> [Exp]
forall a. Semigroup a => a -> a -> a
<> [Exp]
suffix
        keptPrefix :: [Exp]
keptPrefix = [Exp
d | (Exp
d, DimIndex
i) <- [Exp] -> SliceBase Info VName -> [(Exp, DimIndex)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Exp]
prefix SliceBase Info VName
idxs, DimIndex -> Bool
forall {f :: * -> *} {vn}. DimIndexBase f vn -> Bool
keepDim DimIndex
i]
        prefix :: [Exp]
prefix = Int -> [Exp] -> [Exp]
forall a. Int -> [a] -> [a]
take (SliceBase Info VName -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length SliceBase Info VName
idxs) [Exp]
dims
        suffix :: [Exp]
suffix = Int -> [Exp] -> [Exp]
forall a. Int -> [a] -> [a]
drop (SliceBase Info VName -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length SliceBase Info VName
idxs) [Exp]
dims
        keepDim :: DimIndexBase f vn -> Bool
keepDim E.DimFix {} = Bool
False
        keepDim E.DimSlice {} = Bool
True
    indexType StructType
t SliceBase Info VName
_ = StructType
t
internaliseExp Name
desc (E.Attr AttrInfo VName
attr Exp
e SrcLoc
loc) = do
  attr' <- AttrInfo VName -> InternaliseM Attr
internaliseAttr AttrInfo VName
attr
  e' <- local (f attr') $ internaliseExp desc e
  case attr' of
    Attr
"trace" ->
      Text -> [SubExp] -> InternaliseM [SubExp]
traceRes ([Char] -> Text
T.pack ([Char] -> Text) -> [Char] -> Text
forall a b. (a -> b) -> a -> b
$ SrcLoc -> [Char]
forall a. Located a => a -> [Char]
locStr SrcLoc
loc) [SubExp]
e'
    I.AttrComp Name
"trace" [I.AttrName Name
tag] ->
      Text -> [SubExp] -> InternaliseM [SubExp]
traceRes (Name -> Text
nameToText Name
tag) [SubExp]
e'
    Attr
"opaque" ->
      (SubExp -> InternaliseM SubExp)
-> [SubExp] -> InternaliseM [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 InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
desc (Exp SOACS -> InternaliseM SubExp)
-> (SubExp -> Exp SOACS) -> SubExp -> InternaliseM SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BasicOp -> Exp SOACS
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp SOACS)
-> (SubExp -> BasicOp) -> SubExp -> Exp SOACS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. OpaqueOp -> SubExp -> BasicOp
Opaque OpaqueOp
OpaqueNil) [SubExp]
e'
    Attr
"scratch" -> do
      ts <- (SubExp -> InternaliseM (TypeBase Shape NoMode))
-> [SubExp] -> InternaliseM [TypeBase Shape NoMode]
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 SubExp -> InternaliseM (TypeBase Shape NoMode)
forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoMode)
subExpType [SubExp]
e'
      forM (zip ts e') $ \(TypeBase Shape NoMode
t, SubExp
se) ->
        case TypeBase Shape NoMode
t of
          I.Array PrimType
pt Shape
shape NoMode
_ ->
            Name -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
desc (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
forall a b. (a -> b) -> a -> b
$ PrimType -> [SubExp] -> BasicOp
I.Scratch PrimType
pt ([SubExp] -> BasicOp) -> [SubExp] -> BasicOp
forall a b. (a -> b) -> a -> b
$ Shape -> [SubExp]
forall d. ShapeBase d -> [d]
I.shapeDims Shape
shape
          -- Ignore scratch on non-arrays because they are sometimes applied
          -- too promisciously.
          TypeBase Shape NoMode
_ -> SubExp -> InternaliseM SubExp
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SubExp
se
    Attr
"blank" -> do
      ts <- (SubExp -> InternaliseM (TypeBase Shape NoMode))
-> [SubExp] -> InternaliseM [TypeBase Shape NoMode]
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 SubExp -> InternaliseM (TypeBase Shape NoMode)
forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoMode)
subExpType [SubExp]
e'
      forM (zip ts e') $ \(TypeBase Shape NoMode
t, SubExp
se) ->
        case TypeBase Shape NoMode
t of
          I.Array PrimType
pt Shape
shape NoMode
_ ->
            Name -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
desc (Exp SOACS -> InternaliseM SubExp)
-> (PrimValue -> Exp SOACS) -> PrimValue -> InternaliseM SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BasicOp -> Exp SOACS
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp SOACS)
-> (PrimValue -> BasicOp) -> PrimValue -> Exp SOACS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Shape -> SubExp -> BasicOp
I.Replicate Shape
shape (SubExp -> BasicOp)
-> (PrimValue -> SubExp) -> PrimValue -> BasicOp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PrimValue -> SubExp
forall v. IsValue v => v -> SubExp
constant (PrimValue -> InternaliseM SubExp)
-> PrimValue -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
              PrimType -> PrimValue
blankPrimValue PrimType
pt
          I.Prim PrimType
pt ->
            SubExp -> InternaliseM SubExp
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SubExp -> InternaliseM SubExp) -> SubExp -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ PrimValue -> SubExp
forall v. IsValue v => v -> SubExp
constant (PrimValue -> SubExp) -> PrimValue -> SubExp
forall a b. (a -> b) -> a -> b
$ PrimType -> PrimValue
blankPrimValue PrimType
pt
          TypeBase Shape NoMode
_ -> SubExp -> InternaliseM SubExp
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SubExp
se
    I.AttrComp Name
"param" [I.AttrName Name
tag] -> do
      case [SubExp]
e' of
        [SubExp
se] -> do
          se_t <- SubExp -> InternaliseM (TypeBase Shape NoMode)
forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoMode)
I.subExpType SubExp
se
          if se_t == I.Prim int64
            then fmap pure $ letSubExp desc $ I.BasicOp $ I.UserParam tag se
            else pure e'
        [SubExp]
_ -> [SubExp] -> InternaliseM [SubExp]
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [SubExp]
e'
    Attr
_ ->
      [SubExp] -> InternaliseM [SubExp]
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [SubExp]
e'
  where
    traceRes :: Text -> [SubExp] -> InternaliseM [SubExp]
traceRes Text
tag' =
      (SubExp -> InternaliseM SubExp)
-> [SubExp] -> InternaliseM [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 InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
desc (Exp SOACS -> InternaliseM SubExp)
-> (SubExp -> Exp SOACS) -> SubExp -> InternaliseM SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BasicOp -> Exp SOACS
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp SOACS)
-> (SubExp -> BasicOp) -> SubExp -> Exp SOACS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. OpaqueOp -> SubExp -> BasicOp
Opaque (Text -> OpaqueOp
OpaqueTrace Text
tag'))
    f :: Attr -> InternaliseEnv -> InternaliseEnv
f Attr
attr' InternaliseEnv
env
      | Attr
attr' Attr -> Attr -> Bool
forall a. Eq a => a -> a -> Bool
== Attr
"unsafe",
        Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ InternaliseEnv -> Bool
envSafe InternaliseEnv
env =
          InternaliseEnv
env {envDoBoundsChecks = False}
      | Bool
otherwise =
          InternaliseEnv
env {envAttrs = envAttrs env <> oneAttr attr'}
internaliseExp Name
desc (E.Assert Exp
e1 Exp
e2 (Info Text
check) SrcLoc
loc) = SrcLoc -> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a b. Located a => a -> InternaliseM b -> InternaliseM b
locating SrcLoc
loc (InternaliseM [SubExp] -> InternaliseM [SubExp])
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ do
  e1' <- Name -> Exp -> InternaliseM SubExp
internaliseExp1 Name
"assert_cond" Exp
e1
  c <- assert "assert_c" e1' $ errorMsg [ErrorString $ "Assertion is false: " <> check]
  -- Make sure there are some bindings to certify.
  certifying c $ mapM rebind =<< internaliseExp desc e2
  where
    rebind :: SubExp -> m SubExp
rebind SubExp
v = do
      v' <- Name -> m VName
forall (m :: * -> *). MonadFreshNames m => Name -> m VName
newVName Name
"assert_res"
      letBindNames [v'] $ I.BasicOp $ I.SubExp v
      pure $ I.Var v'
internaliseExp Name
_ (E.Constr Name
c [Exp]
es (Info (E.Scalar (E.Sum Map Name [StructType]
fs))) SrcLoc
loc) = SrcLoc -> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a b. Located a => a -> InternaliseM b -> InternaliseM b
locating SrcLoc
loc (InternaliseM [SubExp] -> InternaliseM [SubExp])
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ do
  (ts, constr_map) <- Map Name [StructType]
-> InternaliseM ([TypeBase ExtShape NoMode], [(Name, [Int])])
internaliseSumType (Map Name [StructType]
 -> InternaliseM ([TypeBase ExtShape NoMode], [(Name, [Int])]))
-> Map Name [StructType]
-> InternaliseM ([TypeBase ExtShape NoMode], [(Name, [Int])])
forall a b. (a -> b) -> a -> b
$ ([StructType] -> [StructType])
-> Map Name [StructType] -> Map Name [StructType]
forall a b k. (a -> b) -> Map k a -> Map k b
M.map ((StructType -> StructType) -> [StructType] -> [StructType]
forall a b. (a -> b) -> [a] -> [b]
map StructType -> StructType
forall dim o. TypeBase dim o -> TypeBase dim NoMode
E.toStruct) Map Name [StructType]
fs
  es' <- concat <$> mapM (internaliseExp "payload") es

  let noExt p
_ = SubExp -> f SubExp
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SubExp -> f SubExp) -> SubExp -> f SubExp
forall a b. (a -> b) -> a -> b
$ IntType -> Integer -> SubExp
intConst IntType
Int64 Integer
0
  ts' <- instantiateShapes noExt $ map fromDecl ts

  case lookupWithIndex c constr_map of
    Just (Int
i, [Int]
js)
      | Map Name [StructType] -> Int
forall a. Map Name a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Map Name [StructType]
fs Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1 ->
          Int
-> [TypeBase Shape NoMode]
-> [(Int, SubExp)]
-> InternaliseM [SubExp]
forall {f :: * -> *} {t}.
(Num t, MonadBuilder f, Eq t) =>
t -> [TypeBase Shape NoMode] -> [(t, SubExp)] -> f [SubExp]
clauses Int
0 [TypeBase Shape NoMode]
ts' ([Int] -> [SubExp] -> [(Int, SubExp)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int]
js [SubExp]
es')
      | Bool
otherwise ->
          (IntType -> Integer -> SubExp
intConst IntType
Int8 (Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
i) SubExp -> [SubExp] -> [SubExp]
forall a. a -> [a] -> [a]
:) ([SubExp] -> [SubExp])
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int
-> [TypeBase Shape NoMode]
-> [(Int, SubExp)]
-> InternaliseM [SubExp]
forall {f :: * -> *} {t}.
(Num t, MonadBuilder f, Eq t) =>
t -> [TypeBase Shape NoMode] -> [(t, SubExp)] -> f [SubExp]
clauses Int
0 [TypeBase Shape NoMode]
ts' ([Int] -> [SubExp] -> [(Int, SubExp)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int]
js [SubExp]
es')
    Maybe (Int, [Int])
Nothing ->
      [Char] -> InternaliseM [SubExp]
forall a. HasCallStack => [Char] -> a
error [Char]
"internaliseExp Constr: missing constructor"
  where
    clauses :: t -> [TypeBase Shape NoMode] -> [(t, SubExp)] -> f [SubExp]
clauses t
j (TypeBase Shape NoMode
t : [TypeBase Shape NoMode]
ts) [(t, SubExp)]
js_to_es
      | Just SubExp
e <- t
j t -> [(t, SubExp)] -> Maybe SubExp
forall a b. Eq a => a -> [(a, b)] -> Maybe b
`lookup` [(t, SubExp)]
js_to_es =
          (SubExp
e SubExp -> [SubExp] -> [SubExp]
forall a. a -> [a] -> [a]
:) ([SubExp] -> [SubExp]) -> f [SubExp] -> f [SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> t -> [TypeBase Shape NoMode] -> [(t, SubExp)] -> f [SubExp]
clauses (t
j t -> t -> t
forall a. Num a => a -> a -> a
+ t
1) [TypeBase Shape NoMode]
ts [(t, SubExp)]
js_to_es
      | Bool
otherwise = do
          blank <-
            -- Cannot use eBlank here for arrays, because when doing
            -- equality comparisons on sum types, we end up looking at
            -- the array elements. (#2081) This is a bit of an edge
            -- case, but arrays in sum types are known to be
            -- inefficient.
            Name -> Exp (Rep f) -> f SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"zero"
              (Exp (Rep f) -> f SubExp) -> f (Exp (Rep f)) -> f SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< case TypeBase Shape NoMode
t of
                I.Array {} ->
                  Exp (Rep f) -> f (Exp (Rep f))
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp (Rep f) -> f (Exp (Rep f))) -> Exp (Rep f) -> f (Exp (Rep f))
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep f)
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp (Rep f)) -> BasicOp -> Exp (Rep f)
forall a b. (a -> b) -> a -> b
$ Shape -> SubExp -> BasicOp
Replicate (TypeBase Shape NoMode -> Shape
forall shape o. ArrayShape shape => TypeBase shape o -> shape
I.arrayShape TypeBase Shape NoMode
t) (SubExp -> BasicOp) -> SubExp -> BasicOp
forall a b. (a -> b) -> a -> b
$ PrimValue -> SubExp
I.Constant (PrimValue -> SubExp) -> PrimValue -> SubExp
forall a b. (a -> b) -> a -> b
$ PrimType -> PrimValue
blankPrimValue (PrimType -> PrimValue) -> PrimType -> PrimValue
forall a b. (a -> b) -> a -> b
$ TypeBase Shape NoMode -> PrimType
forall shape o. TypeBase shape o -> PrimType
elemType TypeBase Shape NoMode
t
                TypeBase Shape NoMode
_ -> TypeBase Shape NoMode -> f (Exp (Rep f))
forall (m :: * -> *).
MonadBuilder m =>
TypeBase Shape NoMode -> m (Exp (Rep m))
eBlank TypeBase Shape NoMode
t
          (blank :) <$> clauses (j + 1) ts js_to_es
    clauses t
_ [] [(t, SubExp)]
_ =
      [SubExp] -> f [SubExp]
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure []
internaliseExp Name
_ (E.Constr Name
_ [Exp]
_ (Info StructType
t) SrcLoc
loc) =
  [Char] -> InternaliseM [SubExp]
forall a. HasCallStack => [Char] -> a
error ([Char] -> InternaliseM [SubExp])
-> [Char] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ [Char]
"internaliseExp: constructor with type " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ StructType -> [Char]
forall a. Pretty a => a -> [Char]
prettyString StructType
t [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
" at " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ SrcLoc -> [Char]
forall a. Located a => a -> [Char]
locStr SrcLoc
loc
-- The "interesting" cases are over, now it's mostly boilerplate.

internaliseExp Name
_ (E.Literal PrimValue
v SrcLoc
_) =
  [SubExp] -> InternaliseM [SubExp]
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [PrimValue -> SubExp
I.Constant (PrimValue -> SubExp) -> PrimValue -> SubExp
forall a b. (a -> b) -> a -> b
$ PrimValue -> PrimValue
internalisePrimValue PrimValue
v]
internaliseExp Name
_ (E.IntLit Integer
v (Info StructType
t) SrcLoc
_) =
  case StructType
t of
    E.Scalar (E.Prim (E.Signed IntType
it)) ->
      [SubExp] -> InternaliseM [SubExp]
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [PrimValue -> SubExp
I.Constant (PrimValue -> SubExp) -> PrimValue -> SubExp
forall a b. (a -> b) -> a -> b
$ IntValue -> PrimValue
I.IntValue (IntValue -> PrimValue) -> IntValue -> PrimValue
forall a b. (a -> b) -> a -> b
$ IntType -> Integer -> IntValue
forall int. Integral int => IntType -> int -> IntValue
intValue IntType
it Integer
v]
    E.Scalar (E.Prim (E.Unsigned IntType
it)) ->
      [SubExp] -> InternaliseM [SubExp]
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [PrimValue -> SubExp
I.Constant (PrimValue -> SubExp) -> PrimValue -> SubExp
forall a b. (a -> b) -> a -> b
$ IntValue -> PrimValue
I.IntValue (IntValue -> PrimValue) -> IntValue -> PrimValue
forall a b. (a -> b) -> a -> b
$ IntType -> Integer -> IntValue
forall int. Integral int => IntType -> int -> IntValue
intValue IntType
it Integer
v]
    E.Scalar (E.Prim (E.FloatType FloatType
ft)) ->
      [SubExp] -> InternaliseM [SubExp]
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [PrimValue -> SubExp
I.Constant (PrimValue -> SubExp) -> PrimValue -> SubExp
forall a b. (a -> b) -> a -> b
$ FloatValue -> PrimValue
I.FloatValue (FloatValue -> PrimValue) -> FloatValue -> PrimValue
forall a b. (a -> b) -> a -> b
$ FloatType -> Integer -> FloatValue
forall num. Real num => FloatType -> num -> FloatValue
floatValue FloatType
ft Integer
v]
    StructType
_ -> [Char] -> InternaliseM [SubExp]
forall a. HasCallStack => [Char] -> a
error ([Char] -> InternaliseM [SubExp])
-> [Char] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ [Char]
"internaliseExp: nonsensical type for integer literal: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ StructType -> [Char]
forall a. Pretty a => a -> [Char]
prettyString StructType
t
internaliseExp Name
_ (E.FloatLit Double
v (Info StructType
t) SrcLoc
_) =
  case StructType
t of
    E.Scalar (E.Prim (E.FloatType FloatType
ft)) ->
      [SubExp] -> InternaliseM [SubExp]
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [PrimValue -> SubExp
I.Constant (PrimValue -> SubExp) -> PrimValue -> SubExp
forall a b. (a -> b) -> a -> b
$ FloatValue -> PrimValue
I.FloatValue (FloatValue -> PrimValue) -> FloatValue -> PrimValue
forall a b. (a -> b) -> a -> b
$ FloatType -> Double -> FloatValue
forall num. Real num => FloatType -> num -> FloatValue
floatValue FloatType
ft Double
v]
    StructType
_ -> [Char] -> InternaliseM [SubExp]
forall a. HasCallStack => [Char] -> a
error ([Char] -> InternaliseM [SubExp])
-> [Char] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ [Char]
"internaliseExp: nonsensical type for float literal: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ StructType -> [Char]
forall a. Pretty a => a -> [Char]
prettyString StructType
t
-- Builtin operators are handled specially because they are
-- overloaded.
internaliseExp Name
desc (E.Project Name
k Exp
e (Info StructType
rt) SrcLoc
_) = do
  let i' :: Int
i' = [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ([Int] -> Int) -> ([StructType] -> [Int]) -> [StructType] -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (StructType -> Int) -> [StructType] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map StructType -> Int
forall o. TypeBase Exp o -> Int
internalisedTypeSize ([StructType] -> Int) -> [StructType] -> Int
forall a b. (a -> b) -> a -> b
$
        case Exp -> StructType
E.typeOf Exp
e of
          E.Scalar (Record Map Name StructType
fs) ->
            ((Name, StructType) -> StructType)
-> [(Name, StructType)] -> [StructType]
forall a b. (a -> b) -> [a] -> [b]
map (Name, StructType) -> StructType
forall a b. (a, b) -> b
snd ([(Name, StructType)] -> [StructType])
-> [(Name, StructType)] -> [StructType]
forall a b. (a -> b) -> a -> b
$ ((Name, StructType) -> Bool)
-> [(Name, StructType)] -> [(Name, StructType)]
forall a. (a -> Bool) -> [a] -> [a]
takeWhile ((Name -> Name -> Bool
forall a. Eq a => a -> a -> Bool
/= Name
k) (Name -> Bool)
-> ((Name, StructType) -> Name) -> (Name, StructType) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Name, StructType) -> Name
forall a b. (a, b) -> a
fst) ([(Name, StructType)] -> [(Name, StructType)])
-> [(Name, StructType)] -> [(Name, StructType)]
forall a b. (a -> b) -> a -> b
$ Map Name StructType -> [(Name, StructType)]
forall a. Map Name a -> [(Name, a)]
sortFields Map Name StructType
fs
          StructType
t -> [StructType
t]
  Int -> [SubExp] -> [SubExp]
forall a. Int -> [a] -> [a]
take (StructType -> Int
forall o. TypeBase Exp o -> Int
internalisedTypeSize StructType
rt) ([SubExp] -> [SubExp])
-> ([SubExp] -> [SubExp]) -> [SubExp] -> [SubExp]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> [SubExp] -> [SubExp]
forall a. Int -> [a] -> [a]
drop Int
i' ([SubExp] -> [SubExp])
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Name -> Exp -> InternaliseM [SubExp]
internaliseExp Name
desc Exp
e
internaliseExp Name
_ e :: Exp
e@E.Lambda {} =
  [Char] -> InternaliseM [SubExp]
forall a. HasCallStack => [Char] -> a
error ([Char] -> InternaliseM [SubExp])
-> [Char] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ [Char]
"internaliseExp: Unexpected lambda at " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ SrcLoc -> [Char]
forall a. Located a => a -> [Char]
locStr (Exp -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf Exp
e)
internaliseExp Name
_ e :: Exp
e@E.OpSection {} =
  [Char] -> InternaliseM [SubExp]
forall a. HasCallStack => [Char] -> a
error ([Char] -> InternaliseM [SubExp])
-> [Char] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ [Char]
"internaliseExp: Unexpected operator section at " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ SrcLoc -> [Char]
forall a. Located a => a -> [Char]
locStr (Exp -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf Exp
e)
internaliseExp Name
_ e :: Exp
e@E.OpSectionLeft {} =
  [Char] -> InternaliseM [SubExp]
forall a. HasCallStack => [Char] -> a
error ([Char] -> InternaliseM [SubExp])
-> [Char] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ [Char]
"internaliseExp: Unexpected left operator section at " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ SrcLoc -> [Char]
forall a. Located a => a -> [Char]
locStr (Exp -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf Exp
e)
internaliseExp Name
_ e :: Exp
e@E.OpSectionRight {} =
  [Char] -> InternaliseM [SubExp]
forall a. HasCallStack => [Char] -> a
error ([Char] -> InternaliseM [SubExp])
-> [Char] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ [Char]
"internaliseExp: Unexpected right operator section at " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ SrcLoc -> [Char]
forall a. Located a => a -> [Char]
locStr (Exp -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf Exp
e)
internaliseExp Name
_ e :: Exp
e@E.UpdateSection {} =
  [Char] -> InternaliseM [SubExp]
forall a. HasCallStack => [Char] -> a
error ([Char] -> InternaliseM [SubExp])
-> [Char] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ [Char]
"internaliseExp: Unexpected projection section at " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ SrcLoc -> [Char]
forall a. Located a => a -> [Char]
locStr (Exp -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf Exp
e)

internaliseArg :: Name -> (E.Exp, Maybe VName) -> InternaliseM [SubExp]
internaliseArg :: Name -> (Exp, Maybe VName) -> InternaliseM [SubExp]
internaliseArg Name
desc (Exp
arg, Maybe VName
argdim) = do
  exists <- InternaliseM (Scope SOACS)
forall rep (m :: * -> *). HasScope rep m => m (Scope rep)
askScope
  case argdim of
    Just VName
d | VName
d VName -> Scope SOACS -> Bool
forall k a. Ord k => k -> Map k a -> Bool
`M.member` Scope SOACS
exists -> [SubExp] -> InternaliseM [SubExp]
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [VName -> SubExp
I.Var VName
d]
    Maybe VName
_ -> do
      arg' <- Name -> Exp -> InternaliseM [SubExp]
internaliseExp Name
desc Exp
arg
      case (arg', argdim) of
        ([SubExp
se], Just VName
d) -> do
          [VName] -> Exp (Rep InternaliseM) -> InternaliseM ()
forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [VName
d] (Exp (Rep InternaliseM) -> InternaliseM ())
-> Exp (Rep InternaliseM) -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
SubExp SubExp
se
        ([SubExp], Maybe VName)
_ -> () -> InternaliseM ()
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
      pure arg'

internalisePatLit :: E.PatLit -> E.StructType -> I.PrimValue
internalisePatLit :: PatLit -> StructType -> PrimValue
internalisePatLit (E.PatLitPrim PrimValue
v) StructType
_ =
  PrimValue -> PrimValue
internalisePrimValue PrimValue
v
internalisePatLit (E.PatLitInt Integer
x) (E.Scalar (E.Prim (E.Signed IntType
it))) =
  IntValue -> PrimValue
I.IntValue (IntValue -> PrimValue) -> IntValue -> PrimValue
forall a b. (a -> b) -> a -> b
$ IntType -> Integer -> IntValue
forall int. Integral int => IntType -> int -> IntValue
intValue IntType
it Integer
x
internalisePatLit (E.PatLitInt Integer
x) (E.Scalar (E.Prim (E.Unsigned IntType
it))) =
  IntValue -> PrimValue
I.IntValue (IntValue -> PrimValue) -> IntValue -> PrimValue
forall a b. (a -> b) -> a -> b
$ IntType -> Integer -> IntValue
forall int. Integral int => IntType -> int -> IntValue
intValue IntType
it Integer
x
internalisePatLit (E.PatLitFloat Double
x) (E.Scalar (E.Prim (E.FloatType FloatType
ft))) =
  FloatValue -> PrimValue
I.FloatValue (FloatValue -> PrimValue) -> FloatValue -> PrimValue
forall a b. (a -> b) -> a -> b
$ FloatType -> Double -> FloatValue
forall num. Real num => FloatType -> num -> FloatValue
floatValue FloatType
ft Double
x
internalisePatLit PatLit
l StructType
t =
  [Char] -> PrimValue
forall a. HasCallStack => [Char] -> a
error ([Char] -> PrimValue) -> [Char] -> PrimValue
forall a b. (a -> b) -> a -> b
$ [Char]
"Nonsensical pattern and type: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ (PatLit, StructType) -> [Char]
forall a. Show a => a -> [Char]
show (PatLit
l, StructType
t)

generateCond ::
  E.Pat StructType ->
  [I.SubExp] ->
  InternaliseM ([Maybe I.PrimValue], [I.SubExp])
generateCond :: PatBase Info VName StructType
-> [SubExp] -> InternaliseM ([Maybe PrimValue], [SubExp])
generateCond PatBase Info VName StructType
orig_p [SubExp]
orig_ses = do
  (cmps, pertinent, _) <- PatBase Info VName StructType
-> [SubExp] -> InternaliseM ([Maybe PrimValue], [SubExp], [SubExp])
forall {vn} {a}.
IsName vn =>
PatBase Info vn StructType
-> [a] -> InternaliseM ([Maybe PrimValue], [a], [a])
compares PatBase Info VName StructType
orig_p [SubExp]
orig_ses
  pure (cmps, pertinent)
  where
    compares :: PatBase Info vn StructType
-> [a] -> InternaliseM ([Maybe PrimValue], [a], [a])
compares (E.PatLit PatLit
l (Info StructType
t) SrcLoc
_) (a
se : [a]
ses) =
      ([Maybe PrimValue], [a], [a])
-> InternaliseM ([Maybe PrimValue], [a], [a])
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([PrimValue -> Maybe PrimValue
forall a. a -> Maybe a
Just (PrimValue -> Maybe PrimValue) -> PrimValue -> Maybe PrimValue
forall a b. (a -> b) -> a -> b
$ PatLit -> StructType -> PrimValue
internalisePatLit PatLit
l StructType
t], [a
se], [a]
ses)
    compares (E.PatConstr Name
c (Info (E.Scalar (E.Sum Map Name [StructType]
fs))) [PatBase Info vn StructType]
pats SrcLoc
_) [a]
sum_ses = do
      -- Handle the unary sum type special case.
      let unary :: Bool
unary = Map Name [StructType] -> Int
forall a. Map Name a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Map Name [StructType]
fs Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1
          ses :: [a]
ses = if Bool
unary then [a]
sum_ses else [a] -> [a]
forall a. HasCallStack => [a] -> [a]
tail [a]
sum_ses
      (payload_ts, m) <- Map Name [StructType]
-> InternaliseM ([TypeBase ExtShape NoMode], [(Name, [Int])])
internaliseSumType (Map Name [StructType]
 -> InternaliseM ([TypeBase ExtShape NoMode], [(Name, [Int])]))
-> Map Name [StructType]
-> InternaliseM ([TypeBase ExtShape NoMode], [(Name, [Int])])
forall a b. (a -> b) -> a -> b
$ ([StructType] -> [StructType])
-> Map Name [StructType] -> Map Name [StructType]
forall a b k. (a -> b) -> Map k a -> Map k b
M.map ((StructType -> StructType) -> [StructType] -> [StructType]
forall a b. (a -> b) -> [a] -> [b]
map StructType -> StructType
forall dim o. TypeBase dim o -> TypeBase dim NoMode
toStruct) Map Name [StructType]
fs
      case lookupWithIndex c m of
        Just (Int
tag, [Int]
payload_is) -> do
          let ([a]
payload_ses, [a]
ses') = Int -> [a] -> ([a], [a])
forall a. Int -> [a] -> ([a], [a])
splitAt ([TypeBase ExtShape NoMode] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [TypeBase ExtShape NoMode]
payload_ts) [a]
ses
          (cmps, pertinent, _) <-
            [PatBase Info vn StructType]
-> [a] -> InternaliseM ([Maybe PrimValue], [a], [a])
comparesMany [PatBase Info vn StructType]
pats ([a] -> InternaliseM ([Maybe PrimValue], [a], [a]))
-> [a] -> InternaliseM ([Maybe PrimValue], [a], [a])
forall a b. (a -> b) -> a -> b
$ (Int -> a) -> [Int] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map ([a]
payload_ses [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!!) [Int]
payload_is
          let missingCmps Int
i a
_ =
                case Int
i Int -> [Int] -> Maybe Int
forall a. Eq a => a -> [a] -> Maybe Int
`elemIndex` [Int]
payload_is of
                  Just Int
j -> [Maybe PrimValue]
cmps [Maybe PrimValue] -> Int -> Maybe PrimValue
forall a. HasCallStack => [a] -> Int -> a
!! Int
j
                  Maybe Int
Nothing -> Maybe PrimValue
forall a. Maybe a
Nothing
          pure
            ( [Just (I.IntValue $ intValue Int8 $ toInteger tag) | not unary]
                <> zipWith missingCmps [0 ..] payload_ses,
              pertinent,
              ses'
            )
        Maybe (Int, [Int])
Nothing ->
          [Char] -> InternaliseM ([Maybe PrimValue], [a], [a])
forall a. HasCallStack => [Char] -> a
error [Char]
"generateCond: missing constructor"
    compares (E.PatConstr Name
_ (Info StructType
t) [PatBase Info vn StructType]
_ SrcLoc
_) [a]
_ =
      [Char] -> InternaliseM ([Maybe PrimValue], [a], [a])
forall a. HasCallStack => [Char] -> a
error ([Char] -> InternaliseM ([Maybe PrimValue], [a], [a]))
-> [Char] -> InternaliseM ([Maybe PrimValue], [a], [a])
forall a b. (a -> b) -> a -> b
$ [Char]
"generateCond: PatConstr has nonsensical type: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ StructType -> [Char]
forall a. Pretty a => a -> [Char]
prettyString StructType
t
    compares (E.Id vn
_ Info StructType
t SrcLoc
loc) [a]
ses =
      PatBase Info vn StructType
-> [a] -> InternaliseM ([Maybe PrimValue], [a], [a])
compares (Info StructType -> SrcLoc -> PatBase Info vn StructType
forall (f :: * -> *) vn t. f t -> SrcLoc -> PatBase f vn t
E.Wildcard Info StructType
t SrcLoc
loc) [a]
ses
    compares (E.Wildcard (Info StructType
t) SrcLoc
_) [a]
ses = do
      let ([a]
id_ses, [a]
rest_ses) = Int -> [a] -> ([a], [a])
forall a. Int -> [a] -> ([a], [a])
splitAt (StructType -> Int
forall o. TypeBase Exp o -> Int
internalisedTypeSize (StructType -> Int) -> StructType -> Int
forall a b. (a -> b) -> a -> b
$ StructType -> StructType
forall dim o. TypeBase dim o -> TypeBase dim NoMode
E.toStruct StructType
t) [a]
ses
      ([Maybe PrimValue], [a], [a])
-> InternaliseM ([Maybe PrimValue], [a], [a])
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((a -> Maybe PrimValue) -> [a] -> [Maybe PrimValue]
forall a b. (a -> b) -> [a] -> [b]
map (Maybe PrimValue -> a -> Maybe PrimValue
forall a b. a -> b -> a
const Maybe PrimValue
forall a. Maybe a
Nothing) [a]
id_ses, [a]
id_ses, [a]
rest_ses)
    compares (E.PatParens PatBase Info vn StructType
pat SrcLoc
_) [a]
ses =
      PatBase Info vn StructType
-> [a] -> InternaliseM ([Maybe PrimValue], [a], [a])
compares PatBase Info vn StructType
pat [a]
ses
    compares (E.PatAttr AttrInfo vn
_ PatBase Info vn StructType
pat SrcLoc
_) [a]
ses =
      PatBase Info vn StructType
-> [a] -> InternaliseM ([Maybe PrimValue], [a], [a])
compares PatBase Info vn StructType
pat [a]
ses
    compares (E.TuplePat [] SrcLoc
loc) [a]
ses =
      PatBase Info vn StructType
-> [a] -> InternaliseM ([Maybe PrimValue], [a], [a])
compares (Info StructType -> SrcLoc -> PatBase Info vn StructType
forall (f :: * -> *) vn t. f t -> SrcLoc -> PatBase f vn t
E.Wildcard (StructType -> Info StructType
forall a. a -> Info a
Info (StructType -> Info StructType) -> StructType -> Info StructType
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase Exp NoMode -> StructType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
E.Scalar (ScalarTypeBase Exp NoMode -> StructType)
-> ScalarTypeBase Exp NoMode -> StructType
forall a b. (a -> b) -> a -> b
$ Map Name StructType -> ScalarTypeBase Exp NoMode
forall dim o. Map Name (TypeBase dim o) -> ScalarTypeBase dim o
E.Record Map Name StructType
forall a. Monoid a => a
mempty) SrcLoc
loc) [a]
ses
    compares (E.RecordPat [] SrcLoc
loc) [a]
ses =
      PatBase Info vn StructType
-> [a] -> InternaliseM ([Maybe PrimValue], [a], [a])
compares (Info StructType -> SrcLoc -> PatBase Info vn StructType
forall (f :: * -> *) vn t. f t -> SrcLoc -> PatBase f vn t
E.Wildcard (StructType -> Info StructType
forall a. a -> Info a
Info (StructType -> Info StructType) -> StructType -> Info StructType
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase Exp NoMode -> StructType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
E.Scalar (ScalarTypeBase Exp NoMode -> StructType)
-> ScalarTypeBase Exp NoMode -> StructType
forall a b. (a -> b) -> a -> b
$ Map Name StructType -> ScalarTypeBase Exp NoMode
forall dim o. Map Name (TypeBase dim o) -> ScalarTypeBase dim o
E.Record Map Name StructType
forall a. Monoid a => a
mempty) SrcLoc
loc) [a]
ses
    compares (E.TuplePat [PatBase Info vn StructType]
pats SrcLoc
_) [a]
ses =
      [PatBase Info vn StructType]
-> [a] -> InternaliseM ([Maybe PrimValue], [a], [a])
comparesMany [PatBase Info vn StructType]
pats [a]
ses
    compares (E.RecordPat [(L Name, PatBase Info vn StructType)]
fs SrcLoc
_) [a]
ses =
      [PatBase Info vn StructType]
-> [a] -> InternaliseM ([Maybe PrimValue], [a], [a])
comparesMany (((Name, PatBase Info vn StructType) -> PatBase Info vn StructType)
-> [(Name, PatBase Info vn StructType)]
-> [PatBase Info vn StructType]
forall a b. (a -> b) -> [a] -> [b]
map (Name, PatBase Info vn StructType) -> PatBase Info vn StructType
forall a b. (a, b) -> b
snd ([(Name, PatBase Info vn StructType)]
 -> [PatBase Info vn StructType])
-> [(Name, PatBase Info vn StructType)]
-> [PatBase Info vn StructType]
forall a b. (a -> b) -> a -> b
$ Map Name (PatBase Info vn StructType)
-> [(Name, PatBase Info vn StructType)]
forall a. Map Name a -> [(Name, a)]
E.sortFields (Map Name (PatBase Info vn StructType)
 -> [(Name, PatBase Info vn StructType)])
-> Map Name (PatBase Info vn StructType)
-> [(Name, PatBase Info vn StructType)]
forall a b. (a -> b) -> a -> b
$ [(Name, PatBase Info vn StructType)]
-> Map Name (PatBase Info vn StructType)
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(Name, PatBase Info vn StructType)]
 -> Map Name (PatBase Info vn StructType))
-> [(Name, PatBase Info vn StructType)]
-> Map Name (PatBase Info vn StructType)
forall a b. (a -> b) -> a -> b
$ ((L Name, PatBase Info vn StructType)
 -> (Name, PatBase Info vn StructType))
-> [(L Name, PatBase Info vn StructType)]
-> [(Name, PatBase Info vn StructType)]
forall a b. (a -> b) -> [a] -> [b]
map ((L Name -> Name)
-> (L Name, PatBase Info vn StructType)
-> (Name, PatBase Info vn StructType)
forall a b c. (a -> b) -> (a, c) -> (b, c)
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first L Name -> Name
forall a. L a -> a
unLoc) [(L Name, PatBase Info vn StructType)]
fs) [a]
ses
    compares (E.PatAscription PatBase Info vn StructType
pat TypeExp (ExpBase Info vn) vn
_ SrcLoc
_) [a]
ses =
      PatBase Info vn StructType
-> [a] -> InternaliseM ([Maybe PrimValue], [a], [a])
compares PatBase Info vn StructType
pat [a]
ses
    compares PatBase Info vn StructType
pat [] =
      [Char] -> InternaliseM ([Maybe PrimValue], [a], [a])
forall a. HasCallStack => [Char] -> a
error ([Char] -> InternaliseM ([Maybe PrimValue], [a], [a]))
-> [Char] -> InternaliseM ([Maybe PrimValue], [a], [a])
forall a b. (a -> b) -> a -> b
$ [Char]
"generateCond: No values left for pattern " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ PatBase Info vn StructType -> [Char]
forall a. Pretty a => a -> [Char]
prettyString PatBase Info vn StructType
pat

    comparesMany :: [PatBase Info vn StructType]
-> [a] -> InternaliseM ([Maybe PrimValue], [a], [a])
comparesMany [] [a]
ses = ([Maybe PrimValue], [a], [a])
-> InternaliseM ([Maybe PrimValue], [a], [a])
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([], [], [a]
ses)
    comparesMany (PatBase Info vn StructType
pat : [PatBase Info vn StructType]
pats) [a]
ses = do
      (cmps1, pertinent1, ses') <- PatBase Info vn StructType
-> [a] -> InternaliseM ([Maybe PrimValue], [a], [a])
compares PatBase Info vn StructType
pat [a]
ses
      (cmps2, pertinent2, ses'') <- comparesMany pats ses'
      pure
        ( cmps1 <> cmps2,
          pertinent1 <> pertinent2,
          ses''
        )

internalisePat ::
  Name ->
  [E.SizeBinder VName] ->
  E.Pat StructType ->
  E.Exp ->
  InternaliseM a ->
  InternaliseM a
internalisePat :: forall a.
Name
-> [SizeBinder VName]
-> PatBase Info VName StructType
-> Exp
-> InternaliseM a
-> InternaliseM a
internalisePat Name
desc [SizeBinder VName]
sizes PatBase Info VName StructType
p Exp
e InternaliseM a
m = do
  ses <- Name -> Exp -> InternaliseM [SubExp]
internaliseExp Name
desc' Exp
e
  internalisePat' sizes p ses m
  where
    desc' :: Name
desc' = case PatBase Info VName StructType -> [IdentBase Info VName StructType]
forall (f :: * -> *) vn t. PatBase f vn t -> [IdentBase f vn t]
E.patIdents PatBase Info VName StructType
p of
      [IdentBase Info VName StructType
v] -> VName -> Name
baseName (VName -> Name) -> VName -> Name
forall a b. (a -> b) -> a -> b
$ IdentBase Info VName StructType -> VName
forall {k} (f :: k -> *) vn (t :: k). IdentBase f vn t -> vn
E.identName IdentBase Info VName StructType
v
      [IdentBase Info VName StructType]
_ -> Name
desc

internalisePat' ::
  [E.SizeBinder VName] ->
  E.Pat StructType ->
  [I.SubExp] ->
  InternaliseM a ->
  InternaliseM a
internalisePat' :: forall a.
[SizeBinder VName]
-> PatBase Info VName StructType
-> [SubExp]
-> InternaliseM a
-> InternaliseM a
internalisePat' [SizeBinder VName]
sizes PatBase Info VName StructType
p [SubExp]
ses InternaliseM a
m = do
  ses_ts <- (SubExp -> InternaliseM (TypeBase Shape NoMode))
-> [SubExp] -> InternaliseM [TypeBase Shape NoMode]
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 SubExp -> InternaliseM (TypeBase Shape NoMode)
forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoMode)
subExpType [SubExp]
ses
  stmPat (toParam E.Observe <$> p) ses_ts $ \[VName]
pat_names -> do
    AppRes -> [SubExp] -> InternaliseM ()
bindExtSizes (StructType -> [VName] -> AppRes
AppRes (PatBase Info VName StructType -> StructType
forall d o. Pat (TypeBase d o) -> TypeBase d o
E.patternType PatBase Info VName StructType
p) ((SizeBinder VName -> VName) -> [SizeBinder VName] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map SizeBinder VName -> VName
forall vn. SizeBinder vn -> vn
E.sizeName [SizeBinder VName]
sizes)) [SubExp]
ses
    [(VName, SubExp)]
-> ((VName, SubExp) -> InternaliseM ()) -> InternaliseM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([VName] -> [SubExp] -> [(VName, SubExp)]
forall a b. [a] -> [b] -> [(a, b)]
zip [VName]
pat_names [SubExp]
ses) (((VName, SubExp) -> InternaliseM ()) -> InternaliseM ())
-> ((VName, SubExp) -> InternaliseM ()) -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$ \(VName
v, SubExp
se) ->
      [VName] -> Exp (Rep InternaliseM) -> InternaliseM ()
forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [VName
v] (Exp (Rep InternaliseM) -> InternaliseM ())
-> Exp (Rep InternaliseM) -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
I.SubExp SubExp
se
    InternaliseM a
m

internaliseSlice ::
  [SubExp] ->
  [E.DimIndex] ->
  InternaliseM ([I.DimIndex SubExp], Certs)
internaliseSlice :: [SubExp]
-> SliceBase Info VName -> InternaliseM ([DimIndex SubExp], Certs)
internaliseSlice [SubExp]
dims SliceBase Info VName
idxs = do
  (idxs', oks, parts) <- [(DimIndex SubExp, SubExp, [ErrorMsgPart SubExp])]
-> ([DimIndex SubExp], [SubExp], [[ErrorMsgPart SubExp]])
forall a b c. [(a, b, c)] -> ([a], [b], [c])
unzip3 ([(DimIndex SubExp, SubExp, [ErrorMsgPart SubExp])]
 -> ([DimIndex SubExp], [SubExp], [[ErrorMsgPart SubExp]]))
-> InternaliseM [(DimIndex SubExp, SubExp, [ErrorMsgPart SubExp])]
-> InternaliseM
     ([DimIndex SubExp], [SubExp], [[ErrorMsgPart SubExp]])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (SubExp
 -> DimIndex
 -> InternaliseM (DimIndex SubExp, SubExp, [ErrorMsgPart SubExp]))
-> [SubExp]
-> SliceBase Info VName
-> InternaliseM [(DimIndex SubExp, SubExp, [ErrorMsgPart SubExp])]
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM SubExp
-> DimIndex
-> InternaliseM (DimIndex SubExp, SubExp, [ErrorMsgPart SubExp])
internaliseDimIndex [SubExp]
dims SliceBase Info VName
idxs
  ok <- letSubExp "index_ok" =<< eAll oks
  let msg =
        [ErrorMsgPart SubExp] -> ErrorMsg SubExp
forall a. [ErrorMsgPart a] -> ErrorMsg a
errorMsg ([ErrorMsgPart SubExp] -> ErrorMsg SubExp)
-> [ErrorMsgPart SubExp] -> ErrorMsg SubExp
forall a b. (a -> b) -> a -> b
$
          [ErrorMsgPart SubExp
"Index ["]
            [ErrorMsgPart SubExp]
-> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
forall a. [a] -> [a] -> [a]
++ [ErrorMsgPart SubExp]
-> [[ErrorMsgPart SubExp]] -> [ErrorMsgPart SubExp]
forall a. [a] -> [[a]] -> [a]
intercalate [ErrorMsgPart SubExp
", "] [[ErrorMsgPart SubExp]]
parts
            [ErrorMsgPart SubExp]
-> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
forall a. [a] -> [a] -> [a]
++ [ErrorMsgPart SubExp
"] out of bounds for array of shape ["]
            [ErrorMsgPart SubExp]
-> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
forall a. [a] -> [a] -> [a]
++ ErrorMsgPart SubExp
-> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
forall a. a -> [a] -> [a]
intersperse ErrorMsgPart SubExp
"][" ((SubExp -> ErrorMsgPart SubExp)
-> [SubExp] -> [ErrorMsgPart SubExp]
forall a b. (a -> b) -> [a] -> [b]
map (PrimType -> SubExp -> ErrorMsgPart SubExp
forall a. PrimType -> a -> ErrorMsgPart a
ErrorVal PrimType
int64) ([SubExp] -> [ErrorMsgPart SubExp])
-> [SubExp] -> [ErrorMsgPart SubExp]
forall a b. (a -> b) -> a -> b
$ Int -> [SubExp] -> [SubExp]
forall a. Int -> [a] -> [a]
take (SliceBase Info VName -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length SliceBase Info VName
idxs) [SubExp]
dims)
            [ErrorMsgPart SubExp]
-> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
forall a. [a] -> [a] -> [a]
++ [ErrorMsgPart SubExp
"]."]
  c <- assert "index_certs" ok msg
  pure (idxs', c)

internaliseDimIndex ::
  SubExp ->
  E.DimIndex ->
  InternaliseM (I.DimIndex SubExp, SubExp, [ErrorMsgPart SubExp])
internaliseDimIndex :: SubExp
-> DimIndex
-> InternaliseM (DimIndex SubExp, SubExp, [ErrorMsgPart SubExp])
internaliseDimIndex SubExp
w (E.DimFix Exp
i) = do
  (i', _) <- Name -> Exp -> InternaliseM (SubExp, IntType)
internaliseSizeExp Name
"i" Exp
i
  let lowerBound =
        BasicOp -> Exp SOACS
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp SOACS) -> BasicOp -> Exp SOACS
forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (IntType -> CmpOp
I.CmpSle IntType
I.Int64) (Int64 -> SubExp
forall v. IsValue v => v -> SubExp
I.constant (Int64
0 :: I.Int64)) SubExp
i'
      upperBound =
        BasicOp -> Exp SOACS
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp SOACS) -> BasicOp -> Exp SOACS
forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (IntType -> CmpOp
I.CmpSlt IntType
I.Int64) SubExp
i' SubExp
w
  ok <- letSubExp "bounds_check" =<< eBinOp I.LogAnd (pure lowerBound) (pure upperBound)
  pure (I.DimFix i', ok, [ErrorVal int64 i'])

-- Special-case an important common case that otherwise leads to horrible code.
internaliseDimIndex
  SubExp
w
  ( E.DimSlice
      Maybe Exp
Nothing
      Maybe Exp
Nothing
      (Just (E.Negate (E.IntLit Integer
1 Info StructType
_ SrcLoc
_) SrcLoc
_))
    ) = do
    w_minus_1 <-
      Name -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"w_minus_1" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
        BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
forall a b. (a -> b) -> a -> b
$
          BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
Sub IntType
Int64 Overflow
I.OverflowWrap) SubExp
w SubExp
one
    pure
      ( I.DimSlice w_minus_1 w $ intConst Int64 (-1),
        constant True,
        mempty
      )
    where
      one :: SubExp
one = Int64 -> SubExp
forall v. IsValue v => v -> SubExp
constant (Int64
1 :: Int64)
internaliseDimIndex SubExp
w (E.DimSlice Maybe Exp
i Maybe Exp
j Maybe Exp
s) = do
  s' <- InternaliseM SubExp
-> (Exp -> InternaliseM SubExp) -> Maybe Exp -> InternaliseM SubExp
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (SubExp -> InternaliseM SubExp
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SubExp
one) (((SubExp, IntType) -> SubExp)
-> InternaliseM (SubExp, IntType) -> InternaliseM SubExp
forall a b. (a -> b) -> InternaliseM a -> InternaliseM b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (SubExp, IntType) -> SubExp
forall a b. (a, b) -> a
fst (InternaliseM (SubExp, IntType) -> InternaliseM SubExp)
-> (Exp -> InternaliseM (SubExp, IntType))
-> Exp
-> InternaliseM SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Name -> Exp -> InternaliseM (SubExp, IntType)
internaliseSizeExp Name
"s") Maybe Exp
s
  s_sign <- letSubExp "s_sign" $ BasicOp $ I.UnOp (I.SSignum Int64) s'
  backwards <- letSubExp "backwards" $ I.BasicOp $ I.CmpOp (I.CmpEq int64) s_sign negone
  w_minus_1 <- letSubExp "w_minus_1" $ BasicOp $ I.BinOp (Sub Int64 I.OverflowWrap) w one
  let i_def =
        Name -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"i_def"
          (Exp SOACS -> InternaliseM SubExp)
-> InternaliseM (Exp SOACS) -> InternaliseM SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< InternaliseM (Exp (Rep InternaliseM))
-> InternaliseM (GBody (Rep InternaliseM) SubExpRes)
-> InternaliseM (GBody (Rep InternaliseM) SubExpRes)
-> InternaliseM (Exp (Rep InternaliseM))
forall (m :: * -> *).
(MonadBuilder m, BranchType (Rep m) ~ TypeBase ExtShape NoMode) =>
m (Exp (Rep m))
-> m (Body (Rep m)) -> m (Body (Rep m)) -> m (Exp (Rep m))
eIf
            (SubExp -> InternaliseM (Exp (Rep InternaliseM))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
backwards)
            ([SubExp] -> InternaliseM (GBody (Rep InternaliseM) SubExpRes)
forall (m :: * -> *).
MonadBuilder m =>
[SubExp] -> m (Body (Rep m))
resultBodyM [SubExp
w_minus_1])
            ([SubExp] -> InternaliseM (GBody (Rep InternaliseM) SubExpRes)
forall (m :: * -> *).
MonadBuilder m =>
[SubExp] -> m (Body (Rep m))
resultBodyM [SubExp
zero])
      j_def =
        Name -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"j_def"
          (Exp SOACS -> InternaliseM SubExp)
-> InternaliseM (Exp SOACS) -> InternaliseM SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< InternaliseM (Exp (Rep InternaliseM))
-> InternaliseM (GBody (Rep InternaliseM) SubExpRes)
-> InternaliseM (GBody (Rep InternaliseM) SubExpRes)
-> InternaliseM (Exp (Rep InternaliseM))
forall (m :: * -> *).
(MonadBuilder m, BranchType (Rep m) ~ TypeBase ExtShape NoMode) =>
m (Exp (Rep m))
-> m (Body (Rep m)) -> m (Body (Rep m)) -> m (Exp (Rep m))
eIf
            (SubExp -> InternaliseM (Exp (Rep InternaliseM))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
backwards)
            ([SubExp] -> InternaliseM (GBody (Rep InternaliseM) SubExpRes)
forall (m :: * -> *).
MonadBuilder m =>
[SubExp] -> m (Body (Rep m))
resultBodyM [SubExp
negone])
            ([SubExp] -> InternaliseM (GBody (Rep InternaliseM) SubExpRes)
forall (m :: * -> *).
MonadBuilder m =>
[SubExp] -> m (Body (Rep m))
resultBodyM [SubExp
w])
  i' <- maybe i_def (fmap fst . internaliseSizeExp "i") i
  j' <- maybe j_def (fmap fst . internaliseSizeExp "j") j
  j_m_i <- letSubExp "j_m_i" $ BasicOp $ I.BinOp (Sub Int64 I.OverflowWrap) j' i'
  -- Something like a division-rounding-up, but accomodating negative
  -- operands.
  let divRounding m (Exp (Rep m))
x m (Exp (Rep m))
y =
        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
          (IntType -> Safety -> BinOp
SQuot IntType
Int64 Safety
Safe)
          ( 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
              (IntType -> Overflow -> BinOp
Add IntType
Int64 Overflow
I.OverflowWrap)
              m (Exp (Rep m))
x
              (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 (IntType -> Overflow -> BinOp
Sub IntType
Int64 Overflow
I.OverflowWrap) m (Exp (Rep m))
y (m (Exp (Rep m)) -> m (Exp (Rep m))
forall (m :: * -> *).
MonadBuilder m =>
m (Exp (Rep m)) -> m (Exp (Rep m))
eSignum m (Exp (Rep m))
y))
          )
          m (Exp (Rep m))
y
  n <- letSubExp "n" =<< divRounding (toExp j_m_i) (toExp s')

  zero_stride <- letSubExp "zero_stride" $ I.BasicOp $ I.CmpOp (CmpEq int64) s_sign zero
  nonzero_stride <- letSubExp "nonzero_stride" $ I.BasicOp $ I.UnOp (I.Neg I.Bool) zero_stride

  -- Bounds checks depend on whether we are slicing forwards or
  -- backwards.  If forwards, we must check '0 <= i && i <= j'.  If
  -- backwards, '-1 <= j && j <= i'.  In both cases, we check '0 <=
  -- i+n*s && i+(n-1)*s < w'.  We only check if the slice is nonempty.
  empty_slice <- letSubExp "empty_slice" $ I.BasicOp $ I.CmpOp (CmpEq int64) n zero

  m <- letSubExp "m" $ I.BasicOp $ I.BinOp (Sub Int64 I.OverflowWrap) n one
  m_t_s <- letSubExp "m_t_s" $ I.BasicOp $ I.BinOp (Mul Int64 I.OverflowWrap) m s'
  i_p_m_t_s <- letSubExp "i_p_m_t_s" $ I.BasicOp $ I.BinOp (Add Int64 I.OverflowWrap) i' m_t_s
  zero_leq_i_p_m_t_s <-
    letSubExp "zero_leq_i_p_m_t_s" $
      I.BasicOp $
        I.CmpOp (I.CmpSle Int64) zero i_p_m_t_s
  i_p_m_t_s_leq_w <-
    letSubExp "i_p_m_t_s_leq_w" $
      I.BasicOp $
        I.CmpOp (I.CmpSle Int64) i_p_m_t_s w
  i_p_m_t_s_lth_w <-
    letSubExp "i_p_m_t_s_leq_w" $
      I.BasicOp $
        I.CmpOp (I.CmpSlt Int64) i_p_m_t_s w

  zero_lte_i <- letSubExp "zero_lte_i" $ I.BasicOp $ I.CmpOp (I.CmpSle Int64) zero i'
  i_lte_j <- letSubExp "i_lte_j" $ I.BasicOp $ I.CmpOp (I.CmpSle Int64) i' j'
  forwards_ok <-
    letSubExp "forwards_ok"
      =<< eAll [zero_lte_i, i_lte_j, zero_leq_i_p_m_t_s, i_p_m_t_s_lth_w]

  negone_lte_j <- letSubExp "negone_lte_j" $ I.BasicOp $ I.CmpOp (I.CmpSle Int64) negone j'
  j_lte_i <- letSubExp "j_lte_i" $ I.BasicOp $ I.CmpOp (I.CmpSle Int64) j' i'
  backwards_ok <-
    letSubExp "backwards_ok"
      =<< eAll
        [negone_lte_j, j_lte_i, zero_leq_i_p_m_t_s, i_p_m_t_s_leq_w]

  slice_ok <-
    letSubExp "slice_ok"
      =<< eIf
        (eSubExp backwards)
        (resultBodyM [backwards_ok])
        (resultBodyM [forwards_ok])

  ok_or_empty <-
    letSubExp "ok_or_empty" $
      I.BasicOp $
        I.BinOp I.LogOr empty_slice slice_ok

  acceptable <-
    letSubExp "slice_acceptable" $
      I.BasicOp $
        I.BinOp I.LogAnd nonzero_stride ok_or_empty

  let parts = case (Maybe Exp
i, Maybe Exp
j, Maybe Exp
s) of
        (Maybe Exp
_, Maybe Exp
_, Just {}) ->
          [ ErrorMsgPart SubExp
-> (Exp -> ErrorMsgPart SubExp) -> Maybe Exp -> ErrorMsgPart SubExp
forall b a. b -> (a -> b) -> Maybe a -> b
maybe ErrorMsgPart SubExp
"" (ErrorMsgPart SubExp -> Exp -> ErrorMsgPart SubExp
forall a b. a -> b -> a
const (ErrorMsgPart SubExp -> Exp -> ErrorMsgPart SubExp)
-> ErrorMsgPart SubExp -> Exp -> ErrorMsgPart SubExp
forall a b. (a -> b) -> a -> b
$ PrimType -> SubExp -> ErrorMsgPart SubExp
forall a. PrimType -> a -> ErrorMsgPart a
ErrorVal PrimType
int64 SubExp
i') Maybe Exp
i,
            ErrorMsgPart SubExp
":",
            ErrorMsgPart SubExp
-> (Exp -> ErrorMsgPart SubExp) -> Maybe Exp -> ErrorMsgPart SubExp
forall b a. b -> (a -> b) -> Maybe a -> b
maybe ErrorMsgPart SubExp
"" (ErrorMsgPart SubExp -> Exp -> ErrorMsgPart SubExp
forall a b. a -> b -> a
const (ErrorMsgPart SubExp -> Exp -> ErrorMsgPart SubExp)
-> ErrorMsgPart SubExp -> Exp -> ErrorMsgPart SubExp
forall a b. (a -> b) -> a -> b
$ PrimType -> SubExp -> ErrorMsgPart SubExp
forall a. PrimType -> a -> ErrorMsgPart a
ErrorVal PrimType
int64 SubExp
j') Maybe Exp
j,
            ErrorMsgPart SubExp
":",
            PrimType -> SubExp -> ErrorMsgPart SubExp
forall a. PrimType -> a -> ErrorMsgPart a
ErrorVal PrimType
int64 SubExp
s'
          ]
        (Maybe Exp
_, Just {}, Maybe Exp
_) ->
          [ ErrorMsgPart SubExp
-> (Exp -> ErrorMsgPart SubExp) -> Maybe Exp -> ErrorMsgPart SubExp
forall b a. b -> (a -> b) -> Maybe a -> b
maybe ErrorMsgPart SubExp
"" (ErrorMsgPart SubExp -> Exp -> ErrorMsgPart SubExp
forall a b. a -> b -> a
const (ErrorMsgPart SubExp -> Exp -> ErrorMsgPart SubExp)
-> ErrorMsgPart SubExp -> Exp -> ErrorMsgPart SubExp
forall a b. (a -> b) -> a -> b
$ PrimType -> SubExp -> ErrorMsgPart SubExp
forall a. PrimType -> a -> ErrorMsgPart a
ErrorVal PrimType
int64 SubExp
i') Maybe Exp
i,
            ErrorMsgPart SubExp
":",
            PrimType -> SubExp -> ErrorMsgPart SubExp
forall a. PrimType -> a -> ErrorMsgPart a
ErrorVal PrimType
int64 SubExp
j'
          ]
            [ErrorMsgPart SubExp]
-> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
forall a. [a] -> [a] -> [a]
++ [ErrorMsgPart SubExp]
-> (Exp -> [ErrorMsgPart SubExp])
-> Maybe Exp
-> [ErrorMsgPart SubExp]
forall b a. b -> (a -> b) -> Maybe a -> b
maybe [ErrorMsgPart SubExp]
forall a. Monoid a => a
mempty ([ErrorMsgPart SubExp] -> Exp -> [ErrorMsgPart SubExp]
forall a b. a -> b -> a
const [ErrorMsgPart SubExp
":", PrimType -> SubExp -> ErrorMsgPart SubExp
forall a. PrimType -> a -> ErrorMsgPart a
ErrorVal PrimType
int64 SubExp
s']) Maybe Exp
s
        (Maybe Exp
_, Maybe Exp
Nothing, Maybe Exp
Nothing) ->
          [PrimType -> SubExp -> ErrorMsgPart SubExp
forall a. PrimType -> a -> ErrorMsgPart a
ErrorVal PrimType
int64 SubExp
i', ErrorMsgPart SubExp
":"]
  pure (I.DimSlice i' n s', acceptable, parts)
  where
    zero :: SubExp
zero = Int64 -> SubExp
forall v. IsValue v => v -> SubExp
constant (Int64
0 :: Int64)
    negone :: SubExp
negone = Int64 -> SubExp
forall v. IsValue v => v -> SubExp
constant (-Int64
1 :: Int64)
    one :: SubExp
one = Int64 -> SubExp
forall v. IsValue v => v -> SubExp
constant (Int64
1 :: Int64)

internaliseFoldLambda ::
  E.Exp ->
  [I.Type] ->
  [I.Type] ->
  InternaliseM (I.Lambda SOACS)
internaliseFoldLambda :: Exp
-> [TypeBase Shape NoMode]
-> [TypeBase Shape NoMode]
-> InternaliseM (Lambda SOACS)
internaliseFoldLambda Exp
lam [TypeBase Shape NoMode]
acctypes [TypeBase Shape NoMode]
arrtypes = do
  let rowtypes :: [TypeBase Shape NoMode]
rowtypes = (TypeBase Shape NoMode -> TypeBase Shape NoMode)
-> [TypeBase Shape NoMode] -> [TypeBase Shape NoMode]
forall a b. (a -> b) -> [a] -> [b]
map TypeBase Shape NoMode -> TypeBase Shape NoMode
forall d m. TypeBase (ShapeBase d) m -> TypeBase (ShapeBase d) m
I.rowType [TypeBase Shape NoMode]
arrtypes
  (params, body, rettype) <- Exp
-> [TypeBase Shape NoMode]
-> InternaliseM
     ([LParam SOACS], Body SOACS, [TypeBase Shape NoMode])
internaliseLambda Exp
lam ([TypeBase Shape NoMode]
 -> InternaliseM
      ([LParam SOACS], Body SOACS, [TypeBase Shape NoMode]))
-> [TypeBase Shape NoMode]
-> InternaliseM
     ([LParam SOACS], Body SOACS, [TypeBase Shape NoMode])
forall a b. (a -> b) -> a -> b
$ [TypeBase Shape NoMode]
acctypes [TypeBase Shape NoMode]
-> [TypeBase Shape NoMode] -> [TypeBase Shape NoMode]
forall a. [a] -> [a] -> [a]
++ [TypeBase Shape NoMode]
rowtypes
  let rettype' =
        [ TypeBase Shape NoMode
t TypeBase Shape NoMode -> Shape -> TypeBase Shape NoMode
forall newshape oldshape o.
ArrayShape newshape =>
TypeBase oldshape o -> newshape -> TypeBase newshape o
`I.setArrayShape` TypeBase Shape NoMode -> Shape
forall shape o. ArrayShape shape => TypeBase shape o -> shape
I.arrayShape TypeBase Shape NoMode
shape
        | (TypeBase Shape NoMode
t, TypeBase Shape NoMode
shape) <- [TypeBase Shape NoMode]
-> [TypeBase Shape NoMode]
-> [(TypeBase Shape NoMode, TypeBase Shape NoMode)]
forall a b. [a] -> [b] -> [(a, b)]
zip [TypeBase Shape NoMode]
rettype [TypeBase Shape NoMode]
acctypes
        ]
  -- The result of the body must have the exact same shape as the
  -- initial accumulator.
  mkLambda params $
    ensureResultShape
      (ErrorMsg [ErrorString "shape of result does not match shape of initial value"])
      rettype'
      =<< bodyBind body

internaliseScanOrReduce ::
  Name ->
  Name ->
  (SubExp -> I.Lambda SOACS -> [SubExp] -> [VName] -> InternaliseM (SOAC SOACS)) ->
  (E.Exp, E.Exp, E.Exp) ->
  InternaliseM [SubExp]
internaliseScanOrReduce :: Name
-> Name
-> (SubExp
    -> Lambda SOACS
    -> [SubExp]
    -> [VName]
    -> InternaliseM (SOAC SOACS))
-> (Exp, Exp, Exp)
-> InternaliseM [SubExp]
internaliseScanOrReduce Name
desc Name
what SubExp
-> Lambda SOACS -> [SubExp] -> [VName] -> InternaliseM (SOAC SOACS)
f (Exp
lam, Exp
ne, Exp
arr) = do
  arrs <- Name -> Exp -> InternaliseM [VName]
internaliseExpToVars (Name
what Name -> Name -> Name
forall a. Semigroup a => a -> a -> a
<> Name
"_arr") Exp
arr
  nes <- internaliseExp (what <> "_ne") ne
  nes' <- forM (zip nes arrs) $ \(SubExp
ne', VName
arr') -> do
    rowtype <- Int -> TypeBase Shape NoMode -> TypeBase Shape NoMode
forall d m.
Int -> TypeBase (ShapeBase d) m -> TypeBase (ShapeBase d) m
I.stripArray Int
1 (TypeBase Shape NoMode -> TypeBase Shape NoMode)
-> InternaliseM (TypeBase Shape NoMode)
-> InternaliseM (TypeBase Shape NoMode)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> InternaliseM (TypeBase Shape NoMode)
forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoMode)
lookupType VName
arr'
    ensureShape
      "Row shape of input array does not match shape of neutral element"
      rowtype
      (what <> "_ne_right_shape")
      ne'
  nests <- mapM I.subExpType nes'
  arrts <- mapM lookupType arrs
  lam' <- internaliseFoldLambda lam nests arrts
  w <- arraysSize 0 <$> mapM lookupType arrs
  letValExp' desc . I.Op =<< f w lam' nes' arrs

internaliseHist ::
  Int ->
  Name ->
  E.Exp ->
  E.Exp ->
  E.Exp ->
  E.Exp ->
  E.Exp ->
  E.Exp ->
  InternaliseM [SubExp]
internaliseHist :: Int
-> Name
-> Exp
-> Exp
-> Exp
-> Exp
-> Exp
-> Exp
-> InternaliseM [SubExp]
internaliseHist Int
dim Name
desc Exp
rf Exp
hist Exp
op Exp
ne Exp
buckets Exp
img = do
  rf' <- Name -> Exp -> InternaliseM SubExp
internaliseExp1 Name
"hist_rf" Exp
rf
  ne' <- internaliseExp "hist_ne" ne
  hist' <- internaliseExpToVars "hist_hist" hist
  buckets' <- internaliseExpToVars "hist_buckets" buckets
  img' <- internaliseExpToVars "hist_img" img

  -- reshape neutral element to have same size as the destination array
  ne_shp <- forM (zip ne' hist') $ \(SubExp
n, VName
h) -> do
    rowtype <- Int -> TypeBase Shape NoMode -> TypeBase Shape NoMode
forall d m.
Int -> TypeBase (ShapeBase d) m -> TypeBase (ShapeBase d) m
I.stripArray Int
dim (TypeBase Shape NoMode -> TypeBase Shape NoMode)
-> InternaliseM (TypeBase Shape NoMode)
-> InternaliseM (TypeBase Shape NoMode)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> InternaliseM (TypeBase Shape NoMode)
forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoMode)
lookupType VName
h
    ensureShape
      "Row shape of destination array does not match shape of neutral element"
      rowtype
      "hist_ne_right_shape"
      n
  ne_ts <- mapM I.subExpType ne_shp
  his_ts <- mapM (fmap (I.stripArray (dim - 1)) . lookupType) hist'
  op' <- internaliseFoldLambda op ne_ts his_ts

  -- reshape return type of bucket function to have same size as neutral element
  -- (modulo the index)
  bucket_params <- replicateM dim (newParam "bucket_p" $ I.Prim int64)
  img_params <- mapM (newParam "img_p" . rowType) =<< mapM lookupType img'
  let params = [Param (TypeBase Shape NoMode)]
bucket_params [Param (TypeBase Shape NoMode)]
-> [Param (TypeBase Shape NoMode)]
-> [Param (TypeBase Shape NoMode)]
forall a. [a] -> [a] -> [a]
++ [Param (TypeBase Shape NoMode)]
img_params
      rettype = Int -> TypeBase Shape NoMode -> [TypeBase Shape NoMode]
forall a. Int -> a -> [a]
replicate Int
dim (PrimType -> TypeBase Shape NoMode
forall shape o. PrimType -> TypeBase shape o
I.Prim PrimType
int64) [TypeBase Shape NoMode]
-> [TypeBase Shape NoMode] -> [TypeBase Shape NoMode]
forall a. [a] -> [a] -> [a]
++ [TypeBase Shape NoMode]
ne_ts
      body = Stms SOACS -> [SubExpRes] -> Body SOACS
forall res. IsResult res => Stms SOACS -> [res] -> GBody SOACS res
forall rep res.
(Buildable rep, IsResult res) =>
Stms rep -> [res] -> GBody rep res
mkBody Stms SOACS
forall a. Monoid a => a
mempty ([SubExpRes] -> Body SOACS) -> [SubExpRes] -> Body SOACS
forall a b. (a -> b) -> a -> b
$ [VName] -> [SubExpRes]
varsRes ([VName] -> [SubExpRes]) -> [VName] -> [SubExpRes]
forall a b. (a -> b) -> a -> b
$ (Param (TypeBase Shape NoMode) -> VName)
-> [Param (TypeBase Shape NoMode)] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param (TypeBase Shape NoMode) -> VName
forall dec. Param dec -> VName
I.paramName [Param (TypeBase Shape NoMode)]
params
  lam' <-
    mkLambda params $
      ensureResultShape
        "Row shape of value array does not match row shape of hist target"
        rettype
        =<< bodyBind body

  -- get sizes of histogram and image arrays
  shape_hist <- I.Shape . take dim . I.arrayDims <$> lookupType (head hist')
  w_img <- I.arraySize 0 <$> lookupType (head img')

  letValExp' desc . I.Op $
    I.Hist w_img (buckets' ++ img') [HistOp shape_hist rf' hist' ne_shp op'] lam'

internaliseStreamAcc ::
  Name ->
  E.Exp ->
  Maybe (E.Exp, E.Exp) ->
  E.Exp ->
  E.Exp ->
  InternaliseM [SubExp]
internaliseStreamAcc :: Name
-> Exp -> Maybe (Exp, Exp) -> Exp -> Exp -> InternaliseM [SubExp]
internaliseStreamAcc Name
desc Exp
dest Maybe (Exp, Exp)
op Exp
lam Exp
bs = do
  dest' <- Name -> Exp -> InternaliseM [VName]
internaliseExpToVars Name
"scatter_dest" Exp
dest
  bs' <- internaliseExpToVars "scatter_input" bs

  acc_cert_v <- newVName "acc_cert"
  dest_ts <- mapM lookupType dest'
  let dest_w = Int -> [TypeBase Shape NoMode] -> SubExp
forall o. Int -> [TypeBase Shape o] -> SubExp
arraysSize Int
0 [TypeBase Shape NoMode]
dest_ts
      acc_t = VName -> Shape -> [TypeBase Shape NoMode] -> TypeBase Shape NoMode
forall shape o.
VName -> Shape -> [TypeBase Shape NoMode] -> TypeBase shape o
Acc VName
acc_cert_v ([SubExp] -> Shape
forall d. [d] -> ShapeBase d
I.Shape [SubExp
dest_w]) ((TypeBase Shape NoMode -> TypeBase Shape NoMode)
-> [TypeBase Shape NoMode] -> [TypeBase Shape NoMode]
forall a b. (a -> b) -> [a] -> [b]
map TypeBase Shape NoMode -> TypeBase Shape NoMode
forall d m. TypeBase (ShapeBase d) m -> TypeBase (ShapeBase d) m
rowType [TypeBase Shape NoMode]
dest_ts)
  acc_p <- newParam "acc_p" acc_t
  withacc_lam <- mkLambda [Param mempty acc_cert_v (I.Prim I.Unit), acc_p] $ do
    bs_ts <- mapM lookupType bs'
    lam' <- internaliseLambdaCoerce lam $ map rowType $ paramType acc_p : bs_ts
    let w = Int -> [TypeBase Shape NoMode] -> SubExp
forall o. Int -> [TypeBase Shape o] -> SubExp
arraysSize Int
0 [TypeBase Shape NoMode]
bs_ts
    fmap subExpsRes
      . letValExp' "acc_res"
      . I.Op
      . I.Screma w (I.paramName acc_p : bs')
      =<< I.mapSOAC lam'

  op' <-
    case op of
      Just (Exp
op_lam, Exp
ne) -> do
        ne' <- Name -> Exp -> InternaliseM [SubExp]
internaliseExp Name
"hist_ne" Exp
ne
        ne_ts <- mapM I.subExpType ne'
        (lam_params, lam_body, lam_rettype) <-
          internaliseLambda op_lam $ ne_ts ++ ne_ts
        idxp <- newParam "idx" $ I.Prim int64
        let op_lam' = [LParam SOACS]
-> [TypeBase Shape NoMode] -> Body SOACS -> Lambda SOACS
forall rep t. [LParam rep] -> [t] -> Body rep -> GLambda rep t
I.Lambda (Param (TypeBase Shape NoMode)
idxp Param (TypeBase Shape NoMode)
-> [Param (TypeBase Shape NoMode)]
-> [Param (TypeBase Shape NoMode)]
forall a. a -> [a] -> [a]
: [Param (TypeBase Shape NoMode)]
[LParam SOACS]
lam_params) [TypeBase Shape NoMode]
lam_rettype Body SOACS
lam_body
        pure $ Just (op_lam', ne')
      Maybe (Exp, Exp)
Nothing ->
        Maybe (Lambda SOACS, [SubExp])
-> InternaliseM (Maybe (Lambda SOACS, [SubExp]))
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe (Lambda SOACS, [SubExp])
forall a. Maybe a
Nothing

  destw <- arraysSize 0 <$> mapM lookupType dest'
  fmap (map I.Var) $
    letTupExp desc $
      WithAcc [(I.Shape [destw], dest', op')] withacc_lam

internaliseExp1 :: Name -> E.Exp -> InternaliseM I.SubExp
internaliseExp1 :: Name -> Exp -> InternaliseM SubExp
internaliseExp1 Name
desc Exp
e = do
  vs <- Name -> Exp -> InternaliseM [SubExp]
internaliseExp Name
desc Exp
e
  case vs of
    [SubExp
se] -> SubExp -> InternaliseM SubExp
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SubExp
se
    [SubExp]
_ -> [Char] -> InternaliseM SubExp
forall a. HasCallStack => [Char] -> a
error [Char]
"Internalise.internaliseExp1: was passed not just a single subexpression"

-- | Promote to dimension type as appropriate for the original type.
-- Also return original type.
internaliseSizeExp :: Name -> E.Exp -> InternaliseM (I.SubExp, IntType)
internaliseSizeExp :: Name -> Exp -> InternaliseM (SubExp, IntType)
internaliseSizeExp Name
s Exp
e = do
  e' <- Name -> Exp -> InternaliseM SubExp
internaliseExp1 Name
s Exp
e
  case E.typeOf e of
    E.Scalar (E.Prim (E.Signed IntType
it)) -> (,IntType
it) (SubExp -> (SubExp, IntType))
-> InternaliseM SubExp -> InternaliseM (SubExp, IntType)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IntType -> SubExp -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
IntType -> SubExp -> m SubExp
asIntS IntType
Int64 SubExp
e'
    StructType
_ -> [Char] -> InternaliseM (SubExp, IntType)
forall a. HasCallStack => [Char] -> a
error [Char]
"internaliseSizeExp: bad type"

asVar :: Name -> I.SubExp -> InternaliseM I.VName
asVar :: Name -> SubExp -> InternaliseM VName
asVar Name
desc SubExp
se = Name -> Exp (Rep InternaliseM) -> InternaliseM VName
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m VName
letExp Name
desc (Exp (Rep InternaliseM) -> InternaliseM VName)
-> Exp (Rep InternaliseM) -> InternaliseM VName
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
I.SubExp SubExp
se

internaliseExpToVars :: Name -> E.Exp -> InternaliseM [I.VName]
internaliseExpToVars :: Name -> Exp -> InternaliseM [VName]
internaliseExpToVars Name
desc Exp
e =
  (SubExp -> InternaliseM VName) -> [SubExp] -> InternaliseM [VName]
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 -> SubExp -> InternaliseM VName
asVar Name
desc) ([SubExp] -> InternaliseM [VName])
-> InternaliseM [SubExp] -> InternaliseM [VName]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Name -> Exp -> InternaliseM [SubExp]
internaliseExp Name
desc Exp
e

internaliseOperation ::
  Name ->
  E.Exp ->
  (I.VName -> InternaliseM I.BasicOp) ->
  InternaliseM [I.SubExp]
internaliseOperation :: Name
-> Exp -> (VName -> InternaliseM BasicOp) -> InternaliseM [SubExp]
internaliseOperation Name
s Exp
e VName -> InternaliseM BasicOp
op = do
  vs <- Name -> Exp -> InternaliseM [VName]
internaliseExpToVars Name
s Exp
e
  mapM (letSubExp s . I.BasicOp <=< op) vs

certifyingNonzero ::
  IntType ->
  SubExp ->
  InternaliseM a ->
  InternaliseM a
certifyingNonzero :: forall a. IntType -> SubExp -> InternaliseM a -> InternaliseM a
certifyingNonzero IntType
t SubExp
x InternaliseM a
m = do
  zero <-
    Name -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"zero" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
      BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
forall a b. (a -> b) -> a -> b
$
        CmpOp -> SubExp -> SubExp -> BasicOp
CmpOp (PrimType -> CmpOp
CmpEq (IntType -> PrimType
IntType IntType
t)) SubExp
x (IntType -> Integer -> SubExp
intConst IntType
t Integer
0)
  nonzero <- letSubExp "nonzero" $ I.BasicOp $ UnOp (I.Neg I.Bool) zero
  c <- assert "nonzero_cert" nonzero "division by zero"
  certifying c m

certifyingNonnegative ::
  IntType ->
  SubExp ->
  InternaliseM a ->
  InternaliseM a
certifyingNonnegative :: forall a. IntType -> SubExp -> InternaliseM a -> InternaliseM a
certifyingNonnegative IntType
t SubExp
x InternaliseM a
m = do
  nonnegative <-
    Name -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"nonnegative" (Exp SOACS -> InternaliseM SubExp)
-> (BasicOp -> Exp SOACS) -> BasicOp -> InternaliseM SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BasicOp -> Exp SOACS
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> InternaliseM SubExp) -> BasicOp -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$
      CmpOp -> SubExp -> SubExp -> BasicOp
CmpOp (IntType -> CmpOp
CmpSle IntType
t) (IntType -> Integer -> SubExp
intConst IntType
t Integer
0) SubExp
x
  c <- assert "nonzero_cert" nonnegative "negative exponent"
  certifying c m

internaliseBinOp ::
  Name ->
  E.BinOp ->
  I.SubExp ->
  I.SubExp ->
  E.PrimType ->
  E.PrimType ->
  InternaliseM [I.SubExp]
internaliseBinOp :: Name
-> BinOp
-> SubExp
-> SubExp
-> PrimType
-> PrimType
-> InternaliseM [SubExp]
internaliseBinOp Name
desc BinOp
E.LogAnd SubExp
x SubExp
y PrimType
E.Bool PrimType
_ =
  Name -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp Name
desc BinOp
I.LogAnd SubExp
x SubExp
y
internaliseBinOp Name
desc BinOp
E.LogOr SubExp
x SubExp
y PrimType
E.Bool PrimType
_ =
  Name -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp Name
desc BinOp
I.LogOr SubExp
x SubExp
y
internaliseBinOp Name
desc BinOp
E.Plus SubExp
x SubExp
y (E.Signed IntType
t) PrimType
_ =
  Name -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp Name
desc (IntType -> Overflow -> BinOp
I.Add IntType
t Overflow
I.OverflowWrap) SubExp
x SubExp
y
internaliseBinOp Name
desc BinOp
E.Plus SubExp
x SubExp
y (E.Unsigned IntType
t) PrimType
_ =
  Name -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp Name
desc (IntType -> Overflow -> BinOp
I.Add IntType
t Overflow
I.OverflowWrap) SubExp
x SubExp
y
internaliseBinOp Name
desc BinOp
E.Plus SubExp
x SubExp
y (E.FloatType FloatType
t) PrimType
_ =
  Name -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp Name
desc (FloatType -> BinOp
I.FAdd FloatType
t) SubExp
x SubExp
y
internaliseBinOp Name
desc BinOp
E.Minus SubExp
x SubExp
y (E.Signed IntType
t) PrimType
_ =
  Name -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp Name
desc (IntType -> Overflow -> BinOp
I.Sub IntType
t Overflow
I.OverflowWrap) SubExp
x SubExp
y
internaliseBinOp Name
desc BinOp
E.Minus SubExp
x SubExp
y (E.Unsigned IntType
t) PrimType
_ =
  Name -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp Name
desc (IntType -> Overflow -> BinOp
I.Sub IntType
t Overflow
I.OverflowWrap) SubExp
x SubExp
y
internaliseBinOp Name
desc BinOp
E.Minus SubExp
x SubExp
y (E.FloatType FloatType
t) PrimType
_ =
  Name -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp Name
desc (FloatType -> BinOp
I.FSub FloatType
t) SubExp
x SubExp
y
internaliseBinOp Name
desc BinOp
E.Times SubExp
x SubExp
y (E.Signed IntType
t) PrimType
_ =
  Name -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp Name
desc (IntType -> Overflow -> BinOp
I.Mul IntType
t Overflow
I.OverflowWrap) SubExp
x SubExp
y
internaliseBinOp Name
desc BinOp
E.Times SubExp
x SubExp
y (E.Unsigned IntType
t) PrimType
_ =
  Name -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp Name
desc (IntType -> Overflow -> BinOp
I.Mul IntType
t Overflow
I.OverflowWrap) SubExp
x SubExp
y
internaliseBinOp Name
desc BinOp
E.Times SubExp
x SubExp
y (E.FloatType FloatType
t) PrimType
_ =
  Name -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp Name
desc (FloatType -> BinOp
I.FMul FloatType
t) SubExp
x SubExp
y
internaliseBinOp Name
desc BinOp
E.Divide SubExp
x SubExp
y (E.Signed IntType
t) PrimType
_ =
  IntType -> SubExp -> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a. IntType -> SubExp -> InternaliseM a -> InternaliseM a
certifyingNonzero IntType
t SubExp
y (InternaliseM [SubExp] -> InternaliseM [SubExp])
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$
    Name -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp Name
desc (IntType -> Safety -> BinOp
I.SDiv IntType
t Safety
I.Unsafe) SubExp
x SubExp
y
internaliseBinOp Name
desc BinOp
E.Divide SubExp
x SubExp
y (E.Unsigned IntType
t) PrimType
_ =
  IntType -> SubExp -> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a. IntType -> SubExp -> InternaliseM a -> InternaliseM a
certifyingNonzero IntType
t SubExp
y (InternaliseM [SubExp] -> InternaliseM [SubExp])
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$
    Name -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp Name
desc (IntType -> Safety -> BinOp
I.UDiv IntType
t Safety
I.Unsafe) SubExp
x SubExp
y
internaliseBinOp Name
desc BinOp
E.Divide SubExp
x SubExp
y (E.FloatType FloatType
t) PrimType
_ =
  Name -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp Name
desc (FloatType -> BinOp
I.FDiv FloatType
t) SubExp
x SubExp
y
internaliseBinOp Name
desc BinOp
E.Pow SubExp
x SubExp
y (E.FloatType FloatType
t) PrimType
_ =
  Name -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp Name
desc (FloatType -> BinOp
I.FPow FloatType
t) SubExp
x SubExp
y
internaliseBinOp Name
desc BinOp
E.Pow SubExp
x SubExp
y (E.Signed IntType
t) PrimType
_ =
  IntType -> SubExp -> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a. IntType -> SubExp -> InternaliseM a -> InternaliseM a
certifyingNonnegative IntType
t SubExp
y (InternaliseM [SubExp] -> InternaliseM [SubExp])
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$
    Name -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp Name
desc (IntType -> BinOp
I.Pow IntType
t) SubExp
x SubExp
y
internaliseBinOp Name
desc BinOp
E.Pow SubExp
x SubExp
y (E.Unsigned IntType
t) PrimType
_ =
  Name -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp Name
desc (IntType -> BinOp
I.Pow IntType
t) SubExp
x SubExp
y
internaliseBinOp Name
desc BinOp
E.Mod SubExp
x SubExp
y (E.Signed IntType
t) PrimType
_ =
  IntType -> SubExp -> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a. IntType -> SubExp -> InternaliseM a -> InternaliseM a
certifyingNonzero IntType
t SubExp
y (InternaliseM [SubExp] -> InternaliseM [SubExp])
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$
    Name -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp Name
desc (IntType -> Safety -> BinOp
I.SMod IntType
t Safety
I.Unsafe) SubExp
x SubExp
y
internaliseBinOp Name
desc BinOp
E.Mod SubExp
x SubExp
y (E.Unsigned IntType
t) PrimType
_ =
  IntType -> SubExp -> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a. IntType -> SubExp -> InternaliseM a -> InternaliseM a
certifyingNonzero IntType
t SubExp
y (InternaliseM [SubExp] -> InternaliseM [SubExp])
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$
    Name -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp Name
desc (IntType -> Safety -> BinOp
I.UMod IntType
t Safety
I.Unsafe) SubExp
x SubExp
y
internaliseBinOp Name
desc BinOp
E.Mod SubExp
x SubExp
y (E.FloatType FloatType
t) PrimType
_ =
  Name -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp Name
desc (FloatType -> BinOp
I.FMod FloatType
t) SubExp
x SubExp
y
internaliseBinOp Name
desc BinOp
E.Quot SubExp
x SubExp
y (E.Signed IntType
t) PrimType
_ =
  IntType -> SubExp -> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a. IntType -> SubExp -> InternaliseM a -> InternaliseM a
certifyingNonzero IntType
t SubExp
y (InternaliseM [SubExp] -> InternaliseM [SubExp])
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$
    Name -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp Name
desc (IntType -> Safety -> BinOp
I.SQuot IntType
t Safety
I.Unsafe) SubExp
x SubExp
y
internaliseBinOp Name
desc BinOp
E.Quot SubExp
x SubExp
y (E.Unsigned IntType
t) PrimType
_ =
  IntType -> SubExp -> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a. IntType -> SubExp -> InternaliseM a -> InternaliseM a
certifyingNonzero IntType
t SubExp
y (InternaliseM [SubExp] -> InternaliseM [SubExp])
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$
    Name -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp Name
desc (IntType -> Safety -> BinOp
I.UDiv IntType
t Safety
I.Unsafe) SubExp
x SubExp
y
internaliseBinOp Name
desc BinOp
E.Rem SubExp
x SubExp
y (E.Signed IntType
t) PrimType
_ =
  IntType -> SubExp -> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a. IntType -> SubExp -> InternaliseM a -> InternaliseM a
certifyingNonzero IntType
t SubExp
y (InternaliseM [SubExp] -> InternaliseM [SubExp])
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$
    Name -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp Name
desc (IntType -> Safety -> BinOp
I.SRem IntType
t Safety
I.Unsafe) SubExp
x SubExp
y
internaliseBinOp Name
desc BinOp
E.Rem SubExp
x SubExp
y (E.Unsigned IntType
t) PrimType
_ =
  IntType -> SubExp -> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a. IntType -> SubExp -> InternaliseM a -> InternaliseM a
certifyingNonzero IntType
t SubExp
y (InternaliseM [SubExp] -> InternaliseM [SubExp])
-> InternaliseM [SubExp] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$
    Name -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp Name
desc (IntType -> Safety -> BinOp
I.UMod IntType
t Safety
I.Unsafe) SubExp
x SubExp
y
internaliseBinOp Name
desc BinOp
E.ShiftR SubExp
x SubExp
y (E.Signed IntType
t) PrimType
_ =
  Name -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp Name
desc (IntType -> BinOp
I.AShr IntType
t) SubExp
x SubExp
y
internaliseBinOp Name
desc BinOp
E.ShiftR SubExp
x SubExp
y (E.Unsigned IntType
t) PrimType
_ =
  Name -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp Name
desc (IntType -> BinOp
I.LShr IntType
t) SubExp
x SubExp
y
internaliseBinOp Name
desc BinOp
E.ShiftL SubExp
x SubExp
y (E.Signed IntType
t) PrimType
_ =
  Name -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp Name
desc (IntType -> BinOp
I.Shl IntType
t) SubExp
x SubExp
y
internaliseBinOp Name
desc BinOp
E.ShiftL SubExp
x SubExp
y (E.Unsigned IntType
t) PrimType
_ =
  Name -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp Name
desc (IntType -> BinOp
I.Shl IntType
t) SubExp
x SubExp
y
internaliseBinOp Name
desc BinOp
E.Band SubExp
x SubExp
y (E.Signed IntType
t) PrimType
_ =
  Name -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp Name
desc (IntType -> BinOp
I.And IntType
t) SubExp
x SubExp
y
internaliseBinOp Name
desc BinOp
E.Band SubExp
x SubExp
y (E.Unsigned IntType
t) PrimType
_ =
  Name -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp Name
desc (IntType -> BinOp
I.And IntType
t) SubExp
x SubExp
y
internaliseBinOp Name
desc BinOp
E.Xor SubExp
x SubExp
y (E.Signed IntType
t) PrimType
_ =
  Name -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp Name
desc (IntType -> BinOp
I.Xor IntType
t) SubExp
x SubExp
y
internaliseBinOp Name
desc BinOp
E.Xor SubExp
x SubExp
y (E.Unsigned IntType
t) PrimType
_ =
  Name -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp Name
desc (IntType -> BinOp
I.Xor IntType
t) SubExp
x SubExp
y
internaliseBinOp Name
desc BinOp
E.Bor SubExp
x SubExp
y (E.Signed IntType
t) PrimType
_ =
  Name -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp Name
desc (IntType -> BinOp
I.Or IntType
t) SubExp
x SubExp
y
internaliseBinOp Name
desc BinOp
E.Bor SubExp
x SubExp
y (E.Unsigned IntType
t) PrimType
_ =
  Name -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp Name
desc (IntType -> BinOp
I.Or IntType
t) SubExp
x SubExp
y
internaliseBinOp Name
desc BinOp
E.Equal SubExp
x SubExp
y PrimType
t PrimType
_ =
  Name -> CmpOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleCmpOp Name
desc (PrimType -> CmpOp
I.CmpEq (PrimType -> CmpOp) -> PrimType -> CmpOp
forall a b. (a -> b) -> a -> b
$ PrimType -> PrimType
internalisePrimType PrimType
t) SubExp
x SubExp
y
internaliseBinOp Name
desc BinOp
E.NotEqual SubExp
x SubExp
y PrimType
t PrimType
_ = do
  eq <- Name -> Exp (Rep InternaliseM) -> InternaliseM 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
"true") (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
I.CmpEq (PrimType -> CmpOp) -> PrimType -> CmpOp
forall a b. (a -> b) -> a -> b
$ PrimType -> PrimType
internalisePrimType PrimType
t) SubExp
x SubExp
y
  fmap pure $ letSubExp desc $ I.BasicOp $ I.UnOp (I.Neg I.Bool) eq
internaliseBinOp Name
desc BinOp
E.Less SubExp
x SubExp
y (E.Signed IntType
t) PrimType
_ =
  Name -> CmpOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleCmpOp Name
desc (IntType -> CmpOp
I.CmpSlt IntType
t) SubExp
x SubExp
y
internaliseBinOp Name
desc BinOp
E.Less SubExp
x SubExp
y (E.Unsigned IntType
t) PrimType
_ =
  Name -> CmpOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleCmpOp Name
desc (IntType -> CmpOp
I.CmpUlt IntType
t) SubExp
x SubExp
y
internaliseBinOp Name
desc BinOp
E.Leq SubExp
x SubExp
y (E.Signed IntType
t) PrimType
_ =
  Name -> CmpOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleCmpOp Name
desc (IntType -> CmpOp
I.CmpSle IntType
t) SubExp
x SubExp
y
internaliseBinOp Name
desc BinOp
E.Leq SubExp
x SubExp
y (E.Unsigned IntType
t) PrimType
_ =
  Name -> CmpOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleCmpOp Name
desc (IntType -> CmpOp
I.CmpUle IntType
t) SubExp
x SubExp
y
internaliseBinOp Name
desc BinOp
E.Greater SubExp
x SubExp
y (E.Signed IntType
t) PrimType
_ =
  Name -> CmpOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleCmpOp Name
desc (IntType -> CmpOp
I.CmpSlt IntType
t) SubExp
y SubExp
x -- Note the swapped x and y
internaliseBinOp Name
desc BinOp
E.Greater SubExp
x SubExp
y (E.Unsigned IntType
t) PrimType
_ =
  Name -> CmpOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleCmpOp Name
desc (IntType -> CmpOp
I.CmpUlt IntType
t) SubExp
y SubExp
x -- Note the swapped x and y
internaliseBinOp Name
desc BinOp
E.Geq SubExp
x SubExp
y (E.Signed IntType
t) PrimType
_ =
  Name -> CmpOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleCmpOp Name
desc (IntType -> CmpOp
I.CmpSle IntType
t) SubExp
y SubExp
x -- Note the swapped x and y
internaliseBinOp Name
desc BinOp
E.Geq SubExp
x SubExp
y (E.Unsigned IntType
t) PrimType
_ =
  Name -> CmpOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleCmpOp Name
desc (IntType -> CmpOp
I.CmpUle IntType
t) SubExp
y SubExp
x -- Note the swapped x and y
internaliseBinOp Name
desc BinOp
E.Less SubExp
x SubExp
y (E.FloatType FloatType
t) PrimType
_ =
  Name -> CmpOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleCmpOp Name
desc (FloatType -> CmpOp
I.FCmpLt FloatType
t) SubExp
x SubExp
y
internaliseBinOp Name
desc BinOp
E.Leq SubExp
x SubExp
y (E.FloatType FloatType
t) PrimType
_ =
  Name -> CmpOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleCmpOp Name
desc (FloatType -> CmpOp
I.FCmpLe FloatType
t) SubExp
x SubExp
y
internaliseBinOp Name
desc BinOp
E.Greater SubExp
x SubExp
y (E.FloatType FloatType
t) PrimType
_ =
  Name -> CmpOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleCmpOp Name
desc (FloatType -> CmpOp
I.FCmpLt FloatType
t) SubExp
y SubExp
x -- Note the swapped x and y
internaliseBinOp Name
desc BinOp
E.Geq SubExp
x SubExp
y (E.FloatType FloatType
t) PrimType
_ =
  Name -> CmpOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleCmpOp Name
desc (FloatType -> CmpOp
I.FCmpLe FloatType
t) SubExp
y SubExp
x -- Note the swapped x and y

-- Relational operators for booleans.
internaliseBinOp Name
desc BinOp
E.Less SubExp
x SubExp
y PrimType
E.Bool PrimType
_ =
  Name -> CmpOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleCmpOp Name
desc CmpOp
I.CmpLlt SubExp
x SubExp
y
internaliseBinOp Name
desc BinOp
E.Leq SubExp
x SubExp
y PrimType
E.Bool PrimType
_ =
  Name -> CmpOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleCmpOp Name
desc CmpOp
I.CmpLle SubExp
x SubExp
y
internaliseBinOp Name
desc BinOp
E.Greater SubExp
x SubExp
y PrimType
E.Bool PrimType
_ =
  Name -> CmpOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleCmpOp Name
desc CmpOp
I.CmpLlt SubExp
y SubExp
x -- Note the swapped x and y
internaliseBinOp Name
desc BinOp
E.Geq SubExp
x SubExp
y PrimType
E.Bool PrimType
_ =
  Name -> CmpOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleCmpOp Name
desc CmpOp
I.CmpLle SubExp
y SubExp
x -- Note the swapped x and y
internaliseBinOp Name
_ BinOp
op SubExp
_ SubExp
_ PrimType
t1 PrimType
t2 =
  [Char] -> InternaliseM [SubExp]
forall a. HasCallStack => [Char] -> a
error ([Char] -> InternaliseM [SubExp])
-> [Char] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$
    [Char]
"Invalid binary operator "
      [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ BinOp -> [Char]
forall a. Pretty a => a -> [Char]
prettyString BinOp
op
      [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
" with operand types "
      [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ PrimType -> [Char]
forall a. Pretty a => a -> [Char]
prettyString PrimType
t1
      [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
", "
      [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ PrimType -> [Char]
forall a. Pretty a => a -> [Char]
prettyString PrimType
t2

simpleBinOp ::
  Name ->
  I.BinOp ->
  I.SubExp ->
  I.SubExp ->
  InternaliseM [I.SubExp]
simpleBinOp :: Name -> BinOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleBinOp Name
desc BinOp
bop SubExp
x SubExp
y =
  Name -> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m [SubExp]
letTupExp' Name
desc (Exp (Rep InternaliseM) -> InternaliseM [SubExp])
-> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp BinOp
bop SubExp
x SubExp
y

simpleCmpOp ::
  Name ->
  I.CmpOp ->
  I.SubExp ->
  I.SubExp ->
  InternaliseM [I.SubExp]
simpleCmpOp :: Name -> CmpOp -> SubExp -> SubExp -> InternaliseM [SubExp]
simpleCmpOp Name
desc CmpOp
op SubExp
x SubExp
y =
  Name -> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m [SubExp]
letTupExp' Name
desc (Exp (Rep InternaliseM) -> InternaliseM [SubExp])
-> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp CmpOp
op SubExp
x SubExp
y

data Function
  = FunctionName (E.QualName VName)
  | FunctionHole SrcLoc
  deriving (Int -> Function -> [Char] -> [Char]
[Function] -> [Char] -> [Char]
Function -> [Char]
(Int -> Function -> [Char] -> [Char])
-> (Function -> [Char])
-> ([Function] -> [Char] -> [Char])
-> Show Function
forall a.
(Int -> a -> [Char] -> [Char])
-> (a -> [Char]) -> ([a] -> [Char] -> [Char]) -> Show a
$cshowsPrec :: Int -> Function -> [Char] -> [Char]
showsPrec :: Int -> Function -> [Char] -> [Char]
$cshow :: Function -> [Char]
show :: Function -> [Char]
$cshowList :: [Function] -> [Char] -> [Char]
showList :: [Function] -> [Char] -> [Char]
Show)

findFuncall :: E.AppExp -> (Function, [(E.Exp, Maybe VName)])
findFuncall :: AppExp -> (Function, [(Exp, Maybe VName)])
findFuncall (E.Apply Exp
f NonEmpty (Info (Maybe VName), Exp)
args SrcLoc
_)
  | E.Var QualName VName
fname Info StructType
_ SrcLoc
_ <- Exp
f =
      (QualName VName -> Function
FunctionName QualName VName
fname, ((Info (Maybe VName), Exp) -> (Exp, Maybe VName))
-> [(Info (Maybe VName), Exp)] -> [(Exp, Maybe VName)]
forall a b. (a -> b) -> [a] -> [b]
map (Info (Maybe VName), Exp) -> (Exp, Maybe VName)
forall {b} {a}. (Info b, a) -> (a, b)
onArg ([(Info (Maybe VName), Exp)] -> [(Exp, Maybe VName)])
-> [(Info (Maybe VName), Exp)] -> [(Exp, Maybe VName)]
forall a b. (a -> b) -> a -> b
$ NonEmpty (Info (Maybe VName), Exp) -> [(Info (Maybe VName), Exp)]
forall a. NonEmpty a -> [a]
NE.toList NonEmpty (Info (Maybe VName), Exp)
args)
  | E.Hole (Info StructType
_) SrcLoc
loc <- Exp
f =
      (SrcLoc -> Function
FunctionHole SrcLoc
loc, ((Info (Maybe VName), Exp) -> (Exp, Maybe VName))
-> [(Info (Maybe VName), Exp)] -> [(Exp, Maybe VName)]
forall a b. (a -> b) -> [a] -> [b]
map (Info (Maybe VName), Exp) -> (Exp, Maybe VName)
forall {b} {a}. (Info b, a) -> (a, b)
onArg ([(Info (Maybe VName), Exp)] -> [(Exp, Maybe VName)])
-> [(Info (Maybe VName), Exp)] -> [(Exp, Maybe VName)]
forall a b. (a -> b) -> a -> b
$ NonEmpty (Info (Maybe VName), Exp) -> [(Info (Maybe VName), Exp)]
forall a. NonEmpty a -> [a]
NE.toList NonEmpty (Info (Maybe VName), Exp)
args)
  where
    onArg :: (Info b, a) -> (a, b)
onArg (Info b
argext, a
e) = (a
e, b
argext)
findFuncall AppExp
e =
  [Char] -> (Function, [(Exp, Maybe VName)])
forall a. HasCallStack => [Char] -> a
error ([Char] -> (Function, [(Exp, Maybe VName)]))
-> [Char] -> (Function, [(Exp, Maybe VName)])
forall a b. (a -> b) -> a -> b
$ [Char]
"Invalid function expression in application:\n" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ AppExp -> [Char]
forall a. Pretty a => a -> [Char]
prettyString AppExp
e

-- The type of a body.  Watch out: this only works for the degenerate
-- case where the body does not already return its context.
bodyExtType :: Body SOACS -> InternaliseM [ExtType]
bodyExtType :: Body SOACS -> InternaliseM [TypeBase ExtShape NoMode]
bodyExtType (Body BodyDec SOACS
_ Stms SOACS
stms [SubExpRes]
res) =
  [VName] -> [TypeBase ExtShape NoMode] -> [TypeBase ExtShape NoMode]
existentialiseExtTypes (Scope SOACS -> [VName]
forall k a. Map k a -> [k]
M.keys Scope SOACS
stmsscope) ([TypeBase ExtShape NoMode] -> [TypeBase ExtShape NoMode])
-> ([TypeBase Shape NoMode] -> [TypeBase ExtShape NoMode])
-> [TypeBase Shape NoMode]
-> [TypeBase ExtShape NoMode]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [TypeBase Shape NoMode] -> [TypeBase ExtShape NoMode]
forall o. [TypeBase Shape o] -> [TypeBase ExtShape o]
staticShapes
    ([TypeBase Shape NoMode] -> [TypeBase ExtShape NoMode])
-> InternaliseM [TypeBase Shape NoMode]
-> InternaliseM [TypeBase ExtShape NoMode]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ExtendedScope SOACS InternaliseM [TypeBase Shape NoMode]
-> Scope SOACS -> InternaliseM [TypeBase Shape NoMode]
forall rep (m :: * -> *) a.
ExtendedScope rep m a -> Scope rep -> m a
extendedScope ((SubExpRes
 -> ExtendedScope SOACS InternaliseM (TypeBase Shape NoMode))
-> [SubExpRes]
-> ExtendedScope SOACS InternaliseM [TypeBase Shape NoMode]
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 SOACS InternaliseM (TypeBase Shape NoMode)
forall t (m :: * -> *).
HasScope t m =>
SubExpRes -> m (TypeBase Shape NoMode)
subExpResType [SubExpRes]
res) Scope SOACS
stmsscope
  where
    stmsscope :: Scope SOACS
stmsscope = Stms SOACS -> Scope SOACS
forall rep a. Scoped rep a => a -> Scope rep
scopeOf Stms SOACS
stms

internaliseLambda ::
  E.Exp ->
  [I.Type] ->
  InternaliseM ([I.LParam SOACS], I.Body SOACS, [I.Type])
internaliseLambda :: Exp
-> [TypeBase Shape NoMode]
-> InternaliseM
     ([LParam SOACS], Body SOACS, [TypeBase Shape NoMode])
internaliseLambda (E.Parens Exp
e SrcLoc
_) [TypeBase Shape NoMode]
rowtypes =
  Exp
-> [TypeBase Shape NoMode]
-> InternaliseM
     ([LParam SOACS], Body SOACS, [TypeBase Shape NoMode])
internaliseLambda Exp
e [TypeBase Shape NoMode]
rowtypes
internaliseLambda (E.Lambda [PatBase Info VName ParamType]
params Exp
body Maybe (TypeExp Exp VName)
_ (Info (RetType [VName]
_ TypeBase Exp Freshness
rettype)) SrcLoc
_) [TypeBase Shape NoMode]
rowtypes =
  [PatBase Info VName ParamType]
-> [TypeBase Shape NoMode]
-> ([LParam SOACS]
    -> InternaliseM
         ([LParam SOACS], Body SOACS, [TypeBase Shape NoMode]))
-> InternaliseM
     ([LParam SOACS], Body SOACS, [TypeBase Shape NoMode])
forall a.
[PatBase Info VName ParamType]
-> [TypeBase Shape NoMode]
-> ([LParam SOACS] -> InternaliseM a)
-> InternaliseM a
bindingLambdaParams [PatBase Info VName ParamType]
params [TypeBase Shape NoMode]
rowtypes (([LParam SOACS]
  -> InternaliseM
       ([LParam SOACS], Body SOACS, [TypeBase Shape NoMode]))
 -> InternaliseM
      ([LParam SOACS], Body SOACS, [TypeBase Shape NoMode]))
-> ([LParam SOACS]
    -> InternaliseM
         ([LParam SOACS], Body SOACS, [TypeBase Shape NoMode]))
-> InternaliseM
     ([LParam SOACS], Body SOACS, [TypeBase Shape NoMode])
forall a b. (a -> b) -> a -> b
$ \[LParam SOACS]
params' -> do
    body' <- Name -> Exp -> InternaliseM (Body SOACS)
internaliseBody Name
"lam" Exp
body
    rettype' <- internaliseLambdaReturnType rettype =<< bodyExtType body'
    pure (params', body', rettype')
internaliseLambda Exp
e [TypeBase Shape NoMode]
_ = [Char]
-> InternaliseM
     ([LParam SOACS], Body SOACS, [TypeBase Shape NoMode])
forall a. HasCallStack => [Char] -> a
error ([Char]
 -> InternaliseM
      ([LParam SOACS], Body SOACS, [TypeBase Shape NoMode]))
-> [Char]
-> InternaliseM
     ([LParam SOACS], Body SOACS, [TypeBase Shape NoMode])
forall a b. (a -> b) -> a -> b
$ [Char]
"internaliseLambda: unexpected expression:\n" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Exp -> [Char]
forall a. Pretty a => a -> [Char]
prettyString Exp
e

internaliseLambdaCoerce :: E.Exp -> [Type] -> InternaliseM (I.Lambda SOACS)
internaliseLambdaCoerce :: Exp -> [TypeBase Shape NoMode] -> InternaliseM (Lambda SOACS)
internaliseLambdaCoerce Exp
lam [TypeBase Shape NoMode]
argtypes = do
  (params, body, rettype) <- Exp
-> [TypeBase Shape NoMode]
-> InternaliseM
     ([LParam SOACS], Body SOACS, [TypeBase Shape NoMode])
internaliseLambda Exp
lam [TypeBase Shape NoMode]
argtypes
  mkLambda params $
    ensureResultShape
      (ErrorMsg [ErrorString "unexpected lambda result size"])
      rettype
      =<< bodyBind body

-- | The number of internalised values in the array component of the result of
-- the lambda of a 'flatmap', which returns that array paired with a value of
-- regular type.
flatLambdaIrregulars :: E.Exp -> Int
flatLambdaIrregulars :: Exp -> Int
flatLambdaIrregulars (E.Parens Exp
e SrcLoc
_) = Exp -> Int
flatLambdaIrregulars Exp
e
flatLambdaIrregulars (E.Lambda [PatBase Info VName ParamType]
_ Exp
_ Maybe (TypeExp Exp VName)
_ (Info (E.RetType [VName]
_ TypeBase Exp Freshness
t)) SrcLoc
_)
  | Just (TypeBase Exp Freshness
irreg_t : [TypeBase Exp Freshness]
_) <- TypeBase Exp Freshness -> Maybe [TypeBase Exp Freshness]
forall dim o. TypeBase dim o -> Maybe [TypeBase dim o]
E.isTupleRecord TypeBase Exp Freshness
t =
      StructType -> Int
forall o. TypeBase Exp o -> Int
internalisedTypeSize (StructType -> Int) -> StructType -> Int
forall a b. (a -> b) -> a -> b
$ TypeBase Exp Freshness -> StructType
forall dim o. TypeBase dim o -> TypeBase dim NoMode
E.toStruct TypeBase Exp Freshness
irreg_t
flatLambdaIrregulars Exp
e =
  [Char] -> Int
forall a. HasCallStack => [Char] -> a
error ([Char] -> Int) -> [Char] -> Int
forall a b. (a -> b) -> a -> b
$ [Char]
"flatLambdaIrregulars: unexpected expression:\n" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Exp -> [Char]
forall a. Pretty a => a -> [Char]
prettyString Exp
e

-- | Internalise the lambda of a 'flatmap'. The irregular results (the ones that
-- are concatenated) are distinguished in the 'ExtLambda' by having the
-- existential size as their outermost size, so that size must be a variable -
-- when the lambda produces arrays of some other size, we bind that size and
-- coerce the arrays to it. By the source language type rules, this cannot fail.
internaliseFlatLambda :: E.Exp -> [Type] -> InternaliseM (I.ExtLambda SOACS)
internaliseFlatLambda :: Exp -> [TypeBase Shape NoMode] -> InternaliseM (ExtLambda SOACS)
internaliseFlatLambda Exp
lam [TypeBase Shape NoMode]
argtypes = do
  (params, body, _) <- Exp
-> [TypeBase Shape NoMode]
-> InternaliseM
     ([LParam SOACS], Body SOACS, [TypeBase Shape NoMode])
internaliseLambda Exp
lam [TypeBase Shape NoMode]
argtypes
  (body', ret) <- buildBody . localScope (scopeOfLParams params) $ do
    res <- bodyBind body
    let (irreg_res, reg_res) = splitAt (flatLambdaIrregulars lam) res
    k <- irregularSize irreg_res
    irreg_res' <- mapM (coerceOuter k) irreg_res
    irreg_ts <- mapM (subExpType . resSubExp) irreg_res'
    reg_ts <- mapM (subExpType . resSubExp) reg_res
    pure
      ( subExpRes (I.Var k) : irreg_res' <> reg_res,
        I.Prim int64
          : existentialiseExtTypes [k] (staticShapes irreg_ts)
            <> staticShapes reg_ts
      )
  pure $ I.Lambda params ret body'
  where
    irregularSize :: [SubExpRes] -> InternaliseM VName
irregularSize [] =
      [Char] -> InternaliseM VName
forall a. HasCallStack => [Char] -> a
error [Char]
"internaliseFlatLambda: lambda has no irregular results."
    irregularSize (SubExpRes
r : [SubExpRes]
_) =
      Name -> SubExp -> InternaliseM VName
asVar Name
"flatmap_k" (SubExp -> InternaliseM VName)
-> (TypeBase Shape NoMode -> SubExp)
-> TypeBase Shape NoMode
-> InternaliseM VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> TypeBase Shape NoMode -> SubExp
forall o. Int -> TypeBase Shape o -> SubExp
arraySize Int
0 (TypeBase Shape NoMode -> InternaliseM VName)
-> InternaliseM (TypeBase Shape NoMode) -> InternaliseM VName
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< SubExp -> InternaliseM (TypeBase Shape NoMode)
forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoMode)
subExpType (SubExpRes -> SubExp
resSubExp SubExpRes
r)

    coerceOuter :: VName -> SubExpRes -> m SubExpRes
coerceOuter VName
k r :: SubExpRes
r@(SubExpRes Certs
cs (I.Var VName
v)) = do
      t <- VName -> m (TypeBase Shape NoMode)
forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoMode)
lookupType VName
v
      case arrayDims t of
        SubExp
d : [SubExp]
ds
          | SubExp
d SubExp -> SubExp -> Bool
forall a. Eq a => a -> a -> Bool
/= VName -> SubExp
I.Var VName
k ->
              Certs -> SubExp -> SubExpRes
SubExpRes Certs
cs
                (SubExp -> SubExpRes) -> m SubExp -> m SubExpRes
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
"flatmap_irreg" ([SubExp] -> VName -> Exp (Rep m)
forall rep. [SubExp] -> VName -> Exp rep
shapeCoerce (VName -> SubExp
I.Var VName
k SubExp -> [SubExp] -> [SubExp]
forall a. a -> [a] -> [a]
: [SubExp]
ds) VName
v)
        [SubExp]
_ -> SubExpRes -> m SubExpRes
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SubExpRes
r
    coerceOuter VName
_ SubExpRes
r = SubExpRes -> m SubExpRes
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SubExpRes
r

-- | Overloaded operators are treated here.
isOverloadedFunction ::
  E.QualName VName ->
  Name ->
  Maybe ([(E.StructType, [SubExp])] -> InternaliseM [SubExp])
isOverloadedFunction :: QualName VName
-> Name
-> Maybe ([(StructType, [SubExp])] -> InternaliseM [SubExp])
isOverloadedFunction QualName VName
qname Name
desc = do
  Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> Maybe ()) -> Bool -> Maybe ()
forall a b. (a -> b) -> a -> b
$ VName -> Bool
isIntrinsic (VName -> Bool) -> VName -> Bool
forall a b. (a -> b) -> a -> b
$ QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
qname
  Name -> Maybe ([(StructType, [SubExp])] -> InternaliseM [SubExp])
handle (Name -> Maybe ([(StructType, [SubExp])] -> InternaliseM [SubExp]))
-> Name
-> Maybe ([(StructType, [SubExp])] -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ VName -> Name
baseName (VName -> Name) -> VName -> Name
forall a b. (a -> b) -> a -> b
$ QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
qname
  where
    -- Handle equality and inequality specially, to treat the case of
    -- arrays.
    handle :: Name -> Maybe ([(StructType, [SubExp])] -> InternaliseM [SubExp])
handle Name
op
      | Just SubExp -> InternaliseM [SubExp]
cmp_f <- Name -> Maybe (SubExp -> InternaliseM [SubExp])
isEqlOp Name
op = ([(StructType, [SubExp])] -> InternaliseM [SubExp])
-> Maybe ([(StructType, [SubExp])] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([(StructType, [SubExp])] -> InternaliseM [SubExp])
 -> Maybe ([(StructType, [SubExp])] -> InternaliseM [SubExp]))
-> ([(StructType, [SubExp])] -> InternaliseM [SubExp])
-> Maybe ([(StructType, [SubExp])] -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \[(StructType
_, [SubExp]
xe'), (StructType
_, [SubExp]
ye')] -> do
          rs <- (SubExp -> SubExp -> InternaliseM SubExp)
-> [SubExp] -> [SubExp] -> InternaliseM [SubExp]
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM SubExp -> SubExp -> InternaliseM SubExp
doComparison [SubExp]
xe' [SubExp]
ye'
          cmp_f =<< letSubExp "eq" =<< eAll rs
      where
        isEqlOp :: Name -> Maybe (SubExp -> InternaliseM [SubExp])
isEqlOp Name
"!=" = (SubExp -> InternaliseM [SubExp])
-> Maybe (SubExp -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((SubExp -> InternaliseM [SubExp])
 -> Maybe (SubExp -> InternaliseM [SubExp]))
-> (SubExp -> InternaliseM [SubExp])
-> Maybe (SubExp -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \SubExp
eq ->
          Name -> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m [SubExp]
letTupExp' Name
desc (Exp (Rep InternaliseM) -> InternaliseM [SubExp])
-> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
forall a b. (a -> b) -> a -> b
$ UnOp -> SubExp -> BasicOp
I.UnOp (PrimType -> UnOp
I.Neg PrimType
I.Bool) SubExp
eq
        isEqlOp Name
"==" = (SubExp -> InternaliseM [SubExp])
-> Maybe (SubExp -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((SubExp -> InternaliseM [SubExp])
 -> Maybe (SubExp -> InternaliseM [SubExp]))
-> (SubExp -> InternaliseM [SubExp])
-> Maybe (SubExp -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \SubExp
eq ->
          [SubExp] -> InternaliseM [SubExp]
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [SubExp
eq]
        isEqlOp Name
_ = Maybe (SubExp -> InternaliseM [SubExp])
forall a. Maybe a
Nothing

        doComparison :: SubExp -> SubExp -> InternaliseM SubExp
doComparison SubExp
x SubExp
y = do
          x_t <- SubExp -> InternaliseM (TypeBase Shape NoMode)
forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoMode)
I.subExpType SubExp
x
          y_t <- I.subExpType y
          case x_t of
            I.Prim PrimType
t -> Name -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
desc (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
I.CmpEq PrimType
t) SubExp
x SubExp
y
            TypeBase Shape NoMode
_ -> do
              let x_dims :: [SubExp]
x_dims = TypeBase Shape NoMode -> [SubExp]
forall d m. TypeBase (ShapeBase d) m -> [d]
I.arrayDims TypeBase Shape NoMode
x_t
                  y_dims :: [SubExp]
y_dims = TypeBase Shape NoMode -> [SubExp]
forall d m. TypeBase (ShapeBase d) m -> [d]
I.arrayDims TypeBase Shape NoMode
y_t
              dims_match <- [(SubExp, SubExp)]
-> ((SubExp, SubExp) -> InternaliseM SubExp)
-> InternaliseM [SubExp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM ([SubExp] -> [SubExp] -> [(SubExp, SubExp)]
forall a b. [a] -> [b] -> [(a, b)]
zip [SubExp]
x_dims [SubExp]
y_dims) (((SubExp, SubExp) -> InternaliseM SubExp)
 -> InternaliseM [SubExp])
-> ((SubExp, SubExp) -> InternaliseM SubExp)
-> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ \(SubExp
x_dim, SubExp
y_dim) ->
                Name -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"dim_eq" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
forall a b. (a -> b) -> a -> b
$ CmpOp -> SubExp -> SubExp -> BasicOp
I.CmpOp (PrimType -> CmpOp
I.CmpEq PrimType
int64) SubExp
x_dim SubExp
y_dim
              shapes_match <- letSubExp "shapes_match" =<< eAll dims_match
              let compare_elems_body = Builder SOACS [SubExpRes] -> InternaliseM (Body SOACS)
forall rep (m :: * -> *) somerep res.
(Buildable rep, MonadFreshNames m, HasScope somerep m,
 SameScope somerep rep, IsResult res) =>
Builder rep [res] -> m (GBody rep res)
runBodyBuilder (Builder SOACS [SubExpRes] -> InternaliseM (Body SOACS))
-> Builder SOACS [SubExpRes] -> InternaliseM (Body SOACS)
forall a b. (a -> b) -> a -> b
$ do
                    -- Flatten both x and y.
                    x_num_elems <-
                      Name
-> Exp (Rep (BuilderT SOACS (State VNameSource)))
-> BuilderT SOACS (State VNameSource) SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"x_num_elems"
                        (Exp SOACS -> BuilderT SOACS (State VNameSource) SubExp)
-> BuilderT SOACS (State VNameSource) (Exp SOACS)
-> BuilderT SOACS (State VNameSource) SubExp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< BinOp
-> SubExp
-> [SubExp]
-> BuilderT
     SOACS
     (State VNameSource)
     (Exp (Rep (BuilderT SOACS (State VNameSource))))
forall (m :: * -> *).
MonadBuilder m =>
BinOp -> SubExp -> [SubExp] -> m (Exp (Rep m))
foldBinOp (IntType -> Overflow -> BinOp
I.Mul IntType
Int64 Overflow
I.OverflowUndef) (Int64 -> SubExp
forall v. IsValue v => v -> SubExp
constant (Int64
1 :: Int64)) [SubExp]
x_dims
                    x' <- letExp "x" $ I.BasicOp $ I.SubExp x
                    y' <- letExp "x" $ I.BasicOp $ I.SubExp y
                    x_flat <-
                      letExp "x_flat" . I.BasicOp $
                        I.Reshape x' (reshapeAll (I.arrayShape x_t) (I.Shape [x_num_elems]))
                    y_flat <-
                      letExp "y_flat" . I.BasicOp $
                        I.Reshape y' (reshapeAll (I.arrayShape x_t) (I.Shape [x_num_elems]))

                    -- Compare the elements.
                    cmp_lam <- cmpOpLambda $ I.CmpEq (elemType x_t)
                    cmps <-
                      letExp "cmps"
                        . I.Op
                        . I.Screma x_num_elems [x_flat, y_flat]
                        =<< I.mapSOAC cmp_lam

                    -- Check that all were equal.
                    and_lam <- binOpLambda I.LogAnd I.Bool
                    reduce <- I.reduceSOAC [Reduce Commutative and_lam [constant True]]
                    all_equal <- letSubExp "all_equal" $ I.Op $ I.Screma x_num_elems [cmps] reduce
                    pure $ subExpsRes [all_equal]

              letSubExp "arrays_equal"
                =<< eIf (eSubExp shapes_match) compare_elems_body (resultBodyM [constant False])
    handle Name
name
      | Just BinOp
bop <-
          (BinOp -> Bool) -> [BinOp] -> Maybe BinOp
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find
            ((Name
name Name -> Name -> Bool
forall a. Eq a => a -> a -> Bool
==) (Name -> Bool) -> (BinOp -> Name) -> BinOp -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> Name
nameFromText (Text -> Name) -> (BinOp -> Text) -> BinOp -> Name
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BinOp -> Text
forall a. Pretty a => a -> Text
prettyText)
            [BinOp
forall a. Bounded a => a
minBound .. BinOp
forall a. Bounded a => a
maxBound :: E.BinOp] =
          ([(StructType, [SubExp])] -> InternaliseM [SubExp])
-> Maybe ([(StructType, [SubExp])] -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just (([(StructType, [SubExp])] -> InternaliseM [SubExp])
 -> Maybe ([(StructType, [SubExp])] -> InternaliseM [SubExp]))
-> ([(StructType, [SubExp])] -> InternaliseM [SubExp])
-> Maybe ([(StructType, [SubExp])] -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \[(StructType
x_t, [SubExp
x']), (StructType
y_t, [SubExp
y'])] ->
            case (StructType
x_t, StructType
y_t) of
              (E.Scalar (E.Prim PrimType
t1), E.Scalar (E.Prim PrimType
t2)) ->
                Name
-> BinOp
-> SubExp
-> SubExp
-> PrimType
-> PrimType
-> InternaliseM [SubExp]
internaliseBinOp Name
desc BinOp
bop SubExp
x' SubExp
y' PrimType
t1 PrimType
t2
              (StructType, StructType)
_ -> [Char] -> InternaliseM [SubExp]
forall a. HasCallStack => [Char] -> a
error [Char]
"Futhark.Internalise.internaliseExp: non-primitive type in BinOp."
    handle Name
_ = Maybe ([(StructType, [SubExp])] -> InternaliseM [SubExp])
forall a. Maybe a
Nothing

scatterF :: Int -> E.Exp -> E.Exp -> E.Exp -> Name -> InternaliseM [SubExp]
scatterF :: Int -> Exp -> Exp -> Exp -> Name -> InternaliseM [SubExp]
scatterF Int
rank Exp
dest Exp
is Exp
v Name
desc = do
  is' <- Name -> Exp -> InternaliseM [VName]
internaliseExpToVars Name
"write_arg_i" Exp
is
  v' <- internaliseExpToVars "write_arg_v" v
  dest' <- internaliseExpToVars "write_arg_a" dest
  map I.Var
    <$> doScatter desc rank dest' (is' <> v') (pure . map (I.Var . I.paramName))

-- | Handle intrinsic functions.  These are only allowed to be called
-- in the prelude, and their internalisation may involve inspecting
-- the AST.
isIntrinsicFunction ::
  E.QualName VName ->
  [E.Exp] ->
  Maybe (Name -> InternaliseM [SubExp])
isIntrinsicFunction :: QualName VName -> [Exp] -> Maybe (Name -> InternaliseM [SubExp])
isIntrinsicFunction QualName VName
qname [Exp]
args = do
  Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> Maybe ()) -> Bool -> Maybe ()
forall a b. (a -> b) -> a -> b
$ VName -> Bool
isIntrinsic (VName -> Bool) -> VName -> Bool
forall a b. (a -> b) -> a -> b
$ QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
qname
  let handlers :: [[Exp] -> Name -> Maybe (Name -> InternaliseM [SubExp])]
handlers =
        [ [Exp] -> Name -> Maybe (Name -> InternaliseM [SubExp])
forall {a}.
(Eq a, IsString a) =>
[Exp] -> a -> Maybe (Name -> InternaliseM [SubExp])
handleSign,
          [Exp] -> Name -> Maybe (Name -> InternaliseM [SubExp])
forall {f :: * -> *}.
Applicative f =>
[Exp] -> Name -> Maybe (Name -> InternaliseM (f SubExp))
handleOps,
          [Exp] -> Name -> Maybe (Name -> InternaliseM [SubExp])
forall {a}.
(Eq a, IsString a) =>
[Exp] -> a -> Maybe (Name -> InternaliseM [SubExp])
handleSOACs,
          [Exp] -> Name -> Maybe (Name -> InternaliseM [SubExp])
forall {a}.
(Eq a, IsString a) =>
[Exp] -> a -> Maybe (Name -> InternaliseM [SubExp])
handleAccs,
          [Exp] -> Name -> Maybe (Name -> InternaliseM [SubExp])
forall {a}.
(Eq a, IsString a) =>
[Exp] -> a -> Maybe (Name -> InternaliseM [SubExp])
handleAD,
          [Exp] -> Name -> Maybe (Name -> InternaliseM [SubExp])
forall {p} {t} {t}. p -> Name -> Maybe (t -> t)
handleIO,
          [Exp] -> Name -> Maybe (Name -> InternaliseM [SubExp])
forall {a}.
(IsString a, Eq a) =>
[Exp] -> a -> Maybe (Name -> InternaliseM [SubExp])
handleRest
        ]
  [Maybe (Name -> InternaliseM [SubExp])]
-> Maybe (Name -> InternaliseM [SubExp])
forall (t :: * -> *) (m :: * -> *) a.
(Foldable t, MonadPlus m) =>
t (m a) -> m a
msum [[Exp] -> Name -> Maybe (Name -> InternaliseM [SubExp])
h [Exp]
args (Name -> Maybe (Name -> InternaliseM [SubExp]))
-> Name -> Maybe (Name -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ VName -> Name
baseName (VName -> Name) -> VName -> Name
forall a b. (a -> b) -> a -> b
$ QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
qname | [Exp] -> Name -> Maybe (Name -> InternaliseM [SubExp])
h <- [[Exp] -> Name -> Maybe (Name -> InternaliseM [SubExp])]
handlers]
  where
    handleSign :: [Exp] -> a -> Maybe (Name -> InternaliseM [SubExp])
handleSign [Exp
x] a
"sign_i8" = (Name -> InternaliseM [SubExp])
-> Maybe (Name -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((Name -> InternaliseM [SubExp])
 -> Maybe (Name -> InternaliseM [SubExp]))
-> (Name -> InternaliseM [SubExp])
-> Maybe (Name -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ IntType -> Exp -> Name -> InternaliseM [SubExp]
toSigned IntType
I.Int8 Exp
x
    handleSign [Exp
x] a
"sign_i16" = (Name -> InternaliseM [SubExp])
-> Maybe (Name -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((Name -> InternaliseM [SubExp])
 -> Maybe (Name -> InternaliseM [SubExp]))
-> (Name -> InternaliseM [SubExp])
-> Maybe (Name -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ IntType -> Exp -> Name -> InternaliseM [SubExp]
toSigned IntType
I.Int16 Exp
x
    handleSign [Exp
x] a
"sign_i32" = (Name -> InternaliseM [SubExp])
-> Maybe (Name -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((Name -> InternaliseM [SubExp])
 -> Maybe (Name -> InternaliseM [SubExp]))
-> (Name -> InternaliseM [SubExp])
-> Maybe (Name -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ IntType -> Exp -> Name -> InternaliseM [SubExp]
toSigned IntType
I.Int32 Exp
x
    handleSign [Exp
x] a
"sign_i64" = (Name -> InternaliseM [SubExp])
-> Maybe (Name -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((Name -> InternaliseM [SubExp])
 -> Maybe (Name -> InternaliseM [SubExp]))
-> (Name -> InternaliseM [SubExp])
-> Maybe (Name -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ IntType -> Exp -> Name -> InternaliseM [SubExp]
toSigned IntType
I.Int64 Exp
x
    handleSign [Exp
x] a
"unsign_i8" = (Name -> InternaliseM [SubExp])
-> Maybe (Name -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((Name -> InternaliseM [SubExp])
 -> Maybe (Name -> InternaliseM [SubExp]))
-> (Name -> InternaliseM [SubExp])
-> Maybe (Name -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ IntType -> Exp -> Name -> InternaliseM [SubExp]
toUnsigned IntType
I.Int8 Exp
x
    handleSign [Exp
x] a
"unsign_i16" = (Name -> InternaliseM [SubExp])
-> Maybe (Name -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((Name -> InternaliseM [SubExp])
 -> Maybe (Name -> InternaliseM [SubExp]))
-> (Name -> InternaliseM [SubExp])
-> Maybe (Name -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ IntType -> Exp -> Name -> InternaliseM [SubExp]
toUnsigned IntType
I.Int16 Exp
x
    handleSign [Exp
x] a
"unsign_i32" = (Name -> InternaliseM [SubExp])
-> Maybe (Name -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((Name -> InternaliseM [SubExp])
 -> Maybe (Name -> InternaliseM [SubExp]))
-> (Name -> InternaliseM [SubExp])
-> Maybe (Name -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ IntType -> Exp -> Name -> InternaliseM [SubExp]
toUnsigned IntType
I.Int32 Exp
x
    handleSign [Exp
x] a
"unsign_i64" = (Name -> InternaliseM [SubExp])
-> Maybe (Name -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((Name -> InternaliseM [SubExp])
 -> Maybe (Name -> InternaliseM [SubExp]))
-> (Name -> InternaliseM [SubExp])
-> Maybe (Name -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ IntType -> Exp -> Name -> InternaliseM [SubExp]
toUnsigned IntType
I.Int64 Exp
x
    handleSign [Exp]
_ a
_ = Maybe (Name -> InternaliseM [SubExp])
forall a. Maybe a
Nothing

    handleOps :: [Exp] -> Name -> Maybe (Name -> InternaliseM (f SubExp))
handleOps [Exp
x] Name
s
      | Just UnOp
unop <- (UnOp -> Bool) -> [UnOp] -> Maybe UnOp
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find ((Name -> Name -> Bool
forall a. Eq a => a -> a -> Bool
== Name
s) (Name -> Bool) -> (UnOp -> Name) -> UnOp -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> Name
nameFromText (Text -> Name) -> (UnOp -> Text) -> UnOp -> Name
forall b c a. (b -> c) -> (a -> b) -> a -> c
. UnOp -> Text
forall a. Pretty a => a -> Text
prettyText) [UnOp]
allUnOps = (Name -> InternaliseM (f SubExp))
-> Maybe (Name -> InternaliseM (f SubExp))
forall a. a -> Maybe a
Just ((Name -> InternaliseM (f SubExp))
 -> Maybe (Name -> InternaliseM (f SubExp)))
-> (Name -> InternaliseM (f SubExp))
-> Maybe (Name -> InternaliseM (f SubExp))
forall a b. (a -> b) -> a -> b
$ \Name
desc -> do
          x' <- Name -> Exp -> InternaliseM SubExp
internaliseExp1 Name
"x" Exp
x
          fmap pure $ letSubExp desc $ I.BasicOp $ I.UnOp unop x'
    handleOps [TupLit [Exp
x, Exp
y] SrcLoc
_] Name
s
      | Just BinOp
bop <- (BinOp -> Bool) -> [BinOp] -> Maybe BinOp
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find ((Name -> Name -> Bool
forall a. Eq a => a -> a -> Bool
== Name
s) (Name -> Bool) -> (BinOp -> Name) -> BinOp -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> Name
nameFromText (Text -> Name) -> (BinOp -> Text) -> BinOp -> Name
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BinOp -> Text
forall a. Pretty a => a -> Text
prettyText) [BinOp]
allBinOps = (Name -> InternaliseM (f SubExp))
-> Maybe (Name -> InternaliseM (f SubExp))
forall a. a -> Maybe a
Just ((Name -> InternaliseM (f SubExp))
 -> Maybe (Name -> InternaliseM (f SubExp)))
-> (Name -> InternaliseM (f SubExp))
-> Maybe (Name -> InternaliseM (f SubExp))
forall a b. (a -> b) -> a -> b
$ \Name
desc -> do
          x' <- Name -> Exp -> InternaliseM SubExp
internaliseExp1 Name
"x" Exp
x
          y' <- internaliseExp1 "y" y
          fmap pure $ letSubExp desc $ I.BasicOp $ I.BinOp bop x' y'
      | Just CmpOp
cmp <- (CmpOp -> Bool) -> [CmpOp] -> Maybe CmpOp
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find ((Name -> Name -> Bool
forall a. Eq a => a -> a -> Bool
== Name
s) (Name -> Bool) -> (CmpOp -> Name) -> CmpOp -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> Name
nameFromText (Text -> Name) -> (CmpOp -> Text) -> CmpOp -> Name
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CmpOp -> Text
forall a. Pretty a => a -> Text
prettyText) [CmpOp]
allCmpOps = (Name -> InternaliseM (f SubExp))
-> Maybe (Name -> InternaliseM (f SubExp))
forall a. a -> Maybe a
Just ((Name -> InternaliseM (f SubExp))
 -> Maybe (Name -> InternaliseM (f SubExp)))
-> (Name -> InternaliseM (f SubExp))
-> Maybe (Name -> InternaliseM (f SubExp))
forall a b. (a -> b) -> a -> b
$ \Name
desc -> do
          x' <- Name -> Exp -> InternaliseM SubExp
internaliseExp1 Name
"x" Exp
x
          y' <- internaliseExp1 "y" y
          fmap pure $ letSubExp desc $ I.BasicOp $ I.CmpOp cmp x' y'
    handleOps [Exp
x] Name
s
      | Just ConvOp
conv <- (ConvOp -> Bool) -> [ConvOp] -> Maybe ConvOp
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find ((Name -> Name -> Bool
forall a. Eq a => a -> a -> Bool
== Name
s) (Name -> Bool) -> (ConvOp -> Name) -> ConvOp -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> Name
nameFromText (Text -> Name) -> (ConvOp -> Text) -> ConvOp -> Name
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ConvOp -> Text
forall a. Pretty a => a -> Text
prettyText) [ConvOp]
allConvOps = (Name -> InternaliseM (f SubExp))
-> Maybe (Name -> InternaliseM (f SubExp))
forall a. a -> Maybe a
Just ((Name -> InternaliseM (f SubExp))
 -> Maybe (Name -> InternaliseM (f SubExp)))
-> (Name -> InternaliseM (f SubExp))
-> Maybe (Name -> InternaliseM (f SubExp))
forall a b. (a -> b) -> a -> b
$ \Name
desc -> do
          x' <- Name -> Exp -> InternaliseM SubExp
internaliseExp1 Name
"x" Exp
x
          fmap pure $ letSubExp desc $ I.BasicOp $ I.ConvOp conv x'
    handleOps [Exp]
_ Name
_ = Maybe (Name -> InternaliseM (f SubExp))
forall a. Maybe a
Nothing

    handleSOACs :: [Exp] -> a -> Maybe (Name -> InternaliseM [SubExp])
handleSOACs [Exp
lam, Exp
arr] a
"map" = (Name -> InternaliseM [SubExp])
-> Maybe (Name -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((Name -> InternaliseM [SubExp])
 -> Maybe (Name -> InternaliseM [SubExp]))
-> (Name -> InternaliseM [SubExp])
-> Maybe (Name -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \Name
desc -> do
      arr' <- Name -> Exp -> InternaliseM [VName]
internaliseExpToVars Name
"map_arr" Exp
arr
      arr_ts <- mapM lookupType arr'
      lam' <- internaliseLambdaCoerce lam $ map rowType arr_ts
      let w = Int -> [TypeBase Shape NoMode] -> SubExp
forall o. Int -> [TypeBase Shape o] -> SubExp
arraysSize Int
0 [TypeBase Shape NoMode]
arr_ts
      letTupExp' desc . I.Op . I.Screma w arr' =<< I.mapSOAC lam'
    handleSOACs [Exp
lam, Exp
ne, Exp
arr] a
"reduce" = (Name -> InternaliseM [SubExp])
-> Maybe (Name -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((Name -> InternaliseM [SubExp])
 -> Maybe (Name -> InternaliseM [SubExp]))
-> (Name -> InternaliseM [SubExp])
-> Maybe (Name -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \Name
desc ->
      Name
-> Name
-> (SubExp
    -> Lambda SOACS
    -> [SubExp]
    -> [VName]
    -> InternaliseM (SOAC SOACS))
-> (Exp, Exp, Exp)
-> InternaliseM [SubExp]
internaliseScanOrReduce Name
desc Name
"reduce" SubExp
-> Lambda SOACS -> [SubExp] -> [VName] -> InternaliseM (SOAC SOACS)
forall {rep} {f :: * -> *}.
(Buildable rep, MonadFreshNames f) =>
SubExp -> Lambda rep -> [SubExp] -> [VName] -> f (SOAC rep)
reduce (Exp
lam, Exp
ne, Exp
arr)
      where
        reduce :: SubExp -> Lambda rep -> [SubExp] -> [VName] -> f (SOAC rep)
reduce SubExp
w Lambda rep
red_lam [SubExp]
nes [VName]
arrs =
          SubExp -> [VName] -> ScremaForm rep -> SOAC rep
forall rep. SubExp -> [VName] -> ScremaForm rep -> SOAC rep
I.Screma SubExp
w [VName]
arrs
            (ScremaForm rep -> SOAC rep) -> f (ScremaForm rep) -> f (SOAC rep)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Reduce rep] -> f (ScremaForm rep)
forall rep (m :: * -> *).
(Buildable rep, MonadFreshNames m) =>
[Reduce rep] -> m (ScremaForm rep)
I.reduceSOAC [Commutativity -> Lambda rep -> [SubExp] -> Reduce rep
forall rep. Commutativity -> Lambda rep -> [SubExp] -> Reduce rep
Reduce Commutativity
Noncommutative Lambda rep
red_lam [SubExp]
nes]
    handleSOACs [Exp
lam, Exp
ne, Exp
arr] a
"reduce_comm" = (Name -> InternaliseM [SubExp])
-> Maybe (Name -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((Name -> InternaliseM [SubExp])
 -> Maybe (Name -> InternaliseM [SubExp]))
-> (Name -> InternaliseM [SubExp])
-> Maybe (Name -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \Name
desc ->
      Name
-> Name
-> (SubExp
    -> Lambda SOACS
    -> [SubExp]
    -> [VName]
    -> InternaliseM (SOAC SOACS))
-> (Exp, Exp, Exp)
-> InternaliseM [SubExp]
internaliseScanOrReduce Name
desc Name
"reduce" SubExp
-> Lambda SOACS -> [SubExp] -> [VName] -> InternaliseM (SOAC SOACS)
forall {rep} {f :: * -> *}.
(Buildable rep, MonadFreshNames f) =>
SubExp -> Lambda rep -> [SubExp] -> [VName] -> f (SOAC rep)
reduce (Exp
lam, Exp
ne, Exp
arr)
      where
        reduce :: SubExp -> Lambda rep -> [SubExp] -> [VName] -> f (SOAC rep)
reduce SubExp
w Lambda rep
red_lam [SubExp]
nes [VName]
arrs =
          SubExp -> [VName] -> ScremaForm rep -> SOAC rep
forall rep. SubExp -> [VName] -> ScremaForm rep -> SOAC rep
I.Screma SubExp
w [VName]
arrs
            (ScremaForm rep -> SOAC rep) -> f (ScremaForm rep) -> f (SOAC rep)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Reduce rep] -> f (ScremaForm rep)
forall rep (m :: * -> *).
(Buildable rep, MonadFreshNames m) =>
[Reduce rep] -> m (ScremaForm rep)
I.reduceSOAC [Commutativity -> Lambda rep -> [SubExp] -> Reduce rep
forall rep. Commutativity -> Lambda rep -> [SubExp] -> Reduce rep
Reduce Commutativity
Commutative Lambda rep
red_lam [SubExp]
nes]
    handleSOACs [Exp
lam, Exp
ne, Exp
arr] a
"scan" = (Name -> InternaliseM [SubExp])
-> Maybe (Name -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((Name -> InternaliseM [SubExp])
 -> Maybe (Name -> InternaliseM [SubExp]))
-> (Name -> InternaliseM [SubExp])
-> Maybe (Name -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \Name
desc ->
      Name
-> Name
-> (SubExp
    -> Lambda SOACS
    -> [SubExp]
    -> [VName]
    -> InternaliseM (SOAC SOACS))
-> (Exp, Exp, Exp)
-> InternaliseM [SubExp]
internaliseScanOrReduce Name
desc Name
"scan" SubExp
-> Lambda SOACS -> [SubExp] -> [VName] -> InternaliseM (SOAC SOACS)
forall {rep} {f :: * -> *}.
(Buildable rep, MonadFreshNames f) =>
SubExp -> Lambda rep -> [SubExp] -> [VName] -> f (SOAC rep)
reduce (Exp
lam, Exp
ne, Exp
arr)
      where
        reduce :: SubExp -> Lambda rep -> [SubExp] -> [VName] -> f (SOAC rep)
reduce SubExp
w Lambda rep
scan_lam [SubExp]
nes [VName]
arrs =
          SubExp -> [VName] -> ScremaForm rep -> SOAC rep
forall rep. SubExp -> [VName] -> ScremaForm rep -> SOAC rep
I.Screma SubExp
w [VName]
arrs (ScremaForm rep -> SOAC rep) -> f (ScremaForm rep) -> f (SOAC rep)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Scan rep] -> f (ScremaForm rep)
forall rep (m :: * -> *).
(Buildable rep, MonadFreshNames m) =>
[Scan rep] -> m (ScremaForm rep)
I.scanSOAC [Lambda rep -> [SubExp] -> Scan rep
forall rep. Lambda rep -> [SubExp] -> Scan rep
Scan Lambda rep
scan_lam [SubExp]
nes]
    handleSOACs [Exp
rf, Exp
dest, Exp
op, Exp
ne, Exp
buckets, Exp
img] a
"hist_1d" = (Name -> InternaliseM [SubExp])
-> Maybe (Name -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((Name -> InternaliseM [SubExp])
 -> Maybe (Name -> InternaliseM [SubExp]))
-> (Name -> InternaliseM [SubExp])
-> Maybe (Name -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \Name
desc ->
      Int
-> Name
-> Exp
-> Exp
-> Exp
-> Exp
-> Exp
-> Exp
-> InternaliseM [SubExp]
internaliseHist Int
1 Name
desc Exp
rf Exp
dest Exp
op Exp
ne Exp
buckets Exp
img
    handleSOACs [Exp
rf, Exp
dest, Exp
op, Exp
ne, Exp
buckets, Exp
img] a
"hist_2d" = (Name -> InternaliseM [SubExp])
-> Maybe (Name -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((Name -> InternaliseM [SubExp])
 -> Maybe (Name -> InternaliseM [SubExp]))
-> (Name -> InternaliseM [SubExp])
-> Maybe (Name -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \Name
desc ->
      Int
-> Name
-> Exp
-> Exp
-> Exp
-> Exp
-> Exp
-> Exp
-> InternaliseM [SubExp]
internaliseHist Int
2 Name
desc Exp
rf Exp
dest Exp
op Exp
ne Exp
buckets Exp
img
    handleSOACs [Exp
rf, Exp
dest, Exp
op, Exp
ne, Exp
buckets, Exp
img] a
"hist_3d" = (Name -> InternaliseM [SubExp])
-> Maybe (Name -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((Name -> InternaliseM [SubExp])
 -> Maybe (Name -> InternaliseM [SubExp]))
-> (Name -> InternaliseM [SubExp])
-> Maybe (Name -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \Name
desc ->
      Int
-> Name
-> Exp
-> Exp
-> Exp
-> Exp
-> Exp
-> Exp
-> InternaliseM [SubExp]
internaliseHist Int
3 Name
desc Exp
rf Exp
dest Exp
op Exp
ne Exp
buckets Exp
img
    handleSOACs [Exp
lam, Exp
arr] a
"flatmap" = (Name -> InternaliseM [SubExp])
-> Maybe (Name -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((Name -> InternaliseM [SubExp])
 -> Maybe (Name -> InternaliseM [SubExp]))
-> (Name -> InternaliseM [SubExp])
-> Maybe (Name -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \Name
desc -> do
      arr' <- Name -> Exp -> InternaliseM [VName]
internaliseExpToVars Name
"map_arr" Exp
arr
      arr_ts <- mapM lookupType arr'
      lam' <- internaliseFlatLambda lam $ map rowType arr_ts
      let w = Int -> [TypeBase Shape NoMode] -> SubExp
forall o. Int -> [TypeBase Shape o] -> SubExp
arraysSize Int
0 [TypeBase Shape NoMode]
arr_ts
      letValExp' desc $ I.Op $ FlatMap w arr' lam'
    handleSOACs [Exp]
_ a
_ = Maybe (Name -> InternaliseM [SubExp])
forall a. Maybe a
Nothing

    handleAccs :: [Exp] -> a -> Maybe (Name -> InternaliseM [SubExp])
handleAccs [Exp
dest, Exp
f, Exp
bs] a
"scatter_stream" = (Name -> InternaliseM [SubExp])
-> Maybe (Name -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((Name -> InternaliseM [SubExp])
 -> Maybe (Name -> InternaliseM [SubExp]))
-> (Name -> InternaliseM [SubExp])
-> Maybe (Name -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \Name
desc ->
      Name
-> Exp -> Maybe (Exp, Exp) -> Exp -> Exp -> InternaliseM [SubExp]
internaliseStreamAcc Name
desc Exp
dest Maybe (Exp, Exp)
forall a. Maybe a
Nothing Exp
f Exp
bs
    handleAccs [Exp
dest, Exp
op, Exp
ne, Exp
f, Exp
bs] a
"hist_stream" = (Name -> InternaliseM [SubExp])
-> Maybe (Name -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((Name -> InternaliseM [SubExp])
 -> Maybe (Name -> InternaliseM [SubExp]))
-> (Name -> InternaliseM [SubExp])
-> Maybe (Name -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \Name
desc ->
      Name
-> Exp -> Maybe (Exp, Exp) -> Exp -> Exp -> InternaliseM [SubExp]
internaliseStreamAcc Name
desc Exp
dest ((Exp, Exp) -> Maybe (Exp, Exp)
forall a. a -> Maybe a
Just (Exp
op, Exp
ne)) Exp
f Exp
bs
    handleAccs [Exp
acc, Exp
i, Exp
v] a
"acc_write" = (Name -> InternaliseM [SubExp])
-> Maybe (Name -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((Name -> InternaliseM [SubExp])
 -> Maybe (Name -> InternaliseM [SubExp]))
-> (Name -> InternaliseM [SubExp])
-> Maybe (Name -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \Name
desc -> do
      acc' <- [VName] -> VName
forall a. HasCallStack => [a] -> a
head ([VName] -> VName) -> InternaliseM [VName] -> InternaliseM VName
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Name -> Exp -> InternaliseM [VName]
internaliseExpToVars Name
"acc" Exp
acc
      i' <- internaliseExp1 "acc_i" i
      vs <- internaliseExp "acc_v" v
      fmap pure $ letSubExp desc $ BasicOp $ UpdateAcc Safe acc' [i'] vs
    handleAccs [Exp]
_ a
_ = Maybe (Name -> InternaliseM [SubExp])
forall a. Maybe a
Nothing

    handleAD :: [Exp] -> a -> Maybe (Name -> InternaliseM [SubExp])
handleAD [Exp
f, Exp
x, Exp
v] a
fname
      | a
fname a -> [a] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [a
"jvp2", a
"vjp2", a
"jmp2", a
"mjp2"] = (Name -> InternaliseM [SubExp])
-> Maybe (Name -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((Name -> InternaliseM [SubExp])
 -> Maybe (Name -> InternaliseM [SubExp]))
-> (Name -> InternaliseM [SubExp])
-> Maybe (Name -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \Name
desc -> do
          x' <- Name -> Exp -> InternaliseM [SubExp]
internaliseExp Name
"ad_x" Exp
x
          v' <- internaliseExp "ad_v" v
          x_t <- subExpType $ head x'
          v_t <- subExpType $ head v'
          lam <- internaliseLambdaCoerce f =<< mapM subExpType x'
          fmap (map I.Var) . letTupExp desc . Op $
            case fname of
              a
"jvp2" -> Shape -> [SubExp] -> [SubExp] -> Lambda SOACS -> SOAC SOACS
forall rep. Shape -> [SubExp] -> [SubExp] -> Lambda rep -> SOAC rep
JVP Shape
forall a. Monoid a => a
mempty [SubExp]
x' [SubExp]
v' Lambda SOACS
lam
              a
"vjp2" -> Shape -> [SubExp] -> [SubExp] -> Lambda SOACS -> SOAC SOACS
forall rep. Shape -> [SubExp] -> [SubExp] -> Lambda rep -> SOAC rep
VJP Shape
forall a. Monoid a => a
mempty [SubExp]
x' [SubExp]
v' Lambda SOACS
lam
              a
"jmp2" ->
                Shape -> [SubExp] -> [SubExp] -> Lambda SOACS -> SOAC SOACS
forall rep. Shape -> [SubExp] -> [SubExp] -> Lambda rep -> SOAC rep
JVP (TypeBase Shape NoMode -> TypeBase Shape NoMode -> Shape
forall {shape} {o} {d} {m}.
ArrayShape shape =>
TypeBase shape o -> TypeBase (ShapeBase d) m -> ShapeBase d
vecShape TypeBase Shape NoMode
x_t TypeBase Shape NoMode
v_t) [SubExp]
x' [SubExp]
v' Lambda SOACS
lam
              a
"mjp2" ->
                Shape -> [SubExp] -> [SubExp] -> Lambda SOACS -> SOAC SOACS
forall rep. Shape -> [SubExp] -> [SubExp] -> Lambda rep -> SOAC rep
VJP (TypeBase Shape NoMode -> TypeBase Shape NoMode -> Shape
forall {shape} {o} {d} {m}.
ArrayShape shape =>
TypeBase shape o -> TypeBase (ShapeBase d) m -> ShapeBase d
vecShape ([TypeBase Shape NoMode] -> TypeBase Shape NoMode
forall a. HasCallStack => [a] -> a
head (Lambda SOACS -> [TypeBase Shape NoMode]
forall rep t. GLambda rep t -> [t]
lambdaReturnType Lambda SOACS
lam)) TypeBase Shape NoMode
v_t) [SubExp]
x' [SubExp]
v' Lambda SOACS
lam
              a
_ -> [Char] -> SOAC SOACS
forall a. HasCallStack => [Char] -> a
error [Char]
"handleAD: not supposed to happen."
      where
        vecShape :: TypeBase shape o -> TypeBase (ShapeBase d) m -> ShapeBase d
vecShape TypeBase shape o
t1 TypeBase (ShapeBase d) m
t2 =
          [d] -> ShapeBase d
forall d. [d] -> ShapeBase d
I.Shape ([d] -> ShapeBase d) -> [d] -> ShapeBase d
forall a b. (a -> b) -> a -> b
$ Int -> [d] -> [d]
forall a. Int -> [a] -> [a]
take (TypeBase (ShapeBase d) m -> Int
forall shape o. ArrayShape shape => TypeBase shape o -> Int
I.arrayRank TypeBase (ShapeBase d) m
t2 Int -> Int -> Int
forall a. Num a => a -> a -> a
- TypeBase shape o -> Int
forall shape o. ArrayShape shape => TypeBase shape o -> Int
I.arrayRank TypeBase shape o
t1) (TypeBase (ShapeBase d) m -> [d]
forall d m. TypeBase (ShapeBase d) m -> [d]
I.arrayDims TypeBase (ShapeBase d) m
t2)
    handleAD [Exp
f, Exp
f_adj, Exp
x] a
"with_vjp" = (Name -> InternaliseM [SubExp])
-> Maybe (Name -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((Name -> InternaliseM [SubExp])
 -> Maybe (Name -> InternaliseM [SubExp]))
-> (Name -> InternaliseM [SubExp])
-> Maybe (Name -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \Name
desc -> do
      x' <- Name -> Exp -> InternaliseM [SubExp]
internaliseExp Name
"ad_x" Exp
x
      lam <- internaliseLambdaCoerce f =<< mapM subExpType x'
      lam_adj <-
        internaliseLambdaCoerce f_adj $
          lambdaReturnType lam ++ lambdaReturnType lam
      fmap (map I.Var) . letTupExp desc . Op $ WithVJP x' lam lam_adj
    handleAD [Exp]
_ a
_ = Maybe (Name -> InternaliseM [SubExp])
forall a. Maybe a
Nothing

    handleRest :: [Exp] -> a -> Maybe (Name -> InternaliseM [SubExp])
handleRest [Exp
a, Exp
si, Exp
v] a
"scatter" = (Name -> InternaliseM [SubExp])
-> Maybe (Name -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((Name -> InternaliseM [SubExp])
 -> Maybe (Name -> InternaliseM [SubExp]))
-> (Name -> InternaliseM [SubExp])
-> Maybe (Name -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ Int -> Exp -> Exp -> Exp -> Name -> InternaliseM [SubExp]
scatterF Int
1 Exp
a Exp
si Exp
v
    handleRest [Exp
a, Exp
si, Exp
v] a
"scatter_2d" = (Name -> InternaliseM [SubExp])
-> Maybe (Name -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((Name -> InternaliseM [SubExp])
 -> Maybe (Name -> InternaliseM [SubExp]))
-> (Name -> InternaliseM [SubExp])
-> Maybe (Name -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ Int -> Exp -> Exp -> Exp -> Name -> InternaliseM [SubExp]
scatterF Int
2 Exp
a Exp
si Exp
v
    handleRest [Exp
a, Exp
si, Exp
v] a
"scatter_3d" = (Name -> InternaliseM [SubExp])
-> Maybe (Name -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((Name -> InternaliseM [SubExp])
 -> Maybe (Name -> InternaliseM [SubExp]))
-> (Name -> InternaliseM [SubExp])
-> Maybe (Name -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ Int -> Exp -> Exp -> Exp -> Name -> InternaliseM [SubExp]
scatterF Int
3 Exp
a Exp
si Exp
v
    handleRest [Exp
n, Exp
m, Exp
arr] a
"unflatten" = (Name -> InternaliseM [SubExp])
-> Maybe (Name -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((Name -> InternaliseM [SubExp])
 -> Maybe (Name -> InternaliseM [SubExp]))
-> (Name -> InternaliseM [SubExp])
-> Maybe (Name -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \Name
desc -> do
      arrs <- Name -> Exp -> InternaliseM [VName]
internaliseExpToVars Name
"unflatten_arr" Exp
arr
      n' <- internaliseExp1 "n" n
      m' <- internaliseExp1 "m" m
      -- Each dimension must be nonnegative, and the unflattened
      -- dimension needs to have the same number of elements as the
      -- original dimension.
      old_dim <- I.arraysSize 0 <$> mapM lookupType arrs
      dim_ok <-
        letSubExp "dim_ok" <=< toExp $
          pe64 old_dim .==. pe64 n'
            * pe64 m'
              .&&. pe64 n'
              .>=. 0
              .&&. pe64 m'
              .>=. 0
      dim_ok_cert <-
        assert "dim_ok_cert" dim_ok $
          ErrorMsg
            [ "Cannot unflatten array of shape [",
              ErrorVal int64 old_dim,
              "] to array of shape [",
              ErrorVal int64 n',
              "][",
              ErrorVal int64 m',
              "]"
            ]
      certifying dim_ok_cert $
        forM arrs $ \VName
arr' -> do
          arr_t <- VName -> InternaliseM (TypeBase Shape NoMode)
forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoMode)
lookupType VName
arr'
          letSubExp desc . I.BasicOp $
            I.Reshape arr' $
              reshapeAll (I.arrayShape arr_t) $
                reshapeOuter (I.Shape [n', m']) 1 $
                  I.arrayShape arr_t
    handleRest [Exp
arr] a
"manifest" = (Name -> InternaliseM [SubExp])
-> Maybe (Name -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((Name -> InternaliseM [SubExp])
 -> Maybe (Name -> InternaliseM [SubExp]))
-> (Name -> InternaliseM [SubExp])
-> Maybe (Name -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \Name
desc -> do
      arrs <- Name -> Exp -> InternaliseM [VName]
internaliseExpToVars Name
"flatten_arr" Exp
arr
      forM arrs $ \VName
arr' -> do
        r <- TypeBase Shape NoMode -> Int
forall shape o. ArrayShape shape => TypeBase shape o -> Int
I.arrayRank (TypeBase Shape NoMode -> Int)
-> InternaliseM (TypeBase Shape NoMode) -> InternaliseM Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> InternaliseM (TypeBase Shape NoMode)
forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoMode)
lookupType VName
arr'
        if r == 0
          then pure $ I.Var arr'
          else letSubExp desc $ I.BasicOp $ I.Manifest arr' [0 .. r - 1]
    handleRest [Exp
arr] a
"flatten" = (Name -> InternaliseM [SubExp])
-> Maybe (Name -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((Name -> InternaliseM [SubExp])
 -> Maybe (Name -> InternaliseM [SubExp]))
-> (Name -> InternaliseM [SubExp])
-> Maybe (Name -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \Name
desc -> do
      arrs <- Name -> Exp -> InternaliseM [VName]
internaliseExpToVars Name
"flatten_arr" Exp
arr
      forM arrs $ \VName
arr' -> do
        arr_t <- VName -> InternaliseM (TypeBase Shape NoMode)
forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoMode)
lookupType VName
arr'
        let n = Int -> TypeBase Shape NoMode -> SubExp
forall o. Int -> TypeBase Shape o -> SubExp
arraySize Int
0 TypeBase Shape NoMode
arr_t
            m = Int -> TypeBase Shape NoMode -> SubExp
forall o. Int -> TypeBase Shape o -> SubExp
arraySize Int
1 TypeBase Shape NoMode
arr_t
        k <- letSubExp "flat_dim" $ I.BasicOp $ I.BinOp (Mul Int64 I.OverflowUndef) n m
        letSubExp desc . I.BasicOp $
          I.Reshape arr' $
            reshapeAll (I.arrayShape arr_t) $
              reshapeOuter (I.Shape [k]) 2 $
                I.arrayShape arr_t
    handleRest [Exp
x, Exp
y] a
"concat" = (Name -> InternaliseM [SubExp])
-> Maybe (Name -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((Name -> InternaliseM [SubExp])
 -> Maybe (Name -> InternaliseM [SubExp]))
-> (Name -> InternaliseM [SubExp])
-> Maybe (Name -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \Name
desc -> do
      xs <- Name -> Exp -> InternaliseM [VName]
internaliseExpToVars Name
"concat_x" Exp
x
      ys <- internaliseExpToVars "concat_y" y
      outer_size <- arraysSize 0 <$> mapM lookupType xs
      let sumdims SubExp
xsize SubExp
ysize =
            Name -> Exp (Rep m) -> m SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"conc_tmp" (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
I.BasicOp (BasicOp -> Exp (Rep m)) -> BasicOp -> Exp (Rep m)
forall a b. (a -> b) -> a -> b
$
                BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
I.Add IntType
I.Int64 Overflow
I.OverflowUndef) SubExp
xsize SubExp
ysize
      ressize <-
        foldM sumdims outer_size
          =<< mapM (fmap (arraysSize 0) . mapM lookupType) [ys]

      let conc VName
xarr VName
yarr =
            BasicOp -> Exp SOACS
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp SOACS) -> BasicOp -> Exp SOACS
forall a b. (a -> b) -> a -> b
$ Int -> NonEmpty VName -> SubExp -> BasicOp
I.Concat Int
0 (VName
xarr VName -> [VName] -> NonEmpty VName
forall a. a -> [a] -> NonEmpty a
:| [VName
yarr]) SubExp
ressize
      mapM (letSubExp desc) $ zipWith conc xs ys
    handleRest [Exp
e] a
"transpose" = (Name -> InternaliseM [SubExp])
-> Maybe (Name -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((Name -> InternaliseM [SubExp])
 -> Maybe (Name -> InternaliseM [SubExp]))
-> (Name -> InternaliseM [SubExp])
-> Maybe (Name -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \Name
desc ->
      Name
-> Exp -> (VName -> InternaliseM BasicOp) -> InternaliseM [SubExp]
internaliseOperation Name
desc Exp
e ((VName -> InternaliseM BasicOp) -> InternaliseM [SubExp])
-> (VName -> InternaliseM BasicOp) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ \VName
v -> do
        r <- TypeBase Shape NoMode -> Int
forall shape o. ArrayShape shape => TypeBase shape o -> Int
I.arrayRank (TypeBase Shape NoMode -> Int)
-> InternaliseM (TypeBase Shape NoMode) -> InternaliseM Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> InternaliseM (TypeBase Shape NoMode)
forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (TypeBase Shape NoMode)
lookupType VName
v
        pure $ I.Rearrange v ([1, 0] ++ [2 .. r - 1])
    handleRest [Exp
x, Exp
y] a
"zip" = (Name -> InternaliseM [SubExp])
-> Maybe (Name -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((Name -> InternaliseM [SubExp])
 -> Maybe (Name -> InternaliseM [SubExp]))
-> (Name -> InternaliseM [SubExp])
-> Maybe (Name -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \Name
desc ->
      (VName -> InternaliseM SubExp) -> [VName] -> InternaliseM [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 InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"zip_copy" (Exp SOACS -> InternaliseM SubExp)
-> (VName -> Exp SOACS) -> VName -> InternaliseM SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BasicOp -> Exp SOACS
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp SOACS) -> (VName -> BasicOp) -> VName -> Exp SOACS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Shape -> SubExp -> BasicOp
Replicate Shape
forall a. Monoid a => a
mempty (SubExp -> BasicOp) -> (VName -> SubExp) -> VName -> BasicOp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> SubExp
I.Var)
        ([VName] -> InternaliseM [SubExp])
-> InternaliseM [VName] -> InternaliseM [SubExp]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< ( [VName] -> [VName] -> [VName]
forall a. [a] -> [a] -> [a]
(++)
                ([VName] -> [VName] -> [VName])
-> InternaliseM [VName] -> InternaliseM ([VName] -> [VName])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Name -> Exp -> InternaliseM [VName]
internaliseExpToVars (Name
desc Name -> Name -> Name
forall a. Semigroup a => a -> a -> a
<> Name
"_zip_x") Exp
x
                InternaliseM ([VName] -> [VName])
-> InternaliseM [VName] -> InternaliseM [VName]
forall a b.
InternaliseM (a -> b) -> InternaliseM a -> InternaliseM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Name -> Exp -> InternaliseM [VName]
internaliseExpToVars (Name
desc Name -> Name -> Name
forall a. Semigroup a => a -> a -> a
<> Name
"_zip_y") Exp
y
            )
    handleRest [Exp
x] a
"unzip" = (Name -> InternaliseM [SubExp])
-> Maybe (Name -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((Name -> InternaliseM [SubExp])
 -> Maybe (Name -> InternaliseM [SubExp]))
-> (Name -> InternaliseM [SubExp])
-> Maybe (Name -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \Name
desc ->
      (VName -> InternaliseM SubExp) -> [VName] -> InternaliseM [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 InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
desc (Exp SOACS -> InternaliseM SubExp)
-> (VName -> Exp SOACS) -> VName -> InternaliseM SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BasicOp -> Exp SOACS
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp SOACS) -> (VName -> BasicOp) -> VName -> Exp SOACS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Shape -> SubExp -> BasicOp
Replicate Shape
forall a. Monoid a => a
mempty (SubExp -> BasicOp) -> (VName -> SubExp) -> VName -> BasicOp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> SubExp
I.Var)
        ([VName] -> InternaliseM [SubExp])
-> InternaliseM [VName] -> InternaliseM [SubExp]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Name -> Exp -> InternaliseM [VName]
internaliseExpToVars Name
desc Exp
x
    handleRest [Exp
arr, Exp
offset, Exp
n1, Exp
s1, Exp
n2, Exp
s2] a
"flat_index_2d" = (Name -> InternaliseM [SubExp])
-> Maybe (Name -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((Name -> InternaliseM [SubExp])
 -> Maybe (Name -> InternaliseM [SubExp]))
-> (Name -> InternaliseM [SubExp])
-> Maybe (Name -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \Name
desc -> do
      Name -> Exp -> Exp -> [(Exp, Exp)] -> InternaliseM [SubExp]
flatIndexHelper Name
desc Exp
arr Exp
offset [(Exp
n1, Exp
s1), (Exp
n2, Exp
s2)]
    handleRest [Exp
arr1, Exp
offset, Exp
s1, Exp
s2, Exp
arr2] a
"flat_update_2d" = (Name -> InternaliseM [SubExp])
-> Maybe (Name -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((Name -> InternaliseM [SubExp])
 -> Maybe (Name -> InternaliseM [SubExp]))
-> (Name -> InternaliseM [SubExp])
-> Maybe (Name -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \Name
desc -> do
      Name -> Exp -> Exp -> [Exp] -> Exp -> InternaliseM [SubExp]
flatUpdateHelper Name
desc Exp
arr1 Exp
offset [Exp
s1, Exp
s2] Exp
arr2
    handleRest [Exp
arr, Exp
offset, Exp
n1, Exp
s1, Exp
n2, Exp
s2, Exp
n3, Exp
s3] a
"flat_index_3d" = (Name -> InternaliseM [SubExp])
-> Maybe (Name -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((Name -> InternaliseM [SubExp])
 -> Maybe (Name -> InternaliseM [SubExp]))
-> (Name -> InternaliseM [SubExp])
-> Maybe (Name -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \Name
desc -> do
      Name -> Exp -> Exp -> [(Exp, Exp)] -> InternaliseM [SubExp]
flatIndexHelper Name
desc Exp
arr Exp
offset [(Exp
n1, Exp
s1), (Exp
n2, Exp
s2), (Exp
n3, Exp
s3)]
    handleRest [Exp
arr1, Exp
offset, Exp
s1, Exp
s2, Exp
s3, Exp
arr2] a
"flat_update_3d" = (Name -> InternaliseM [SubExp])
-> Maybe (Name -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((Name -> InternaliseM [SubExp])
 -> Maybe (Name -> InternaliseM [SubExp]))
-> (Name -> InternaliseM [SubExp])
-> Maybe (Name -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \Name
desc -> do
      Name -> Exp -> Exp -> [Exp] -> Exp -> InternaliseM [SubExp]
flatUpdateHelper Name
desc Exp
arr1 Exp
offset [Exp
s1, Exp
s2, Exp
s3] Exp
arr2
    handleRest [Exp
arr, Exp
offset, Exp
n1, Exp
s1, Exp
n2, Exp
s2, Exp
n3, Exp
s3, Exp
n4, Exp
s4] a
"flat_index_4d" = (Name -> InternaliseM [SubExp])
-> Maybe (Name -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((Name -> InternaliseM [SubExp])
 -> Maybe (Name -> InternaliseM [SubExp]))
-> (Name -> InternaliseM [SubExp])
-> Maybe (Name -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \Name
desc -> do
      Name -> Exp -> Exp -> [(Exp, Exp)] -> InternaliseM [SubExp]
flatIndexHelper Name
desc Exp
arr Exp
offset [(Exp
n1, Exp
s1), (Exp
n2, Exp
s2), (Exp
n3, Exp
s3), (Exp
n4, Exp
s4)]
    handleRest [Exp
arr1, Exp
offset, Exp
s1, Exp
s2, Exp
s3, Exp
s4, Exp
arr2] a
"flat_update_4d" = (Name -> InternaliseM [SubExp])
-> Maybe (Name -> InternaliseM [SubExp])
forall a. a -> Maybe a
Just ((Name -> InternaliseM [SubExp])
 -> Maybe (Name -> InternaliseM [SubExp]))
-> (Name -> InternaliseM [SubExp])
-> Maybe (Name -> InternaliseM [SubExp])
forall a b. (a -> b) -> a -> b
$ \Name
desc -> do
      Name -> Exp -> Exp -> [Exp] -> Exp -> InternaliseM [SubExp]
flatUpdateHelper Name
desc Exp
arr1 Exp
offset [Exp
s1, Exp
s2, Exp
s3, Exp
s4] Exp
arr2
    handleRest [Exp]
_ a
_ = Maybe (Name -> InternaliseM [SubExp])
forall a. Maybe a
Nothing

    handleIO :: p -> Name -> Maybe (t -> t)
handleIO p
_ Name
s
      | Just Text
f <- Text -> Text -> Maybe Text
T.stripPrefix Text
"io_" (Name -> Text
nameToText Name
s) = (t -> t) -> Maybe (t -> t)
forall a. a -> Maybe a
Just ((t -> t) -> Maybe (t -> t)) -> (t -> t) -> Maybe (t -> t)
forall a b. (a -> b) -> a -> b
$ \t
_ ->
          Text -> t
forall a. Text -> a
compilerLimitation (Text -> t) -> Text -> t
forall a b. (a -> b) -> a -> b
$ Text
"io." Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
f Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" is only supported in interpreted code."
    handleIO p
_ Name
_ = Maybe (t -> t)
forall a. Maybe a
Nothing

    toSigned :: IntType -> Exp -> Name -> InternaliseM [SubExp]
toSigned IntType
int_to Exp
e Name
desc = do
      e' <- Name -> Exp -> InternaliseM SubExp
internaliseExp1 Name
"trunc_arg" Exp
e
      case E.typeOf e of
        E.Scalar (E.Prim PrimType
E.Bool) ->
          Name -> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m [SubExp]
letTupExp' Name
desc
            (Exp SOACS -> InternaliseM [SubExp])
-> InternaliseM (Exp SOACS) -> InternaliseM [SubExp]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< InternaliseM (Exp (Rep InternaliseM))
-> InternaliseM (GBody (Rep InternaliseM) SubExpRes)
-> InternaliseM (GBody (Rep InternaliseM) SubExpRes)
-> InternaliseM (Exp (Rep InternaliseM))
forall (m :: * -> *).
(MonadBuilder m, BranchType (Rep m) ~ TypeBase ExtShape NoMode) =>
m (Exp (Rep m))
-> m (Body (Rep m)) -> m (Body (Rep m)) -> m (Exp (Rep m))
eIf
              (SubExp -> InternaliseM (Exp (Rep InternaliseM))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
e')
              ([SubExp] -> InternaliseM (GBody (Rep InternaliseM) SubExpRes)
forall (m :: * -> *).
MonadBuilder m =>
[SubExp] -> m (Body (Rep m))
resultBodyM [IntType -> Integer -> SubExp
intConst IntType
int_to Integer
1])
              ([SubExp] -> InternaliseM (GBody (Rep InternaliseM) SubExpRes)
forall (m :: * -> *).
MonadBuilder m =>
[SubExp] -> m (Body (Rep m))
resultBodyM [IntType -> Integer -> SubExp
intConst IntType
int_to Integer
0])
        E.Scalar (E.Prim (E.Signed IntType
int_from)) ->
          Name -> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m [SubExp]
letTupExp' Name
desc (Exp (Rep InternaliseM) -> InternaliseM [SubExp])
-> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
forall a b. (a -> b) -> a -> b
$ ConvOp -> SubExp -> BasicOp
I.ConvOp (IntType -> IntType -> ConvOp
I.SExt IntType
int_from IntType
int_to) SubExp
e'
        E.Scalar (E.Prim (E.Unsigned IntType
int_from)) ->
          Name -> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m [SubExp]
letTupExp' Name
desc (Exp (Rep InternaliseM) -> InternaliseM [SubExp])
-> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
forall a b. (a -> b) -> a -> b
$ ConvOp -> SubExp -> BasicOp
I.ConvOp (IntType -> IntType -> ConvOp
I.ZExt IntType
int_from IntType
int_to) SubExp
e'
        E.Scalar (E.Prim (E.FloatType FloatType
float_from)) ->
          Name -> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m [SubExp]
letTupExp' Name
desc (Exp (Rep InternaliseM) -> InternaliseM [SubExp])
-> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
forall a b. (a -> b) -> a -> b
$ ConvOp -> SubExp -> BasicOp
I.ConvOp (FloatType -> IntType -> ConvOp
I.FPToSI FloatType
float_from IntType
int_to) SubExp
e'
        StructType
_ -> [Char] -> InternaliseM [SubExp]
forall a. HasCallStack => [Char] -> a
error [Char]
"Futhark.Internalise: non-numeric type in ToSigned"

    toUnsigned :: IntType -> Exp -> Name -> InternaliseM [SubExp]
toUnsigned IntType
int_to Exp
e Name
desc = do
      e' <- Name -> Exp -> InternaliseM SubExp
internaliseExp1 Name
"trunc_arg" Exp
e
      case E.typeOf e of
        E.Scalar (E.Prim PrimType
E.Bool) ->
          Name -> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m [SubExp]
letTupExp' Name
desc
            (Exp SOACS -> InternaliseM [SubExp])
-> InternaliseM (Exp SOACS) -> InternaliseM [SubExp]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< InternaliseM (Exp (Rep InternaliseM))
-> InternaliseM (GBody (Rep InternaliseM) SubExpRes)
-> InternaliseM (GBody (Rep InternaliseM) SubExpRes)
-> InternaliseM (Exp (Rep InternaliseM))
forall (m :: * -> *).
(MonadBuilder m, BranchType (Rep m) ~ TypeBase ExtShape NoMode) =>
m (Exp (Rep m))
-> m (Body (Rep m)) -> m (Body (Rep m)) -> m (Exp (Rep m))
eIf
              (SubExp -> InternaliseM (Exp (Rep InternaliseM))
forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
eSubExp SubExp
e')
              ([SubExp] -> InternaliseM (GBody (Rep InternaliseM) SubExpRes)
forall (m :: * -> *).
MonadBuilder m =>
[SubExp] -> m (Body (Rep m))
resultBodyM [IntType -> Integer -> SubExp
intConst IntType
int_to Integer
1])
              ([SubExp] -> InternaliseM (GBody (Rep InternaliseM) SubExpRes)
forall (m :: * -> *).
MonadBuilder m =>
[SubExp] -> m (Body (Rep m))
resultBodyM [IntType -> Integer -> SubExp
intConst IntType
int_to Integer
0])
        E.Scalar (E.Prim (E.Signed IntType
int_from)) ->
          Name -> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m [SubExp]
letTupExp' Name
desc (Exp (Rep InternaliseM) -> InternaliseM [SubExp])
-> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
forall a b. (a -> b) -> a -> b
$ ConvOp -> SubExp -> BasicOp
I.ConvOp (IntType -> IntType -> ConvOp
I.ZExt IntType
int_from IntType
int_to) SubExp
e'
        E.Scalar (E.Prim (E.Unsigned IntType
int_from)) ->
          Name -> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m [SubExp]
letTupExp' Name
desc (Exp (Rep InternaliseM) -> InternaliseM [SubExp])
-> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
forall a b. (a -> b) -> a -> b
$ ConvOp -> SubExp -> BasicOp
I.ConvOp (IntType -> IntType -> ConvOp
I.ZExt IntType
int_from IntType
int_to) SubExp
e'
        E.Scalar (E.Prim (E.FloatType FloatType
float_from)) ->
          Name -> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m [SubExp]
letTupExp' Name
desc (Exp (Rep InternaliseM) -> InternaliseM [SubExp])
-> Exp (Rep InternaliseM) -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
forall a b. (a -> b) -> a -> b
$ ConvOp -> SubExp -> BasicOp
I.ConvOp (FloatType -> IntType -> ConvOp
I.FPToUI FloatType
float_from IntType
int_to) SubExp
e'
        StructType
_ -> [Char] -> InternaliseM [SubExp]
forall a. HasCallStack => [Char] -> a
error [Char]
"Futhark.Internalise.internaliseExp: non-numeric type in ToUnsigned"

flatIndexHelper :: Name -> E.Exp -> E.Exp -> [(E.Exp, E.Exp)] -> InternaliseM [SubExp]
flatIndexHelper :: Name -> Exp -> Exp -> [(Exp, Exp)] -> InternaliseM [SubExp]
flatIndexHelper Name
desc Exp
arr Exp
offset [(Exp, Exp)]
slices = do
  arrs <- Name -> Exp -> InternaliseM [VName]
internaliseExpToVars Name
"arr" Exp
arr
  offset' <- internaliseExp1 "offset" offset
  old_dim <- I.arraysSize 0 <$> mapM lookupType arrs
  offset_inbounds_down <- letSubExp "offset_inbounds_down" $ I.BasicOp $ I.CmpOp (I.CmpUle Int64) (intConst Int64 0) offset'
  offset_inbounds_up <- letSubExp "offset_inbounds_up" $ I.BasicOp $ I.CmpOp (I.CmpUlt Int64) offset' old_dim
  slices' <-
    mapM
      ( \(Exp
n, Exp
s) -> do
          n' <- Name -> Exp -> InternaliseM SubExp
internaliseExp1 Name
"n" Exp
n
          s' <- internaliseExp1 "s" s
          pure (n', s')
      )
      slices
  (min_bound, max_bound) <-
    foldM
      ( \(SubExp
lower, SubExp
upper) (SubExp
n, SubExp
s) -> do
          n_m1 <- Name -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"span" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
I.Sub IntType
Int64 Overflow
I.OverflowUndef) SubExp
n (IntType -> Integer -> SubExp
intConst IntType
Int64 Integer
1)
          spn <- letSubExp "span" $ I.BasicOp $ I.BinOp (I.Mul Int64 I.OverflowUndef) n_m1 s

          span_and_lower <- letSubExp "span_and_lower" $ I.BasicOp $ I.BinOp (I.Add Int64 I.OverflowUndef) spn lower
          span_and_upper <- letSubExp "span_and_upper" $ I.BasicOp $ I.BinOp (I.Add Int64 I.OverflowUndef) spn upper

          lower' <- letSubExp "minimum" $ I.BasicOp $ I.BinOp (I.UMin Int64) span_and_lower lower
          upper' <- letSubExp "maximum" $ I.BasicOp $ I.BinOp (I.UMax Int64) span_and_upper upper

          pure (lower', upper')
      )
      (offset', offset')
      slices'
  min_in_bounds <- letSubExp "min_in_bounds" $ I.BasicOp $ I.CmpOp (I.CmpUle Int64) (intConst Int64 0) min_bound
  max_in_bounds <- letSubExp "max_in_bounds" $ I.BasicOp $ I.CmpOp (I.CmpUlt Int64) max_bound old_dim

  all_bounds <-
    foldM
      (\SubExp
x SubExp
y -> Name -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"inBounds" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp BinOp
I.LogAnd SubExp
x SubExp
y)
      offset_inbounds_down
      [offset_inbounds_up, min_in_bounds, max_in_bounds]

  c <-
    assert "bounds_cert" all_bounds $
      ErrorMsg [ErrorString $ "Flat slice out of bounds: " <> prettyText old_dim <> " and " <> prettyText slices']
  let slice = SubExp -> [FlatDimIndex SubExp] -> FlatSlice SubExp
forall d. d -> [FlatDimIndex d] -> FlatSlice d
I.FlatSlice SubExp
offset' ([FlatDimIndex SubExp] -> FlatSlice SubExp)
-> [FlatDimIndex SubExp] -> FlatSlice SubExp
forall a b. (a -> b) -> a -> b
$ ((SubExp, SubExp) -> FlatDimIndex SubExp)
-> [(SubExp, SubExp)] -> [FlatDimIndex SubExp]
forall a b. (a -> b) -> [a] -> [b]
map ((SubExp -> SubExp -> FlatDimIndex SubExp)
-> (SubExp, SubExp) -> FlatDimIndex SubExp
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry SubExp -> SubExp -> FlatDimIndex SubExp
forall d. d -> d -> FlatDimIndex d
FlatDimIndex) [(SubExp, SubExp)]
slices'
  certifying c $
    forM arrs $ \VName
arr' ->
      Name -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
desc (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
forall a b. (a -> b) -> a -> b
$ VName -> FlatSlice SubExp -> BasicOp
I.FlatIndex VName
arr' FlatSlice SubExp
slice

flatUpdateHelper :: Name -> E.Exp -> E.Exp -> [E.Exp] -> E.Exp -> InternaliseM [SubExp]
flatUpdateHelper :: Name -> Exp -> Exp -> [Exp] -> Exp -> InternaliseM [SubExp]
flatUpdateHelper Name
desc Exp
arr1 Exp
offset [Exp]
slices Exp
arr2 = do
  arrs1 <- Name -> Exp -> InternaliseM [VName]
internaliseExpToVars Name
"arr" Exp
arr1
  offset' <- internaliseExp1 "offset" offset
  old_dim <- I.arraysSize 0 <$> mapM lookupType arrs1
  offset_inbounds_down <- letSubExp "offset_inbounds_down" $ I.BasicOp $ I.CmpOp (I.CmpUle Int64) (intConst Int64 0) offset'
  offset_inbounds_up <- letSubExp "offset_inbounds_up" $ I.BasicOp $ I.CmpOp (I.CmpUlt Int64) offset' old_dim
  arrs2 <- internaliseExpToVars "arr" arr2
  ts <- mapM lookupType arrs2
  slices' <-
    mapM
      ( \(Exp
s, Int
i) -> do
          s' <- Name -> Exp -> InternaliseM SubExp
internaliseExp1 Name
"s" Exp
s
          let n = Int -> [TypeBase Shape NoMode] -> SubExp
forall o. Int -> [TypeBase Shape o] -> SubExp
arraysSize Int
i [TypeBase Shape NoMode]
ts
          pure (n, s')
      )
      $ zip slices [0 ..]
  (min_bound, max_bound) <-
    foldM
      ( \(SubExp
lower, SubExp
upper) (SubExp
n, SubExp
s) -> do
          n_m1 <- Name -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"span" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp (IntType -> Overflow -> BinOp
I.Sub IntType
Int64 Overflow
I.OverflowUndef) SubExp
n (IntType -> Integer -> SubExp
intConst IntType
Int64 Integer
1)
          spn <- letSubExp "span" $ I.BasicOp $ I.BinOp (I.Mul Int64 I.OverflowUndef) n_m1 s

          span_and_lower <- letSubExp "span_and_lower" $ I.BasicOp $ I.BinOp (I.Add Int64 I.OverflowUndef) spn lower
          span_and_upper <- letSubExp "span_and_upper" $ I.BasicOp $ I.BinOp (I.Add Int64 I.OverflowUndef) spn upper

          lower' <- letSubExp "minimum" $ I.BasicOp $ I.BinOp (I.UMin Int64) span_and_lower lower
          upper' <- letSubExp "maximum" $ I.BasicOp $ I.BinOp (I.UMax Int64) span_and_upper upper

          pure (lower', upper')
      )
      (offset', offset')
      slices'
  min_in_bounds <- letSubExp "min_in_bounds" $ I.BasicOp $ I.CmpOp (I.CmpUle Int64) (intConst Int64 0) min_bound
  max_in_bounds <- letSubExp "max_in_bounds" $ I.BasicOp $ I.CmpOp (I.CmpUlt Int64) max_bound old_dim

  all_bounds <-
    foldM
      (\SubExp
x SubExp
y -> Name -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
"inBounds" (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
forall a b. (a -> b) -> a -> b
$ BinOp -> SubExp -> SubExp -> BasicOp
I.BinOp BinOp
I.LogAnd SubExp
x SubExp
y)
      offset_inbounds_down
      [offset_inbounds_up, min_in_bounds, max_in_bounds]

  c <-
    assert "bounds_cert" all_bounds $
      ErrorMsg [ErrorString $ "Flat slice out of bounds: " <> prettyText old_dim <> " and " <> prettyText slices']
  let slice = SubExp -> [FlatDimIndex SubExp] -> FlatSlice SubExp
forall d. d -> [FlatDimIndex d] -> FlatSlice d
I.FlatSlice SubExp
offset' ([FlatDimIndex SubExp] -> FlatSlice SubExp)
-> [FlatDimIndex SubExp] -> FlatSlice SubExp
forall a b. (a -> b) -> a -> b
$ ((SubExp, SubExp) -> FlatDimIndex SubExp)
-> [(SubExp, SubExp)] -> [FlatDimIndex SubExp]
forall a b. (a -> b) -> [a] -> [b]
map ((SubExp -> SubExp -> FlatDimIndex SubExp)
-> (SubExp, SubExp) -> FlatDimIndex SubExp
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry SubExp -> SubExp -> FlatDimIndex SubExp
forall d. d -> d -> FlatDimIndex d
FlatDimIndex) [(SubExp, SubExp)]
slices'
  certifying c $
    forM (zip arrs1 arrs2) $ \(VName
arr1', VName
arr2') ->
      Name -> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall (m :: * -> *).
MonadBuilder m =>
Name -> Exp (Rep m) -> m SubExp
letSubExp Name
desc (Exp (Rep InternaliseM) -> InternaliseM SubExp)
-> Exp (Rep InternaliseM) -> InternaliseM SubExp
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
I.BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
forall a b. (a -> b) -> a -> b
$ VName -> FlatSlice SubExp -> VName -> BasicOp
I.FlatUpdate VName
arr1' FlatSlice SubExp
slice VName
arr2'

funcall ::
  Name ->
  QualName VName ->
  [SubExp] ->
  InternaliseM [SubExp]
funcall :: Name -> QualName VName -> [SubExp] -> InternaliseM [SubExp]
funcall Name
desc (QualName [VName]
_ VName
fname) [SubExp]
args = do
  (shapes, value_paramts, fun_params, rettype_fun) <- VName -> InternaliseM FunInfo
lookupFunction VName
fname
  argts <- mapM subExpType args

  shapeargs <- argShapes shapes fun_params argts
  let diets =
        Int -> Diet -> [Diet]
forall a. Int -> a -> [a]
replicate ([SubExp] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [SubExp]
shapeargs) Diet
I.Observe
          [Diet] -> [Diet] -> [Diet]
forall a. [a] -> [a] -> [a]
++ (DeclType -> Diet) -> [DeclType] -> [Diet]
forall a b. (a -> b) -> [a] -> [b]
map DeclType -> Diet
forall shape. TypeBase shape Diet -> Diet
I.diet [DeclType]
value_paramts
  args' <-
    ensureArgShapes
      "function arguments of wrong shape"
      (map I.paramName fun_params)
      (map I.paramType fun_params)
      (shapeargs ++ args)
  argts' <- mapM subExpType args'
  case rettype_fun $ zip args' argts' of
    Maybe [(TypeBase ExtShape NoMode, RetAls)]
Nothing ->
      [Char] -> InternaliseM [SubExp]
forall a. HasCallStack => [Char] -> a
error ([Char] -> InternaliseM [SubExp])
-> ([[Char]] -> [Char]) -> [[Char]] -> InternaliseM [SubExp]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [[Char]] -> [Char]
unlines ([[Char]] -> InternaliseM [SubExp])
-> [[Char]] -> InternaliseM [SubExp]
forall a b. (a -> b) -> a -> b
$
        [ [Char]
"Cannot apply "
            [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> VName -> [Char]
forall a. Pretty a => a -> [Char]
prettyString VName
fname
            [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> [Char]
" to "
            [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> Int -> [Char]
forall a. Show a => a -> [Char]
show ([SubExp] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [SubExp]
args')
            [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> [Char]
" arguments",
          [Char]
" " [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> [SubExp] -> [Char]
forall a. Pretty a => a -> [Char]
prettyString [SubExp]
args',
          [Char]
"of types",
          [Char]
" " [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> [TypeBase Shape NoMode] -> [Char]
forall a. Pretty a => a -> [Char]
prettyString [TypeBase Shape NoMode]
argts',
          [Char]
"Function has " [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> Int -> [Char]
forall a. Show a => a -> [Char]
show ([Param DeclType] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Param DeclType]
fun_params) [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> [Char]
" parameters",
          [Char]
" " [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> [Param DeclType] -> [Char]
forall a. Pretty a => a -> [Char]
prettyString [Param DeclType]
fun_params
        ]
    Just [(TypeBase ExtShape NoMode, RetAls)]
ts -> do
      safety <- InternaliseM Safety
askSafety
      attrs <- asks envAttrs
      attributing attrs . letValExp' desc $
        I.Apply (internaliseFunName fname) (zip args' diets) ts safety

-- Bind existential names defined by an expression, based on the
-- concrete values that expression evaluated to.  This most
-- importantly should be done after function calls, but also
-- everything else that can produce existentials in the source
-- language.
bindExtSizes :: AppRes -> [SubExp] -> InternaliseM ()
bindExtSizes :: AppRes -> [SubExp] -> InternaliseM ()
bindExtSizes (AppRes StructType
ret [VName]
retext) [SubExp]
ses = do
  let ts :: [TypeBase ExtShape NoMode]
ts = (Tree (TypeBase ExtShape NoMode) -> [TypeBase ExtShape NoMode])
-> [Tree (TypeBase ExtShape NoMode)] -> [TypeBase ExtShape NoMode]
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Tree (TypeBase ExtShape NoMode) -> [TypeBase ExtShape NoMode]
forall a. Free [] a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList ([Tree (TypeBase ExtShape NoMode)] -> [TypeBase ExtShape NoMode])
-> [Tree (TypeBase ExtShape NoMode)] -> [TypeBase ExtShape NoMode]
forall a b. (a -> b) -> a -> b
$ StructType -> [Tree (TypeBase ExtShape NoMode)]
internaliseType (StructType -> [Tree (TypeBase ExtShape NoMode)])
-> StructType -> [Tree (TypeBase ExtShape NoMode)]
forall a b. (a -> b) -> a -> b
$ StructType -> StructType
forall dim o. TypeBase dim o -> TypeBase dim NoMode
E.toStruct StructType
ret
  ses_ts <- (SubExp -> InternaliseM (TypeBase Shape NoMode))
-> [SubExp] -> InternaliseM [TypeBase Shape NoMode]
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 SubExp -> InternaliseM (TypeBase Shape NoMode)
forall t (m :: * -> *).
HasScope t m =>
SubExp -> m (TypeBase Shape NoMode)
subExpType [SubExp]
ses

  let combine TypeBase ExtShape NoMode
t1 TypeBase Shape NoMode
t2 =
        [Map VName SubExp] -> Map VName SubExp
forall a. Monoid a => [a] -> a
mconcat ([Map VName SubExp] -> Map VName SubExp)
-> [Map VName SubExp] -> Map VName SubExp
forall a b. (a -> b) -> a -> b
$ (Ext SubExp -> SubExp -> Map VName SubExp)
-> [Ext SubExp] -> [SubExp] -> [Map VName SubExp]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Ext SubExp -> SubExp -> Map VName SubExp
combine' (TypeBase ExtShape NoMode -> [Ext SubExp]
forall d m. TypeBase (ShapeBase d) m -> [d]
arrayDims TypeBase ExtShape NoMode
t1) (TypeBase Shape NoMode -> [SubExp]
forall d m. TypeBase (ShapeBase d) m -> [d]
arrayDims TypeBase Shape NoMode
t2)
      combine' (I.Free (I.Var VName
v)) SubExp
se
        | VName
v VName -> [VName] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [VName]
retext = VName -> SubExp -> Map VName SubExp
forall k a. k -> a -> Map k a
M.singleton VName
v SubExp
se
      combine' Ext SubExp
_ SubExp
_ = Map VName SubExp
forall a. Monoid a => a
mempty

  forM_ (M.toList $ mconcat $ zipWith combine ts ses_ts) $ \(VName
v, SubExp
se) ->
    [VName] -> Exp (Rep InternaliseM) -> InternaliseM ()
forall (m :: * -> *).
MonadBuilder m =>
[VName] -> Exp (Rep m) -> m ()
letBindNames [VName
v] (Exp (Rep InternaliseM) -> InternaliseM ())
-> Exp (Rep InternaliseM) -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp (Rep InternaliseM)
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp (Rep InternaliseM))
-> BasicOp -> Exp (Rep InternaliseM)
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
SubExp SubExp
se

askSafety :: InternaliseM Safety
askSafety :: InternaliseM Safety
askSafety = do
  check <- (InternaliseEnv -> Bool) -> InternaliseM Bool
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks InternaliseEnv -> Bool
envDoBoundsChecks
  pure $ if check then I.Safe else I.Unsafe

sizeExpForError :: E.Size -> InternaliseM [ErrorMsgPart SubExp]
sizeExpForError :: Exp -> InternaliseM [ErrorMsgPart SubExp]
sizeExpForError Exp
e
  | Just Int
_ <- Exp -> Maybe Int
isAnySize Exp
e = [ErrorMsgPart SubExp] -> InternaliseM [ErrorMsgPart SubExp]
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [ErrorMsgPart SubExp
"[]"]
  | Bool
otherwise = do
      e' <- Name -> Exp -> InternaliseM SubExp
internaliseExp1 Name
"size" Exp
e
      pure ["[", ErrorVal int64 e', "]"]

typeExpForError :: E.TypeBase Size o -> InternaliseM [ErrorMsgPart SubExp]
typeExpForError :: forall o. TypeBase Exp o -> InternaliseM [ErrorMsgPart SubExp]
typeExpForError (E.Scalar (E.Prim PrimType
t)) = [ErrorMsgPart SubExp] -> InternaliseM [ErrorMsgPart SubExp]
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [Text -> ErrorMsgPart SubExp
forall a. Text -> ErrorMsgPart a
ErrorString (Text -> ErrorMsgPart SubExp) -> Text -> ErrorMsgPart SubExp
forall a b. (a -> b) -> a -> b
$ PrimType -> Text
forall a. Pretty a => a -> Text
prettyText PrimType
t]
typeExpForError (E.Scalar (E.TypeVar o
_ QualName VName
v [TypeArg Exp]
args)) = do
  args' <- [[ErrorMsgPart SubExp]] -> [ErrorMsgPart SubExp]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[ErrorMsgPart SubExp]] -> [ErrorMsgPart SubExp])
-> InternaliseM [[ErrorMsgPart SubExp]]
-> InternaliseM [ErrorMsgPart SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (TypeArg Exp -> InternaliseM [ErrorMsgPart SubExp])
-> [TypeArg Exp] -> InternaliseM [[ErrorMsgPart 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 TypeArg Exp -> InternaliseM [ErrorMsgPart SubExp]
onArg [TypeArg Exp]
args
  pure $ intersperse " " $ ErrorString (prettyText v) : args'
  where
    onArg :: TypeArg Exp -> InternaliseM [ErrorMsgPart SubExp]
onArg (TypeArgDim Exp
d) = Exp -> InternaliseM [ErrorMsgPart SubExp]
sizeExpForError Exp
d
    onArg (TypeArgType StructType
t) = StructType -> InternaliseM [ErrorMsgPart SubExp]
forall o. TypeBase Exp o -> InternaliseM [ErrorMsgPart SubExp]
typeExpForError StructType
t
typeExpForError (E.Scalar (E.Record Map Name (TypeBase Exp o)
fs))
  | Just [TypeBase Exp o]
ts <- Map Name (TypeBase Exp o) -> Maybe [TypeBase Exp o]
forall a. Map Name a -> Maybe [a]
E.areTupleFields Map Name (TypeBase Exp o)
fs = do
      ts' <- (TypeBase Exp o -> InternaliseM [ErrorMsgPart SubExp])
-> [TypeBase Exp o] -> InternaliseM [[ErrorMsgPart 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 TypeBase Exp o -> InternaliseM [ErrorMsgPart SubExp]
forall o. TypeBase Exp o -> InternaliseM [ErrorMsgPart SubExp]
typeExpForError [TypeBase Exp o]
ts
      pure $ ["("] ++ intercalate [", "] ts' ++ [")"]
  | Bool
otherwise = do
      fs' <- ((Name, TypeBase Exp o) -> InternaliseM [ErrorMsgPart SubExp])
-> [(Name, TypeBase Exp o)] -> InternaliseM [[ErrorMsgPart 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, TypeBase Exp o) -> InternaliseM [ErrorMsgPart SubExp]
forall {a} {o}.
Pretty a =>
(a, TypeBase Exp o) -> InternaliseM [ErrorMsgPart SubExp]
onField ([(Name, TypeBase Exp o)] -> InternaliseM [[ErrorMsgPart SubExp]])
-> [(Name, TypeBase Exp o)] -> InternaliseM [[ErrorMsgPart SubExp]]
forall a b. (a -> b) -> a -> b
$ Map Name (TypeBase Exp o) -> [(Name, TypeBase Exp o)]
forall k a. Map k a -> [(k, a)]
M.toList Map Name (TypeBase Exp o)
fs
      pure $ ["{"] ++ intercalate [", "] fs' ++ ["}"]
  where
    onField :: (a, TypeBase Exp o) -> InternaliseM [ErrorMsgPart SubExp]
onField (a
k, TypeBase Exp o
te) =
      (Text -> ErrorMsgPart SubExp
forall a. Text -> ErrorMsgPart a
ErrorString (a -> Text
forall a. Pretty a => a -> Text
prettyText a
k Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
": ") ErrorMsgPart SubExp
-> [ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp]
forall a. a -> [a] -> [a]
:) ([ErrorMsgPart SubExp] -> [ErrorMsgPart SubExp])
-> InternaliseM [ErrorMsgPart SubExp]
-> InternaliseM [ErrorMsgPart SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TypeBase Exp o -> InternaliseM [ErrorMsgPart SubExp]
forall o. TypeBase Exp o -> InternaliseM [ErrorMsgPart SubExp]
typeExpForError TypeBase Exp o
te
typeExpForError (E.Array o
_ Shape Exp
shape ScalarTypeBase Exp NoMode
et) = do
  shape' <- [[ErrorMsgPart SubExp]] -> [ErrorMsgPart SubExp]
forall a. Monoid a => [a] -> a
mconcat ([[ErrorMsgPart SubExp]] -> [ErrorMsgPart SubExp])
-> InternaliseM [[ErrorMsgPart SubExp]]
-> InternaliseM [ErrorMsgPart SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Exp -> InternaliseM [ErrorMsgPart SubExp])
-> [Exp] -> InternaliseM [[ErrorMsgPart 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 Exp -> InternaliseM [ErrorMsgPart SubExp]
sizeExpForError (Shape Exp -> [Exp]
forall dim. Shape dim -> [dim]
E.shapeDims Shape Exp
shape)
  et' <- typeExpForError $ Scalar et
  pure $ shape' ++ et'
typeExpForError (E.Scalar (E.Sum Map Name [TypeBase Exp o]
cs)) = do
  cs' <- ((Name, [TypeBase Exp o]) -> InternaliseM [ErrorMsgPart SubExp])
-> [(Name, [TypeBase Exp o])]
-> InternaliseM [[ErrorMsgPart 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, [TypeBase Exp o]) -> InternaliseM [ErrorMsgPart SubExp]
forall {a} {o}.
Pretty a =>
(a, [TypeBase Exp o]) -> InternaliseM [ErrorMsgPart SubExp]
onConstructor ([(Name, [TypeBase Exp o])]
 -> InternaliseM [[ErrorMsgPart SubExp]])
-> [(Name, [TypeBase Exp o])]
-> InternaliseM [[ErrorMsgPart SubExp]]
forall a b. (a -> b) -> a -> b
$ Map Name [TypeBase Exp o] -> [(Name, [TypeBase Exp o])]
forall k a. Map k a -> [(k, a)]
M.toList Map Name [TypeBase Exp o]
cs
  pure $ intercalate [" | "] cs'
  where
    onConstructor :: (a, [TypeBase Exp o]) -> InternaliseM [ErrorMsgPart SubExp]
onConstructor (a
c, [TypeBase Exp o]
ts) = do
      ts' <- (TypeBase Exp o -> InternaliseM [ErrorMsgPart SubExp])
-> [TypeBase Exp o] -> InternaliseM [[ErrorMsgPart 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 TypeBase Exp o -> InternaliseM [ErrorMsgPart SubExp]
forall o. TypeBase Exp o -> InternaliseM [ErrorMsgPart SubExp]
typeExpForError [TypeBase Exp o]
ts
      pure $ ErrorString ("#" <> prettyText c <> " ") : intercalate [" "] ts'
typeExpForError (E.Scalar Arrow {}) = [ErrorMsgPart SubExp] -> InternaliseM [ErrorMsgPart SubExp]
forall a. a -> InternaliseM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [ErrorMsgPart SubExp
"#<fun>"]

-- A smart constructor that compacts neighbouring literals for easier
-- reading in the IR.
errorMsg :: [ErrorMsgPart a] -> ErrorMsg a
errorMsg :: forall a. [ErrorMsgPart a] -> ErrorMsg a
errorMsg = [ErrorMsgPart a] -> ErrorMsg a
forall a. [ErrorMsgPart a] -> ErrorMsg a
ErrorMsg ([ErrorMsgPart a] -> ErrorMsg a)
-> ([ErrorMsgPart a] -> [ErrorMsgPart a])
-> [ErrorMsgPart a]
-> ErrorMsg a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [ErrorMsgPart a] -> [ErrorMsgPart a]
forall {a}. [ErrorMsgPart a] -> [ErrorMsgPart a]
compact
  where
    compact :: [ErrorMsgPart a] -> [ErrorMsgPart a]
compact [] = []
    compact (ErrorString Text
x : ErrorString Text
y : [ErrorMsgPart a]
parts) =
      [ErrorMsgPart a] -> [ErrorMsgPart a]
compact (Text -> ErrorMsgPart a
forall a. Text -> ErrorMsgPart a
ErrorString (Text
x Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
y) ErrorMsgPart a -> [ErrorMsgPart a] -> [ErrorMsgPart a]
forall a. a -> [a] -> [a]
: [ErrorMsgPart a]
parts)
    compact (ErrorMsgPart a
x : [ErrorMsgPart a]
y) = ErrorMsgPart a
x ErrorMsgPart a -> [ErrorMsgPart a] -> [ErrorMsgPart a]
forall a. a -> [a] -> [a]
: [ErrorMsgPart a] -> [ErrorMsgPart a]
compact [ErrorMsgPart a]
y