{-# LANGUAGE FlexibleContexts #-}

module Futhark.Internalise.Lambdas
  ( InternaliseLambda,
    internaliseMapLambda,
    internaliseStreamMapLambda,
    internaliseFoldLambda,
    internaliseStreamLambda,
    internalisePartitionLambda,
  )
where

import Futhark.IR.SOACS as I
import Futhark.Internalise.AccurateSizes
import Futhark.Internalise.Monad
import Language.Futhark as E

-- | A function for internalising lambdas.
type InternaliseLambda =
  E.Exp -> [I.Type] -> InternaliseM ([I.LParam], I.Body, [I.Type])

internaliseMapLambda ::
  InternaliseLambda ->
  E.Exp ->
  [I.SubExp] ->
  InternaliseM I.Lambda
internaliseMapLambda :: InternaliseLambda -> Exp -> [SubExp] -> InternaliseM Lambda
internaliseMapLambda InternaliseLambda
internaliseLambda Exp
lam [SubExp]
args = do
  [Type]
argtypes <- (SubExp -> InternaliseM Type) -> [SubExp] -> InternaliseM [Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM SubExp -> InternaliseM Type
forall t (m :: * -> *). HasScope t m => SubExp -> m Type
I.subExpType [SubExp]
args
  let rowtypes :: [Type]
rowtypes = (Type -> Type) -> [Type] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map Type -> Type
forall shape u.
ArrayShape shape =>
TypeBase shape u -> TypeBase shape u
I.rowType [Type]
argtypes
  ([Param Type]
params, Body
body, [Type]
rettype) <- InternaliseLambda
internaliseLambda Exp
lam [Type]
rowtypes
  Body
body' <-
    Scope SOACS -> InternaliseM Body -> InternaliseM Body
forall lore (m :: * -> *) a.
LocalScope lore m =>
Scope lore -> m a -> m a
localScope ([Param Type] -> Scope SOACS
forall lore dec.
(LParamInfo lore ~ dec) =>
[Param dec] -> Scope lore
scopeOfLParams [Param Type]
params) (InternaliseM Body -> InternaliseM Body)
-> InternaliseM Body -> InternaliseM Body
forall a b. (a -> b) -> a -> b
$
      ErrorMsg SubExp -> SrcLoc -> [Type] -> Body -> InternaliseM Body
ensureResultShape
        ([ErrorMsgPart SubExp] -> ErrorMsg SubExp
forall a. [ErrorMsgPart a] -> ErrorMsg a
ErrorMsg [String -> ErrorMsgPart SubExp
forall a. String -> ErrorMsgPart a
ErrorString String
"not all iterations produce same shape"])
        (Exp -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf Exp
lam)
        [Type]
rettype
        Body
body
  Lambda -> InternaliseM Lambda
forall (m :: * -> *) a. Monad m => a -> m a
return (Lambda -> InternaliseM Lambda) -> Lambda -> InternaliseM Lambda
forall a b. (a -> b) -> a -> b
$ [LParam SOACS] -> Body -> [Type] -> Lambda
forall lore. [LParam lore] -> BodyT lore -> [Type] -> LambdaT lore
I.Lambda [Param Type]
[LParam SOACS]
params Body
body' [Type]
rettype

internaliseStreamMapLambda ::
  InternaliseLambda ->
  E.Exp ->
  [I.SubExp] ->
  InternaliseM I.Lambda
internaliseStreamMapLambda :: InternaliseLambda -> Exp -> [SubExp] -> InternaliseM Lambda
internaliseStreamMapLambda InternaliseLambda
internaliseLambda Exp
lam [SubExp]
args = do
  VName
chunk_size <- String -> InternaliseM VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newVName String
"chunk_size"
  let chunk_param :: Param (TypeBase shape u)
chunk_param = VName -> TypeBase shape u -> Param (TypeBase shape u)
forall dec. VName -> dec -> Param dec
I.Param VName
chunk_size (PrimType -> TypeBase shape u
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int64)
      outer :: TypeBase (ShapeBase SubExp) u -> TypeBase (ShapeBase SubExp) u
outer = (TypeBase (ShapeBase SubExp) u
-> SubExp -> TypeBase (ShapeBase SubExp) u
forall d u.
ArrayShape (ShapeBase d) =>
TypeBase (ShapeBase d) u -> d -> TypeBase (ShapeBase d) u
`setOuterSize` VName -> SubExp
I.Var VName
chunk_size)
  Scope SOACS -> InternaliseM Lambda -> InternaliseM Lambda
forall lore (m :: * -> *) a.
LocalScope lore m =>
Scope lore -> m a -> m a
localScope ([Param Type] -> Scope SOACS
forall lore dec.
(LParamInfo lore ~ dec) =>
[Param dec] -> Scope lore
scopeOfLParams [Param Type
forall shape u. Param (TypeBase shape u)
chunk_param]) (InternaliseM Lambda -> InternaliseM Lambda)
-> InternaliseM Lambda -> InternaliseM Lambda
forall a b. (a -> b) -> a -> b
$ do
    [Type]
argtypes <- (SubExp -> InternaliseM Type) -> [SubExp] -> InternaliseM [Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM SubExp -> InternaliseM Type
forall t (m :: * -> *). HasScope t m => SubExp -> m Type
I.subExpType [SubExp]
args
    ([Param Type]
lam_params, Body
orig_body, [Type]
rettype) <-
      InternaliseLambda
internaliseLambda Exp
lam ([Type] -> InternaliseM ([LParam SOACS], Body, [Type]))
-> [Type] -> InternaliseM ([LParam SOACS], Body, [Type])
forall a b. (a -> b) -> a -> b
$ PrimType -> Type
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int64 Type -> [Type] -> [Type]
forall a. a -> [a] -> [a]
: (Type -> Type) -> [Type] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map Type -> Type
forall u.
TypeBase (ShapeBase SubExp) u -> TypeBase (ShapeBase SubExp) u
outer [Type]
argtypes
    let Param Type
orig_chunk_param : [Param Type]
params = [Param Type]
lam_params
    Body
body <- Binder SOACS Body -> InternaliseM Body
forall lore (m :: * -> *) somelore.
(Bindable lore, MonadFreshNames m, HasScope somelore m,
 SameScope somelore lore) =>
Binder lore (Body lore) -> m (Body lore)
runBodyBinder (Binder SOACS Body -> InternaliseM Body)
-> Binder SOACS Body -> InternaliseM Body
forall a b. (a -> b) -> a -> b
$ do
      [VName]
-> Exp (Lore (BinderT SOACS (State VNameSource)))
-> BinderT SOACS (State VNameSource) ()
forall (m :: * -> *).
MonadBinder m =>
[VName] -> Exp (Lore m) -> m ()
letBindNames [Param Type -> VName
forall dec. Param dec -> VName
paramName Param Type
orig_chunk_param] (Exp (Lore (BinderT SOACS (State VNameSource)))
 -> BinderT SOACS (State VNameSource) ())
-> Exp (Lore (BinderT SOACS (State VNameSource)))
-> BinderT SOACS (State VNameSource) ()
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
I.SubExp (SubExp -> BasicOp) -> SubExp -> BasicOp
forall a b. (a -> b) -> a -> b
$ VName -> SubExp
I.Var VName
chunk_size
      Body -> Binder SOACS Body
forall (m :: * -> *) a. Monad m => a -> m a
return Body
orig_body
    Body
body' <- Scope SOACS -> InternaliseM Body -> InternaliseM Body
forall lore (m :: * -> *) a.
LocalScope lore m =>
Scope lore -> m a -> m a
localScope ([Param Type] -> Scope SOACS
forall lore dec.
(LParamInfo lore ~ dec) =>
[Param dec] -> Scope lore
scopeOfLParams [Param Type]
params) (InternaliseM Body -> InternaliseM Body)
-> InternaliseM Body -> InternaliseM Body
forall a b. (a -> b) -> a -> b
$
      InternaliseM (Body (Lore InternaliseM))
-> InternaliseM (Body (Lore InternaliseM))
forall (m :: * -> *).
MonadBinder m =>
m (Body (Lore m)) -> m (Body (Lore m))
insertStmsM (InternaliseM (Body (Lore InternaliseM))
 -> InternaliseM (Body (Lore InternaliseM)))
-> InternaliseM (Body (Lore InternaliseM))
-> InternaliseM (Body (Lore InternaliseM))
forall a b. (a -> b) -> a -> b
$ do
        [VName] -> Exp (Lore InternaliseM) -> InternaliseM ()
forall (m :: * -> *).
MonadBinder m =>
[VName] -> Exp (Lore m) -> m ()
letBindNames [Param Type -> VName
forall dec. Param dec -> VName
paramName Param Type
orig_chunk_param] (Exp (Lore InternaliseM) -> InternaliseM ())
-> Exp (Lore InternaliseM) -> InternaliseM ()
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
I.SubExp (SubExp -> BasicOp) -> SubExp -> BasicOp
forall a b. (a -> b) -> a -> b
$ VName -> SubExp
I.Var VName
chunk_size
        ErrorMsg SubExp -> SrcLoc -> [Type] -> Body -> InternaliseM Body
ensureResultShape
          ([ErrorMsgPart SubExp] -> ErrorMsg SubExp
forall a. [ErrorMsgPart a] -> ErrorMsg a
ErrorMsg [String -> ErrorMsgPart SubExp
forall a. String -> ErrorMsgPart a
ErrorString String
"not all iterations produce same shape"])
          (Exp -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf Exp
lam)
          ((Type -> Type) -> [Type] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map Type -> Type
forall u.
TypeBase (ShapeBase SubExp) u -> TypeBase (ShapeBase SubExp) u
outer [Type]
rettype)
          Body
body
    Lambda -> InternaliseM Lambda
forall (m :: * -> *) a. Monad m => a -> m a
return (Lambda -> InternaliseM Lambda) -> Lambda -> InternaliseM Lambda
forall a b. (a -> b) -> a -> b
$ [LParam SOACS] -> Body -> [Type] -> Lambda
forall lore. [LParam lore] -> BodyT lore -> [Type] -> LambdaT lore
I.Lambda (Param Type
forall shape u. Param (TypeBase shape u)
chunk_param Param Type -> [Param Type] -> [Param Type]
forall a. a -> [a] -> [a]
: [Param Type]
params) Body
body' ((Type -> Type) -> [Type] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map Type -> Type
forall u.
TypeBase (ShapeBase SubExp) u -> TypeBase (ShapeBase SubExp) u
outer [Type]
rettype)

internaliseFoldLambda ::
  InternaliseLambda ->
  E.Exp ->
  [I.Type] ->
  [I.Type] ->
  InternaliseM I.Lambda
internaliseFoldLambda :: InternaliseLambda -> Exp -> [Type] -> [Type] -> InternaliseM Lambda
internaliseFoldLambda InternaliseLambda
internaliseLambda Exp
lam [Type]
acctypes [Type]
arrtypes = do
  let rowtypes :: [Type]
rowtypes = (Type -> Type) -> [Type] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map Type -> Type
forall shape u.
ArrayShape shape =>
TypeBase shape u -> TypeBase shape u
I.rowType [Type]
arrtypes
  ([Param Type]
params, Body
body, [Type]
rettype) <- InternaliseLambda
internaliseLambda Exp
lam ([Type] -> InternaliseM ([LParam SOACS], Body, [Type]))
-> [Type] -> InternaliseM ([LParam SOACS], Body, [Type])
forall a b. (a -> b) -> a -> b
$ [Type]
acctypes [Type] -> [Type] -> [Type]
forall a. [a] -> [a] -> [a]
++ [Type]
rowtypes
  let rettype' :: [Type]
rettype' =
        [ Type
t Type -> ShapeBase SubExp -> Type
forall newshape oldshape u.
ArrayShape newshape =>
TypeBase oldshape u -> newshape -> TypeBase newshape u
`I.setArrayShape` Type -> ShapeBase SubExp
forall shape u. ArrayShape shape => TypeBase shape u -> shape
I.arrayShape Type
shape
          | (Type
t, Type
shape) <- [Type] -> [Type] -> [(Type, Type)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Type]
rettype [Type]
acctypes
        ]
  -- The result of the body must have the exact same shape as the
  -- initial accumulator.  We accomplish this with an assertion and
  -- reshape().
  Body
body' <-
    Scope SOACS -> InternaliseM Body -> InternaliseM Body
forall lore (m :: * -> *) a.
LocalScope lore m =>
Scope lore -> m a -> m a
localScope ([Param Type] -> Scope SOACS
forall lore dec.
(LParamInfo lore ~ dec) =>
[Param dec] -> Scope lore
scopeOfLParams [Param Type]
params) (InternaliseM Body -> InternaliseM Body)
-> InternaliseM Body -> InternaliseM Body
forall a b. (a -> b) -> a -> b
$
      ErrorMsg SubExp -> SrcLoc -> [Type] -> Body -> InternaliseM Body
ensureResultShape
        ([ErrorMsgPart SubExp] -> ErrorMsg SubExp
forall a. [ErrorMsgPart a] -> ErrorMsg a
ErrorMsg [String -> ErrorMsgPart SubExp
forall a. String -> ErrorMsgPart a
ErrorString String
"shape of result does not match shape of initial value"])
        (Exp -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf Exp
lam)
        [Type]
rettype'
        Body
body
  Lambda -> InternaliseM Lambda
forall (m :: * -> *) a. Monad m => a -> m a
return (Lambda -> InternaliseM Lambda) -> Lambda -> InternaliseM Lambda
forall a b. (a -> b) -> a -> b
$ [LParam SOACS] -> Body -> [Type] -> Lambda
forall lore. [LParam lore] -> BodyT lore -> [Type] -> LambdaT lore
I.Lambda [Param Type]
[LParam SOACS]
params Body
body' [Type]
rettype'

internaliseStreamLambda ::
  InternaliseLambda ->
  E.Exp ->
  [I.Type] ->
  InternaliseM ([LParam], Body)
internaliseStreamLambda :: InternaliseLambda
-> Exp -> [Type] -> InternaliseM ([LParam SOACS], Body)
internaliseStreamLambda InternaliseLambda
internaliseLambda Exp
lam [Type]
rowts = do
  VName
chunk_size <- String -> InternaliseM VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newVName String
"chunk_size"
  let chunk_param :: Param (TypeBase shape u)
chunk_param = VName -> TypeBase shape u -> Param (TypeBase shape u)
forall dec. VName -> dec -> Param dec
I.Param VName
chunk_size (TypeBase shape u -> Param (TypeBase shape u))
-> TypeBase shape u -> Param (TypeBase shape u)
forall a b. (a -> b) -> a -> b
$ PrimType -> TypeBase shape u
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int64
      chunktypes :: [Type]
chunktypes = (Type -> Type) -> [Type] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map (Type -> SubExp -> Type
forall d.
ArrayShape (ShapeBase d) =>
TypeBase (ShapeBase d) NoUniqueness
-> d -> TypeBase (ShapeBase d) NoUniqueness
`arrayOfRow` VName -> SubExp
I.Var VName
chunk_size) [Type]
rowts
  Scope SOACS
-> InternaliseM ([Param Type], Body)
-> InternaliseM ([Param Type], Body)
forall lore (m :: * -> *) a.
LocalScope lore m =>
Scope lore -> m a -> m a
localScope ([Param Type] -> Scope SOACS
forall lore dec.
(LParamInfo lore ~ dec) =>
[Param dec] -> Scope lore
scopeOfLParams [Param Type
forall shape u. Param (TypeBase shape u)
chunk_param]) (InternaliseM ([Param Type], Body)
 -> InternaliseM ([Param Type], Body))
-> InternaliseM ([Param Type], Body)
-> InternaliseM ([Param Type], Body)
forall a b. (a -> b) -> a -> b
$ do
    ([Param Type]
lam_params, Body
orig_body, [Type]
_) <-
      InternaliseLambda
internaliseLambda Exp
lam ([Type] -> InternaliseM ([LParam SOACS], Body, [Type]))
-> [Type] -> InternaliseM ([LParam SOACS], Body, [Type])
forall a b. (a -> b) -> a -> b
$ PrimType -> Type
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int64 Type -> [Type] -> [Type]
forall a. a -> [a] -> [a]
: [Type]
chunktypes
    let Param Type
orig_chunk_param : [Param Type]
params = [Param Type]
lam_params
    Body
body <- Binder SOACS Body -> InternaliseM Body
forall lore (m :: * -> *) somelore.
(Bindable lore, MonadFreshNames m, HasScope somelore m,
 SameScope somelore lore) =>
Binder lore (Body lore) -> m (Body lore)
runBodyBinder (Binder SOACS Body -> InternaliseM Body)
-> Binder SOACS Body -> InternaliseM Body
forall a b. (a -> b) -> a -> b
$ do
      [VName]
-> Exp (Lore (BinderT SOACS (State VNameSource)))
-> BinderT SOACS (State VNameSource) ()
forall (m :: * -> *).
MonadBinder m =>
[VName] -> Exp (Lore m) -> m ()
letBindNames [Param Type -> VName
forall dec. Param dec -> VName
paramName Param Type
orig_chunk_param] (Exp (Lore (BinderT SOACS (State VNameSource)))
 -> BinderT SOACS (State VNameSource) ())
-> Exp (Lore (BinderT SOACS (State VNameSource)))
-> BinderT SOACS (State VNameSource) ()
forall a b. (a -> b) -> a -> b
$ BasicOp -> ExpT SOACS
forall lore. BasicOp -> ExpT lore
I.BasicOp (BasicOp -> ExpT SOACS) -> BasicOp -> ExpT SOACS
forall a b. (a -> b) -> a -> b
$ SubExp -> BasicOp
I.SubExp (SubExp -> BasicOp) -> SubExp -> BasicOp
forall a b. (a -> b) -> a -> b
$ VName -> SubExp
I.Var VName
chunk_size
      Body -> Binder SOACS Body
forall (m :: * -> *) a. Monad m => a -> m a
return Body
orig_body
    ([Param Type], Body) -> InternaliseM ([Param Type], Body)
forall (m :: * -> *) a. Monad m => a -> m a
return (Param Type
forall shape u. Param (TypeBase shape u)
chunk_param Param Type -> [Param Type] -> [Param Type]
forall a. a -> [a] -> [a]
: [Param Type]
params, Body
body)

-- Given @k@ lambdas, this will return a lambda that returns an
-- (k+2)-element tuple of integers.  The first element is the
-- equivalence class ID in the range [0,k].  The remaining are all zero
-- except for possibly one element.
internalisePartitionLambda ::
  InternaliseLambda ->
  Int ->
  E.Exp ->
  [I.SubExp] ->
  InternaliseM I.Lambda
internalisePartitionLambda :: InternaliseLambda -> Int -> Exp -> [SubExp] -> InternaliseM Lambda
internalisePartitionLambda InternaliseLambda
internaliseLambda Int
k Exp
lam [SubExp]
args = do
  [Type]
argtypes <- (SubExp -> InternaliseM Type) -> [SubExp] -> InternaliseM [Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM SubExp -> InternaliseM Type
forall t (m :: * -> *). HasScope t m => SubExp -> m Type
I.subExpType [SubExp]
args
  let rowtypes :: [Type]
rowtypes = (Type -> Type) -> [Type] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map Type -> Type
forall shape u.
ArrayShape shape =>
TypeBase shape u -> TypeBase shape u
I.rowType [Type]
argtypes
  ([Param Type]
params, Body
body, [Type]
_) <- InternaliseLambda
internaliseLambda Exp
lam [Type]
rowtypes
  Body
body' <-
    Scope SOACS -> InternaliseM Body -> InternaliseM Body
forall lore (m :: * -> *) a.
LocalScope lore m =>
Scope lore -> m a -> m a
localScope ([Param Type] -> Scope SOACS
forall lore dec.
(LParamInfo lore ~ dec) =>
[Param dec] -> Scope lore
scopeOfLParams [Param Type]
params) (InternaliseM Body -> InternaliseM Body)
-> InternaliseM Body -> InternaliseM Body
forall a b. (a -> b) -> a -> b
$
      Body -> InternaliseM Body
lambdaWithIncrement Body
body
  Lambda -> InternaliseM Lambda
forall (m :: * -> *) a. Monad m => a -> m a
return (Lambda -> InternaliseM Lambda) -> Lambda -> InternaliseM Lambda
forall a b. (a -> b) -> a -> b
$ [LParam SOACS] -> Body -> [Type] -> Lambda
forall lore. [LParam lore] -> BodyT lore -> [Type] -> LambdaT lore
I.Lambda [Param Type]
[LParam SOACS]
params Body
body' [Type]
forall shape u. [TypeBase shape u]
rettype
  where
    rettype :: [TypeBase shape u]
rettype = Int -> TypeBase shape u -> [TypeBase shape u]
forall a. Int -> a -> [a]
replicate (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
2) (TypeBase shape u -> [TypeBase shape u])
-> TypeBase shape u -> [TypeBase shape u]
forall a b. (a -> b) -> a -> b
$ PrimType -> TypeBase shape u
forall shape u. PrimType -> TypeBase shape u
I.Prim PrimType
int64
    result :: Int -> [SubExp]
result Int
i =
      (Int64 -> SubExp) -> [Int64] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map Int64 -> SubExp
forall v. IsValue v => v -> SubExp
constant ([Int64] -> [SubExp]) -> [Int64] -> [SubExp]
forall a b. (a -> b) -> a -> b
$
        Int -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
i Int64 -> [Int64] -> [Int64]
forall a. a -> [a] -> [a]
:
        (Int -> Int64 -> [Int64]
forall a. Int -> a -> [a]
replicate Int
i Int64
0 [Int64] -> [Int64] -> [Int64]
forall a. [a] -> [a] -> [a]
++ [Int64
1 :: Int64] [Int64] -> [Int64] -> [Int64]
forall a. [a] -> [a] -> [a]
++ Int -> Int64 -> [Int64]
forall a. Int -> a -> [a]
replicate (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
i) Int64
0)

    mkResult :: SubExp -> Int -> m [SubExp]
mkResult SubExp
_ Int
i | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
k = [SubExp] -> m [SubExp]
forall (m :: * -> *) a. Monad m => a -> m a
return ([SubExp] -> m [SubExp]) -> [SubExp] -> m [SubExp]
forall a b. (a -> b) -> a -> b
$ Int -> [SubExp]
result Int
i
    mkResult SubExp
eq_class Int
i = do
      SubExp
is_i <-
        String -> Exp (Lore m) -> m SubExp
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m SubExp
letSubExp String
"is_i" (Exp (Lore m) -> m SubExp) -> Exp (Lore m) -> m SubExp
forall a b. (a -> b) -> a -> b
$
          BasicOp -> Exp (Lore m)
forall lore. BasicOp -> ExpT lore
BasicOp (BasicOp -> Exp (Lore m)) -> BasicOp -> Exp (Lore m)
forall a b. (a -> b) -> a -> b
$
            CmpOp -> SubExp -> SubExp -> BasicOp
CmpOp (PrimType -> CmpOp
CmpEq PrimType
int64) SubExp
eq_class (SubExp -> BasicOp) -> SubExp -> BasicOp
forall a b. (a -> b) -> a -> b
$
              IntType -> Integer -> SubExp
intConst IntType
Int64 (Integer -> SubExp) -> Integer -> SubExp
forall a b. (a -> b) -> a -> b
$ Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
i
      ([VName] -> [SubExp]) -> m [VName] -> m [SubExp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((VName -> SubExp) -> [VName] -> [SubExp]
forall a b. (a -> b) -> [a] -> [b]
map VName -> SubExp
I.Var) (m [VName] -> m [SubExp])
-> (Exp (Lore m) -> m [VName]) -> Exp (Lore m) -> m [SubExp]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> Exp (Lore m) -> m [VName]
forall (m :: * -> *).
MonadBinder m =>
String -> Exp (Lore m) -> m [VName]
letTupExp String
"part_res"
        (Exp (Lore m) -> m [SubExp]) -> m (Exp (Lore m)) -> m [SubExp]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< m (Exp (Lore m))
-> m (Body (Lore m)) -> m (Body (Lore m)) -> m (Exp (Lore m))
forall (m :: * -> *).
(MonadBinder m, BranchType (Lore m) ~ ExtType) =>
m (Exp (Lore m))
-> m (Body (Lore m)) -> m (Body (Lore m)) -> m (Exp (Lore m))
eIf
          (SubExp -> m (Exp (Lore m))
forall (m :: * -> *). MonadBinder m => SubExp -> m (Exp (Lore m))
eSubExp SubExp
is_i)
          (Body (Lore m) -> m (Body (Lore m))
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Body (Lore m) -> m (Body (Lore m)))
-> Body (Lore m) -> m (Body (Lore m))
forall a b. (a -> b) -> a -> b
$ [SubExp] -> Body (Lore m)
forall lore. Bindable lore => [SubExp] -> Body lore
resultBody ([SubExp] -> Body (Lore m)) -> [SubExp] -> Body (Lore m)
forall a b. (a -> b) -> a -> b
$ Int -> [SubExp]
result Int
i)
          ([SubExp] -> Body (Lore m)
forall lore. Bindable lore => [SubExp] -> Body lore
resultBody ([SubExp] -> Body (Lore m)) -> m [SubExp] -> m (Body (Lore m))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SubExp -> Int -> m [SubExp]
mkResult SubExp
eq_class (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1))

    lambdaWithIncrement :: I.Body -> InternaliseM I.Body
    lambdaWithIncrement :: Body -> InternaliseM Body
lambdaWithIncrement Body
lam_body = Binder SOACS Body -> InternaliseM Body
forall lore (m :: * -> *) somelore.
(Bindable lore, MonadFreshNames m, HasScope somelore m,
 SameScope somelore lore) =>
Binder lore (Body lore) -> m (Body lore)
runBodyBinder (Binder SOACS Body -> InternaliseM Body)
-> Binder SOACS Body -> InternaliseM Body
forall a b. (a -> b) -> a -> b
$ do
      SubExp
eq_class <- [SubExp] -> SubExp
forall a. [a] -> a
head ([SubExp] -> SubExp)
-> BinderT SOACS (State VNameSource) [SubExp]
-> BinderT SOACS (State VNameSource) SubExp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Body (Lore (BinderT SOACS (State VNameSource)))
-> BinderT SOACS (State VNameSource) [SubExp]
forall (m :: * -> *). MonadBinder m => Body (Lore m) -> m [SubExp]
bodyBind Body (Lore (BinderT SOACS (State VNameSource)))
Body
lam_body
      [SubExp] -> Body
forall lore. Bindable lore => [SubExp] -> Body lore
resultBody ([SubExp] -> Body)
-> BinderT SOACS (State VNameSource) [SubExp] -> Binder SOACS Body
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SubExp -> Int -> BinderT SOACS (State VNameSource) [SubExp]
forall (m :: * -> *).
(MonadBinder m, Bindable (Lore m)) =>
SubExp -> Int -> m [SubExp]
mkResult SubExp
eq_class Int
0