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

    import Data.Maybe (listToMaybe)
    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 SOACS], I.Body SOACS, [I.Type])

    internaliseFoldLambda ::
      InternaliseLambda ->
      E.Exp ->
      [I.Type] ->
      [I.Type] ->
      InternaliseM (I.Lambda SOACS)
    internaliseFoldLambda :: InternaliseLambda
    -> Exp -> [Type] -> [Type] -> InternaliseM (Lambda SOACS)
    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 u.
    TypeBase (ShapeBase SubExp) u -> TypeBase (ShapeBase SubExp) u
    I.rowType [Type]
    arrtypes
      ([Param Type]
    params, Body SOACS
    body, [Type]
    rettype) <- InternaliseLambda
    internaliseLambda Exp
    lam ([Type] -> InternaliseM ([LParam SOACS], Body SOACS, [Type]))
    -> [Type] -> InternaliseM ([LParam SOACS], Body SOACS, [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.
      [LParam (Rep InternaliseM)]
    -> InternaliseM Result -> InternaliseM (Lambda (Rep InternaliseM))
    forall (m :: * -> *).
    MonadBuilder m =>
    [LParam (Rep m)] -> m Result -> m (Lambda (Rep m))
    mkLambda [Param Type]
    [LParam (Rep InternaliseM)]
    params (InternaliseM Result -> InternaliseM (Lambda (Rep InternaliseM)))
    -> InternaliseM Result -> InternaliseM (Lambda (Rep InternaliseM))
    forall a b. (a -> b) -> a -> b
    $
        ErrorMsg SubExp -> [Type] -> Result -> InternaliseM Result
    ensureResultShape
          ([ErrorMsgPart SubExp] -> ErrorMsg SubExp
    forall a. [ErrorMsgPart a] -> ErrorMsg a
    ErrorMsg [Text -> ErrorMsgPart SubExp
    forall a. Text -> ErrorMsgPart a
    ErrorString Text
    "shape of result does not match shape of initial value"])
          [Type]
    rettype'
          (Result -> InternaliseM Result)
    -> InternaliseM Result -> InternaliseM Result
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
    =<< Body (Rep InternaliseM) -> InternaliseM Result
    forall (m :: * -> *). MonadBuilder m => Body (Rep m) -> m Result
    bodyBind Body (Rep InternaliseM)
    Body SOACS
    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 SOACS)
    internalisePartitionLambda :: InternaliseLambda
    -> Int -> Exp -> [SubExp] -> InternaliseM (Lambda SOACS)
    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)
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [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 u.
    TypeBase (ShapeBase SubExp) u -> TypeBase (ShapeBase SubExp) u
    I.rowType [Type]
    argtypes
      ([Param Type]
    params, Body SOACS
    body, [Type]
    _) <- InternaliseLambda
    internaliseLambda Exp
    lam [Type]
    rowtypes
      Body SOACS
    body' <-
        Scope SOACS
    -> InternaliseM (Body SOACS) -> InternaliseM (Body SOACS)
    forall a. Scope SOACS -> InternaliseM a -> InternaliseM a
    forall rep (m :: * -> *) a.
    LocalScope rep m =>
    Scope rep -> m a -> m a
    localScope ([Param Type] -> Scope SOACS
    forall rep dec. (LParamInfo rep ~ dec) => [Param dec] -> Scope rep
    scopeOfLParams [Param Type]
    params) (InternaliseM (Body SOACS) -> InternaliseM (Body SOACS))
    -> InternaliseM (Body SOACS) -> InternaliseM (Body SOACS)
    forall a b. (a -> b) -> a -> b
    $
          Body SOACS -> InternaliseM (Body SOACS)
    lambdaWithIncrement Body SOACS
    body
      Lambda SOACS -> InternaliseM (Lambda SOACS)
    forall a. a -> InternaliseM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Lambda SOACS -> InternaliseM (Lambda SOACS))
    -> Lambda SOACS -> InternaliseM (Lambda SOACS)
    forall a b. (a -> b) -> a -> b
    $ [LParam SOACS] -> [Type] -> Body SOACS -> Lambda SOACS
    forall rep. [LParam rep] -> [Type] -> Body rep -> Lambda rep
    I.Lambda [Param Type]
    [LParam SOACS]
    params [Type]
    forall {shape} {u}. [TypeBase shape u]
    rettype Body SOACS
    body'
      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 -> f [SubExp]
    mkResult SubExp
    _ Int
    i | Int
    i Int -> Int -> Bool
    forall a. Ord a => a -> a -> Bool
    >= Int
    k = [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
    $ Int -> [SubExp]
    result Int
    i
        mkResult SubExp
    eq_class Int
    i = do
          SubExp
    is_i <-
            String -> Exp (Rep f) -> f SubExp
    forall (m :: * -> *).
    MonadBuilder m =>
    String -> Exp (Rep m) -> m SubExp
    letSubExp String
    "is_i" (Exp (Rep f) -> f SubExp) -> Exp (Rep f) -> f SubExp
    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
    $
                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
          String -> Exp (Rep f) -> f [SubExp]
    forall (m :: * -> *).
    MonadBuilder m =>
    String -> Exp (Rep m) -> m [SubExp]
    letTupExp' String
    "part_res"
            (Exp (Rep f) -> f [SubExp]) -> f (Exp (Rep f)) -> f [SubExp]
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
    =<< f (Exp (Rep f))
    -> f (Body (Rep f)) -> f (Body (Rep f)) -> f (Exp (Rep f))
    forall (m :: * -> *).
    (MonadBuilder m, BranchType (Rep m) ~ ExtType) =>
    m (Exp (Rep m))
    -> m (Body (Rep m)) -> m (Body (Rep m)) -> m (Exp (Rep m))
    eIf
              (SubExp -> f (Exp (Rep f))
    forall (m :: * -> *). MonadBuilder m => SubExp -> m (Exp (Rep m))
    eSubExp SubExp
    is_i)
              (Body (Rep f) -> f (Body (Rep f))
    forall a. a -> f a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Body (Rep f) -> f (Body (Rep f)))
    -> Body (Rep f) -> f (Body (Rep f))
    forall a b. (a -> b) -> a -> b
    $ [SubExp] -> Body (Rep f)
    forall rep. Buildable rep => [SubExp] -> Body rep
    resultBody ([SubExp] -> Body (Rep f)) -> [SubExp] -> Body (Rep f)
    forall a b. (a -> b) -> a -> b
    $ Int -> [SubExp]
    result Int
    i)
              ([SubExp] -> Body (Rep f)
    forall rep. Buildable rep => [SubExp] -> Body rep
    resultBody ([SubExp] -> Body (Rep f)) -> f [SubExp] -> f (Body (Rep f))
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> SubExp -> Int -> f [SubExp]
    mkResult SubExp
    eq_class (Int
    i Int -> Int -> Int
    forall a. Num a => a -> a -> a
    + Int
    1))

        lambdaWithIncrement :: I.Body SOACS -> InternaliseM (I.Body SOACS)
        lambdaWithIncrement :: Body SOACS -> InternaliseM (Body SOACS)
    lambdaWithIncrement Body SOACS
    lam_body = Builder SOACS Result -> InternaliseM (Body SOACS)
    forall rep (m :: * -> *) somerep.
    (Buildable rep, MonadFreshNames m, HasScope somerep m,
     SameScope somerep rep) =>
    Builder rep Result -> m (Body rep)
    runBodyBuilder (Builder SOACS Result -> InternaliseM (Body SOACS))
    -> Builder SOACS Result -> InternaliseM (Body SOACS)
    forall a b. (a -> b) -> a -> b
    $ do
          SubExp
    eq_class <-
            SubExp -> (SubExpRes -> SubExp) -> Maybe SubExpRes -> SubExp
    forall b a. b -> (a -> b) -> Maybe a -> b
    maybe (IntType -> Integer -> SubExp
    intConst IntType
    Int64 Integer
    0) SubExpRes -> SubExp
    resSubExp (Maybe SubExpRes -> SubExp)
    -> (Result -> Maybe SubExpRes) -> Result -> SubExp
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Result -> Maybe SubExpRes
    forall a. [a] -> Maybe a
    listToMaybe (Result -> SubExp)
    -> Builder SOACS Result
    -> BuilderT SOACS (State VNameSource) SubExp
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Body (Rep (BuilderT SOACS (State VNameSource)))
    -> Builder SOACS Result
    forall (m :: * -> *). MonadBuilder m => Body (Rep m) -> m Result
    bodyBind Body (Rep (BuilderT SOACS (State VNameSource)))
    Body SOACS
    lam_body
          [SubExp] -> Result
    subExpsRes ([SubExp] -> Result)
    -> BuilderT SOACS (State VNameSource) [SubExp]
    -> Builder SOACS Result
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> SubExp -> Int -> BuilderT SOACS (State VNameSource) [SubExp]
    forall {f :: * -> *}.
    (MonadBuilder f, Buildable (Rep f)) =>
    SubExp -> Int -> f [SubExp]
    mkResult SubExp
    eq_class Int
    0